Thinking Inside the Mask: In-Place Prompting in Diffusion LLMs

📄 arXiv: 2508.10736v2 📥 PDF

作者: Xiangqi Jin, Yuxuan Wang, Yifeng Gao, Zichen Wen, Biqing Qi, Dongrui Liu, Linfeng Zhang

分类: cs.CL

发布日期: 2025-08-14 (更新: 2025-10-11)

💡 一句话要点

提出ICE框架，利用扩散LLM的In-Place Prompting提升推理性能并加速计算。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散语言模型 In-Place Prompting 链式思考 提前退出 迭代优化

📋 核心要点

  1. 传统LLM的前缀式prompting限制了双向信息流动，影响推理能力。
  2. ICE框架利用dLLM的迭代优化和双向注意力，实现In-Place Prompting。
  3. ICE通过提前退出机制显著降低计算成本，并在多个任务上取得性能提升。

📝 摘要（中文）

大型语言模型（LLMs）取得了显著成功，但其前缀式提示范式和顺序生成过程在双向信息利用方面存在局限性。扩散大型语言模型（dLLMs）通过其双向注意力机制和迭代优化过程，为更灵活的In-Place Prompting策略提供了新的机会。我们提出了ICE（In-Place Chain-of-Thought Prompting with Early Exit），这是一个新颖的框架，专门为dLLMs将前缀式提示转换为In-Place Prompting。ICE在迭代优化期间将In-Place Prompt直接集成到掩码token位置中，并采用置信度感知的提前退出机制来显著降低计算开销。大量实验表明了ICE的有效性，在GSM8K上实现了高达17.29%的准确率提升和4.12倍的加速，在MMLU上实现了高达276.67倍的加速，同时保持了具有竞争力的性能。

🔬 方法详解

问题定义：现有的大型语言模型主要采用前缀式提示（prefix-only prompting）范式，这种方式限制了模型在生成过程中对上下文信息的双向利用。尤其是在需要复杂推理的任务中，模型难以充分利用中间步骤的信息进行迭代优化。此外，顺序生成过程也导致计算效率较低，难以满足实际应用的需求。

核心思路：ICE的核心思路是将传统的prefix-only prompting转换为In-Place Prompting，即在模型生成过程中，直接在被掩码的token位置插入提示信息，并利用扩散语言模型的双向注意力机制，使模型能够充分利用这些提示信息进行迭代优化。同时，引入置信度感知的提前退出机制，在模型达到一定置信度时提前停止迭代，从而降低计算开销。

技术框架：ICE框架主要包含两个关键组成部分：In-Place Prompting和置信度感知的提前退出机制。首先，将输入问题和In-Place Prompt一起输入到扩散语言模型中，模型通过迭代优化过程逐步生成答案。在每次迭代中，模型都会根据In-Place Prompt更新其内部状态，并生成新的token。其次，在每次迭代后，模型会计算当前生成答案的置信度，如果置信度达到预设的阈值，则提前退出迭代过程，输出最终答案。

关键创新：ICE最重要的技术创新点在于将In-Place Prompting引入到扩散语言模型中。与传统的前缀式提示相比，In-Place Prompting能够更有效地利用上下文信息，提高模型的推理能力。此外，置信度感知的提前退出机制能够显著降低计算开销，提高模型的效率。

关键设计：ICE的关键设计包括In-Place Prompt的格式和位置选择，以及置信度阈值的设定。In-Place Prompt通常采用自然语言描述，例如“Let's think step by step”。Prompt的位置选择通常是根据任务的特点进行调整，例如，在数学推理任务中，可以将Prompt插入到每一步计算之后。置信度阈值的设定需要根据实际情况进行调整，以在性能和效率之间取得平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ICE框架在GSM8K数据集上实现了高达17.29%的准确率提升，并加速了4.12倍。在MMLU数据集上，ICE实现了高达276.67倍的加速，同时保持了具有竞争力的性能。这些结果表明，ICE框架能够显著提升LLM的推理能力和计算效率。

🎯 应用场景

ICE框架具有广泛的应用前景，可应用于数学推理、常识推理、代码生成等多个领域。通过提升LLM的推理能力和计算效率，ICE能够帮助解决更复杂的实际问题，例如智能客服、自动化编程、科学研究等。未来，ICE有望成为新一代LLM的重要组成部分，推动人工智能技术的发展。

📄 摘要（原文）

Despite large language models (LLMs) have achieved remarkable success, their prefix-only prompting paradigm and sequential generation process offer limited flexibility for bidirectional information. Diffusion large language models (dLLMs) present new opportunities through their bidirectional attention mechanisms and iterative refinement processes, enabling more flexible in-place prompting strategies. We introduce ICE (In-Place Chain-of-Thought Prompting with Early Exit), a novel framework that transforms prefix-only prompting into in-place prompting specifically designed for dLLMs. ICE integrates in-place prompts directly within masked token positions during iterative refinement and employs a confidence-aware early exit mechanism to significantly reduce computational overhead. Extensive experiments demonstrate ICE's effectiveness, achieving up to 17.29% accuracy improvement with 4.12$\times$ speedup on GSM8K, and up to 276.67$\times$ acceleration on MMLU while maintaining competitive performance.