A2D: Any-Order, Any-Step Safety Alignment for Diffusion Language Models

📄 arXiv: 2509.23286v2 📥 PDF

作者: Wonje Jeung, Sangyeon Yoon, Yoonjun Cho, Dongjae Jeon, Sangwoo Shin, Hyesoo Hong, Albert No

分类: cs.CL, cs.AI

发布日期: 2025-09-27 (更新: 2026-02-03)

备注: Accepted at ICLR 2026. Code and models are available at https://ai-isl.github.io/A2D

💡 一句话要点

A2D：针对扩散语言模型的任意顺序、任意步骤安全对齐方法

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

关键词: 扩散语言模型 安全对齐 token级别防御 预填充攻击 有害内容检测

📋 核心要点

  1. 扩散语言模型易受攻击，有害内容可能出现在任意位置，现有防御方法难以应对。
  2. A2D通过token级别的安全对齐，使模型在检测到有害内容时输出[EOS]拒绝信号。
  3. 实验表明，A2D能有效防御各种攻击，显著降低有害内容生成率，并实现快速安全终止。

📝 摘要（中文）

扩散大型语言模型(dLLMs)支持任意顺序的生成，但也因此扩大了攻击面：有害内容可能出现在任意位置，并且基于模板的预填充攻击（如DIJA）会绕过响应级别的拒绝。我们提出了A2D（任意顺序、任意步骤防御），这是一种token级别的对齐方法，它使dLLMs在出现有害内容时发出[EOS]拒绝信号。通过在随机掩码下直接在token级别对齐安全性，A2D在各种条件下实现了对任意解码顺序和任意步骤预填充攻击的鲁棒性。它还支持实时监控：dLLMs可以开始响应，但如果出现不安全的延续，则会自动终止。在安全基准测试中，A2D始终如一地阻止有害输出的生成，将DIJA的成功率从80%以上降低到接近于零（在LLaDA-8B-Instruct上为1.3%，在Dream-v0-Instruct-7B上为0.0%），并且阈值化的[EOS]概率允许提前拒绝，从而实现高达19.3倍的更快安全终止。

🔬 方法详解

问题定义：扩散语言模型（dLLMs）虽然具有灵活的生成能力，但也面临着严重的安全风险。传统的防御方法通常在响应级别进行判断，容易被基于模板的预填充攻击（如DIJA）绕过。此外，由于dLLMs支持任意顺序的生成，有害内容可能出现在生成的任何位置，这使得传统的防御方法难以有效应对。

核心思路：A2D的核心思路是在token级别进行安全对齐。具体来说，A2D训练模型在生成有害token时，立即输出[EOS]（end-of-sequence）信号，从而终止有害内容的生成。这种token级别的防御机制能够有效地应对任意顺序生成和预填充攻击，因为无论有害内容出现在哪个位置，模型都能及时检测并终止生成。

技术框架：A2D的技术框架主要包括以下几个步骤：1) 数据准备：构建包含有害内容和安全内容的训练数据集。2) 随机掩码：在训练过程中，对输入序列进行随机掩码，模拟各种可能的生成顺序。3) Token级别对齐：训练模型预测每个token的安全概率，并对有害token进行惩罚，促使模型输出[EOS]信号。4) 阈值化：设置[EOS]概率的阈值，当模型输出的[EOS]概率超过阈值时，立即终止生成。

关键创新：A2D的关键创新在于其token级别的安全对齐机制。与传统的响应级别防御方法相比，A2D能够更精细地控制生成过程，及时阻止有害内容的生成。此外，A2D采用随机掩码技术，增强了模型对各种生成顺序的鲁棒性。

关键设计：A2D的关键设计包括：1) 损失函数：使用交叉熵损失函数来训练模型预测每个token的安全概率。对于有害token，增加损失权重，促使模型输出[EOS]信号。2) [EOS]阈值：设置合适的[EOS]概率阈值，以平衡安全性和生成质量。阈值过高可能导致模型过度拒绝，阈值过低则可能无法有效阻止有害内容的生成。3) 随机掩码策略：采用不同的随机掩码策略，例如随机删除、随机替换等，以增强模型的鲁棒性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

A2D在安全基准测试中表现出色，能够有效防御DIJA等预填充攻击。在LLaDA-8B-Instruct上，A2D将DIJA的成功率从80%以上降低到1.3%，在Dream-v0-Instruct-7B上更是降低到0.0%。此外，A2D的阈值化[EOS]概率允许提前拒绝，从而实现高达19.3倍的更快安全终止。

🎯 应用场景

A2D可应用于各种基于扩散语言模型的应用场景，例如聊天机器人、文本生成、代码生成等。通过A2D的安全对齐，可以有效防止模型生成有害、不当或具有攻击性的内容，提高用户体验，并降低潜在的法律和道德风险。此外，A2D的实时监控功能可以帮助开发者及时发现和修复模型中的安全漏洞。

📄 摘要（原文）

Diffusion large language models (dLLMs) enable any-order generation, but this flexibility enlarges the attack surface: harmful spans may appear at arbitrary positions, and template-based prefilling attacks such as DIJA bypass response-level refusals. We introduce A2D (Any-Order, Any-Step Defense), a token-level alignment method that aligns dLLMs to emit an [EOS] refusal signal whenever harmful content arises. By aligning safety directly at the token-level under randomized masking, A2D achieves robustness to both any-decoding-order and any-step prefilling attacks under various conditions. It also enables real-time monitoring: dLLMs may begin a response but automatically terminate if unsafe continuation emerges. On safety benchmarks, A2D consistently prevents the generation of harmful outputs, slashing DIJA success rates from over 80% to near-zero (1.3% on LLaDA-8B-Instruct, 0.0% on Dream-v0-Instruct-7B), and thresholded [EOS] probabilities allow early rejection, yielding up to 19.3x faster safe termination.