Sinkhorn 双随机注意力比 softmax 更能保持秩吗？

秩坍缩意味着 token 不再可区分

自注意力层会让每个 token 与所有其他 token 混合。标准实现中注意力矩阵是行随机的：每一行之和为 1，因此可以读作一个概率分布，表示某个查询 token 对各 token 的关注强度。秩坍缩指的是，反复施加这类层会把所有 token 向量推向同一个值，于是序列携带的区分性信息越来越少。论文用残差来度量这一现象，即表示中减去各 token 共享的均值分量后剩下的部分；残差范数越小，表示越接近均匀。

这与熵坍缩是两种不同的失效模式。论文单独把熵坍缩描述为注意力分布高度集中，使每个查询只关注少数几个 token。秩坍缩关心的是混合之后 token 表示的多样性；熵坍缩关心的是注意力权重有多尖锐。论文的实验针对前者，其核心比较是两种注意力矩阵归一化方式，而不是不同的温度设置。

第二种归一化是双随机：行和列之和都为 1。论文用 Sinkhorn 算法替代 softmax 来得到它，该算法交替对行和列重新缩放。作者报告，实验中注意力矩阵在 50 次 Sinkhorn 迭代后实际上已变为双随机，所有实验都采用这一数值。论文给出的动机是，双随机注意力起到类似熵正则化的作用，比只对行归一化更均匀地分散注意力。

理论与实验回答的是不同问题

论文区分了两个容易混淆的论断。理论部分分析的是纯自注意力网络，即堆叠的注意力层，前馈层和跳跃连接都被关闭。实验部分测量的是完整训练好的 Transformer 内部的秩衰减，其中每个块包含一个自注意力层、一个前馈层、跳跃连接和层归一化。因此，关于纯网络的结论不会自动描述训练好的 Transformer，论文也把两者分开处理。

在实验内部，作者还测量了两个不同的量。一个是沿层间路径的注意力矩阵乘积的秩衰减，另一个是自注意力输出本身的秩衰减。两者相关但并不等同，论文把它们放在不同小节报告，而不是用一个替代另一个。

实验设置覆盖三个数据集。文本分类使用 AG's News，文章分为四类，采用 24 层 Transformer 编码器，模型维度 128，四个注意力头，最大序列长度 128。图像分类使用 MNIST 和 Cats and Dogs，采用 24 层、模型维度 128、四个头的 Vision Transformer；MNIST 图像切成 4×4 的 patch，Cats and Dogs 切成 16×16 的 patch。所有情况下都在相同架构下比较 softmax 注意力与 Sinkhorn 注意力。

训练好的模型中 Sinkhorn 注意力衰减更慢

在训练好的架构中，随着深度增加，Sinkhorn 归一化下测得的残差范数始终高于 softmax，这是论文最主要的实证结果：相对于标准行随机注意力，双随机注意力缓解了秩坍缩。作者称这一效应在秩坍缩尚不严重的浅层、以及存在跳跃连接时最为明显。

跳跃连接这一条件并非细节。论文指出，跳跃连接对缓解秩坍缩至关重要，这与先前关于 softmax 注意力的工作一致；实验也改变了是否保留跳跃连接和前馈层。因此，Sinkhorn 归一化所报告的优势与特定架构设置绑定，而不只取决于深度。

测量结果以批次上的均值报告，阴影区域表示一个标准差，第零层的值对应第一个注意力层之前的输入嵌入。这一基线对理解曲线很重要：比较的是初始残差逐层保留了多少，而不是某一层的绝对表示质量。

纯自注意力仍会坍缩到秩一

理论结果给出的是一个上界，而不是豁免。对于使用 Sinkhorn 归一化的单头单层注意力，输出残差范数被一个常数因子乘以输入残差范数的三次方所界定，该常数涉及序列长度、查询与键的维度以及权重矩阵的谱范数。堆叠 L 个这样的层会让指数相乘，得到随输入残差的 3^L 次方缩放的界。

这个指数是关键数字。由于指数随深度以幂次增长而非线性增长，残差随深度以双指数速度收敛到零，输出趋近秩一矩阵。论文指出，同样的双指数行为此前已在 softmax 注意力中确立，因此 Sinkhorn 归一化改变的是常数和实证表现，而不是纯网络中的这一渐近结果。

该上界在特定范围内成立：未归一化的注意力 logits 矩阵的谱范数必须远小于 1，这使作者可以在零附近用线性化局部近似 Sinkhorn 算子。论文明确指出这是一种近似，并把界的紧致性列为局限，指出这些估计有可能被收紧。该结果应被理解为关于这一范围的陈述，而不是对任意注意力幅度的普遍定理。

随机矩阵乘积中 Sinkhorn 没有优势

一个对照实验厘清了实证优势的含义与边界。作者生成随机矩阵，用 softmax 或 Sinkhorn 对每个乘积归一化，并测量同一残差量随路径深度的变化，每个深度对 100 条采样注意力路径取平均。当乘积中的矩阵互不相关时，两种归一化表现出基本相同的秩衰减。

当随机一部分因子被替换为单位矩阵以模拟跳跃连接时，也出现同样的零结果，模拟中替换了一半矩阵。因此，训练好的 Transformer 中观察到的 Sinkhorn 优势并不是双随机矩阵在乘法下的普遍性质。

作者给出了一个解释，并明确将其标为解释：Transformer 中沿路径的注意力矩阵彼此相关，因为某一层的注意力矩阵以上一层的输出为输入。他们还观察到，随机随机矩阵的乘积比从训练好的 Transformer 中提取的注意力矩阵衰减更慢，并将其归因于这种输入依赖产生的级联效应放大了衰减。Sinkhorn 与 softmax 注意力矩阵如何相互作用仍留待未来工作，因此训练模型优势背后的机制并未被这篇论文确定。

Sinkhorn 优势改变判断的适用条件

论文支持的实践区分是有条件的。如果你在带跳跃连接的完整 Transformer 内比较归一化方案，证据指向 Sinkhorn 双随机注意力是衰减更慢的选项，且差距在网络早期最大。如果你考虑的是去掉跳跃连接和前馈层的精简注意力层堆叠，理论表明两种归一化最终都落到秩一，换归一化并不能挽救表示。

随机矩阵对照增加了第二个条件：优势取决于训练模型中注意力矩阵相关且依赖输入的结构，而不只取决于双随机性。声称 Sinkhorn 归一化普遍比 softmax 更好地保持秩，会超出证据范围，因为论文自己的不相关测试显示两者持平。

两项局限界定了实证结果能走多远。模型按完整 Transformer 训练，只在推理时关闭前馈层以隔离纯自注意力行为，因此从头训练纯自注意力网络仍未测试。理论界依赖小 logits 范围内对 Sinkhorn 算子的一阶近似。论文自己列出的未来工作优先事项是比较 Sinkhorn 与 softmax 注意力的 Lipschitz 性质，这关系到鲁棒性和训练稳定性，而不只是秩。

因此值得带走的判断是一个分情况结论，而不是一句定论：在带跳跃连接的训练好的 Transformer 浅层，Sinkhorn 归一化换来更慢的秩衰减；而在纯自注意力中，坍缩到秩一的双指数行为在更换归一化后依然存在。

本文引用的文献