이산 확산을 통한 LLM의 무손실 속도 향상 구현

초록

대규모 언어 모델(LLM)은 그 성공의 상당 부분을 다음 토큰 예측(next-token prediction, NTP) 덕분에 얻지만, 자기회귀(autoregressive, AR) 구조는 느리고 순차적인 토큰 생성을 요구한다. 이 병목을 해결하고자 우리는 확산으로 증강된 LLM(diffusion-augmented LLM)이라는 새로운 모델 부류를 제안한다. 이 모델은 AR 모델 분포를 정의하면서도 확산(diffusion)을 사용하여 해당 분포에서 여러 토큰을 병렬로 추출한다. 우리는 이 모델들의 매개변수를 두 집합으로 분리한다. 하나는 표준 NTP 목적 함수로 학습되는 AR 가중치이고, 다른 하나는 여러 토큰을 동시에 생성하도록 학습되는 경량 확산 가중치이다. 확산 가중치는 기존 LLM 훈련 파이프라인에 무시할 만한 오버헤드만을 추가하는 간단한 확산 증류(Diffusion Distillation) 단계를 통해 학습된다. 또한 고정된 컨텍스트 길이에서 무손실 가속과 추론 시간 스케일링을 가능하게 하는 샘플러 계열인 Ψ-Spec을 소개한다. 추측적 디코딩(speculative decoding)과 달리, 우리 방법은 별도의 초안 모델이 필요하지 않다. 확산 LLM(d-LLM)과 달리, 이 방법은 기본 AR 모델의 품질을 희생하지 않고 생성을 가속화한다. 그 결과로 얻은 모델인 Uno는 처음부터 학습하거나 기존 오픈 가중치 AR LLM을 증강하여 구축할 수 있다. Uno는 평가된 모든 배치 크기에서 최고 수준의 추측적 디코딩 방법보다 높은 처리량을 달성하며, 장치가 지원하는 가장 큰 배치 크기에서도 기본 AR 모델 대비 최대 3배의 속도 향상을 제공한다. 특히, 우리의 8B Uno 모델은 에이전트 기반 도구 사용, 코딩, 장문맥 추론 등 평가된 모든 벤치마크에서 선도적인 오픈 d-LLM인 26B DiffusionGemma와 독점 모델인 Mercury 2를 능가한다. 코드와 체크포인트는 https://s-sahoo.github.io/uno/ 에서 공개한다.

English

Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation. To overcome this bottleneck, we introduce diffusion-augmented LLMs, a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution. We decouple the parameters of these models into two sets: AR weights, trained using the standard NTP objective, and lightweight diffusion weights, trained to generate multiple tokens simultaneously. The diffusion weights are learned through a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines. We also introduce Ψ-Spec, a family of samplers that enables lossless acceleration and inference-time scaling at a fixed context length. Unlike speculative decoding, our method requires no separate draft model. Unlike diffusion LLMs (d-LLMs), it accelerates generation without sacrificing the quality of the underlying AR model. The resulting models, called Uno, can be trained from scratch or built by augmenting existing open-weight AR LLMs. Uno achieves higher throughput than leading speculative-decoding methods at every evaluated batch size and delivers up to 3times speedups over the base AR model, including at the largest batch size supported by the device. Notably, our 8B Uno model outperforms the leading open d-LLM, the 26B DiffusionGemma, and the proprietary Mercury 2 across all evaluated benchmarks in agentic tool use, coding, and long-context reasoning. We release code and checkpoints at: https://s-sahoo.github.io/uno/