逐字生成非最优？试试逐「块」生成！Block Diffusion打通了自回归与扩散

去年初，OpenAI 的视频生成模型 Sora 带火了扩散模型。

如今，扩散模型被广泛用于生成图像和视频，并在生成文本或生物序列等离散数据方面变得越来越有效。从技术上讲，与自回归模型相比，扩散模型具有加速生成和提高模型输出可控性的潜力。 

目前，离散扩散模型目前面临至少三个限制。首先，在聊天系统等应用中，模型必须生成任意长度的输出序列（例如对用户问题的回答）。但是，大多数最新的扩散架构仅能生成固定长度的向量。其次，离散扩散模型在生成过程中使用双向上下文，因此无法使用 KV 缓存重用以前的计算，这会降低推理效率。第三，以困惑度等标准指标衡量的离散扩散模型，质量落后于自回归方法，进一步限制了其适用性。

本文中，来自 Cornell Tech、斯坦福大学、Cohere 的研究者提出通过 块离散去噪扩散语言模型（Block Discrete Denoising Diffusion Language Models，BD3-LMs）来解决以上限制，该模型在扩散和自回归模型之间进行插值。

具体来讲，块扩散模型（也是半自回归模型）定义了离散随机变量块的自回归概率分布，而给定先前块的条件概率由离散去噪扩散模型指定。

• 论文标题：Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

• 论文地址：https://arxiv.org/pdf/2503.09573

• 项目主页：https://m-arriola.com/bd3lms/

下图为 Block Diffusion 与自回归、扩散模型的生成效果对比：

研究者表示，开发有效的 BD3-LM 面临以下两个挑战：一是使用神经网络的一次标准前向传递无法有效地计算块扩散模型的训练目标，需要开发专门的算法。二是扩散目标梯度的高方差阻碍了训练，导致 BD3-LM 即使在块大小为 1 的情况下（当两个模型等效时）也表现不佳。

因此，研究者推导出梯度方差的估计量，并证明它是自回归和扩散之间困惑度差距的关键因素。然后，他们提出了自定义噪声过程，以实现最小化梯度方差并进一步缩小困惑度差距。

实验部分，研究者在多个语言建模基准上评估了 BD3-LM，并证明它们能够生成任意长度的序列，包括超出其训练上下文的长度。此外，BD3-LM 在离散扩散模型中实现了新的 SOTA 困惑度。与对嵌入进行高斯扩散的替代半自回归方法相比，本文离散方法实现了易于处理的似然估计，并在少一个数量级生成步骤的情况下，生成的样本在困惑度方面得到了改进。

论文一作 Marianne Arriola 发推称，扩散语言模型在并行文本生成领域正在崛起，但与自回归模型相比，它们存在质量、固定长度限制和缺乏 KV 缓存等问题。本文 Block Diffusion 将自回归和扩散模型结合了起来，实现了两全其美。

BD3-LMs 模型概览

研究者结合建模范式，从自回归模型中获得更好的似然估计和灵活的长度生成，并从扩散模型中获得了快速的并行生成效果。

块扩散似然

研究者提出了一个建模框架，该框架对 token 块进行自回归建模，并在每个块内执行扩散操作。他们对长度为 L′ 的 B 个块进行似然分解，如下所示：

每个 pθ(x^b|x^<b) 都使用包含 L′个 token 的块上的离散扩散 ELBO 进行建模，并通过优化以下似然边界来获得原则性学习目标 L_BD (x,θ)：

研究者使用简单的离散扩散参数化对每个块的似然进行建模，最终目标是对交叉熵项进行加权总和：

高效的训练与采样算法

简单来说，研究者想要通过在一个 loop 中应用 B 次来计算 logits。不过，他们只需要两次前向传递。第一次传递分别预计算完整序列 x 的键和值 K^1:B、V^1:B，在第二次前向传递中使用 同时计算所有块的去噪预测。

为了从 BD3-LM 中采样，研究者以先前采样的块为条件，一次生成一个块。生成块后，他们缓存其键和值，类似于 AR。同时在每个块的 T 个采样步下，使用任何扩散采样流程 SAMPLE 中进行采样。来从条件分布 pθ

算法 1（块扩散训练）和算法 2（块扩散采样）分别如下图（左）和（右）所示。

                                    ^{BD3-LM 训练和采样算法。}

理解扩散模型与自回归模型之间的似然差距

案例研究：单 Token 生成

该研究中的块扩散参数化在期望上等同于自回归负对数似然 (NLL)，特别是在 L′=1 的极限情况下。令人惊讶的是，当在 LM1B 数据集上训练两种模型时，研究发现块扩散模型 (L′=1) 与自回归模型之间存在两点困惑度差距。研究确定扩散目标的高训练方差是导致这一困惑度差距的原因。

                                     ^{在离散扩散 ELBO 下进行训练时，存在高方差。}

高方差训练导致的扩散差距

直观来说，如果采样的掩码率 过低，重构 x 会变得容易，这不能提供有用的学习信号。如果掩码全部内容，最优的重构就是数据分布中每个标记的边际概率，这很容易学习，同样也没有用处。

研究需要找到能够最小化扩散目标引起的训练方差，并进一步减少困惑度差距的噪声调度方案。

基于数据的低方差训练噪声调度

为了避免导致高方差训练的掩码率，研究者在「裁剪的』掩码率 下来训练 BD3-LMs。通过降低训练方差，研究者在均匀采样的掩码率评估下改善了似然度。

由于最佳掩码率可能会根据块大小 L′的不同而变化，他们在训练期间自适应地学习 β,ω。在实践中，研究者在每个验证步骤后（经过 5K 次梯度更新）使用网格搜索来优化 。

在下文中，研究者展示了针对每个块大小优化噪声调度可以减少损失估计器的方差，并与其他替代调度方案相比实现最佳困惑度。

实验结果

似然评估

BD3-LMs 在扩散模型中实现了最先进的似然水平。研究表明，通过调整块长度 L′，BD3-LMs 可以在扩散和自回归似然之间实现插值。

                                                    ^{在 OWT 上测试针对 262B 标记训练的模型的困惑度 (PPL; ↓)。}

任意长度序列生成

许多现有扩散语言模型的一个主要缺点是，它们无法生成超过训练时选择的输出上下文长度的完整文档。例如，OpenWebText 包含最长达 131K tokens 的文档，而离散扩散模型 SEDD（Lou 等人）仅限于生成 1024 tokens。研究表明，BD3-LMs 能够通过解码任意数量的块来生成可变长度的文档。

                                                       ^{从在 OWT 上训练的模型中抽样 500 个文档得出的生成长度统计信息。}

研究者评估了 BD3-LMs 在变长序列上的生成质量，使用相同数量的生成步骤（NFEs）比较了所有方法。他们用 GPT2-Large 模型测量生成序列的困惑度。结果表明，与之前所有的扩散方法相比，BD3-LMs 实现了最佳的生成困惑度。

^{300 个可变长度样本的生成困惑度 （Gen. PPL;↓） 和功能评估次数 （NFE;↓）。所有模型都在 OWT 上进行训练，上下文长度为 L = 1024，并使用核采样。}

对于 MDLM，研究者使用了其分块解码技术（该技术不同于 BD3-LMs 中的分块扩散训练）处理 L=2048 的序列。研究者还与 SSD-LM（Han 等人提出）进行了比较，后者是一种替代性的分块自回归方法（也称为半自回归），它对词嵌入执行高斯扩散，但无法进行似然估计。该研究的离散方法使用比其他方法少一个数量级的生成步骤，产生了具有更好生成困惑度的样本。

更多细节请参阅原论文。