CACTUS: 제약 조건 기반 수락 스펙큘레이티브 샘플링을 통한 자동 회귀 디코딩 가속화

대형 언어 모델의 추론 속도를 높이기 위해 작은 모델을 활용하는 Speculative Sampling 기법의 한계를 극복했습니다. 기존 방식이 속도를 위해 모델의 원래 답변 분포를 왜곡하던 문제를 수학적 최적화로 해결하여, 품질 저하 없이 더 빠른 추론이 가능함을 입증했습니다.

Transformer 기반 LLM은 토큰을 하나씩 생성할 때마다 수천억 개의 파라미터를 메모리에서 읽어야 하므로 속도가 매우 느리다. 이를 해결하기 위해 작은 '초안 모델(Draft Model)'이 여러 토큰을 먼저 제안하고 큰 '검증 모델(Verifier)'이 이를 한 번에 승인하는 Speculative Sampling이 사용된다. 하지만 기존 방식은 검증 모델의 확률 분포를 완벽하게 따라야 한다는 제약 때문에 정답에 가깝지만 확률이 조금 낮은 토큰들을 불필요하게 거절하여 속도 향상에 한계가 있었다.

Cactus는 이 문제를 '분포의 차이(Divergence)'를 일정 수준 이하로 유지하면서 '수락 확률'을 가장 높이는 수학적 최적화 문제로 풀었다. 즉, 검증 모델의 원래 의도를 크게 벗어나지 않는 선에서 초안 모델이 제안한 토큰을 최대한 더 많이 받아들이도록 설계했다. 이는 마치 엄격한 검토관이 핵심 내용만 틀리지 않는다면 사소한 표현 차이는 눈감아주어 업무 속도를 높이는 것과 유사한 원리이다.

결과적으로 Cactus는 검증 모델의 지식과 논리력을 그대로 유지하면서도, 초안 모델과의 협업 효율을 극대화하여 전체적인 생성 속도를 획기적으로 높였다. 특히 복잡한 추론이 필요한 작업에서 기존 방식들이 겪던 성능 저하 문제 없이 안정적인 가속을 제공한다.

Speculative Sampling을 제약 최적화 프레임워크로 공식화했다. 목표 함수는 초안 모델 p가 제안한 토큰 n에 대한 수락률 ϕ(n)을 극대화하는 것이며, 제약 조건은 목표 분포 h와 검증 모델 분포 q 사이의 f-divergence(Df(h||q))를 하이퍼파라미터 δ 이하로 유지하는 것이다.

최적의 목표 분포 h를 찾기 위해 f-divergence로 KL Divergence를 선택했다. [입력값으로 검증 모델의 토큰 확률 q(n)과 허용 오차 δ를 입력하여] → [2차 테일러 전개를 통해 근사 해를 계산하여] → [수정된 목표 확률 γ를 얻고] → [이 값이 초안 모델의 토큰을 수락할지 결정하는 기준이 된다]. 구체적으로 γ* = min{q(n) + sqrt(2δq(n)(1-q(n))), 1}로 계산된다.

알고리즘 구현 시 검증 모델의 전체 어휘 사전에 접근할 필요 없이 초안 모델이 제안한 특정 토큰의 확률값만 읽으면 되므로 메모리 접근 오버헤드를 최소화했다. 이는 대규모 어휘 사전을 가진 최신 LLM 환경에서 효율적인 추론을 가능하게 한다.

Qwen 3 8B 및 14B 모델을 검증 모델로 사용한 실험에서 Cactus는 기존 Speculative Sampling(SpS) 대비 유의미한 가속을 달성했다. GSM8K 벤치마크에서 SpS의 평균 수락 길이(AL)가 4.49일 때 Cactus는 5.72까지 높였으며, 이는 약 37%의 거절 토큰 감소 효과를 의미한다.

품질 측면에서 Cactus는 Typical Acceptance Sampling(TAS)과 달리 성능 저하가 거의 없었다. 특히 난이도가 높은 GPQA 벤치마크에서 TAS는 정확도가 42.93에서 38.89로 급락한 반면, Cactus는 40.01~45.46의 점수를 기록하며 검증 모델의 성능을 보존하거나 오히려 향상시키는 결과를 보였다.

실제 실행 시간(Wall-time) 측정 결과, A100 GPU 환경에서 vLLM 프레임워크를 사용할 때 검증 모델 단독 추론 대비 최대 1.9배의 속도 향상을 기록했다. 이는 동일한 하드웨어 자원으로 약 2배 더 많은 텍스트를 생성할 수 있음을 시사한다.

Cactus는 f-divergence 제약 하에서 수락률을 최적화하는 닫힌 형태의 해(closed-form solution)를 제공한다. 이론적으로 Cactus의 유효 분포 halg와 검증 모델 분포 q 사이의 거리는 Df(halg||q) ≤ min{Γ(δ), Df(p||q)}를 만족하며, 여기서 Γ(δ)는 δ에 대해 연속적이고 비감소하는 함수임을 증명했다.

기존의 TAS(Typical Acceptance Sampling)가 크로스 엔트로피를 제약 조건으로 사용하여 분포가 결정론적으로 붕괴(collapse)될 위험이 있는 것과 달리, Cactus는 KL Divergence를 직접 제어함으로써 검증 모델의 엔트로피와 다양성을 보존한다. 구현 측면에서는 Taylor 근사를 통해 초월 함수 계산 부담을 줄여 런타임 오버헤드를 무시할 수 있는 수준으로 낮췄다.

실험에서는 Qwen 3, Gemma 2, DeepSeek R1, Llama 3.1 등 다양한 모델 아키텍처와 크기 조합에 대해 일관된 성능 향상을 확인했다. 이는 Cactus가 특정 모델의 특성에 의존하지 않는 범용적인 디코딩 가속 기법임을 뒷받침한다.

본 논문은 수락 규칙의 최적화에 집중하고 있으며, 초안 모델의 품질 자체를 개선하거나 KV 캐시 관리와 같은 시스템 레벨의 최적화와 결합했을 때의 시너지 효과에 대해서는 심도 있게 다루지 않았다. 또한 하이퍼파라미터 δ의 최적값이 작업의 종류나 모델 조합에 따라 달라질 수 있다는 점이 명시되었다.