Context-Aware Initialization for Reducing Generative Path Length in Diffusion Language Models
作者: Tongyuan Miao, Gary Huang, Kai Jun Han, Annie Jiang
分类: cs.CL, cs.AI, cs.LG
发布日期: 2025-12-22
💡 一句话要点
提出上下文感知初始化方法,缩短扩散语言模型生成路径,加速推理。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 上下文感知初始化 生成路径缩短 快速推理 重掩码机制
📋 核心要点
- 扩散语言模型推理速度慢,主要瓶颈在于需要多次迭代去噪,将随机噪声转化为连贯文本。
- 提出上下文感知初始化方法,利用轻量级辅助模型注入prompt条件先验,使初始状态更接近目标分布。
- 实验表明,该方法能显著减少去噪迭代次数(约35%),但naive warm-starting可能降低最终精度。
📝 摘要(中文)
扩散语言模型(DLLMs)实现了完全并行的token解码,但由于需要多次去噪迭代才能将无信息的全掩码初始化转化为连贯的文本,因此在推理时通常不实用。现有的大多数加速方法都侧重于通过改进的求解器或采样策略来更有效地遍历此生成轨迹。本文提出了一种互补的视角:通过上下文感知的初始化,从更接近目标分布的位置开始,从而缩短轨迹本身。本文提出了一种无需训练的接口,该接口将来自轻量级辅助模型的prompt条件先验注入到扩散初始化中,并通过两种机制实例化:离散token注入和表示级嵌入插值。由于注入的先验可能不完善,并且仅使用mask解码可能会过早地提交,因此本文还引入了一种基于置信度的简单重掩码机制,作为一种先验怀疑的形式。在GSM8K上的初步证据表明,上下文感知的初始化可以大大减少去噪迭代次数(在本文的设置中减少约35%的函数评估),同时也暴露了一个关键的开放挑战:相对于强大的扩散基线,naive的warm-starting会降低最终的准确性。本文使用这些发现来激发围绕校准、修订机制和表示对齐的研究议程,以实现可靠的warm-started扩散解码。
🔬 方法详解
问题定义:扩散语言模型(DLLMs)虽然具有并行解码的优势,但推理速度慢,主要原因是需要从完全随机的噪声(或全掩码)状态开始,经过多次迭代去噪才能生成连贯的文本。现有的加速方法主要集中在优化去噪过程,例如改进求解器或采样策略,但没有直接解决生成路径过长的问题。
核心思路:本文的核心思路是通过上下文感知初始化来缩短生成路径。具体来说,就是利用prompt信息,通过一个轻量级的辅助模型,预测一个更接近目标分布的初始状态,从而减少后续去噪迭代的次数。这样做的目的是让扩散过程从一个“更好”的起点开始,而不是从完全随机的状态开始。
技术框架:该方法的核心是一个无需训练的接口,用于将prompt条件先验注入到扩散模型的初始化状态中。主要包含以下几个步骤:1) 使用一个轻量级的辅助模型(例如,一个小的Transformer模型)根据prompt预测初始状态;2) 将预测的初始状态以某种方式注入到扩散模型的初始状态中;3) 使用扩散模型进行去噪迭代,生成最终的文本。为了避免辅助模型预测错误导致的问题,还引入了一个基于置信度的重掩码机制。
关键创新:该方法的主要创新在于提出了上下文感知初始化的概念,并提供了一种无需训练的接口来实现这一概念。与现有的加速方法不同,该方法不是优化去噪过程,而是直接缩短生成路径。此外,提出的重掩码机制可以有效地缓解辅助模型预测错误带来的负面影响。
关键设计:该方法使用了两种不同的方式来注入先验信息:离散token注入和表示级嵌入插值。离散token注入直接将辅助模型预测的token替换扩散模型的初始状态中的token。表示级嵌入插值则将辅助模型预测的嵌入向量与扩散模型的初始嵌入向量进行插值。重掩码机制根据辅助模型预测的置信度,对初始状态中的token进行重新掩码,以避免过早地提交到错误的预测。
🖼️ 关键图片
📊 实验亮点
实验结果表明,上下文感知初始化方法可以在GSM8K数据集上减少约35%的函数评估次数,从而显著加速扩散模型的推理过程。然而,naive的warm-starting可能会降低最终的准确性,这表明需要进一步研究校准、修订机制和表示对齐等问题。
🎯 应用场景
该研究成果可应用于各种需要快速文本生成的场景,例如对话系统、机器翻译、文本摘要等。通过减少扩散模型的推理迭代次数,可以显著提高生成速度,降低计算成本,从而使扩散模型更适用于实际应用。
📄 摘要(原文)
Diffusion Large Language Models (DLLMs) enable fully parallel token decoding but often remain impractical at inference time due to the many denoising iterations required to refine an information-free, fully masked initialization into coherent text. Most existing acceleration methods focus on traversing this generative trajectory more efficiently via improved solvers or sampling strategies. We advance a complementary perspective: shorten the trajectory itself by starting closer to the target distribution through context-aware initialization. We propose a training-free interface that injects prompt-conditioned priors from a lightweight auxiliary model into the diffusion initialization, and instantiate it with two mechanisms: discrete token injection and representation-level embedding interpolation. Because injected priors can be imperfect and unmask-only decoding can over-commit early, we also introduce a simple confidence-based remasking mechanism as a form of prior skepticism. Preliminary evidence on GSM8K suggests that context-aware initialization can substantially reduce denoising iterations (about 35\% fewer function evaluations in our setting), while also exposing a key open challenge: naive warm-starting can degrade final accuracy relative to strong diffusion baselines. We use these findings to motivate a research agenda around calibration, revision mechanisms, and representation alignment for reliable warm-started diffusion decoding.