摘要：本文质疑”softmax 注意力的强性能源于它在输入上产生概率分布”这一观点，转而主张 softmax 的有效性在于它对注意力矩阵 Frobenius 范数的隐式正则，从而稳定训练。据此探索替代激活——特别是多项式——以达到类似的正则效果。理论分析表明某些多项式可作为 softmax 的有效替代，尽管违反了 softmax 通常的正性、归一化、稀疏性，仍能在多种 transformer 应用上取得强性能。大量实验支持这些发现，为注意力机制提供了新视角。

1. Intro

Softmax self-attention 满足三条性质：

注意力权重非负，
行归一化使权重和为 1(可解释为概率)，
稀疏性，促使聚焦少数相关 token。

这些性质被普遍认为是有效注意力建模与可解释性的关键，但这一看法主要建立在经验之上、缺乏理论支撑。尽管有人研究过替代激活，softmax 仍占主导，主要因其经验性能与可解释性。

本文提出问题：

Do attention mechanisms in transformers require non-negativity， normalization， and sparsity for effective performance?

并给出新的视角：softmax 的有效性不内禀地源于这三条性质，而源于训练中对注意力矩阵 Frobenius 范数的隐式正则。据此，简单的多项式激活(违反三条性质)仍能通过诱导类似的正则效果，达到与 softmax 相当或更优的性能。这给出了一个与传统概率视角根本不同的解释。

Contributions：

理论分析:softmax attention 隐式正则注意力矩阵的 Frobenius 范数，挑战”非负/归一化/稀疏是其成功主因”的假设。
提出多项式激活作为替代，证明其能在不遵守 softmax 约束的情况下诱导类似正则效果，在图像分类、目标检测、实例分割、文本分类、物理建模上取得有竞争力的性能。

ReLU，ReLU + $\frac{1}{N}$ ，softmax的泰勒展开等。作者认为本文的独特之处在于：建立了注意力矩阵 Frobenius 范数与输入序列长度之间的明确理论联系 ，并据此设计了打破三条 softmax 性质的多项式激活。

3. Preliminaries and Notation

定义Transformer block：

T : R^{N \times D} \to R^{N \times D}, T (x) = F (A (x) + x)

其中 $F$ 是带残差的 FFN， $A$ 是 attention head。

Attention head：

q = QX, k = K X, v = V X, Q, K \in R^{D \times d}, V \in R^{D \times M}

定义 $A (X) = ϕ (S (q, k)) v$ ， $S$ 是相似度变换， $ϕ$ 是激活。最常用 $S (q, k) = q k^{⊤}$ self-attention)，最常用 $ϕ$ 是 softmax，于是

A (X) = softmax (\frac{q k ^{⊤}}{d}) v = softmax (\frac{X Q K ^{⊤} X ^{⊤}}{d}) X V .

Softmax梯度：对 $A \in R^{N \times N}$ ，定义梯度线性映射 $\nabla softmax (A) := J_{softmax} (A)^{⊤}$ 。约定 Frobenius 范数记 $∥ \cdot ∥_{F}$ 。

4. Theoretical Analysis

Frobenius的定义：

对 $A = (a_{ij}) \in R^{m \times n}$ ，定义

∣∣ A ∣ ∣_{F} = i = 1 \sum m j = 1 \sum n a_{ij}^{2}

且有 $∣∣ A ∣ ∣_{F}^{2} = tr (A^{⊤} A) = tr (A A^{⊤}) = \sum_{k} σ_{k}^{2}$ 。

Frobenius 范数等于奇异值的L2范数，因此有 $∣∣ A ∣ ∣_{2} \leq ∣∣ A ∣ ∣_{F} \leq rank (A) ∣∣ A ∣ ∣_{2}$ 。且有：

$⟨ A, B ⟩_{F} = tr (A^{⊤} B)$
$∣∣ U A V ∣ ∣_{F} = ∣∣ A ∣ ∣_{F}$
$∣∣ A B ∣ ∣_{F} \leq ∣∣ A ∣ ∣_{F} ∣∣ B ∣ ∣_{F}, ∣∣ A B ∣ ∣_{F} \leq ∣∣ A ∣ ∣_{2} ∣∣ B ∣ ∣_{F}$

4.1. Implicit regulatization of Softmax

本节证明 softmax 对 self-attention 矩阵 Frobenius 范数的控制_随 token 长度次线性增长_，并且 softmax 关于该矩阵的梯度也有类似规律性。前人从 Lipschitz 常数角度分析过 softmax 的规律性，本文的新意是直接把 Frobenius 范数规律性与 token 长度挂钩，从而给出”注意力激活该如何随 token 长度缩放才能在(基于梯度下降的)训练中保持稳定”的洞见。

Theorem 4.1:

∥ softmax (A) ∥_{F} \leq N, ∥\nabla softmax (A) ∥_{F} \leq 2 N .

证明：(1) $∥ softmax (A) ∥_{F} \leq N$

记 $softmax (A)$ 第 $(k, l)$ 元为 $F_{k l} = \frac{e ^{a_{k l}}}{\sum _{j} e ^{a_{k j}}}$ 。每一行 $k$ 是概率向量 $F_{k l} \geq 0, \sum_{l} F_{k l} = 1$ 。按定义，有：

∥ softmax (A) ∥_{F}^{2} = k = 1 \sum N l = 1 \sum N F_{k l}^{2} = k = 1 \sum N (l = 1 \sum N F_{k l}^{2}) .

由于“平方和小于等于和的平方”可展开，展开后交叉项 $\sum_{l \neq = l^{'}} x_{l} x_{l^{'}} \geq 0$ ，

l \sum F_{k l}^{2} \leq (l \sum F_{k l})^{2} = 1^{2} = 1.

$N$ 行各贡献 $\leq 1$ ，故 $∥ softmax (A) ∥_{F}^{2} \leq N$ ，开方得 $N$ 。

注意到这个界的松紧由稀疏度决定，行向量越接近one-hot的时候， $\sum_{l} F_{k l}^{2} \to 1$ ，取上界；越均匀， $F_{k l}^{2} \to 1/ N$ ，整体 $∣∣ \cdot ∣ ∣_{F} \to 1$ ，因此实际范围应当是 $[1, N]$ 。¹

(2) $∥\nabla softmax (A) ∥_{F} \leq 2 N .$

softmax 逐行独立作用，输出 $F_{k l}$ 只依赖输入第 $k$ 行，故整张 Jacobian是按行分块对角的，共 $N$ 个块，每块是单行 softmax 的 Jacobian。又因 $\nabla softmax = J^{⊤}$ 而 $∥ J^{⊤} ∥_{F} = ∥ J ∥_{F}$ ，只需算 Jacobian。单行(概率 $p = (p_{1}, \dots, p_{N})$ )的 Jacobian 是

J^{(row)} = diag (p) - p p^{⊤}, J_{j l}^{(row)} = p_{j} (δ_{j l} - p_{l}) .

于是

∥\nabla softmax (A) ∥_{F}^{2} = rows \sum diag (p) - p p^{⊤}_{F}^{2} .

把单个输出 $F_{ij}$ 关于本行输入的梯度记为 $\nabla F_{ij}$ ，分量 $(\nabla F_{ij}) l = p_{j} (δ j l - p_{l})$ ，然后用 $∥ v ∥_{2}^{2} \leq ∥ v ∥_{1}^{2}$ 放成 $L_{1}$ :

∥\nabla F_{ij} ∥^{2} \leq (l \sum ∣ (\nabla F_{ij})_{l} ∣)^{2} = p_{j}^{2} (∣1 - p_{j} ∣ + l \neq = j \sum p_{l})^{2}

因行和为 1， $∣1 - p_{j} ∣ = 1 - p_{j}$ 且 $\sum_{l \neq = j} p_{l} = 1 - p_{j}$ ，两者相加等于 $2 (1 - p_{j})$ 代入并再用 $(1 - p_{j})^{2} \leq (1 - p_{j})$ :

∥\nabla F_{ij} ∥^{2} \leq 4 p_{j}^{2} (1 - p_{j})^{2} \leq 4 p_{j}^{2} (1 - p_{j}) .

对 $j$ 求和， $\sum_{j} 4 p_{j}^{2} (1 - p_{j}) \leq 4 \sum_{j} p_{j}^{2} \leq 4$ 。故每行 $\leq 4$ ， $N$ 行 $\leq 4 N$ ，开方 $2 N$ 。

不过实际上注意到，

diag (p) - p p^{⊤}_{F}^{2} = 对角 j \sum (p_{j} - p_{j}^{2})^{2} + 非对角 j \neq = l \sum (p_{j} p_{l})^{2} .

令 $S = \sum_{j} p_{j}^{2}$ ，展开整理(对角项给 $\sum p_{j}^{2} - 2 \sum p_{j}^{3} + \sum p_{j}^{4}$ ，非对角项给 $S^{2} - \sum p_{j}^{4}$ ):

diag (p) - p p^{⊤}_{F}^{2} = S - 2 j \sum p_{j}^{3} + S^{2} .

由柯西不等式， $(\sum p_{j}^{2})^{2} = (\sum p_{j}^{1/2} \cdot p_{j}^{3/2})^{2} \leq (\sum p_{j}) (\sum p_{j}^{3}) = \sum p_{j}^{3}$ ，即 $\sum p_{j}^{3} \geq S^{2}$ 。代入：

S - 2 \sum p_{j}^{3} + S^{2} \leq S - 2 S^{2} + S^{2} = S (1 - S) \leq \frac{1}{4},

最大值 $\frac{1}{4}$ 在 $S = \frac{1}{2}$ 取到（例如某行是两个 $\frac{1}{2}$ 的均匀分布； $N = 2, p = (\frac{1}{2}, \frac{1}{2})$ 时块 $= [1/4 - 1/4 - 1/4 1/4], ∥ \cdot ∥_{F}^{2} = 4 \cdot \frac{1}{16} = \frac{1}{4}$ ,精确取等)。于是每行 $\leq \frac{1}{4}$ , $N$ 行 $\leq \frac{N}{4}$ ,

∥\nabla softmax (A) ∥_{F} \leq \frac{1}{2} N .

所以论文的常数 2 比可达的最优常数 $\frac{1}{2}$ 松了 4 倍。

4.2. Polynomial activations for self-attention

既然 softmax 隐式正则了 Frobenius 范数,现在证明对特定多项式激活做缩放,可以在期望意义上达到类似效果。

Theorem 4.2.

设 $X \in R^{N \times D}, Q, K \in R^{D \times d}$ 为 i.i.d. 高斯, $X \sim N (0, σ_{x}), Q, K \sim N (0, σ_{t})$ ²。则对 $p \geq 1$ ,有：

E [(\frac{X Q K ^{T} X ^{T}}{d})^{p}_{F}] \leq O (N) .

证明：当 $p = 1$ 时， $c_{11} = \sum_{i = 1}^{D} a_{1 i} b_{i 1}$ ，则有

E [c_{11}^{2}] = E [i, j \sum a_{1 i} b_{i 1} a_{1 j} b_{j 1}] = i = j i \sum E [a_{1 i}^{2}] E [b_{i 1}^{2}] + i \neq = j i \neq = j \sum E [a_{1 i}] E [b_{i 1}] E [a_{1 j}] E [b_{j 1}] .

对角项等于 $σ_{1}^{2} σ_{2}^{2}$ ，非对角项因a,b均值零、互相独立而为0，故

E [c_{11}^{2}] = D σ_{1}^{2} σ_{2}^{2} .

每个元同理,共 $N^{2}$ 个, $E ∥ C ∥_{F}^{2} = N^{2} D σ_{1}^{2} σ_{2}^{2}$ ,琴生不等式得

E ∥ C ∥_{F} \leq N D σ_{1} σ_{2} = O (N)

当 $p > 1$ 时，要算 $E [c_{11}^{2}] = E [s^{2 p}], s = \sum_{i = 1}^{D} X_{i}$ ，其中 $X_{i} := a_{1 i} b_{i 1}$ ， i.i.d. 均值零、 $\mathrm{Var}(X_i)=\sigma_1^2\sigma_2^2$$、且$ \mathbb E[X_i^{\text{odd}}]=0 $(奇次矩$ =\mathbb E[a^{\text{odd}}]\mathbb E[b^{\text{odd}}]=0$)。

做多项式展开：

s^{2 p} = (m_{1}, \dots, m_{D}) \sum m_{i} = 2 p \sum \frac{( 2 p )!}{m _{1} ! \dots m _{D} !} i \prod X_{i}^{m_{i}}

取期望，注意到含奇数 $m_{i}$ 的项因为 $E [X_{i}^{odd}] = 0$ 消去了。按 $D$ 的阶看,主导项是 $2 p$ 拆成 $p$ 对 $m_{i} = 2$ 落在 $p$ 个不同下标上(计数 $(p D) = O (D^{p})$ ,这是最多不同因子、计数最高的情形;一个 $m_{i} = 4$ 加 $(p - 2)$ 对只用 $p - 1$ 个下标 $\to O (D^{p - 1})$ ,低阶)。主导项:

E [s^{2 p}] \approx (p D) \cdot \frac{( 2 p )!}{2 ^{p}} \cdot (σ_{1}^{2} σ_{2}^{2})^{p}

其中 $\frac{( 2 p )!}{2 ^{p}}$ 是把 $p$ 个 $m_{i} = 2$ 的多项式系数( $2! = 2$ , $p$ 个)合进去。化简 $(p D) \frac{( 2 p )!}{2 ^{p}} = \frac{D !}{( D - p )!} \cdot \frac{( 2 p )!}{p ! 2 ^{p}}$ ,再用 $(2 p)!! = 2^{p} p!$ 与 $(2 p)! = (2 p)!! (2 p - 1)!!$ 得 $\frac{( 2 p )!}{p ! 2 ^{p}} = \frac{( 2 p )!}{( 2 p )!!} = (2 p - 1)!!$ 。注意到 $\frac{D !}{( D - p )!} = D (D - 1) \dots (D - p + 1) \leq D^{p}$ 于是

E [c_{11}^{2}] \approx \frac{D !}{( D - p )!} (2 p - 1)!! σ_{1}^{2 p} σ_{2}^{2 p} \leq D^{p} (2 p - 1)!! σ_{1}^{2 p} σ_{2}^{2 p}

这是 $N$ 无关的常数。 $N^{2}$ 个元求和: $E ∥ C ∥_{F}^{2} \leq N^{2} D^{p} (2 p - 1)!! σ_{1}^{2 p} σ_{2}^{2 p}$ ，琴生得

E ∥ C ∥_{F} \leq N D^{p} (2 p - 1)!! σ_{1}^{p} σ_{2}^{p} = O (N) .

注意到有个关于 $p$ 的常数 $(2 p - 1)!! D^{p} σ^{2 p}$ 随着 $p$ 指数增长，跟正文里面提到的 $p$ 太大了就无法训练能相互印证。

Theorem 4.3.

E \frac{1}{N} (\frac{X Q K ^{T} X ^{T}}{d})^{p}_{F} \leq O (N) .

显然的缩放，乘 $1/ N$ 即 $O (N)$ 。

Theorem 4.4

记 $M_{ij} = x_{i}^{⊤} Q K^{⊤} x_{j}$ ( $x_{i}$ 为 $X$ 第 $i$ 行, $D$ -维), $Q, K \in R^{D \times d}, x \sim N (0, σ_{x}^{2})$ 、 $K$ 元方差记 $σ_{w}^{2}$ 。要界 $E \frac{\partial}{\partial Q} M^{\circ p}_{F}$ ( $M^{\circ p}$ 逐元素幂)。整张梯度是 4-张量(输出 $N \times N$ ,参数 $D \times d$ ),其 Frobenius 范数平方 $= \sum i, j \frac{\partial M _{ij}^{p}}{\partial Q}_{F}^{2}$ 。

$p = 1$ 时有：

M_{ij} = k, l \sum x_{ik} Q_{k l} (K^{⊤} x_{j}) l, (K^{⊤} x_{j}) l = m \sum K_{m l} x_{j m},

所以

\frac{\partial M _{ij}}{\partial Q _{k l}} = x_{ik} m \sum x_{j m} K_{m l} ⟺ \frac{\partial M _{ij}}{\partial Q} = x_{i} x_{j}^{⊤} K \in R^{D \times d} .

于是

\frac{\partial M _{ij}}{\partial Q} F^{2} = \sum k = 1^{D} l = 1 \sum d x_{ik}^{2} (m \sum x_{j m} K_{m l})^{2} .

展开内平方：

(m \sum x_{j m} K_{m l})^{2} = m \sum x_{j m}^{2} K_{m l}^{2} + m \neq = n \sum x_{j m} x_{j n} K_{m l} K_{n l}

取期望,交叉项含 $E [K_{m l} K_{n l}] = 0$ ( $m \neq = n$ )而消失,剩

E \frac{\partial M _{ij}}{\partial Q} F^{2} = \sum k, l, m E [x_{ik}^{2} x_{j m}^{2}] σ_{w}^{2} .

对全部 $i, j$ 求和：

E \frac{\partial M}{\partial Q} F^{2} = σ_{w}^{2} \sum i, j, k, l, m E [x_{ik}^{2} x_{j m}^{2}] .

$i = j$ 且 $k = m$ : $E [x_{ik}^{4}] = 3 σ_{x}^{4}$ ,共 $N \cdot D \cdot d$ 个。³
其余: $E [x_{ik}^{2} x_{j m}^{2}] = σ_{x}^{4}$ , $N^{2} D^{2} d - N D d$ 个。

合计

E \frac{\partial M}{\partial Q}_{F}^{2} = σ_{x}^{4} σ_{w}^{2} d (N^{2} D^{2} + 2 N D) \approx N^{2} D^{2} d σ_{x}^{4} σ_{w}^{2} .

琴生得 $E ∥ \partial M / \partial Q ∥_{F} \leq N D d σ_{x}^{2} σ_{w} = O (N)$ 。

再考虑 $p > 1$ 的情况，链式法则:

\frac{\partial M _{ij}^{p}}{\partial Q} = p M_{ij}^{p - 1} \frac{\partial M _{ij}}{\partial Q} = p M_{ij}^{p - 1} x_{i} x_{j}^{⊤} K,

\frac{\partial M _{ij}^{p}}{\partial Q} F^{2} = p^{2} M ij^{2 p - 2} k, l \sum (x_{ik} m \sum x_{j m} K_{m l})^{2} .

记 $z_{ij} := M_{ij} = \sum_{k, m} x_{ik} x_{j m} a_{k m}, a_{k m} = \sum_{l} Q_{k l} K_{m l} = (Q K^{⊤})_{k m}$ _。_则要算

E \frac{\partial M ^{\circ p}}{\partial Q} F^{2} = p^{2} \sum i, j E [z ij^{2 p - 2} k, l \sum (x_{ik} \sum_{m} x_{j m} K_{m l})^{2}] .

用与 Thm 4.2 同样的”奇次矩消、偶次配对取 $D$ 的 leading order”机制,论文给出主导项:

E \frac{\partial M ^{\circ p}}{\partial Q} F^{2} \approx N^{2} D^{2 p} d^{p} (2 p - 3)!! σ_{x}^{4 p} σ_{w}^{4 p - 2} .

这里 $(2 p - 3)!!$ 是 $z^{2 p - 2}$ 因子配对的高斯计数（与 $2 p$ 次矩出 $(2 p - 1)!!$ 同理, $(2 p - 2)$ 阶出 $(2 p - 3)!!$ ）。 $σ_{x}^{4 p} σ_{w}^{4 p - 2}$ 是数 $x$ 与 $K$ 因子总幂次的结果( $z$ 携带 $x^{2} \cdot (Q K)$ , $2 p - 2$ 次再乘 $∥ \cdot ∥^{2}$ 的 $x^{2} x^{2} K^{2}$ )。

关键是 $N$ 依赖恰为 $N^{2}$ ,与 $p$ 无关——它纯粹来自外层 $\sum i, j$ 的 $N^{2}$ 个输出位置,每个位置的期望梯度能量 $N$ 无关。

缩放得 $E ∥ \partial M^{\circ p} / \partial Q ∥_{F} \leq O (N)$ ,缩放 $1/ N$ 得 $O (N)$ 。

理论推完发现这文章没中，ICLR2026得分2222，一下子就不想看下去了。感觉实验和理论都不是很好。

也解释了Fig2里面softmax的范数曲线实际上远低于预期的数值，因为它实际上没有工作在one hot like的上界附近。 ↩
其实感觉这里有点不对，因为 $X Q$ 是两个高斯矩阵的乘积，显然不是一个高斯矩阵，并且一行内的元素共享 $x_{i}$ ，所以相互之间也是相关的。论文里面的Remark4.6只说了初始化的时候成立但没讨论这个问题。 ↩
论文里面的证明在此处疑似有问题，感觉它的 $i = j, m \neq = k$ 的地方是不是重复算了？ ↩

Rethinking Attention: Polynomial Alternatives to Softmax in Transformers