Vertex-Softmax：通过精确 Softmax 优化实现紧密的 Transformer 验证

# 通过精确 Softmax 优化实现 Transformer 的紧密验证
代码：https://github.com/navidrezazad/VertexSoftmax
来源：https://arxiv.org/html/2605.10974
Navid Rezazadeh & Arash Gholami Davoodi³
加州大学尔湾分校，[email protected]；卡内基梅隆大学，[email protected]

###### 摘要

Transformer 注意力的认证验证需要在预 softmax 分数的区间约束下界定 softmax 函数。现有的验证器独立于下游目标来松弛 softmax，从而留下了可避免的松弛度。我们证明了该分数框（score-box）问题的精确最优解在约束框的顶点处取得，并建立了一个阈值结构定理，表明在对目标系数排序后，最优解仅存在于线性数量的候选者中，从而提出了具有序列长度对数线性复杂度的 Vertex-Softmax 原语。我们进一步证明了一个形式化的最优性结果，表明 Vertex-Softmax 是仅从分数区间中可获得的紧密且健全的上界，并精确刻画了进一步改进所需增加的额外结构（分数相关性、分数-值耦合）。Vertex-Softmax 被集成到一个具有形式化健全性保证的 CROWN（基于凸松弛的最坏情况神经元优化）风格验证器中，在 MNIST、Fashion-MNIST 和 CIFAR-10 注意力模型上显著提高了认证率，并大幅收紧了下界，同时在成本仅为 baselines 一小部分的情况下，一致地匹配或击败了 alpha-CROWN 和分支定界基线。

## 1 引言

前馈网络的认证验证器已发展成为实用的工具。Transformer 仍然顽固地更难处理——而 softmax 注意力就是原因所在。归一化的指数映射耦合了所有 token，现有的验证器必须以牺牲证书紧密性为代价来对其进行松弛：每一个因松弛过松而无法验证的输入，都是验证器无法认证的输入。缩小这一差距，在不牺牲可扩展性的情况下使注意力验证更加紧密，是这项工作的动机。

表 1：softmax 接口处验证方法的比较。计时列给出了第 5 节（https://arxiv.org/html/2605.10974#S5）中选定小规模注意力实验的每次试验代表性的墙钟秒数；它不是全块混合运行时间，后者在附录表 11（https://arxiv.org/html/2605.10974#A5.T11）中单独分解。†仅在区间抽象之前。‡通过优化的图松弛。§不作为局部分数框原语；通过域分割获取相关性。

CROWN/LiRPA 风格的验证器在大规模上通过神经网络传播仿射界（Zhang 等，2018（https://arxiv.org/html/2605.10974#bib.bib1）；Xue 等，2020（https://arxiv.org/html/2605.10974#bib.bib2）），但注意力引入了归一化的指数映射

$$
a(s)=\operatorname{softmax}(s), \quad a_{j}(s)=\frac{\exp(s_{j})}{\sum_{r=1}^{K}\exp(s_{r})}. \tag{1}
$$

在一次分数界传递后，验证器通常只知道一个独立的分数框

$$
\mathcal{B}=\prod_{j=1}^{K}[\ell_{j},u_{j}] \tag{2}
$$

以及由值侧界或下游边际诱导的方向 $c\in\mathbb{R}^{K}$。因此，验证器在每个注意力行的任务就简化为仅受这些独立分数区间约束，在 softmax 输出上优化线性目标：

$$
\min_{s\in\mathcal{B}}c^{\top}\operatorname{softmax}(s). \tag{3}
$$

通用的 softmax 松弛首先松弛向量映射 $s\mapsto\operatorname{softmax}(s)$，然后再与 $c$ 收缩。我们则直接求解这个依赖于方向的标量问题。

结果是精确且微小的。连续分数框上的最小值在约束框的顶点处取得。直觉是，权衡质量与成本的优化器总是可以将每个坐标推向端点而不会增加目标值。对 $c$ 排序后，最优顶点将上端点赋予最小的 $m$ 个系数，将下端点赋予其余部分，对于某个阈值 $m\in\{0,\dots,K\}$。朴素搜索检查 $2^{K}$ 个顶点。我们表明，对目标系数排序将其压缩为恰好 $K+1$ 个候选者——一次排序和一次前缀和扫描。我们称此原语为 **Vertex-Softmax**。

这也确定了分数框接口的信息面。任何仅以 $(c,\ell,u)$ 为输入的声音下界过程（sound lower-bound procedure）不能返回比式 (3)（https://arxiv.org/html/2605.10974#S1.E3）中精确最小值更大的值。更紧密的证书必须使用在此接口之前丢弃的信息，例如分数相关性、分数-值耦合，或比独立框更小的可达分数集。

#### 贡献。

- •我们证明了式 (3)（https://arxiv.org/html/2605.10974#S1.E3）的顶点精确性和 $K+1$ 阈值精确性，为加权 softmax 框问题提供了一个 $O(K\log K)$ 的精确求解器。
- •我们推导了一个分数框信息最优性陈述：Vertex-Softmax 是仅从独立分数区间中可获得的紧密且健全的上界。
- •我们将该原语集成到 Vertex-CROWN 中，这是一种健全 CROWN 风格的注意力验证器，并在 MNIST 和 Fashion-MNIST 补丁注意力模型上展示了巨大的认证率改进，在选定的注意力块上与 alpha-CROWN 和分支定界基线相比具有竞争力或改进的证书，并在带有 CROWN 后缀混合的注意力-残差-MLP 模型上实现了全块增益。

#### 定位。

基于仿射界传播的可扩展神经网络验证器，包括 CROWN、auto_LiRPA 和 $\alpha,\beta$-CROWN（Zhang 等，2018（https://arxiv.org/html/2605.10974#bib.bib1）；Xue 等，2020（https://arxiv.org/html/2605.10974#bib.bib2）；Wang 等，2021（https://arxiv.org/html/2605.10974#bib.bib3）），已扩展到 Transformer 注意力层（Shi 等，2020（https://arxiv.org/html/2605.10974#bib.bib4）；Bonar 等，2021（https://arxiv.org/html/2605.10974#bib.bib5））。在此领域内，最近的聚焦于 softmax 的松弛收紧了验证器内部使用的向量值 softmax 界。Wei 等（2023）（https://arxiv.org/html/2605.10974#bib.bib7）的凸界和 Zhang 等（2024）（https://arxiv.org/html/2605.10974#bib.bib8）的 GaLileo 松弛具有代表性；两者都是松弛 softmax 而不是精确求解依赖于方向的分数框目标。

在精度-成本光谱的另一端，MILP、MIQCP、SMT 和非线性分支定界等精确方法（Tjeng 等，2019（https://arxiv.org/html/2605.10974#bib.bib9）；Katz 等，2019（https://arxiv.org/html/2605.10974#bib.bib10）；Shi 等，2025（https://arxiv.org/html/2605.10974#bib.bib12））可以在小规模实例上给出完整答案，但无法扩展到稠密 softmax 注意力。我们的贡献介于这两个极端之间：Vertex-Softmax 在 $O(K\log K)$ 时间内精确求解反复出现的分数框子问题，补充了基于松弛和基于搜索的方法。关于相关工作的更广泛讨论，包括 LLM 规模的统计和运行时监控方法，见附录 A（https://arxiv.org/html/2605.10974#A1）。精确性声明局限于分数框接口；验证器松散的其余来源在第 6 节（https://arxiv.org/html/2605.10974#S6）中讨论。表 1（https://arxiv.org/html/2605.10974#S1.T1）总结了现状。

## 2 问题设置

我们现在形式化 CROWN 风格分数界传递与 softmax 层之间的接口，并定义 Vertex-Softmax 求解的标量优化问题。表 2（https://arxiv.org/html/2605.10974#S2.T2）收集了关键符号。

表 2：全文使用的关键符号。

设 $\mathcal{X}$ 为一个输入框，$m:\mathcal{X}\to\mathbb{R}$ 为一个标量边际。认证意味着证明对于所有 $x\in\mathcal{X}$ 有 $m(x)\geq 0$，通常通过计算 $\min_{x\in\mathcal{X}}m(x)$ 的一个健全下界来实现。对于一个注意力行，写

$$
s_{i}(x)\in\mathbb{R}^{K}, \quad a_{i}(x)=\operatorname{softmax}(s_{i}(x)). \tag{4}
$$

当反向界传递到达该行时，此接口处的下游贡献由固定向量 $c_{i}\in\mathbb{R}^{K}$ 总结，因此行证书需要 $c_{i}^{\top}\operatorname{softmax}(s_{i}(x))$ 的下界。

同一次传递提供健全的逐元素分数界

$$
\ell_{ij}\leq s_{ij}(x)\leq u_{ij}, \quad j=1,\dots,K, \quad x\in\mathcal{X}, \tag{5}
$$

形成积区间 $\mathcal{B}_{i}=\prod_{j=1}^{K}[\ell_{ij},u_{ij}]$。省略行索引，定义

$$
F_{c}(s)=c^{\top}\operatorname{softmax}(s)=\frac{\sum_{j=1}^{K}c_{j}e^{s_{j}}}{\sum_{j=1}^{K}e^{s_{j}}}, \quad L_{\mathrm{box}}(c,\ell,u)=\min_{s\in\mathcal{B}}F_{c}(s). \tag{6}
$$

匹配的上界为 $U_{\mathrm{box}}(c,\ell,u)=-L_{\mathrm{box}}(-c,\ell,u)$，因此我们仅陈述下界。全程 $K\geq 1$，区间有限，允许退化区间，且 $c$ 相对于 $s$ 固定。式 (6)（https://arxiv.org/html/2605.10974#S2.E6）对于积框是精确的，但对于可达集 $\{s(x):x\in\mathcal{X}\}\subseteq\mathcal{B}$ 不一定精确。

## 3 精确分数框 Softmax 优化

设 $y_{j}=e^{s_{j}}$，$L_{j}=e^{\ell_{j}}$，且 $U_{j}=e^{u_{j}}$。映射 $s\mapsto y$ 将 $\mathcal{B}$ 双射到 $\mathcal{B}_{y}=\prod_{j=1}^{K}[L_{j},U_{j}]$，且

$$
F_{c}(s)=R(y)\coloneqq\frac{\sum_{j=1}^{K}c_{j}y_{j}}{\sum_{j=1}^{K}y_{j}}. \tag{7}
$$

因此，softmax 问题是正框上的有界线性分式优化问题。框上线性分式规划的顶点性质是经典的（Dinkelbach，1967（https://arxiv.org/html/2605.10974#bib.bib21））；此处的贡献是 $K+1$ 阈值简化及其在注意力验证中的集成。完整证明见附录 B（https://arxiv.org/html/2605.10974#A2）。

###### 定理 3.1（分数框顶点精确性）。

设 $K\geq 1$，$c\in\mathbb{R}^{K}$，且区间有限 $\ell_{j}\leq u_{j}$。则

$$
\min_{s\in\mathcal{B}}F_{c}(s)=\min_{v\in\operatorname{Vert}(\mathcal{B})}F_{c}(v), \tag{8}
$$

且最大值成立类似的等式。

单坐标原因很简单：固定除 $y_{i}$ 外的所有坐标，比值形式为 $(c_{i}y_{i}+A)/(y_{i}+D)$，其导数具有恒定符号。因此，每个坐标都可以推向端点而不会增加目标值。

将系数升序排序，$c_{(1)}\leq\cdots\leq c_{(K)}$，并相应地重新索引 $L,U$。对于 $m\in\{0,\dots,K\}$ 定义

$$
y^{(m)}_{(j)}=\begin{cases} U_{(j)}, & j\leq m, \\ L_{(j)}, & j> m, \end{cases} \tag{9}
$$

以及

$$
\tau_{m}=R(y^{(m)})=\frac{\sum_{j\leq m}c_{(j)}U_{(j)}+\sum_{j> m}c_{(j)}L_{(j)}}{\sum_{j\leq m}U_{(j)}+\sum_{j> m}L_{(j)}}. \tag{10}
$$

###### 定理 3.2（阈值精确性）。

对于每个 $c\in\mathbb{R}^{K}$ 和每个有限框 $\mathcal{B}$，

$$
L_{\mathrm{box}}(c,\ell,u)=\min_{s\in\mathcal{B}}F_{c}(s)=\min_{m=0,\dots,K}\tau_{m}. \tag{11}
$$

因此，$L_{\mathrm{box}}(c,\ell,u)$ 可以在 $O(K\log K)$ 时间内评估。

直觉是，优化器权衡质量与成本：为了最小化加权 softmax，它在具有最小 $c$ 系数的坐标上最大化指数化的分数（使它们尽可能重），并在最昂贵的坐标上最小化分数。因为 $c$ 已排序，“重-便宜”和“轻-昂贵”坐标之间的最优分割始终是一个连续的阈值。

算法 1 Vertex-Softmax 阈值求解器
1: 输入：系数 $c$，分数下界/上界 $\ell, u$
2: 排序索引使得 $c_{(1)}\leq\cdots\leq c_{(K)}$
3: 设 $a=\max_{j}u_{j}$
4: 计算 $L_{j}=\exp(\ell_{j}-a)$ 和 $U_{j}=\exp(u_{j}-a)$
5: 计算 $U_{(j)}$ 和 $c_{(j)}U_{(j)}$ 的前缀和
6: 计算 $L_{(j)}$ 和 $c_{(j)}L_{(j)}$ 的后缀和
7: 评估式 (10)（https://arxiv.org/html/2605.10974#S3.E10）中的所有比值 $\tau_{m}$
8: 返回 $\min_{m}\tau_{m}$

算法 1（https://arxiv.org/html/2605.10974#alg1）中的稳定性移位通过相同的正因子重新缩放每个分子和分母，因此它保留所有比值并避免溢出。产生证明的区间算术评估路径描述在附录 D（https://arxiv.org/html/2605.10974#A4）中。

###### 推论 3.3（分数框最优性）。

设 $G(c,\ell,u)$ 为任何仅依赖于 $c$ 和独立区间的下界过程。若

$$
G(c,\ell,u)\leq F_{c}(s) \quad \forall s\in\mathcal{B}, \tag{12}
$$

则

$$
G(c,\ell,u)\leq L_{\mathrm{box}}(c,\ell,u). \tag{13}
$$

因此，Vertex-Softmax 耗尽了分数框接口处可用的信息：任何更紧密的证书必须利用分数相关性、分数-值耦合，或比独立框更小的可达分数集。

## 4 Vertex-CROWN 注意力界

上一节提供了给定固定系数和分数区间的单个 softmax 行的精确求解器。为了将其转变为工作验证器，我们需要从 CROWN/LiRPA 反向传递中提取这些量，并将行级证书组合成端到端的边际界。Vertex-CROWN 是由此产生的接口。图 1（https://arxiv.org/html/2605.10974#S4.F1）说明了 Vertex-Softmax 原语在整体验证流水线中的位置。

1. 输入扰动集 $\mathcal{X}$，$x'\in[x-\epsilon, x+\epsilon]$
2. 注意力块仿射界，预激活和中间界
3. 每行分数框和值侧系数下界，$s_{r}\in[\ell_{r},u_{r}]$，$z_{r}^{L}\leq z_{r}(x)$
4. Vertex-Softmax 行求解器，$\min/\max c^{\top}\operatorname{softmax}(s_{r})$，$s_{r}\in[\ell_{r},u_{r}]$
5. 在固定系数分数框接口处精确，认证的注意力行输出界，紧密的行级下界/上界
6. 与剩余验证器行、头、残差和 MLP 界组合
7. 类别边际 $m_{t}(x')=f_{y}(x')-f_{t}(x')$ 的最终下界；若 $>0$ 则认证

图 1：Vertex-Softmax 集成到验证流水线中。输入扰动传播到每个注意力行的独立分数框和值侧系数下界。Vertex-Softmax 在此行级分数框接口处计算精确的方向性 softmax 界。