Rethinking Attention: Polynomial Alternatives to Softmax in Transformers

原文：arXiv:2410.18613 · Saratchandran et al., 2024

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

1. Intro

Softmax self-attention 满足三条性质：

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

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

本文提出问题：

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

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

Contributions：

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

2. Related Works

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

3. Preliminaries and Notation

定义 Transformer block：

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

Attention head：

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

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

4. Theoretical Analysis

Frobenius 的定义：

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

且有$||A||^2_F=\text{tr}(A^\top A)=\text{tr}(AA^\top)=\sum_k\sigma_k^2$。

Frobenius 范数等于奇异值的 L2 范数，因此有$||A||_2\leq ||A||_F\leq \sqrt{\text{rank}(A)}||A||_2$。且有：

1. $\langle A,B\rangle_F=\text{tr}(A^\top B)$
2. $||UAV||_F=||A||_F$
3. $||AB||_F\leq ||A||_F||B||_F, ||AB||_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:

证明：(1)$\|\mathrm{softmax}(A)\|_F\le\sqrt N$

记$\mathrm{softmax}(A)$第$(k,l)$元为$F_{kl}=\dfrac{e^{a_{kl}}}{\sum_{j}e^{a_{kj}}}$。每一行$k$是概率向量$F_{kl}\ge 0, \sum_l F_{kl}=1$。按定义，有：

由于 “平方和小于等于和的平方” 可展开，展开后交叉项$\sum_{l\ne l'}x_lx_{l'}\ge 0$，

$N$行各贡献$\le 1$，故$\|\mathrm{softmax}(A)\|_F^2\le N$，开方得$\sqrt N$。

(2)$\|\nabla\mathrm{softmax}(A)\|_F\le 2\sqrt N.$

softmax 逐行独立作用，输出$F_{kl}$只依赖输入第$k$行，故整张 Jacobian 是按行分块对角的，共$N$个块，每块是单行 softmax 的 Jacobian。又因$\nabla\mathrm{softmax}=J^\top$ 而$\|J^\top\|_F=\|J\|_F$，只需算 Jacobian。单行 (概率$p=(p_1,\dots,p_N)$) 的 Jacobian 是

于是

把单个输出$F_{ij}$关于本行输入的梯度记为$\nabla F_{ij}$，分量$(\nabla F_{ij})l=p_j(\delta{jl}-p_l)$，然后用 $\|v\|_2^2\le\|v\|_1^2$放成$L_1$:

因行和为 1，$|1-p_j|=1-p_j$ 且$\sum_{l\ne j}p_l=1-p_j$，两者相加等于$2(1-p_j)$代入并再用$(1-p_j)^2\le(1-p_j)$:

对$j$求和，$\sum_j 4p_j^2(1-p_j)\le 4\sum_j p_j^2\le 4$。故每行$\le 4$，$N$行$\le 4N$，开方$2\sqrt N$。

不过实际上注意到，

令$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$):

由柯西不等式，$\big(\sum p_j^2\big)^2=\big(\sum p_j^{1/2}\cdot p_j^{3/2}\big)^2\le\big(\sum p_j\big)\big(\sum p_j^3\big)=\sum p_j^3$，即$\sum p_j^3\ge S^2$。代入：

最大值$\tfrac14$在$S=\tfrac12$取到（例如某行是两个$\tfrac12$的均匀分布；$N=2,p=(\tfrac12,\tfrac12)$时块 $=\big[\begin{smallmatrix}1/4&-1/4\\-1/4&1/4\end{smallmatrix}\big], \|\cdot\|_F^2=4\cdot\tfrac1{16}=\tfrac14$, 精确取等)。于是每行$\le\tfrac14$, $N$行$\le\tfrac N4$,

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

4.2. Polynomial activations for self-attention

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

Theorem 4.2.

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

对角项等于$\sigma_1^2\sigma_2^2$，非对角项因 a,b 均值零、互相独立而为 0，故

每个元同理, 共$N^2$个,$\mathbb E\|C\|_F^2=N^2D\sigma_1^2\sigma_2^2$, 琴生不等式得

当$p>1$时，要算$\mathbb E[c_{11}^2]=\mathbb E[s^{2p}],s=\sum_{i=1}^D X_i$，其中$X_i:=a_{1i}b_{i1}$， 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$)。

做多项式展开：

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

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

这是$N$无关的常数。$N^2$个元求和:$\mathbb E\|C\|_F^2\le N^2 D^p(2p-1)!!\,\sigma_1^{2p}\sigma_2^{2p}$，琴生得

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

Theorem 4.3.

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

Theorem 4.4

记$M_{ij}=x_i^\top QK^\top x_j$($x_i$为$X$第$i$行,$D$- 维),$Q,K\in\mathbb R^{D\times d}, x\sim\mathcal N(0,\sigma_x^2)$、$K$元方差记 $\sigma_w^2$。要界$\mathbb E\big\|\frac{\partial}{\partial Q}M^{\circ p}\big\|_F$($M^{\circ p}$逐元素幂)。整张梯度是 4 - 张量 (输出 $N\times N$, 参数$D\times d$), 其 Frobenius 范数平方 $=\sum{i,j}\big\|\frac{\partial M_{ij}^p}{\partial Q}\big\|_F^2$。

$p=1$时有：

所以

于是

展开内平方：

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

对全部$i,j$求和：

合计

琴生得$\mathbb E\|\partial M/\partial Q\|_F\le ND\sqrt d\,\sigma_x^2\sigma_w=O(N)$。

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

记 $z_{ij}:=M_{ij}=\sum_{k,m}x_{ik}x_{jm}a_{km}, a_{km}=\sum_l Q_{kl}K_{ml}=(QK^\top)_{km}$_。_则要算

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

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

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

缩放得 $\mathbb E\|\partial M^{\circ p}/\partial Q\|_F\le O(N)$, 缩放 $1/\sqrt N$ 得 $O(\sqrt N)$。

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