Self-Distillation for Multi-Token Prediction

이 논문은 다중 토큰 예측 (MTP) 의 추론 효율성을 높이기 위해 제안된 자기 증류 기법 (MTP-D) 과 루프 확장 전략을 통해 MTP 헤드의 수용률을 크게 향상시키고 추론 속도를 획기적으로 개선하는 방법을 제시합니다.

핵심

🚀 문제: "한 번에 한 마디만 하는 게 너무 느려!"

지금까지의 거대 언어 모델 (LLM) 은 매우 똑똑하지만, 말을 할 때 '한 번에 한 글자 (또는 단어) 만' 내뱉는 방식을 썼습니다.

• 비유: 친구가 "안녕하세요, 오늘 날씨가 정말 좋네요"라고 말하려고 할 때, "안" -> "녕" -> "하" -> "세" -> "요" 순서로 한 글자씩 끊어서 말하는 것과 같습니다.
• 문제점: 글이 길어질수록 이 과정이 너무 오래 걸려서, 실제로 쓸 때 답답합니다.

💡 기존 해결책: "여러 명이서 미리 예측해 보기 (MTP)"

연구자들은 "그럼 한 번에 여러 글자를 미리 예측해서 동시에 말해보자!"라고 생각했습니다. 이를 **다중 토큰 예측 (MTP)**이라고 합니다.

• 비유: 친구가 "안녕하세요"를 말할 때, "안녕"을 말한 순간에 "하세요"까지 미리 입에 담고 있는 것처럼, 한 번에 4~5 글자씩 미리 내뱉는 것입니다.
• 하지만: 이 방식에는 큰 문제가 있었습니다.
  1. 예측이 잘 안 맞음: 미리 내뱉은 글자가 틀릴 확률이 높아서, 매번 "아, 틀렸네" 하고 다시 처음부터 확인해야 했습니다. (수용률 낮음)
  2. 훈련이 어려움: 본인이 (메인 두뇌) 는 잘하는데, 미리 예측하는 부대신들 (MTP 헤드) 은 너무 서툴러서 함께 훈련시키기 힘들었습니다.

✨ 이 논문의 해결책: "MTP-D (스스로 가르치는 속기왕)"

이 논문은 두 가지 핵심 아이디어로 이 문제를 해결했습니다.

1. "스스로를 스승으로 삼아 가르치기 (자기 증류)"

기존 방식은 부대신들이 혼자서 헛다리를 짚는 경우가 많았습니다. 그래서 본인 (메인 두뇌) 이 부대신들에게 "내가 생각한 가장 유력한 답 10,000 가지만 골라봐"라고 가르쳐 준 것입니다.

• 비유:
  • 기존: 부하 직원들이 "아마 이거겠지?"라고 막연히 추측하다가 자주 틀림.
  • 이 방법: 상사 (메인 두뇌) 가 "이 10,000 개 중에 정답이 있을 거야"라고 핵심 힌트만 줌. 부하 직원들은 그 힌트만 보고 집중해서 훈련함.
  • 결과: 부하 직원들이 훨씬 똑똑해져서, 상사의 말을 거의 그대로 따라갈 수 있게 됨 (수용률 7.5% 상승).

2. "레고 블록처럼 반복해서 늘리기 (루프 확장)"

부하 직원들이 4 명만 있는 게 아쉽다면, 이미 훈련된 4 명을 복사해서 8 명, 16 명으로 늘리는 방법을 제안했습니다.

• 비유:
  • 처음에 4 명을 훈련시켰다면, 그 4 명의 능력을 그대로 복사해서 새로운 4 명을 데려와서 "너희는 이 능력을 바탕으로 더 발전해"라고 계속 훈련시킴.
  • 결과: 훈련 비용은 거의 들지 않으면서, 한 번에 내뱉을 수 있는 글자 수를 4 개에서 16 개까지 늘릴 수 있게 됨.

🏆 성과: "속도가 2 배 이상 빨라졌다!"

이 방법을 적용한 결과:

• 정확도 유지: 본래의 똑똑함 (메인 두뇌) 은 그대로 유지하면서, 부하 직원들의 실력이 대폭 향상됨.
• 속도 향상:
  • 부하 직원이 4 명일 때: 약 23% 더 빠름.
  • 부하 직원을 16 명까지 늘렸을 때: 약 220% 더 빠름 (기존보다 3 배 이상).
  • 즉, 같은 작업을 하더라도 3 배 빨리 끝낼 수 있게 된 것입니다.

📝 한 줄 요약

**"거대 AI 가 말을 할 때, 한 번에 한 글자씩 끄집어내는 대신, 가장 유력한 답만 골라 미리 16 개씩 내뱉도록 훈련시켜 속도를 3 배로 높인 방법"**을 개발했습니다.

이 기술은 앞으로 우리가 AI 와 대화할 때, 기다림 없이 훨씬 더 빠르게 답변을 받을 수 있게 해줄 것입니다.

기술 요약

논문 요약: Self-Distillation for Multi-Token Prediction (MTP-D)

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 이 확장됨에 따라 추론 효율성은 주요 병목 현상이 되었습니다. 기존의 **Next-Token Prediction **(NTP) 방식은 토큰 단위로 순차적으로 생성하므로 지연 시간 (latency) 과 계산 비용이 높습니다. 이를 해결하기 위해 **Multi-Token Prediction **(MTP)이 제안되었으나, 기존 MTP 접근법 (예: DeepSeek-V3 의 캐스케이드 아키텍처) 은 다음과 같은 두 가지 주요 과제를 안고 있습니다.

1. **MTP 헤드의 낮은 수용률 **(Limited Acceptance Rates) MTP 헤드가 예측한 토큰이 메인 헤드에 의해 검증될 때 거절되는 비율이 높습니다. 특히 여러 개의 MTP 헤드를 사용할 경우, 개별 헤드의 수용률이 낮으면 누적 수용률 (Cumulative Acceptance Rate) 이 기하급수적으로 감소하여 추론 가속화 효과가 크게 떨어집니다.
2. **다중 헤드의 공동 학습 어려움 **(Difficulty in Joint Training) 메인 헤드의 성능을 유지하면서 여러 MTP 헤드를 동시에 학습시키는 것은 '씨새우 효과 (seesaw effect)'로 인해 어렵습니다. MTP 헤드를 추가할수록 메인 헤드의 성능이 저하되거나 학습이 불안정해지는 문제가 발생합니다.

2. 방법론 (Methodology)

저자들은 MTP-D라는 새로운 자기 증류 (Self-Distillation) 프레임워크와 이를 확장하는 Looped Extension 전략을 제안했습니다.

2.1 MTP-D: 자기 증류 기반 학습

메인 헤드의 성능을 해치지 않으면서 MTP 헤드의 수용률을 높이기 위해 Pre-training 단계에서 다음과 같은 기법을 적용합니다.

• **Gradient-Detached **(기울기 분리) 메인 헤드의 로짓 (logits) 에서 MTP 헤드로의 지식 증류를 수행할 때, 메인 헤드로의 역전파 (backpropagation) 를 차단합니다 (stop-gradient). 이를 통해 MTP 헤드의 학습이 메인 헤드의 최적화를 방해하는 것을 방지합니다.
• TopN-Logits-Selected Distillation: 전체 어휘 (Vocabulary) 에 대한 증류는 계산 비용이 크고 저확률 토큰의 노이즈를 유발합니다. 대신 메인 헤드의 로짓 중 상위 TopN(예: 10,000 개) 의 토큰만 선택하여 MTP 헤드를 지도합니다. 이는 계산 효율성을 높이고 수치적 안정성을 보장합니다.
• **손실 함수 **(Loss Function) 최종 MTP 손실은 두 가지로 구성됩니다.
  1. **Cross-Entropy Loss **(LCE) MTP 헤드가 정답 토큰을 예측하도록 하는 기본 손실.
  2. **KL Divergence Loss **(LKL) 메인 헤드의 TopN 로짓 분포를 MTP 헤드가 따르도록 하는 증류 손실.
  • 전체 손실: $L_{mtp} = L_{CE} + L_{KL}$

2.2 Looped Extension Strategy (연속 전학습 기반 확장)

학습된 MTP 헤드를 기반으로 더 많은 헤드를 추가하여 확장하는 전략입니다.

• 그룹화 및 초기화: 학습된 $m$개의 MTP 헤드를 그룹으로 간주하고, 이 그룹의 가중치를 복사하여 새로운 $m$개의 MTP 헤드를 초기화합니다.
• **연속 전학습 **(Continued Pre-training) 확장된 MTP 헤드를 추가적인 데이터 (70B 토큰) 로 학습시키되, 메인 모델과 기존에 학습된 MTP 헤드는 **동결 **(Frozen) 시킵니다.
• 효과: 그룹 내 MTP 헤드 간의 상관관계와 증류로 인한 분포 일관성을 활용하여 적은 데이터로도 효율적으로 헤드를 확장하고 추론 속도를 높입니다.

3. 주요 기여 (Key Contributions)

1. MTP-D 프레임워크 제안: 메인 헤드의 성능을 유지하면서 MTP 헤드의 수용률을 획기적으로 높이는 효율적인 자기 증류 방법론을 제안했습니다.
2. 비용 효율적인 확장 전략: 학습된 MTP 헤드를 그룹 단위로 반복적으로 확장 (Looped Extension) 하는 전략을 도입하여, 추가 학습 비용 없이도 MTP 헤드의 수를 4 개에서 16 개까지 확장할 수 있음을 입증했습니다.
3. 실험적 검증 및 통찰: 7 개의 벤치마크를 통한 광범위한 실험을 통해 MTP-D 와 확장 전략의 유효성을 입증하고, MTP 의 확장성에 대한 중요한 통찰 (예: 그룹화 확장 효과, 데이터 크기 한계 등) 을 도출했습니다.

4. 실험 결과 (Results)

2B Dense 및 10B MoE 모델을 기반으로 한 실험 결과는 다음과 같습니다.

• 수용률 및 속도 향상:
  • 4 개의 MTP 헤드를 사용할 경우, MTP-D 는 기존 MTP 대비 7.5% 높은 누적 수용률을 달성했습니다.
  • 이는 22.9% 의 추론 속도 향상으로 이어졌습니다.
  • 단일 헤드 (1-head) 기준 MTP-D 는 약 14% 의 속도 향상을 보였습니다.
• **확장성 **(Scalability)
  • Looped Extension을 통해 4 개에서 16 개로 헤드를 확장했을 때, 추가적인 속도 향상 (35.1%) 을 달성하여 전체적으로 220.4%(1-head MTP 대비) 의 속도 향상을 기록했습니다.
  • Training-Free 설정 (추가 학습 없이 확장) 에서도 기존 MTP 는 3 번째 헤드에서 수용률이 0.6% 로 급락했으나, MTP-D 는 26.7% 를 유지하여 확장성이 우수함을 입증했습니다.
• 메인 헤드 성능: MTP-D 를 적용해도 메인 헤드의 정확도 (Accuracy) 는 기존 모델과 비교할 때 유의미한 저하가 없거나 오히려 소폭 향상되었습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 추론 효율성을 극대화하기 위한 실용적인 솔루션을 제시합니다.

• 실용성: 복잡한 하이퍼파라미터 튜닝 없이도 기존 아키텍처에 쉽게 적용 가능한 '간단하고 효과적인' 방법론을 제시하여 산업계에서의 MTP 도입 장벽을 낮춥니다.
• 확장성: MTP 헤드의 수를 늘리는 것이 항상 성능 저하를 가져오는 것이 아니라, 적절한 증류와 확장 전략을 통해 16 개 이상의 헤드까지 효과적으로 확장 가능함을 보여주었습니다.
• 미래 전망: 본 연구는 MTP 기반의 초고속 추론 시스템 구축에 중요한 기반을 제공하며, 향후 더 큰 규모의 모델과 다양한 후학습 (Post-training) 시나리오에서의 적용 가능성을 열어줍니다.

결론적으로, MTP-D는 MTP 헤드의 수용률 문제를 해결하고, Looped Extension은 이를 경제적으로 확장하여 LLM 의 추론 속도를 획기적으로 개선하는 강력한 방법론입니다.