여러 토큰을 동시에 예측해 LLM 성능과 속도를 개선

한줄 요약

LLM 학습 시 다음 토큰 하나 대신 여러 토큰을 동시에 예측하도록 하면 샘플 효율이 높아지고, 추론 속도도 개선된다.

풀어야 하는 문제

기존 next-token prediction은 각 위치에서 단 하나의 토큰만 예측하므로, 장기 의존성 포착과 효율적 학습에 한계가 있다. 특히 코드 생성과 같은 구조적 출력에서 성능 병목이 발생한다.

접근 방법

공유 모델 trunk 위에 n개의 독립적인 출력 헤드를 두어, 각 위치에서 다음 n개 토큰을 동시에 예측하도록 학습한다. 이는 보조 학습 목표로 사용되며, 추론 시에는 첫 번째 헤드만 사용하거나, 여러 헤드를 활용해 병렬 디코딩으로 속도를 높일 수 있다.

결과·기여

13B 파라미터 모델에서 HumanEval 12%, MBPP 17% 성능 향상. 4-토큰 예측 모델은 추론 시 최대 3배 빠름. 작은 알고리즘 태스크에서 induction head 발달과 알고리즘 추론 능력 향상이 관찰됨.