Test-Time Speculation

본 논문은 타겟 모델의 검증 신호를 활용하여 추론 중 드래프트 모델을 지속적으로 적응시키는 온라인 증류 방법인 테스트 시간 추측 (TTS) 을 소개함으로써 기존 추측기들의 긴 시퀀스에서의 성능 저하를 극복하고 수용 길이를 크게 향상시킵니다.

핵심

"테스트 시간 추측 (Test-Time Speculation)"에 대한 설명을 쉬운 언어와 창의적인 비유로 풀어냅니다.

큰 문제: 길을 잃는 "달리는 차"

당신이 천재이지만 생각이 느린 작가 (타겟 모델) 와 함께 매우 긴 이야기 (소설 같은) 를 쓰려고 한다고 상상해 보세요. 시간을 절약하기 위해, 작가가 읽기도 전에 다음 몇 문장을 미리 추측하도록 빠르고 에너지 넘치는 인턴 (드래프트 모델) 을 고용합니다.

AI 세계에서는 이를 "추측적 디코딩 (Speculative Decoding)"이라고 부릅니다. 인턴이 한 단락을 추측하면 작가가 빠르게 검토합니다. 인턴이 맞으면 작가는 "잘했다!"라고 말하고 넘어가며, 그 단어를 처음부터 쓰는 어려운 작업을 건너뜁니다. 인턴이 틀리면 작가는 멈추고 실수를 수정한 뒤 다시 시작해야 합니다.

문제점:
이 논문은 이러한 "인턴"들이 훈련되는 방식에 중대한 결함이 있음을 발견했습니다.

• 훈련: 인턴들은 짧은 이야기 (트윗이나 짧은 이메일 같은) 로 훈련받습니다. 그들은 200 단어짜리 문장에서 다음 단어를 추측하는 데 뛰어납니다.
• 현실: 실제 세계에서는 사람들이 AI 에게 수천 단어에 달하는 긴 보고서, 코드, 또는 이야기를 작성하도록 요청합니다.

이야기가 길어질수록 인턴은 혼란에 빠지기 시작합니다. 짧은 문장만으로 훈련받았기 때문에 텍스트가 길어질수록 "생각의 흐름"을 잃게 됩니다. 그들은 긴 맥락에 맞지 않는 단어를 추측하기 시작합니다.

• 결과: 작가는 인턴의 추측을 거의 모두 거부해야 합니다. 시간을 절약하는 대신, 작가가 인턴을 수정하느라 끊임없이 멈추기 때문에 과정이 느려집니다. 논문은 이를 "수용 길이 (Acceptance Length)"가 거의 1 로 떨어지는 현상 (인턴이 사실상 쓸모없다는 뜻) 이라고 부릅니다.

해결책: "테스트 시간 추측 (TTS)"

저자들은 **테스트 시간 추측 (Test-Time Speculation, TTS)**이라는 영리한 해결책을 제안합니다. 매번 새로운 인턴을 고용하는 대신, 같은 인턴에게 일하는 동안 적응하는 법을 가르치는 것입니다.

비유: 실시간 코칭 세션
인턴이 이야기를 쓰고 작가가 이를 검토한다고 상상해 보세요.

1. 구식 방식: 인턴이 10 단어를 추측하면 작가가 이를 검토합니다. 틀리면 작가가 수정하고 넘어갑니다. 인턴은 다음 문장을 위해 도움이 될 만한 방식으로 왜 틀렸는지 알려주지 않기 때문에 실수로부터 아무것도 배우지 못합니다.
2. TTS 방식: 작가가 인턴의 작업을 검토할 때마다 단순히 "맞다" 또는 "틀리다"라고 말하지 않습니다. 작가는 그 순간을 인턴에게 미니 레슨을 제공하는 기회로 활용합니다.
   • 작가는 말합니다. "너는 '고양이'라고 추측했지만, 이 특정 긴 이야기에서는 '개'라는 단어가 맞아야 해. 내가 사용한 정확한 확률 분포를 보여줄게."
   • 인턴은 즉각 이 구체적인 교훈을 바탕으로 뇌 (내부 수학) 를 업데이트합니다.
   • 이제 인턴이 다음 단어 세트를 추측할 때, 조금 더 똑똑해지고 작가의 현재 기분과 이야기의 긴 역사에 더 잘 맞춰집니다.

왜 이것이 특별한가요?
보통 모델을 더 좋게 만들려면 며칠 동안 멈추고 재훈련해야 합니다. TTS 는 이야기가 쓰이는 동안 이를 즉시 수행합니다. 작가가 어차피 해야 하는 "검증" 단계를 무료 훈련 신호로 활용하는 것입니다. 이는 교사와 대화하며 새로운 언어를 배우는 학생과 같습니다. 교사가 실시간으로 수정해 주므로 대화 끝에는 유창해집니다.

결과: 더 길어질수록 빨라짐

이 논문은 수학 문제 해결, 코드 작성, 과학 질문 답변과 같은 어려운 작업에서 다양한 유형의 "작가 (AI 모델)"와 "인턴 (추측기)"을 대상으로 이를 테스트했습니다.

• 개선: TTS 를 사용하면 "인턴"들이 이야기가 길어질수록 올바른 단어를 추측하는 능력이 훨씬 향상되었습니다.
• 수치: 평균적으로 시스템이 인턴의 추측을 41% 더 수용했습니다. 어떤 경우에는 이전 최선 방법보다 72% 더 뛰어났습니다.
• 추세: 텍스트가 길어질수록 TTS 는 더 잘 작동합니다. 다른 방법들은 수천 단어 이후에 실패하는 반면, TTS 는 인턴이 실시간으로 계속 학습하고 적응하기 때문에 생성이 계속됨에 따라 실제로 더 정확해집니다.

요약

이전 방법들은 100 미터 달리기만 잘하는 빠른 주자를 고용하는 것과 같습니다. 마라톤을 달리게 하면 그들은 쓰러집니다.

테스트 시간 추측은 그 주자 옆을 따라가며 매 단계마다 수정 사항과 전략 조정을 속삭여 주는 코치를 제공하는 것과 같습니다. 주자는 덜 지치고 올바른 길에 머무르며, 전체 팀은 마라톤을 훨씬 더 빠르게 완주합니다.

이 논문은 생성 과정에서 AI 가 "직무 중 학습"을 하도록 함으로써 매우 긴 문서를 작성할 때조차 AI 를 빠르고 효율적으로 유지할 수 있음을 증명합니다.

기술 요약

기술 요약: 테스트 시간 추측 (TTS)

1. 문제 정의

이 논문은 현재 최첨단 추측 디코딩 방법들 (DFlash, EAGLE-3, PARD 등) 이 긴 응답 작업에 적용될 때 직면하는 치명적인 한계를 지적합니다. 추측 디코딩은 빠른 '초안 (draft)' 모델을 사용하여 토큰을 생성하고 느린 '타겟' 모델이 이를 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화하지만, 그 효율성은 주로 라운드당 타겟 모델이 수용하는 연속된 초안 토큰 수인 **수용 길이 (acceptance length)**에 크게 의존합니다.

저자들은 기존 추측기들의 수용 길이가 생성 길이가 증가함에 따라 현저히 저하된다는 것을 관찰했습니다. 단 몇 천 개의 출력 토큰 내에서 수용 길이는 종종 1 에 가까운 값 (예: 1.1) 으로 떨어지며, 이는 사실상 속도 향상 효과를 완전히 무효화합니다. 이러한 저하가 발생하는 이유는 최첨단 추측기들이 일반적으로 2K 토큰 이하의 짧은 시퀀스에 오프라인으로 훈련되어 있기 때문입니다. 이로 인해 추론 시 훨씬 더 긴 시퀀스 (예: 20K~32K 토큰) 에서 타겟 모델을 근사하도록 강요받을 때 분포 불일치가 발생합니다. 생성이 진행됨에 따라 초안 모델의 예측은 타겟 모델의 점점 더 확신에 찬 분포와 멀어지게 되어 빈번한 거절이 발생합니다.

2. 방법론: 테스트 시간 추측 (TTS)

이러한 분포 불일치를 해결하기 위해 저자들은 추론 과정에서 실시간으로 초안 모델을 적응시키는 온라인 증류 접근법인 **테스트 시간 추측 (TTS)**을 제안합니다.

핵심 통찰

핵심적인 깨달음은 표준 추측 디코딩 검증 단계가 추가 비용 없이 적응에 필요한 감독 신호를 이미 생성한다는 점입니다. 매 라운드마다 타겟 모델은 초안 토큰에 대한 전체 확률 분포를 계산합니다. TTS 는 이를 활용하여 다음을 처리합니다:

• 타겟 모델을 "교사 (Teacher)"로 간주합니다.
• 초안 모델을 "학생 (Student)"으로 간주합니다.
• 검증된 초안 토큰을 증류 훈련 샘플로 간주합니다.

알고리즘

TTS 는 생성과 모델 업데이트를 교차하여 수행합니다. 각 추측 라운드의 과정은 다음과 같습니다:

1. 초안 작성 (Drafting): 현재 초안 모델 ($q_t$) 이 $C$개의 토큰으로 구성된 캔버스를 생성합니다.
2. 검증 (Verification): 타겟 모델 ($p$) 이 표준 거절 샘플링을 통해 수용 길이 ($\tau$) 를 결정하면서 단일 순전파로 캔버스를 평가합니다.
3. 증류 손실 (Distillation Loss): 다음 라운드 전에 초안 모델이 증류 손실 함수에 대한 단일 경사 하강 단계를 사용하여 업데이트됩니다:
   $$L_t(q) = \tilde{KL}(p \parallel q) + \lambda \tilde{KL}(q_t \parallel q)$$
   • 첫 번째 항은 캔버스에 대한 타겟 분포와 새로운 초안 분포 간의 쿨백 - 라이블러 (KL) 발산을 근사합니다.
   • 두 번째 항은 초안이 이전 상태 ($q_t$) 에서 너무 멀어지는 것을 방지하는 정규화 구성 요소입니다.
   • 위치 의존적 가중치 ($w_k$) 가 적용되어 캔버스 내의 초기 토큰을 우선시합니다.
4. 업데이트: 초안 모델 파라미터가 업데이트됩니다 ($q_{t+\tau} \leftarrow q_t - \eta \nabla L_t$).

시스템 최적화

개선된 수용 길이와 경사 업데이트의 지연 오버헤드 사이의 균형을 관리하기 위해 TTS 는 다음을 활용합니다:

• 스트라이드 업데이트 (Strided Updates): 매 라운드마다가 아닌 $S$라운드마다 경사 업데이트를 수행하여 계산 비용을 분산합니다.
• 비동기 파이프라이닝 (Asynchronous Pipelining): 업데이트를 후속 $S-1$개의 생성 라운드와 병렬로 실행되는 전용 CUDA 스트림으로 오프로드하여 임계 경로에서의 지연을 숨깁니다.

3. 주요 기여

1. 저하 진단: 저자들은 초안의 짧은 시퀀스 훈련 분포와 긴 시퀀스 추론 분포 간의 불일치로 인해 현재 추측기들의 효과가 생성 길이에 따라 저하됨을 입증합니다.
2. TTS 프레임워크: 검증 단계를 감독 신호로 활용하여 오프라인 재훈련 없이 추론 시 초안 모델을 적응시키는 온라인 증류 방법인 테스트 시간 추측을 제안합니다.
3. 종합 평가: 이 방법은 Qwen-3, Qwen-3.5, Llama3.1 계열을 포함한 다섯 개의 최첨단 모델과 AIME, LiveCodeBench, GPQA 를 포함한 여덟 개의 다양한 벤치마크에서 평가되어 일관된 개선을 보여줍니다.
4. 시스템 통합: 저자들은 커널 미분과 CUDA 그래프 동기화와 같은 시스템 수준의 과제를 해결하기 위해 SGLang 추론 프레임워크 내에서 TTS 를 구현합니다.

4. 실험 결과

• 수용 길이 개선: TTS 는 DFlash 대비 평균 수용 길이를 최대 72%, 평균 41% 개선하며, EAGLE-3 대비 최대 67% (평균 34%) 개선합니다.
• 길이에 따른 확장: TTS 의 이점은 생성 길이에 따라 확장됩니다. 예를 들어, AIME 2024 데이터셋에서 DFlash 대비 개선 폭은 처음 010K 토큰 구간에서 15% 에서 2030K 토큰 구간에서는 **183%**로 증가합니다.
• 처리량: 빈번한 업데이트 (스트라이드 $S=1$) 가 수용 길이를 최대화하지만, 스트라이드 $S=5$는 적응 빈도와 업데이트 오버헤드를 균형 있게 조절하여 DFlash 대비 최대 1.71$\times$의 처리량 속도 향상을 달성합니다.
• 일반화: TTS 는 다양한 모델 크기 (4B~122B) 및 아키텍처 (Dense 및 MoE) 에서 효과적이며, 특히 짧은 컨텍스트 (예: 2K 컨텍스트의 EAGLE-3) 에서 훈련된 추측기를 훨씬 더 큰 컨텍스트 창을 가진 타겟에 적용할 때 이를 보완합니다.

5. 중요성과 주장

이 논문은 TTS 가 실제 세계의 긴 응답 시나리오에서 추측 디코딩의 한계를 근본적으로 해결한다고 주장합니다. 생성 과정 중에 초안 모델을 적응시킴으로써 TTS 는 훈련과 추론 분포 간의 격차를 해소하여, 수만 개의 토큰에 걸친 출력에서도 추측 디코딩이 효과적으로 유지되도록 보장합니다.

저자들은 TTS 가 도메인별 버퍼에 의존하는 이전 온라인 방법과 달리 요청 스트림 구조에 대한 가정을 필요로 하지 않으며, 기존 공개 최첨단 추측기 위에서 직접 작동한다고 강조합니다. 이는 긴 형식의 생성 (예: 코드, 추론, 콘텐츠 제작) 이 주류인 생산 환경에서 높은 추론 처리량을 유지하기 위한 실용적인 솔루션입니다. 이 연구는 LLM 애플리케이션이 더 긴 컨텍스트 창으로 이동함에 따라 추측 디코딩을 실행 가능하게 유지하기 위한 필수적인 진전으로 제시됩니다.