Unlocking Lossless Speedups in LLMs via Discrete Diffusion

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/

한 줄 요약

Uno는 확률 확산을 활용해 대형 언어 모델의 속도를 3배까지 향상시키는 새로운 구조를 제시한다.

핵심 기여도

확률 확산과 자동회귀 모델을 결합한 diffusion-augmented LLMs 구조 제안.
Diffusion Distillation을 통해 기존 NTP 기반 모델에 가벼운 확산 가중치를 학습.
Ψ-Spec 샘플러를 도입해 고정 컨텍스트 길이에서 손실 없는 가속화 가능.
8B Uno 모델이 26B DiffusionGemma와 Mercury 2를 벤치마크에서 상회.

핵심 아이디어

기존 자동회귀 언어 모델은 토큰 생성이 순차적이므로 느리다는 한계가 있다. 이 논문은 확률 확산 프로세스를 도입해 여러 토큰을 병렬 생성함으로써 속도를 향상시키는 새로운 접근법을 제시한다. 확산 가중치는 기존 NTP 기반 AR 가중치와 별도로 학습되며, Diffusion Distillation이라는 간단한 단계를 통해 추가 학습 비용 없이 확산 파라미터를 얻는다. Ψ-Spec 샘플러는 기존 추론 파이프라인에 최소한의 수정만으로 손실 없는 속도 향상을 가능하게 한다. 이는 기존 추론 가속 기법인 speculative decoding과 달리 별도의 draft 모델이 필요하지 않다는 점에서 차별화된다.

기술적 접근법

**모델 구조**: AR 가중치와 확산 가중치로 분리된 diffusion-augmented LLMs.
**학습 방법**: NTP 기반 AR 가중치는 기존 방식으로 학습, 확산 가중치는 Diffusion Distillation을 통해 학습.
**샘플러**: Ψ-Spec 샘플러를 사용해 고정 컨텍스트 길이에서 병렬 토큰 생성.
**모델 구현**: Uno 모델은 기존 오픈소스 AR LLM을 확장하거나 처음부터 학습 가능.

주요 결과

Uno는 기존 추론 가속 방법 대비 모든 배치 크기에서 더 높은 처리량을 달성.
기존 AR 모델 대비 최대 3배의 속도 향상.
8B Uno 모델이 DiffusionGemma 26B와 Mercury 2를 벤치마크에서 상회.
agentic tool use, coding, long-context reasoning에서 우수한 성능.

의의 및 한계

Uno는 확률 확산을 활용한 새로운 추론 가속 기법으로, 기존 AR 모델의 성능을 유지하면서 속도를 향상시킨다. 기존 추론 가속 기법과 달리 별도 모델이 필요하지 않아 실용성이 높다. 그러나 확산 기반 추론의 복잡성은 일부 상황에서 메모리나 계산 자원을 더 요구할 수 있다. 또한, 확산 파라미터 학습의 효과는 모델 크기와 데이터셋에 따라 달라질 수 있다.

실용적 활용

Uno는 대규모 언어 모델의 실시간 추론이 필요한 산업, 예를 들어 챗봇, 코드 생성, 장문 분석 등에서 활용 가능하다. 기존 오픈소스 모델을 확장해 사용할 수 있어 연구 및 개발 환경에서 즉각적으로 적용할 수 있다.