注意力机制的谱分析与秩崩溃

概述

注意力机制是 Transformer 的核心，但其数学性质长期缺乏严格分析。最新研究（Mind the Gap, 2024）使用随机矩阵理论揭示了注意力机制中一个惊人的现象：秩崩溃（Rank Collapse）。¹

本文件深入分析这一现象的数学本质，揭示为什么某些架构改进（如 Differential Transformer）能够有效提升性能。

---

谱分析基础

随机矩阵理论回顾

设 $A\in {\mathbb{R}}^{n\times n}$ 为对称矩阵，其特征值满足：

$${\lambda }_1\ge {\lambda }_2\ge \cdots \ge {\lambda }_n$$

谱范数：$\Vert A{\Vert }_2=|{\lambda }_1|$（最大特征值的绝对值）

稳定秩（Stable Rank）：

$$\mathrm{sr}(A)=\frac{\Vert A{\Vert }_F^2}{\Vert A{\Vert }_2^2}=\frac{{\sum }_i{\lambda }_i^2}{{\lambda }_1^2}$$

稳定秩衡量矩阵的”有效秩”，对噪声和扰动更鲁棒。

关键引理：扰动矩阵的谱

定理（Weyl 不等式扰动界）：

设 $M=A+B$，其中 $A$ 是对称矩阵，$B$ 是扰动。若 $\Vert B\Vert \le ϵ$，则：

$$|{\lambda }_i(M)-{\lambda }_i(A)|\le ϵ$$

---

Softmax 注意力矩阵的谱结构

定义

设输入矩阵 $X\in {\mathbb{R}}^{T\times d}$，其中 $T$ 是序列长度，$d$ 是维度。

注意力分数矩阵：

$$S=\frac{Q{K}^T}{\sqrt{d}}=X{W}_Q{W}_K^T{X}^T/\sqrt{d}$$

注意力矩阵：

$$A=\mathrm{softmax}(S)\in {\mathbb{R}}^{T\times T}$$

注意：$A$ 的每一行是一个概率分布，元素非负且行和为 1。

均匀矩阵

设 $1\in {\mathbb{R}}^T$ 是全 1 向量，均匀矩阵定义为：

$$U=\frac{1}{T}11^T\in {\mathbb{R}}^{T\times T}$$

$U$ 的性质：

• ${\lambda }_1(U)=1$（ Perron-Frobenius 特征值）
• ${\lambda }_2(U)=\cdots ={\lambda }_T(U)=0$（秩为 1）

注意力矩阵的分解

核心发现：当输入 $X$ 的列向量两两正交时，注意力矩阵可以分解为：

$$A(X_0)=U+A^{\perp }$$

其中 $A^{\perp }$ 是与 $U$ 正交的部分，满足 $1^T{A}^{\perp }=0$。

---

秩崩溃定理

主要定理（Mind the Gap, 2024）

定理 1：Softmax 注意力的谱间隙

设 $X_0\in {\mathbb{R}}^{T\times d}$ 的列向量两两正交且范数均为 1，则注意力矩阵 $A(X_0)$ 满足：

$${\lambda }_1(A)=1$$

且

$$\underset{T\to \infty }{\lim }{s}_2(\sqrt{T}A)=2\sigma$$

其中 $s_2(\cdot )$ 是第二奇异值，$\sigma$ 是与初始化方差相关的常数。

定理 2：稳定秩崩溃

$$\mathrm{sr}(A)\to 1\text{当}T\to \infty$$

这意味着随着序列长度增加，有效秩趋向于 1！

数学证明概述

步骤 1：分解注意力矩阵

将 $A$ 分解为均匀部分和扰动部分：

$$A=U+A^{\perp }$$

其中 $U=\frac{1}{T}11^T$。

步骤 2：分析 $A^{\perp }$ 的奇异值

对于 $\sqrt{T}{A}^{\perp }$，考虑其渐近谱分布。使用随机矩阵理论中的半圆定律，可以证明：

$$\underset{T\to \infty }{\lim }\frac{1}{T}\sum _{i=1}^T{\delta }_{{\lambda }_i(\sqrt{T}{A}^{\perp })}\to \text{Semi-circle distribution}$$

步骤 3：应用扰动理论

$A=U+\frac{1}{\sqrt{T}}(\sqrt{T}{A}^{\perp })$

这是一个秩 1 扰动问题。秩 1 扰动的效果是将主特征值从 $O(1/\sqrt{T})$ “推出”到 1。

直观解释

注意力矩阵的谱分布（示意）

T 较小时:                    T → ∞ 时:
    λ                          λ
    ↑                          ↑
    │    ●●●●                  │      ●
    │   ●    ●                 │    ●   ●
    │  ●      ●                │   ●     ●
    │ ●        ●               │  ●       ●
    │●●          ●             │ ●         ●
    └────────────→             └────────────→
    λ₁=1  奇异值分布           λ₁=1  只有一个有效方向

---

秩崩溃的物理意义

奇异值塌缩

当 $T$ 很大时，$\sqrt{T}A$ 的奇异值满足：

$$s_1(\sqrt{T}A)=\sqrt{T}$$ $$s_2(\sqrt{T}A)\to 2\sigma$$

这意味着存在一个主导方向（对应均匀注意力），其他方向被压制。

对应什么？

${\lambda }_1(A)=1$ 对应的特征向量是 $1$（全 1 向量），表示均匀注意力模式。

换句话说，所有 token 的表示趋向于收敛到同一个方向！

宽度方向的崩溃

定理 3：宽度方向的秩崩溃

考虑宽度为 $d$ 的 Transformer 层，定义

$$\Sigma =\frac{1}{T}\sum _{t=1}^T{h}_t{h}_t^T$$

其中 $h_t$ 是第 $t$ 个 token 的隐藏表示。

当 $T\to \infty$ 时：

$$\mathrm{sr}(\Sigma )\to 1$$

即表示的协方差矩阵退化为秩 1！

深度方向的加速崩溃

定理 4：深度加速效应

设第 $\ell$ 层之后，token 表示的稳定秩满足：

$$\mathrm{sr}({\Sigma }^{(\ell )})=\Theta (T^{1-4\ell })$$

这意味着：

• 第 1 层：$\mathrm{sr}\sim T^{-3}$
• 第 2 层：$\mathrm{sr}\sim T^{-7}$
• …

崩溃速度随深度指数加速！

---

秩崩溃的负面影响

1. 表示退化

所有 token 的表示趋向于相同，失去了区分不同位置的能力。

2. 梯度问题

定理 5：梯度爆炸

秩崩溃导致梯度范数指数增长：

$$\Vert \frac{\partial \mathcal{L}}{\partial W^{(1)}}\Vert \approx T^{L-1}$$

其中 $L$ 是层数。

3. 信息瓶颈

主导奇异值对应的方向成为信息传递的”瓶颈”，其他方向的信息被丢失。

4. 优化困难

秩崩溃使得优化 landscape 变得崎岖，收敛到好的解变得更加困难。

---

解决秩崩溃的方法

方法 1：移除均匀矩阵（Differential Transformer）

核心思想：减去均匀注意力来移除谱异常值。

$$A^{\prime }=\mathrm{softmax}(S)-\frac{1}{T}11^T$$

这正是 Differential Transformer（DiT）的做法。²

def differential_attention(Q, K, V, mask=None):
    """
    Differential Attention
    
    A' = softmax(QK^T / sqrt(d)) - (1/T) 11^T
    """
    d_k = Q.shape[-1]
    T = Q.shape[-2]
    
    # 标准注意力
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    attn = F.softmax(scores, dim=-1)
    
    # 减去均匀矩阵
    uniform = torch.ones(T, T, device=Q.device) / T
    attn_differential = attn - uniform
    
    # 重新归一化
    attn_differential = F.softmax(attn_differential * math.sqrt(d_k), dim=-1)
    
    return torch.matmul(attn_differential, V)

方法 2：Centered Self-Attention

Ali 等人（2023） 提出的方法：

$$A^{\prime }=\mathrm{softmax}(S)-\frac{1}{T}11^T$$

直接移除 $T^{-1}11^T$ 项。

方法 3：谱归一化

对注意力矩阵进行谱归一化：

$$A^{\prime }=\frac{A}{\Vert A{\Vert }_2}$$

这限制了主导奇异值，防止秩崩溃。

方法 4：正则化

添加谱正则化项到损失函数：

$${\mathcal{L}}_{\text{reg}}=\mathcal{L}+\lambda \cdot \mathrm{sr}(\Sigma )$$

鼓励模型保持较高的稳定秩。

---

理论验证

实验设置

import torch
import torch.nn as nn
import numpy as np
 
def compute_stable_rank(A):
    """计算矩阵的稳定秩"""
    # 确保 A 是对称的
    A = (A + A.T) / 2
    
    # 特征值分解
    eigenvalues = torch.linalg.eigvalsh(A)
    
    # 稳定秩
    stable_rank = (eigenvalues**2).sum() / (eigenvalues[0]**2)
    return stable_rank.item()
 
def analyze_attention_rank(X, num_layers=12, num_heads=12):
    """
    分析 Transformer 层的秩崩溃
    """
    results = {
        'layer': [],
        'stable_rank': [],
        'attention_entropy': []
    }
    
    d_model = X.shape[-1]
    num_heads = num_heads
    d_k = d_model // num_heads
    
    # 初始化投影矩阵
    W_q = nn.Linear(d_model, d_model)
    W_k = nn.Linear(d_model, d_model)
    W_v = nn.Linear(d_model, d_model)
    
    for layer in range(num_layers):
        # Q, K, V
        Q = W_q(X).reshape(X.shape[0], X.shape[1], num_heads, d_k)
        K = W_k(X).reshape(X.shape[0], X.shape[1], num_heads, d_k)
        V = W_v(X).reshape(X.shape[0], X.shape[1], num_heads, d_k)
        
        # 注意力矩阵
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
        attn = F.softmax(scores, dim=-1)
        
        # 分析秩
        for h in range(num_heads):
            stable_rank = compute_stable_rank(attn[0, :, h, :])
            results['stable_rank'].append(stable_rank)
            results['layer'].append(layer)
        
        # 更新表示
        X = torch.matmul(attn, V).reshape(X.shape[0], X.shape[1], d_model)
    
    return results
 
# 实验
T = 512  # 序列长度
d = 512  # 维度
X = torch.randn(1, T, d) / math.sqrt(d)  # 正交初始化
 
results = analyze_attention_rank(X)
print(f"初始稳定秩: {results['stable_rank'][0]:.4f}")
print(f"最终稳定秩: {results['stable_rank'][-1]:.4f}")

预期结果

Layer 0:  stable_rank ≈ 0.85  (轻微崩溃)
Layer 1:  stable_rank ≈ 0.72
Layer 2:  stable_rank ≈ 0.58
...
Layer 11: stable_rank ≈ 0.12  (严重崩溃)

---

与其他理论的关系

1. 与 Neural Tangent Kernel 的联系

当 $T\to \infty$ 时，Transformer 的训练动态可以用 NTK 描述。

秩崩溃意味着 NTK 矩阵趋向于低秩，失去了捕捉复杂依赖关系的能力。

2. 与特征值浓度分析的联系

ICML 2024 的研究发现：³

“Self-attention networks localize when QK-eigenspectrum concentrates”

这与秩崩溃理论是一致的：特征值集中（concentrate）导致注意力局部化（localize）。

3. 与低秩近似的联系

秩崩溃使得注意力矩阵可以被很好地用低秩矩阵近似：

$$A\approx u{v}^T+{\sigma }_2\cdot \text{(次要方向)}$$

这解释了为什么各种低秩压缩方法能够有效工作。

---

实践建议

1. 架构设计

• 考虑使用 Differential Attention 机制
• 添加显式的去中心化（de-centering）操作
• 使用多查询注意力（MQA）时注意秩崩溃风险

2. 训练策略

• 序列长度预热：从短序列开始，逐渐增加到目标长度
• 谱正则化：惩罚低秩表示
• 注意力 dropout：随机置零部分注意力连接，增加鲁棒性

3. 模型分析

• 定期监控稳定秩指标
• 分析注意力熵：均匀注意力熵最大，崩溃后熵降低
• 可视化特征值分布：观察是否出现明显的秩 1 趋势

---

数学附录

A. 随机矩阵理论关键结果

半圆定律

设 $A\in {\mathbb{R}}^{n\times n}$ 为 Wigner 矩阵，其谱分布趋向于区间 $[-2\sigma ,2\sigma ]$ 上的半圆分布：

$$\rho (\lambda )=\frac{1}{2\pi {\sigma }^2}\sqrt{4{\sigma }^2-{\lambda }^2}$$

Marchenko-Pastur 分布

对于样本协方差矩阵，谱分布由 Marchenko-Pastur 定律描述。

B. 注意力矩阵的渐近分析

设 $T,d\to \infty$ 且 $T/d\to \gamma$（固定比例）。

定理（随机矩阵注意力谱）：

$\sqrt{T}A$ 的奇异值分布趋向于确定性的自由加乘卷积分布，其支撑集为 $[a,b]$。

---

参考

---

相关阅读

• Transformer 与注意力机制 — 注意力机制基础
• 注意力矩阵低秩压缩 — 基于谱分析的压缩方法
• 矩阵范数与神经网络 — 谱归一化等技术
• LSTM 与状态空间对偶性 — SSM 与注意力机制的统一视角

Footnotes

1. Mind the Gap team. “Mind the Gap: Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers”. arXiv:2410.07799, 2024. ↩

2. Ye, T., et al. (2024). “Differential Transformer”. NeurIPS. ↩

3. Bao, et al. (2024). “Self-attention Networks Localize When QK-eigenspectrum Concentrates”. ICML. ↩