通过离散扩散解锁 LLM 中的无损加速

大语言模型 (LLM) 的成功很大程度上归功于下一个词元预测 (NTP)，但其自回归 (AR) 结构需要缓慢、连续的词元生成。为了克服这个瓶颈，我们引入了扩散增强型 LLM，这是一类新的模型，它定义了 AR 模型分布，同时使用扩散从该分布中并行绘制多个标记。我们将这些模型的参数解耦为两组：使用标准 NTP 目标进行训练的 AR 权重，以及经过训练以同时生成多个词元的轻量级扩散权重。扩散权重是通过简单的扩散 蒸馏 阶段学习的，该阶段给现有 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/