특수 토큰 하나로 3배 가속: LLM에 멀티 토큰 예측 내장

연구 개요

에이전틱(Agentic) AI 워크플로가 긴 추론 사슬로 인한 비용·지연을 키우는 가운데, 메릴랜드대·로렌스 리버모어 국립연구소·컬럼비아대·TogetherAI 연구팀은 모델 가중치에 직접 3배 스루풋 향상을 ‘구워 넣는’ 방법을 제안했다. 스펙큘러티브 디코딩처럼 별도 초안 모델을 띄우지 않고, 기존 임베딩 행렬의 미사용 슬롯을 라는 특수 토큰으로 재활용해 병렬 생성을 가능하게 한다.

왜 다음 토큰 예측이 한계인가

다음 토큰 예측(NTP)은 한 번의 포워드 패스에 한 토큰씩 생성해, 수천 토큰을 뽑을 때 스루풋 상한과 비용 병목이 커진다. 특히 사고의 연쇄(CoT)를 길게 펼치는 추론형 모델은 응답 전 수천 토큰을 생성해 체감 지연시간이 크다. 연구팀은 전체 스루풋 최적화 대신 “단일 사용자의 지연시간”을 줄이는 데 초점을 두어, 배치 없이도 GPU를 가득 활용하도록 설계를 전환했다.

제안: 멀티 토큰 예측(MTP)

멀티 토큰 예측(MTP)은 한 번의 패스로 여러 토큰 블록을 동시에 예측하도록 훈련하는 대안 패러다임이다. 다만 표준 지도학습으로는 각 위치의 확률만 따로 맞추어 문법 불일치(예: ‘lion bamboo’)나 원거리 반복(‘...the the the...’)이 생기기 쉽다.

자기 증류 기반 MTP 훈련

연구팀은 학생-교사(self-distillation) 스킴을 도입했다. 학생 모델이 결정론적으로 멀티 토큰 블록을 제안하면, 강한 NTP 교사 모델이 그 블록의 가능도·응집도를 평가해 손실을 부여한다. 이 온폴리시 강화학습 유사 피드백은 고정 데이터가 아닌 학생의 실제 출력에 기반해 동적으로 주어지므로, 문법 불일치와 퇴행적 반복을 함께 억제한다.

스펙큘러티브 디코딩과의 차이

스펙큘러티브 디코딩이나 확산 기반 LLM도 지연 단축에 유용하지만, 초안 모델 운영·검증에 추가 인프라와 연산이 든다. 반면 멀티 토큰 예측(MTP)은 특수 토큰 추가만으로 서빙이 단순하며, MoE·윈도우드 어텐션·SSM 레이어 등 내부 구현을 건드리지 않는다.

추론 최적화: ConfAdapt

멀티 토큰 예측이 항상 정확한 것은 아니므로, 연구팀은 적응형 디코딩 전략 ConfAdapt를 제안했다. 각 스텝에서 90% 같은 신뢰도 임계값을 평가하고, 생성된 블록 중 고신뢰 토큰만 수용한다. 구조적이거나 예측 가능성이 높은 구간은 한 번에 큰 덩어리를 받아 대폭 가속하고, 어려운 구간은 단일 토큰 패스로 집중 처리해 품질을 지킨다.

실험 결과와 전이 성능

Llama‑3.1‑8B‑Magpie와 Qwen3‑4B‑Instruct‑2507에 MetaMathQA로 MTP 튜닝을 적용한 결과, ConfAdapt 설정에서 약 3배 속도를 달성했다. 정확도 저하는 Llama가 3% 미만, Qwen이 약 7%였다. 더 공격적으로는 5배 가속도 가능하나 손실이 커진다. 가속 이득은 예측 가능성이 높은 업무에서 특히 크며, 수학·추론뿐 아니라 창작·요약 같은 개방형 작업으로도 전이되었다.

배포와 통합

아키텍처 변경 사항

건드리는 것은 특수 토큰() 추가뿐이다. 멀티 토큰 예측(MTP)을 적용해도 내부 모듈(MoE, 윈도우드 어텐션, SSM 등)은 그대로여서 기존 모델을 재빌드하지 않고도 적응시킬 수 있다.

시스템 통합 고려사항

vLLM·SGLang 통합 시 배칭 전략과 KV 캐시 처리 방식에 약간의 조정이 필요하지만 일회성 투자에 가깝다. 연구팀은 허깅페이스에 튜닝된 모델을 공개했으며, MTP 프레임워크 코드도 곧 공개 예정이다.

권장 운영 전략

도메인 전이가 관찰되었지만, 특수 산업 도메인에는 해당 데이터의 프롬프트로 멀티 토큰 예측(MTP) 재튜닝을 권장한다. 초기 검증은 숫자 세기나 구문 반복 같은 토이 프롬프트로 ConfAdapt의 가속 효과를 눈으로 확인한 뒤, 배포 도메인 샘플로 미세 조정하는 것이 좋다.

실제 적용 예시

멀티 토큰 예측 도입 전 필수 체크리스트

목표 지연시간·정확도 허용오차 정의, 대상 작업의 예측 가능성(엔트로피) 분석, 배치·KV 캐시 정책 점검, 안전장치(고신뢰 임계값·최대 블록 길이) 설정, 도메인 프롬프트로 소규모 MTP 적응 튜닝 계획 수립.

멀티 토큰 예측 적용 프로세스 단계별 안내

(1) 기존 모델에 토큰 추가 → (2) 자기 증류로 MTP 튜닝 → (3) ConfAdapt 임계값 스윕으로 속도/정확도 스위트스팟 탐색 → (4) vLLM·SGLang에서 배치·KV 캐시 조정 → (5) 실서비스 트래픽 섀도우 테스트 후 점진 배포.