BlockGen: Flexible Blockwise Sequence Modeling with Hybrid Samplers

📄 arXiv: 2606.02241v1 📥 PDF

作者: Justin Deschenaux, Caglar Gulcehre

分类: cs.LG

发布日期: 2026-06-01

🔗 代码/项目: GITHUB

💡 一句话要点

BlockGen：一种使用混合采样器的灵活分块序列建模方法

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 序列建模 扩散模型 分块生成 自回归模型 采样策略 自然语言生成

📋 核心要点

  1. 现有扩散模型在离散序列生成中存在不足，尤其是在校正器设计和分块生成策略上。
  2. BlockGen通过混合采样器和分块生成，实现了更灵活的序列建模，并引入AR信息预测-校正采样。
  3. 实验表明，在不同采样策略和数据集上，BlockGen的掩码扩散和均匀扩散变体表现出不同的性能优势。

📝 摘要（中文）

均匀状态扩散框架是否是离散扩散更强大的范式？最近的研究表明可能是这样。结合预测-校正采样器，均匀状态扩散模型（USDMs）比掩码扩散模型（MDMs）产生更高质量的样本，并且在下游任务中USDMs等于或优于MDMs，即使它们表现出更高的困惑度。然而，仍有两个问题未解决。首先，现有工作比较均匀扩散和掩码扩散时，使用了无信息的校正器，这些校正器在随机位置重新注入噪声，而不是针对最有可能错误的token。其次，先前的工作比较了完整序列扩散模型，因此我们不知道当token逐块生成时，是否会得出相同的结论。为了解决这些问题，我们引入了BlockGen，一种分块序列模型，我们用掩码扩散和均匀扩散实例化它。BlockGen在混合块大小上进行训练，并且其似然比具有固定块大小的模型更精细地在AR和纯扩散之间进行插值。BlockGen支持AR信息预测-校正采样（ARPC），它结合了AR和扩散预测来重新生成不太可能的token，而无需辅助验证器。在祖先采样下，均匀扩散在逐块设置中优于掩码扩散，尤其是在少步态的情况下。在ARPC下，差距在高NFE时缩小并反转。在GSM8K上，块大小为16时，MDMs的准确率略高于USDMs，并且我们在OpenWebText上的生成困惑度中观察到类似的趋势。代码可在https://github.com/jdeschena/blockgen 找到。

🔬 方法详解

问题定义：论文旨在解决离散序列生成任务中，现有扩散模型（尤其是掩码扩散模型MDMs）和均匀状态扩散模型（USDMs）在性能上的差异和局限性。现有方法，如全序列扩散模型，无法有效利用自回归（AR）信息进行校正，并且在分块生成场景下的表现未知。此外，现有校正器通常采用无信息的方式注入噪声，效率较低。

核心思路：论文的核心思路是提出一种名为BlockGen的分块序列模型，该模型允许使用混合块大小进行训练，从而更精细地在自回归和纯扩散之间进行插值。同时，引入AR信息预测-校正采样（ARPC），利用自回归模型的预测信息来指导扩散模型的校正过程，从而更有效地重新生成不太可能的token。

技术框架：BlockGen的整体框架包括以下几个关键部分：1) 分块序列生成：将输入序列分割成多个块，并逐块进行生成。2) 混合块大小训练：模型在不同大小的块上进行训练，以提高泛化能力。3) 掩码扩散和均匀扩散：BlockGen支持两种扩散模式，分别是掩码扩散和均匀扩散。4) AR信息预测-校正采样（ARPC）：结合自回归模型的预测信息和扩散模型的校正过程，提高生成质量。

关键创新：论文的关键创新在于：1) 提出BlockGen分块序列模型，允许灵活的分块大小和混合训练。2) 引入ARPC采样策略，将自回归模型的预测信息融入到扩散模型的校正过程中，从而更有效地纠正错误。3) 对比了掩码扩散和均匀扩散在分块生成场景下的性能差异，并分析了不同采样策略对性能的影响。

关键设计：BlockGen的关键设计包括：1) 块大小的选择：模型在训练时使用混合块大小，以提高泛化能力。2) ARPC采样策略：ARPC结合了自回归模型的预测概率和扩散模型的噪声预测，用于指导校正过程。具体来说，ARPC会优先重新生成自回归模型预测概率较低的token。3) 损失函数：BlockGen使用标准的扩散模型损失函数进行训练，同时考虑了分块生成和混合块大小的影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在GSM8K数据集上，块大小为16时，MDMs的准确率略高于USDMs。在OpenWebText数据集上的生成困惑度也观察到类似的趋势。ARPC采样策略在高NFE时可以缩小甚至反转掩码扩散和均匀扩散之间的差距。这些结果表明，BlockGen的分块生成和ARPC采样策略可以有效提升序列生成模型的性能。

🎯 应用场景

BlockGen具有广泛的应用前景，包括自然语言生成、代码生成、图像生成等领域。通过灵活的分块生成和AR信息校正，BlockGen可以生成更高质量、更连贯的序列数据。该研究对于提升生成模型的性能和效率具有重要意义，并可能推动相关领域的发展。

📄 摘要（原文）

Is the uniform-state diffusion framework a more powerful paradigm for discrete diffusion? Recent studies indicate that this may be the case. In combination with predictor-corrector samplers, uniform-state diffusion models (USDMs) produce samples of higher-quality than masked diffusion models (MDMs), and USDMs equal or outperform MDMs in downstream tasks, even though they exhibit greater perplexity. Two issues remain unresolved. First, existing work compares uniform and masked diffusion with un-informed correctors that re-inject noise at random positions, rather than targeting tokens most likely to be wrong. Second, prior work compares full-sequence diffusion models, so we do not know whether the same conclusion holds when tokens are generated block by block. To address these issues, we introduce BlockGen, a blockwise sequence model that we instantiate with both masked and uniform diffusion. BlockGen trains on a mixture of block sizes and its likelihood interpolates between AR and pure diffusion more finely than models with a fixed block size. BlockGen enables AR-informed predictor-corrector sampling (ARPC), which combines AR and diffusion predictions to re-generate unlikely tokens without an auxiliary verifier. Under ancestral sampling, uniform outperforms masked in the block-by-block setting, especially in the few-step regime. Under ARPC, the gap closes and reverses at high NFE. With block size $16$ on GSM8K, MDMs reach slightly higher accuracy than USDMs, and we observe a similar trend in Generative Perplexity on OpenWebText. Find our code at https://github.com/jdeschena/blockgen.