다중 토큰 예측(MTP, 모델이 이후의 여러 토큰을 추가로 예측하도록 하는 기법)은 이미 일반적인 사전 훈련 레시피로, 다운스트림 성능을 향상시킬 뿐만 아니라 추측 디코딩(speculative decoding)을 위한 초안 헤드로도 활용됩니다. 저자들은 MTP의 교차 엔트로피 그래디언트를 정확히 두 부분으로 분해했습니다. 하나는 MTP 헤드를 모델 자체의 다음 단어 분포로 끌어당기는 KL 자기 증류(KL self-distillation) 항(유용하며 추측 디코딩을 향상시킴)이고, 다른 하나는 신호 대 잡음비(SNR)가 낮아 학습을 방해하고 규모가 커질수록 더 심각해지는 잔차(residual)입니다. Soft-MTP는 '스루-스루 그래디언트 라우팅(straight-through gradient routing)'을 사용하여 KL 항만 유지하고, 기존 교차 엔트로피 계산을 재사용하며, 전체 어휘의 로짓(logits) 행렬이 필요 없고, 융합된 교차 엔트로피 커널과 호환됩니다.
좋아요 9 · 코드 스타 수 6
이유
- 변경 사항이 매우 미미함: 표준 MTP에 대해 코드 한 줄로 구현되며, 매개변수를 추가하지 않고, 저자들은 측정 가능한 추가 오버헤드가 없다고 주장합니다. 또한 DeepSeek식 및 Meta식 MTP를 모두 지원합니다.
- 저자들은 167M에서 1.2B까지의 모든 규모에서 다음 단어 퍼플렉시티(perplexity)가 낮아지고, 추측 디코딩이 빨라지며, 제로샷 정확도가 높아진다고 주장합니다. 1.2B 미세 조정 후 수학적 및 코드 추론 능력이 향상되었으며, 대부분의 지표에서 KL을 직접 최소화하는 방법과 동등하거나 더 나은 성능을 보입니다.
- 재현 가능: 네 가지 규모 × 세 가지 목표(NTP, MTP, Soft-MTP)로 총 12개의 체크포인트가 Hugging Face에 게시되어 있으며, 훈련 구성 및 추측 디코딩 평가 스크립트가 모두 포함되어 있습니다. tests/ 디렉터리에서는 float64에서 그래디언트가 명시적 KL과 일치함을 검증할 수 있습니다. MIT 라이선스.
사용법
- 코드: git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5 (사전 토큰화된 FineWeb-Edu, 각 구성에 대해 31.5B 토큰을 훈련하며, 최소 4개의 샤드가 필요하고, 조립 시 약 120 GB의 디스크 공간이 필요합니다). Muon 옵티마이저는 PyTorch 2.9 이상을 요구합니다.
- 훈련: torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp (대조군으로 hard_mtp / direct_kl / ntp로 변경 가능).
- 기존 모델 확인: AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True), 토크나이저는 원본 gpt2를 사용하며, 추측 디코딩 평가는 python -m src.eval_spec --ckpts …를 사용합니다.
- 자체 훈련 코드로 이식하려면 src/models/loss.py의 combined_mtp_loss()를 참조하세요.
리스크
- 공개된 모델은 모두 1.2B 이하의 연구용 소형 모델로, FineWeb-Edu에서만 훈련되었으며, 바로 사용할 수 있는 대화형 모델이 아닙니다. 더 큰 규모에서도 동일한 효과가 있는지 검증되지 않았습니다.
- 가치는 주로 사전 훈련 단계에 있으며, 이점을 얻으려면 처음부터 직접 훈련해야 합니다. 훈련 평가는 기본적으로 Weights & Biases와 연동됩니다.
- 초록에는 추측 디코딩 가속의 구체적인 배수가 명시되어 있지 않으며, 효과는 저자들이 보고한 것입니다. README에서는 벽시계(wall-clock) 가속이 GPU에 의존하므로, 비교는 동일한 머신에서 실행해야 한다고 경고합니다.
열기 · 코드 · Hugging Face 논문 페이지