[논문리뷰] Set Block Decoding is a Language Model Inference Accelerator

링크: 논문 PDF로 바로 열기

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

핵심 연구 목표

본 논문은 대규모 언어 모델(LLM) 추론, 특히 디코딩 단계에서 발생하는 높은 계산 및 메모리 비용 문제에 초점을 맞춥니다. 이러한 문제를 해결하여 LLM의 실용적인 배포를 가속화하고, **Next Token Prediction (NTP)**과 **Masked Token Prediction (MATP)**을 단일 아키텍처 내에 통합하는 유연한 패러다임인 **Set Block Decoding (SBD)**을 제안합니다. 목표는 정확도 손실 없이 상당한 속도 향상을 달성하고, 기존 모델 아키텍처 변경이나 추가 훈련 하이퍼파라미터 없이도 KV-caching과 호환되도록 하는 것입니다.

핵심 방법론

**Set Block Decoding (SBD)**은 기존 NTP Transformer 아키텍처를 미세 조정하여 NTP와 MATP를 동시에 지원합니다. 훈련 시에는 k개의 미래 토큰 블록을 예측하도록 학습하며, 이 중 일부는 마스크 토큰('m')으로 가려집니다. 이때 이전 토큰에는 인과적 어텐션을, 미래 블록 내 토큰에는 양방향 어텐션을 적용합니다. 추론 단계에서는 Entropy Bounded (EB) Sampler를 활용하여 순차적이지 않은 여러 미래 토큰을 병렬로 샘플링합니다. 이 방식은 Llama-3.1 8B 및 Qwen-3 8B 모델을 미세 조정하여 구현되었으며, 표준 NTP 훈련과 동일한 데이터 및 하이퍼파라미터를 사용했습니다.

주요 결과

SBD는 LiveCodeBench-V6를 포함한 다양한 벤치마크에서 기존 NTP 훈련과 동등한 성능(정확도)을 유지하면서, 생성에 필요한 모델 포워드 패스(NFE) 수를 3-5배 감소시켰습니다. 예를 들어, Llama-3.1 8B 모델에서는 3.0x NFE 속도 향상을, Qwen-3 8B 모델에서는 3.2x NFE 속도 향상을 달성했습니다. 특히, EB-Sampler의 γ (감마) 하이퍼파라미터를 조절하여 속도-정확도 트레이드오프를 제어할 수 있음을 보여주었습니다. 또한, 훈련 시 NTP 손실 항의 포함이 모델의 자기회귀(autoregressive) 능력 유지에 필수적임을 확인했습니다.

AI 실무자를 위한 시사점

SBD는 대규모 언어 모델의 추론 속도를 혁신적으로 가속화할 수 있는 실용적인 방법론을 제공합니다. 이는 복잡한 아키텍처 변경이나 추가적인 드래프트 모델 없이 기존 NTP 모델을 효율적으로 활용하여 3-5배의 포워드 패스 감소를 가능하게 합니다. 따라서 LLM을 활용하는 실시간 애플리케이션의 지연 시간을 줄이고 처리량을 늘리는 데 크게 기여할 수 있습니다. 기존 모델의 KV-caching 호환성을 유지하며, 미세 조정을 통해 쉽게 통합될 수 있다는 점에서 실제 AI 시스템에 적용하기 용이합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.