TASER：面向几何驱动鲁棒性的任务感知Stein正则化

# TASER: 面向几何鲁棒性的任务感知斯坦因正则化

来源: https://arxiv.org/html/2605.30601
Michał Kozyra 牛津大学统计系, 英国 michal\.kozyra@seh\.ox\.ac\.uk & Gesine Reinert 牛津大学统计系, 英国 reinert@stats\.ox\.ac\.uk

###### 摘要

现代深度网络在分布偏移和对抗性扰动下仍然脆弱，这通常是由于过度或结构不当的输入敏感性所致。我们提出 **TASER**（任务感知斯坦因正则化），这是一个基于朗之万斯坦因算子的训练时正则化框架。通过惩罚训练分布下的逐点斯坦因残差，TASER 鼓励预测器与数据密度之间的几何相容性，从而诱导出各向异性、数据感知的平滑性。我们提供了斯坦因正则化与降低一阶偏移敏感性之间的理论联系，开发了与现代架构兼容的可扩展实现变体，并在回归和视觉基准测试中展示了改进的鲁棒性和稳定性。在 CIFAR-10 实验中，TASER 在未引起统计上显著的干净精度下降的情况下，持续改进了现有训练方法的对抗鲁棒性。

## 1 引言

深度神经网络在分布内任务上表现出色，但在分布偏移和对抗性扰动下仍然脆弱 (Hendrycks and Dietterich,2019 (https://arxiv.org/html/2605.30601#bib.bib41); Goodfellow et al.,2015 (https://arxiv.org/html/2605.30601#bib.bib13); Madry et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib14))。这两种现象背后的一个核心失效模式是**输入灵敏度错位**：预测器对相对于数据分布而言微小的扰动表现出巨大响应，同时可能对语义上有意义的变异反应不足 (Tsipras et al.,2019 (https://arxiv.org/html/2605.30601#bib.bib28))。在对抗性设置中，这表现为输入空间中存在这样一些方向，沿这些方向的微小扰动会引发模型输出的巨大变化 (Goodfellow et al.,2015 (https://arxiv.org/html/2605.30601#bib.bib13))。在分布偏移中，当测试输入以结构化的方式偏离训练分布时，这会导致泛化性能下降 (Hendrycks and Dietterich,2019 (https://arxiv.org/html/2605.30601#bib.bib41))。

大量工作通过正则化模型灵敏度来解决这一问题。经典方法包括权重衰减和谱约束 (Loshchilov and Hutter,2019 (https://arxiv.org/html/2605.30601#bib.bib32); Miyato et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib33))，而更直接的方法则惩罚梯度或强制 Lipschitz 界 (Jakubovitz and Giryes,2018 (https://arxiv.org/html/2605.30601#bib.bib34); Cisse et al.,2017 (https://arxiv.org/html/2605.30601#bib.bib35))。对抗训练进一步通过针对最坏情况扰动进行优化来寻求鲁棒性 (Madry et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib14); Zhang et al.,2019 (https://arxiv.org/html/2605.30601#bib.bib21))。尽管取得了成功，这些方法有一个共同的局限性：它们统一地或根据固定的范数约束对待输入空间中的所有方向。特别地，它们没有明确地纳入**训练分布的几何结构**。因此，它们可能会压制语义上有意义的方向上的灵敏度，同时未能充分控制将输入移出高概率区域的方向上的灵敏度。

这项工作引入了一种不同的方法：**根据数据分布的几何结构来正则化模型行为**。我们的出发点是斯坦因方法，它提供了一系列算子，通过均值零恒等式 \(\mathbb{E}_p[\mathcal{T}_p f] = 0\) 来刻画概率分布 (Stein,1972 (https://arxiv.org/html/2605.30601#bib.bib3); Ley et al.,2017 (https://arxiv.org/html/2605.30601#bib.bib4))。对于分布 \(p\)，其得分函数为 \(s_p(x) = \nabla \log p(x)\)，朗之万斯坦因算子

\[
\mathcal{L}_p f(x) = \Delta f(x) + s_p(x)^\top \nabla f(x) \quad (1)
\]

通过曲率和方向导数项的组合来编码 \(p\) 的局部几何结构。其中 \(\Delta f(x) = \mathrm{tr}(\nabla^2 f(x))\) 是拉普拉斯算子，\(\nabla\) 是梯度算子，\(\top\) 表示转置。对于每个固定的函数 \(f\)，我们称 \(r_f(x) = \mathcal{L}_p f(x)\) 为在 \(x\) 处的（逐点）**斯坦因残差**。

参阅图注
图 1：各向同性与几何感知平滑性。标准正则化器（左）对模型灵敏度施加统一的惩罚，平等对待所有输入方向。TASER（右）诱导出一个与数据流形对齐的各向异性平滑包络：沿流形的灵敏度基本不受约束，而沿与得分场 \(\nabla \log p(x)\) 对齐的离流形方向上的灵敏度则受到强烈惩罚。

我们提出 **TASER**（任务感知斯坦因正则化），一种在训练时惩罚逐点斯坦因残差的正则化框架：

\[
\mathcal{L}_{\mathrm{total}}(\theta) = \mathcal{L}_{\mathrm{task}}(\theta) + \lambda \, \mathbb{E}_{X \sim p}\big[(\mathcal{L}_p f_\theta(X))^2\big]. \quad (2)
\]

与在输入空间上统一作用的传统正则化器不同，TASER 施加的约束明确由分布 \(p\) 塑造。特别地，\(\mathcal{L}_p\) 中的得分加权项 \(s_p(x)^\top \nabla f(x)\) 惩罚了数据密度变化最快的方向上的灵敏度，而拉普拉斯项则全局地控制曲率。共同地，它们强制一种**几何感知的平滑性**，使模型灵敏度与数据结构对齐。

在高维数据集中在低维结构附近的环境中，这种视角尤为自然 (Fefferman et al.,2016 (https://arxiv.org/html/2605.30601#bib.bib31))。在这种情形下，最陡峭的密度变化方向往往与高概率质量区域正交，而 TASER 在无需显式流形估计的情况下抑制沿这些方向的灵敏度。这为降低离分布灵敏度提供了一种原则性的机制，而离分布灵敏度正是对抗鲁棒性脆弱的关键驱动因素 (Fawzi et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib29); Gilmer et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib30))。

从理论角度看，TASER 具有直接的鲁棒性解释。斯坦因残差决定了在数据分布平滑扰动下模型的一阶响应。特别地，对于形如 \(q_\varepsilon(x) \propto p(x) e^{\varepsilon h(x)}\) 的指数倾斜，在 \(q_\varepsilon\) 下 \(\mathcal{L}_p f\) 的期望与斯坦因残差和扰动 \(h\) 之间的协方差成比例。因此，最小化 \(\mathcal{L}_p f\) 的方差直接限制了对一大类分布偏移的一阶敏感性。

TASER 易于实现且适用范围广泛。它仅需要访问输入梯度和得分场的估计，后者可以从现代扩散或得分匹配模型中获得 (Ho et al.,2020 (https://arxiv.org/html/2605.30601#bib.bib9); Song et al.,2021 (https://arxiv.org/html/2605.30601#bib.bib10))。该方法与架构和任务无关，并且可以与包括对抗训练在内的现有训练流程结合。

#### 贡献。

本文做出以下贡献：

- • 我们引入 TASER，一个基于斯坦因算子的正则化框架，用于强制对模型灵敏度施加几何感知约束。
- • 我们证明 TASER 惩罚与数据分布对齐的方向导数，为各向同性梯度正则化提供了一种原则性的替代方案。
- • 我们建立了斯坦因残差最小化与降低分布扰动下一阶敏感性之间的理论联系。
- • 我们证明 TASER 通过抑制将输入移出高密度区域方向上的灵敏度，为改善对抗鲁棒性提供了一种自然机制。

更广泛地说，TASER 将斯坦因算子重新定位为**训练**工具，并在生成式建模（通过得分估计）和判别式鲁棒性之间架起了一座桥梁。

## 2 相关工作

#### 正则化模型灵敏度。

控制神经网络对其输入的灵敏度是提高鲁棒性和泛化能力的核心主题。经典方法如权重衰减和谱约束 (Loshchilov and Hutter,2019 (https://arxiv.org/html/2605.30601#bib.bib32); Miyato et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib33)) 通过参数范数间接限制灵敏度，而更直接的方法则惩罚输入梯度，例如通过雅可比范数正则化 (Jakubovitz and Giryes,2018 (https://arxiv.org/html/2605.30601#bib.bib34); Cisse et al.,2017 (https://arxiv.org/html/2605.30601#bib.bib35))。这些技术强制预测器在环境输入空间中具有平滑性，并且通常与底层数据分布无关。因此，它们对所有方向施加统一的约束，而不区分与数据分布一致的变异和对应于不可能或离分布扰动的变异。

#### 对抗训练与鲁棒优化。

对抗训练和鲁棒优化方法通过明确优化模型在指定范数球内的最坏情况扰动下的性能来解决灵敏度问题 (Madry et al.,2018 (https://arxiv.org/html/2605.30601#bib.bib14); Goodfellow et al.,2015 (https://arxiv.org/html/2605.30601#bib.bib13))。像 TRADES 及其相关公式等扩展进一步探索了鲁棒性与准确性之间的权衡 (Zhang et al.,2019 (https://arxiv.org/html/2605.30601#bib.bib21))。尽管这些方法展示了强大的经验鲁棒性，但它们需要解决一个挑战性的内部最大化问题，并且依赖于扰动集的选择，最常见的是由 \(\ell_p\) 范数定义。这种依赖性可能限制泛化能力，因为鲁棒性通常与训练期间看到的特定扰动类别绑定。此外，此类公式并未明确编码数据分布的几何结构，并且可能对与典型数据变异无关的方向进行过度正则化。

#### 基于得分的模型与扩散模型。

基于得分和扩散模型为在高维空间中估计得分场 \(\nabla \log p(x)\) 提供了可扩展的方法 (Ho et al.,2020 (https://arxiv.org/html/2605.30601#bib.bib9); Song et al.,2021 (https://arxiv.org/html/2605.30601#bib.bib10))。这些模型主要用于生成式建模，其中得分定义了一个驱动从噪声到数据分布的随机过程的向量场。除了生成之外，得分场还编码了数据分布的局部几何信息，捕获了最陡峭密度变化的方向。这种表示在生成式建模和几何正则化之间提供了自然的桥梁。

#### 合成数据与基于扩散的鲁棒性。

(Gowal et al.,2021 (https://arxiv.org/html/2605.30601#bib.bib12); Nie et al.,2022 (https://arxiv.org/html/2605.30601#bib.bib11)) 探索了使用生成模型通过使用合成数据增强训练或在潜在空间或生成空间中进行对抗训练来提高鲁棒性。这些方法利用学习到的数据分布来产生更真实的扰动或使用多样化的样本来丰富训练集。基于扩散模型的近期工作使用生成先验进行对抗净化或样本生成 (Nie et al.,2022 (https://arxiv.org/html/2605.30601#bib.bib11))。

#### 机器学习中的斯坦因方法。

斯坦因算子已广泛用于机器学习的拟合优度检验、样本质量评估和基于核的差异度量，相关综述见 (Liu et al.,2026 (https://arxiv.org/html/2605.30601#bib.bib2))。这些方法利用恒等式 \(\mathbb{E}_p[\mathcal{T}_p f] = 0\) 来构建检测与目标分布偏离的统计量。最近，基于斯坦因的量已被探索作为检测分布偏移和模型误设的诊断工具 (Kozyra and Reinert,2026 (https://arxiv.org/html/2605.30601#bib.bib1))。然而，它们的用途主要在事后设置中，即在训练后评估算子，而不是用来塑造训练过程本身。

#### 总结。

总之，现有的鲁棒性方法要么强制统一的平滑性，要么针对预定义的扰动集进行优化，而最近的生成式方法则依赖于昂贵且通常是针对特定问题的流程。基于得分模型提供了一种连续、局部且具有表达力和可扩展性的数据几何表示。本文建立在这些进展的基础上，利用斯坦因算子作为训练时正则化机制，通过得分场直接将模型灵敏度与训练分布的几何结构耦合，同时保持与现有训练方法的简单模块化集成。

## 3 方法：任务感知斯坦因正则化 (TASER)

### 3.1 问题设定与斯坦因公式

考虑一个监督学习问题，输入 \(x \in \mathbb{R}^d\) 来自分布 \(p\)，目标为 \(y\)。设 \(f_\theta: \mathbb{R}^d \to \mathbb{R}^m\) 表示一个模型。目标是学习 \(f_\theta\)，使其能够泛化到训练分布之外，并在结构化输入扰动下保持稳定。

一个核心困难在于，标准训练目标未对预测器相对于数据分布几何结构的行为施加任何约束。特别是，梯度 \(\nabla f_\theta(x)\) 可能与密度 \(p(x)\) 快速变化的方向对齐，导致在将输入移出高概率区域的扰动下输出发生巨大变化。

为了解决这个问题，我们引入一个基于斯坦因算子的正则化原则。设 \(p\) 是一个具有可微（可能未归一化）密度和得分函数 \(s_p(x) = \nabla \log p(x)\) 的分布。对于如 (1) 中所定义的朗之万斯坦因算子 \(\mathcal{L}_p f(x) = \Delta f(x) + s_p(x)^\top \nabla f(x)\)，在标准正则性条件下，斯坦因恒等式 \(\mathbb{E}_{X \sim p}[\mathcal{L}_p f(X)] = 0\) 成立（见附录 B (https://arxiv.org/html/2605.30601#A2)）。TASER 将此恒等式用作训练原则，在样本级别惩罚与之的偏离：

\[
\mathcal{L}_{\mathrm{total}}(\theta) = \mathcal{L}_{\mathrm{task}}(\theta) + \lambda \, \mathbb{E}_{X \sim p}\big[(\mathcal{L}_p f_\theta(X))^2\big]. \quad (3)
\]

由于 \(\mathbb{E}_p[\mathcal{L}_p f] = 0\)，该惩罚对应于训练分布下（逐点）斯坦因残差的方差

\[
r_f(x) = \mathcal{L}_p f(x) =
\]