透過離散擴散解鎖大型語言模型的無損加速

摘要

大型語言模型（Large Language Models, LLMs）的成功在很大程度上歸功於下一詞元預測（next-token prediction, NTP），但其自迴歸（autoregressive, AR）結構需要緩慢且依序地生成詞元。為克服此瓶頸，我們提出了擴散增強的大型語言模型（diffusion-augmented LLMs），這是一類在定義 AR 模型分佈的同時，利用擴散從該分佈中並行抽取多個詞元的新模型。我們將這些模型的參數解耦為兩組：使用標準 NTP 目標訓練的 AR 權重，以及訓練用於同時生成多個詞元的輕量擴散權重。擴散權重透過一個簡單的擴散蒸餾（Diffusion Distillation）階段來學習，該階段對現有 LLM 訓練流程所增加的額外開銷可忽略不計。我們也引入了 Ψ-Spec，這是一個取樣器家族，可在固定上下文長度下實現無損加速與推論時的擴展。與推測性解碼不同，我們的方法不需要獨立的草稿模型。與擴散式 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/