Itai Gat Heli Ben-Hamu Marton Havasi Daniel Haziza Jeremy Reizenstein Gabriel Synnaeve David Lopez-Paz Brian Karrer Yaron Lipman

초록

자기회귀적 다음 토큰 예측 언어 모델은 강력한 능력을 제공하지만, 특히 디코딩 단계에서 발생하는 높은 계산 및 메모리 비용으로 인해 실용적 적용에 큰 도전에 직면해 있다. 본 연구에서는 표준 다음 토큰 예측(NTP, next token prediction)과 마스크된 토큰 예측(MATP, masked autoregressive token prediction)을 단일 아키텍처 내에 통합함으로써 생성 속도를 향상시키는 간단하고 유연한 패러다임인 Set Block Decoding(SBD)을 제안한다. SBD는 모델이 연속적이지 않은 여러 미래 토큰을 병렬로 샘플링할 수 있도록 하며, 이는 기존 가속화 기법과의 핵심적인 차이점이다. 이러한 유연성 덕분에 이산 확산 이론에서 개발된 고급 해법을 활용할 수 있어 정확도를 희생하지 않고도 상당한 속도 향상을 달성할 수 있다. SBD는 아키텍처 변경이나 추가 학습 하이퍼파라미터가 필요 없으며, 정확한 KV 캐싱과 호환되며, 기존의 다음 토큰 예측 모델을 미세조정(fine-tuning)하는 방식으로 구현 가능하다. Llama-3.1 8B 및 Qwen-3 8B 모델을 미세조정한 결과, SBD는 동일한 성능을 유지하면서 생성에 필요한 전방 계산(FW pass) 횟수를 3~5배 감소시킴을 입증하였다.