[Paper] Self-Distillation을 통한 멀티 토큰 예측

Source: arXiv - 2602.06019v1

Overview

연구팀은 사전 학습된 자동 회귀 언어 모델(예: GPT‑2, LLaMA)을 다중 토큰 예측기로 전환하는 놀라울 정도로 간단한 방법을 제안한다. 이 방법은 별도의 “예측자” 모델을 추가하거나 복잡한 디코딩 파이프라인을 구축하지 않고, 온라인 자체 증류 목표를 사용해 원래 모델이 한 번의 순전파로 여러 미래 토큰을 예측하도록 학습한다. 그 결과, 모델의 아키텍처나 배포 코드를 변경하지 않으면서도 추론 속도가 최대 3배 빨라지고, 수학 추론 벤치마크(GSM8K)에서 정확도는 5 % 미만 감소한다.

핵심 기여

• 다중 토큰 예측을 위한 자체 증류: 사전 학습된 모델이 단일 다음 토큰이 아니라 짧은 토큰 시퀀스를 출력하도록 가르치는 온라인 증류 손실을 도입합니다.
• 무변경 배포: 최종 모델은 원래 단일 토큰 모델과 동일한 체크포인트와 추론 코드를 사용합니다—보조 검증기, 예측기, 혹은 맞춤 런타임이 필요 없습니다.
• 경험적 속도‑정밀도 트레이드‑오프: GSM8K에서 5 % 미만의 상대 정확도 손실로 3배 이상의 디코딩 속도 향상을 보여주며, 추측 디코딩과 순수 단일 토큰 생성 사이의 격차를 메웁니다.
• 광범위한 적용 가능성: 크기나 사전 학습 데이터와 무관하게 모든 자동 회귀 언어 모델에서 작동하므로 기존 서비스에 바로 적용할 수 있는 업그레이드입니다.

방법론

1. Baseline model – 동결된 사전학습된 자동회귀 언어 모델(예: 디코더‑전용 트랜스포머)부터 시작합니다.
2. Online teacher – 학습 중에 동일한 모델이 일반적인 단일 토큰 모드로 실행되어 다음 k 토큰에 대한 “교사” 예측을 생성합니다 (k는 4 또는 8과 같은 작은 정수).
3. Student head – 최종 은닉 상태에 가벼운 추가 헤드를 연결하여 k 토큰을 한 번에 직접 출력하도록 학습합니다.
4. Distillation loss – 학생의 로짓이 모든 k 위치에서 교사의 로짓(교차 엔트로피)과 차이날 경우 패널티를 부여하고, 첫 번째 토큰에 대한 원래 언어 모델링 손실은 유지합니다.
5. Curriculum – 학습이 진행됨에 따라 k 값을 점진적으로 증가시켜 모델이 불안정해지지 않으면서 더 긴 예측 범위에 적응하도록 합니다.
6. Inference – 실행 시 모델은 새로운 헤드를 호출해 k 토큰을 출력하고, 컨텍스트를 앞으로 이동시켜 반복합니다. 이를 통해 다중 토큰 추측이 올바른지 확인하는 별도의 검증기가 필요 없게 됩니다.

교사와 학생이 동일한 네트워크이므로 이 과정은 self‑distilling이며 온라인으로 수행할 수 있습니다(별도의 교사 모델이나 데이터셋 생성 단계가 필요 없음).

결과 및 발견

• 속도: 디코딩 속도는 k에 따라 대략 선형적으로 증가하지만, k = 8을 초과하면 정확도 감소가 더 뚜렷해집니다.
• 품질: 약간의 정확도 손실은 초기 토큰 오류가 다중‑토큰 블록을 통해 전파되는 “드리프트” 현상 때문이며, 대부분의 오류는 이후 블록에서 복구될 수 있습니다.
• 호환성: 동일한 접근 방식을 GPT‑2‑medium 및 1.3 B LLaMA 체크포인트에 적용했으며, 아키텍처 변경 없이 일관된 속도 향상을 확인했습니다.

실용적 시사점

• 더 빠른 API: 클라우드 제공자는 자체 증류 손실을 사용해 모델을 미세 조정함으로써 기존 LLM 엔드포인트의 처리량을 향상시킬 수 있습니다—추가 서버나 맞춤형 추론 커널이 필요 없습니다.
• 비용 절감: 생성된 토큰당 전방 패스 횟수를 줄이면 GPU/TPU 사용량이 직접 감소하여, 챗봇, 코드 자동완성 등 대량 애플리케이션의 추론 비용을 낮출 수 있습니다.
• 엣지 배포: 컴퓨팅 자원이 제한된 디바이스(모바일, IoT)에서도 추론 단계당 여러 토큰을 출력함으로써 더 큰 모델을 보다 빠르게 실행할 수 있어, 디바이스 내에서 사용할 수 있는 모델 크기의 한계를 확장합니다.
• 간소화된 파이프라인: 추측 디코딩과 달리 별도의 “검증자” 모델을 유지하거나 추측‑후‑백업 로직을 조정할 필요가 없어 엔지니어링 오버헤드와 잠재적 버그를 줄일 수 있습니다.

제한 사항 및 향후 연구

• 오류 전파: 모델이 블록 내 초기 토큰을 잘못 예측하면 그 실수가 블록의 나머지 부분에 영향을 미쳐 가끔 낮은 품질의 출력이 급증할 수 있습니다.
• 고정 블록 크기: 현재 방법은 정적인 k를 사용합니다; 신뢰도에 기반한 적응형 블록 길이는 정확도와 속도 간의 균형을 개선할 수 있습니다.
• 벤치마크 범위: 실험은 GSM8K(수학 추론)에 초점을 맞추고 있습니다; 일반성을 확인하려면 대화, 코드 생성, 장문 텍스트 등 더 폭넓은 평가가 필요합니다.
• 학습 오버헤드: 자체 증류 미세조정은 순수 추론 전용 배포에 비해 추가 연산을 요구하지만, 사전 학습 비용에 비하면 그 정도는 적습니다.

향후 연구 방향으로는 블록 내에서 신뢰도 기반 조기 중단을 통합하고, 자체 증류를 양자화 또는 프루닝과 결합하여 더 엄격한 지연 시간 예산을 달성하며, 멀티모달 확장(예: 비전‑언어 모델) 탐색이 포함됩니다.

저자

• John Kirchenbauer
• Abhimanyu Hans
• Brian Bartoldson
• Micah Goldblum
• Ashwinee Panda
• Tom Goldstein

논문 정보

• arXiv ID: 2602.06019v1
• 분류: cs.CL, cs.LG
• 출판일: 2026년 2월 5일
• PDF: PDF 다운로드