Unlocking Lossless Speedups in LLMs via Discrete Diffusion
cs.LG
Submitted: 2026-09-03
Updated: 2026-09-03
Comments: Code and Checkpoints at https://s-sahoo.github.io/uno/
Code: https://github.com/GeeeekExplorer/nano-vllm
Project page: https://nvidia.github.io/TensorRT-LLM/developer-guide/perf-benchmarking.html
License: http://creativecommons.org/licenses/by/4.0/
The gist: Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation.
Terminology
Abstract
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 3 times 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/
Sources
- Program Synthesis with Large Language Models
- MathArena: Evaluating LLMs on Uncontaminated Math Competitions
- $\tau^2$-Bench: Evaluating Conversational Agents in a Dual-Control Environment
- LLaDA2.1: Speeding Up Text Diffusion via Token Editing
- Accelerating Large Language Model Decoding with Speculative Sampling
- DFlash: Block Diffusion for Flash Speculative Decoding
- Beyond Benchmarks: MathArena as an Evaluation Platform for Mathematics with LLMs
- DiffusionGemma Technical Report
- Set Block Decoding is a Language Model Inference Accelerator
- GLM-5: from Vibe Coding to Agentic Engineering
- Better & Faster Large Language Models via Multi-token Prediction
- The Llama 3 Herd of Models
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- OpenThoughts: Data Recipes for Reasoning Models
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- Training Compute-Optimal Large Language Models
- Fast and Accurate Causal Parallel Decoding using Jacobi Forcing
- AA-Omniscience: Evaluating Cross-Domain Knowledge Reliability in Large Language Models
- EfficientRollout: System-Aware Self-Speculative Decoding for RL Rollouts
- Kimi K3: Open Frontier Intelligence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks