文本扩散模型中的安全感知去噪器

# 面向文本扩散模型的安全感知去噪器

**来源:** https://arxiv.org/html/2605.08116

###### 摘要

近期关于文本扩散模型的研究为自回归生成提供了一种充满前景的替代方案，但对其安全性的控制仍属研究不足。现有的安全方法主要针对自回归模型，通常依赖事后过滤或推理时干预。这些方法在有效应对文本扩散模型的安全风险方面显得力不从心。我们提出了**安全感知去噪器（Safety-Aware Denoiser, SAD）**，这是一种用于文本扩散模型的安全引导框架。SAD 修改了迭代去噪过程，使得最终去噪步骤生成的文本样本被引导至文本空间中**可证明安全**的区域。这种推理时方法可以将安全约束整合到去噪器中，避免了重新训练底层扩散模型的高昂计算成本，并实现了灵活、轻量级的安全引导。我们使用 SAD 评估了生成文本的安全性，涵盖危害分类、记忆化（memorization）和越狱（jailbreak）等方面。实验结果表明，SAD 在保持生成质量、多样性和流畅性的同时，显著减少了不安全生成，优于现有方法。这些结果证明，我们在去噪过程中的安全引导为在文本扩散模型中实施安全机制提供了一种有效且可扩展的机制。代码详见 https://github.com/ammanyusuf/SAD。机器学习，ICML

## 1 引言

文本扩散模型（Text Diffusion Models, TDMs）（Austin 等人，2021 (https://arxiv.org/html/2605.08116#bib.bib2); Sahoo 等人，2024 (https://arxiv.org/html/2605.08116#bib.bib1); Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)）取得了 rapid 进展，其性能优于同等规模的主流自回归（Autoregressive, AR）大语言模型（LLMs）。例如，LLaDA（Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)）在数学推理和中文理解任务上超越了 LLaMA3-8B（Grattafiori 等人，2024 (https://arxiv.org/html/2605.08116#bib.bib11)）。与 AR-LLMs 的顺序逐 token 生成方式相比，TDMs 通过并行生成和迭代 refine，将完全掩码的序列逐步转化为文本输出，这使得其推理速度更快，且在文本生成方面比 AR-LLMs 更具灵活性。

虽然 TDMs 潜力巨大，但其安全影响 largely understudied，其部署引发了安全担忧，包括生成有毒内容、记忆敏感数据以及对越狱攻击的脆弱性。大多数专注于生成文本毒性和危害的现有安全基准仅在 AR-LLMs 上测试（Gehman 等人，2020 (https://arxiv.org/html/2605.08116#bib.bib13); Hartvigsen 等人，2022 (https://arxiv.org/html/2605.08116#bib.bib15); Jie 等人，2023 (https://arxiv.org/html/2605.08116#bib.bib17)）。此外，现有的安全技术——如事后过滤、拒绝采样或推理时解码约束—— largely 围绕自回归解码设计，不能直接移植到扩散设置中。此外，绝大多数越狱攻击针对 AR-LLMs，只有极少数近期针对 TDM 的越狱利用 TDM 中的并行去噪特性（Li 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib12); Wen 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib8); Zhang 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib31)）。这些论文指出，TDMs 能够以**比**可比 AR 模型**更快**的速度生成有害内容，从而加速不安全输出的出现。这些发现强调，TDMs 构成了强大的攻击面，主要是因为其安全弱点尚未被现有的对齐技术充分缓解，且成功的利用可以快速产生有害文本。这激起了对专门针对扩散范式的防御机制的需求。

解决这一问题的一个直接方案是使用安全目标对文本扩散模型进行重新训练或微调。也可以考虑条件扩散模型，给定明确的安全/不安全标签，从安全分布中采样（Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)）。然而，随着模型规模不断扩大并日益以冻结检查点（frozen checkpoints）形式共享，这些方法计算昂贵且缺乏灵活性。对于条件 TDMs，获取安全数据集本身既昂贵又不灵活，因为安全要求在不断演变。因此，迫切需要轻量级的、推理时的安全机制，能够无缝集成到文本扩散过程中，而无需通过修改或重新训练底层模型来增加能耗。

在本工作中，我们提出了**安全感知去噪器（SAD）**，这是一种原则性的安全引导框架，利用了 Kim 等人（2025 (https://arxiv.org/html/2605.08116#bib.bib4)）提出的安全去噪器。SAD 的核心思想是直接修改迭代去噪过程，调整每一步去噪，使得生成的样本逐步被引导至文本空间中**可证明安全**的区域。与仅作用于最终输出的事后过滤方法不同，SAD 影响**整个**生成轨迹，允许安全约束以更稳定、有效的方式塑造最终文本。值得注意的是，SAD 是**免训练（training-free）**的，无需重新训练扩散模型。安全约束被整合到去噪器本身中，在保留原始模型生成能力的同时，实现灵活且模块化的安全引导。这种设计使 SAD 计算高效、易于部署，并与未来文本扩散架构的进步兼容。

SAD 的有限参考设计使从业者能够策划一小部分已知的不安全示例（如有毒数据集、版权文档或训练数据），而不是详尽列举所有安全内容。我们的消融实验显示，$N \approx 500-1000$ 个参考样本已能饱和安全增益，使得 SAD 即使在大规模下也切实可行。我们在多个安全维度上对 SAD 进行了全面的实证评估，包括危害分类合规性、减少记忆化以及抵御越狱攻击的鲁棒性。我们的结果表明，SAD 在保持生成质量、多样性和流畅性的同时，显著减少了不安全生成，突显了专门为基于扩散的文本生成设计安全机制的重要性。

## 2 背景

我们首先描述两个文本扩散模型的核心成分：MDLM（Sahoo 等人，2024 (https://arxiv.org/html/2605.08116#bib.bib1)）和 LLaDA（Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)），我们将它们用作文本扩散模型。然后，我们描述**安全去噪器**（Kim 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib4)），该方法最初为基于扩散的图像生成提出，我们将其采用并修改，用于使用文本扩散模型进行安全文本生成。

### 2.1 文本扩散模型

类似于用于图像生成的连续状态空间扩散模型，TDMs 也由前向和反向马尔可夫过程及隐变量组成。前向过程定义为 $p(\mathbf{x}_{1:T}|\mathbf{x}_0)=\prod_{t=1}^{T}q(\mathbf{x}_t|\mathbf{x}_{t-1})$，其中 $\mathbf{x}_{1:T}$ 上的联合分布由遵循条件概率 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$ 的一阶马尔可夫过程定义。该条件概率通常假设已知且固定，具有离散随机变量（一组 tokens）上的类别分布，确定逐渐的加噪过程，将干净的数据点 $\mathbf{x}_0$ 转化为损坏的隐变量 $\mathbf{x}_t$。反向马尔可夫过程 $p_{\boldsymbol{\theta}}(\mathbf{x}_{0:T})=p(\mathbf{x}_T)\prod_{t=1}^{T}p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_t)$ 也假设为类别分布，但未知并由参数为 $\boldsymbol{\theta}$ 的神经网络模型参数化。确切参数化因模型而异。早期工作（Austin 等人，2021 (https://arxiv.org/html/2605.08116#bib.bib2)）使用带有称为 `[MASK]` 的吸收状态的转移矩阵，展示了比其他形式转移矩阵一致优越的性能。由于这个原因，掩码（吸收状态）扩散已成为主流文本扩散模型的骨干，如掩码扩散语言模型（MDLM）（Sahoo 等人，2024 (https://arxiv.org/html/2605.08116#bib.bib1)）和大型语言扩散与掩码（LLaDA）（Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)）。这些模型通过最大化变分下界或其近似值进行训练。一旦训练完成，它们产生整个响应草稿，然后遵循其各自的反向马尔可夫过程逐步改进。与 AR-LLMs 相比，这承诺了惊人的推理速度。此外，与限制提示必须放置位置的 AR-LLMs 不同，TDMs 允许提示位于任意位置。例如，TDMs 非常适合诸如重写论文中段落的任务或重构代码块。

### 2.2 安全去噪器

安全去噪器（Kim 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib4)）是一种为安全图像生成提出的免训练安全引导方法。安全去噪器通过利用否定集（例如，不安全图像、版权数据或私有数据）来避免数据分布的特定区域，而无需重新训练或微调模型。安全去噪器假设训练数据 $\mathcal{D}$ 可以划分为不重叠的安全/不安全集合，即，$\mathcal{D}=\mathcal{D}_{\text{safe}}\cup\mathcal{D}_{\text{unsafe}}$，其中安全集远大于不安全集：$|\mathcal{D}_{\text{safe}}|\gg|\mathcal{D}_{\text{unsafe}}|$。使用指示函数，$\mathbb{1}_{\text{safe}}(\mathbf{x}_0)$，如果 $\mathbf{x}_0$ 是安全的（即，$\mathbf{x}_0\in\mathcal{D}_{\text{safe}}$）则取值为 $1$，否则为 $0$。类似地，$\mathbb{1}_{\text{unsafe}}(\mathbf{x}_0)$ 如果 $\mathbf{x}_0$ 是不安全的（即，$\mathbf{x}_0\in\mathcal{D}_{\text{unsafe}}$）则取值为 $1$，否则为 $0$。这些指示函数构成 unity 的划分，结果为：

$$
1=\mathbb{1}_{\text{safe}}(\mathbf{x}_0)+\mathbb{1}_{\text{unsafe}}(\mathbf{x}_0) \quad \forall \mathbf{x}_0\in\text{supp}(p_{\text{data}})
$$

有了这个划分，安全分布 $p_{\text{safe}}(\mathbf{x})$ 的非归一化密度定义为 $\mathbb{1}_{\text{safe}}(\mathbf{x})p_{\text{data}}(\mathbf{x})$。类似地，不安全分布 $p_{\text{unsafe}}(\mathbf{x})$ 的非归一化密度定义为 $\mathbb{1}_{\text{unsafe}}(\mathbf{x})p_{\text{data}}(\mathbf{x})$。这些量定义了相应的条件期望（去噪器）。安全去噪器定义为

$$
\mathbb{E}_{\mathcal{D}_{\text{safe}}}[\mathbf{x}_0|\mathbf{x}_t]=\int_{\mathbf{x}_0}\frac{p_{\text{safe}}(\mathbf{x}_0)q_{t}(\mathbf{x}_t|\mathbf{x}_0)}{p_{\text{safe},t}(\mathbf{x}_t)}\mathrm{d}\mathbf{x}_0,
$$

其中 $p_{\text{safe},t}(\mathbf{x}_t)$ 是从安全分布开始的扩散过程（在 $t$ 时）的边缘分布。类似地，不安全去噪器通过用不安全分布替换安全分布来定义。设 $\mathbb{E}_{\mathcal{D}}[\mathbf{x}_0|\mathbf{x}_t]$ 表示模型的数据去噪器。则以下关系成立：

###### 定理 2.1（（Kim 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib4)）中的定理 3.2。安全 vs. 数据/不安全去噪器）

存在一个非负权重 $\beta^*(\mathbf{x}_t)$——后验似然关于 $\mathbf{x}_t$ 起源于不安全集合单调——使得

$$
\begin{aligned}
\mathbb{E}_{\mathcal{D}_{\text{safe}}}[\mathbf{x}_0|\mathbf{x}_t] &= \mathbb{E}_{\mathcal{D}}[\mathbf{x}_0|\mathbf{x}_t] \\
&\quad + \beta^*(\mathbf{x}_t) \, \big( \mathbb{E}_{\mathcal{D}}[\mathbf{x}_0|\mathbf{x}_t] - \mathbb{E}_{\mathcal{D}_{\text{unsafe}}}[\mathbf{x}_0|\mathbf{x}_t] \big).
\end{aligned}
$$

定理 2.1 (https://arxiv.org/html/2605.08116#S2.Thmthm1) 指出，通过从数据去噪器中减去“不安全”分量，其中 $\beta^*$ 适应当前状态的不安全程度，我们可以可证明地生成来自安全去噪器的样本。详见（Kim 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib4)）补充材料 A 节中该定理的证明。

在实践中，预训练的数据去噪器给出 $\mathbb{E}_{\mathcal{D}}[\mathbf{x}_0|\mathbf{x}_t]$ 的值。因此，我们需要从安全去噪器中采样的唯二量是不安全去噪器 $\mathbb{E}_{\mathcal{D}_{\text{unsafe}}}[\mathbf{x}_0|\mathbf{x}_t]$ 和 $\beta^*$。在连续状态的扩散模型的情况下，前向损坏密度 $q_{t}(\mathbf{x}_t|\mathbf{y}_i)$ 通常由多元高斯定义。当然，当考虑用于文本数据的具有离散状态的扩散模型时，这一数量会发生变化。下文描述了我们如何将安全去噪器应用于上述文本扩散模型，从而得出用于安全文本生成的**安全感知去噪器**。

## 3 方法

在本文中，我们展示了将安全去噪器应用于两个文本扩散模型 MDLM 和 LLaDA（Sahoo 等人，2024 (https://arxiv.org/html/2605.08116#bib.bib1); Nie 等人，2025 (https://arxiv.org/html/2605.08116#bib.bib3)）的应用，作为迈向文本扩散中安全生成的第一步。这并不意味着我们的方法局限于这两个模型。通过将去噪器表达式适配到特定的前向和反向转移概率，可以将 SAD 直接扩展到不同类型的 TDMs。

### 3.1 TDMs 中的前向/反向转移核

MDLM 和 LLaDA 均假设 2.1 节 (https://arxiv.org/html/2605.08116#S2.SS1) 中描述的一阶前向和反向马尔可夫过程。在前向过程中，在每个加噪步骤 $t$，输入 $\mathbf{x}_0$ 以某种概率转移到掩码状态 $\mathbf{m}$，这导致前向过程的边缘由词汇表 $V$ 上表示 token 概率的 $K$ 个 one-hot 编码二元随机变量的以下类别分布给出，即，$\mathbf{x}_t, \mathbf{x}_0 \in V$...