Transformer数学基础

Transformer¹ 是现代深度学习的核心架构。从数学视角看，它是一个矩阵化的多层非线性映射，通过自注意力机制实现全局信息交互。本文从第一性原理出发，建立Transformer的完整数学框架，重点放在线性代数视角——因为这是连接Transformer理论与现代谱分析、表达力界、训练动力学的桥梁。

---

1. 整体架构

一个标准Transformer Block 由两个子层构成：

$$\text{Block}(x)=x+\text{FFN}(\text{LN}(x+\text{Attn}(\text{LN}(x))))$$

其中：

• $\text{Attn}$：多头自注意力
• $\text{FFN}$：逐位置前馈网络
• $\text{LN}$：层归一化（现代多用 RMSNorm）

输入：$X\in {\mathbb{R}}^{n\times d}$，$n$ 为序列长度，$d$ 为模型维度。

1.1 矩阵表示的必要性

为什么 Transformer 天然适合矩阵表示？

1. 并行性：矩阵乘法可在 GPU 上高效并行
2. 线性代数工具：可借用 SVD、谱分析、条件数等工具研究性质
3. 几何直觉：每个矩阵可视为某种”投影”或”变换”

---

2. 自注意力的严格推导

2.1 从内积到注意力

起点：给定查询 $q$ 和键 $k$，如何度量它们的相似度？

最简单：内积 $q^{T}k$，但存在尺度问题。

问题演示：

• 若 $q,k\sim \mathcal{N}(0,1)$，则 $q^{T}k\sim \mathcal{N}(0,d)$
• 当 $d$ 大时，内积方差也大，导致 softmax 进入饱和区

解决方案：缩放因子 $\sqrt{d}$

$$\alpha (q,k)=\frac{q^{T}k}{\sqrt{d}}$$

缩放点积注意力：

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$

其中 $Q\in {\mathbb{R}}^{n\times d_k},K\in {\mathbb{R}}^{n\times d_k},V\in {\mathbb{R}}^{n\times d_v}$。

2.2 从零推导（NumPy实现）

import numpy as np
 
def softmax(x, axis=-1):
    x_max = np.max(x, axis=axis, keepdims=True)
    e = np.exp(x - x_max)
    return e / np.sum(e, axis=axis, keepdims=True)
 
def self_attention(X, W_q, W_k, W_v):
    """
    X: (n, d) 输入
    W_q, W_k, W_v: (d, d_k) 投影矩阵
    """
    n, d = X.shape
    Q = X @ W_q           # (n, d_k)
    K = X @ W_k           # (n, d_k)
    V = X @ W_v           # (n, d_v)
    
    # 注意力分数
    scores = Q @ K.T      # (n, n)
    scores_scaled = scores / np.sqrt(d_k)
    
    # 因果掩码（解码器）
    mask = np.triu(np.ones((n, n)) * -1e9, k=1)
    scores_scaled = scores_scaled + mask
    
    # 注意力权重
    attn = softmax(scores_scaled, axis=-1)  # (n, n)
    
    # 输出
    output = attn @ V     # (n, d_v)
    return output, attn
 
 
# 单元测试
np.random.seed(42)
n, d, d_k = 5, 8, 8
X = np.random.randn(n, d)
W_q = np.random.randn(d, d_k) * 0.01
W_k = np.random.randn(d, d_k) * 0.01
W_v = np.random.randn(d, d_k) * 0.01
 
out, attn = self_attention(X, W_q, W_k, W_v)
print(f"输出形状: {out.shape}")  # (5, 8)
print(f"注意力矩阵每行和: {attn.sum(axis=-1)}")  # 接近1

2.3 PyTorch 模块化实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
 
 
class ScaledDotProductAttention(nn.Module):
    """缩放点积注意力"""
    
    def __init__(self, d_k, dropout=0.1):
        super().__init__()
        self.d_k = d_k
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        # Q: (B, h, n, d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        output = torch.matmul(attn, V)
        
        return output, attn
 
 
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        self.attention = ScaledDotProductAttention(self.d_k, dropout)
    
    def forward(self, query, key, value, mask=None):
        B = query.size(0)
        
        # 线性投影并分头
        Q = self.W_q(query).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        if mask is not None:
            mask = mask.unsqueeze(1)  # 广播到所有头
        
        # 计算注意力
        x, attn_weights = self.attention(Q, K, V, mask)
        
        # 合并多头
        x = x.transpose(1, 2).contiguous().view(B, -1, self.d_model)
        
        return self.W_o(x), attn_weights

---

3. 自注意力的矩阵视角

3.1 注意力矩阵作为线性映射

定义注意力矩阵 $A\in {\mathbb{R}}^{n\times n}$：

$$A=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$$

重要性质：

1. 行随机性：每行 ${\sum }_j{A}_{ij}=1$（softmax 保证）
2. 非负性：$A_{ij}\ge 0$
3. 行向量的单纯形：每行属于概率单纯形 ${\Delta }^{n-1}$

输出可重写为：

$$\text{Output}=A\cdot V$$

即输出是输入 $V$ 的凸组合（行方向）。

3.2 注意力矩阵的低秩结构

核心问题：注意力矩阵的秩是多少？

答案：理论上满秩（$n$），但实践中表现出显著的低秩结构。

经验观察：

# 测量训练后Transformer的注意力矩阵秩
def effective_rank(matrix, threshold=0.99):
    """有效秩：保留 threshold 方差所需的最少奇异值"""
    s = np.linalg.svd(matrix, compute_uv=False)
    s = s[s > 1e-10]
    cumsum = np.cumsum(s) / np.sum(s)
    return np.searchsorted(cumsum, threshold) + 1
 
# 实际Transformer中
# 经验发现：注意力矩阵的有效秩通常远小于序列长度 n
# 例如：n=2048 时，有效秩可能仅 50-200

理论解释：

1. Sanky diagram 注意力常呈现少量 token 主导
2. Attention Sink：初始 token 占据大量注意力
3. 归纳偏置：自然语言的稀疏依赖

3.3 谱分析视角

注意力矩阵的特征值分解：

$$A=\sum _{i=1}^n{\lambda }_i{u}_i{v}_i^T$$

其中 ${\lambda }_i$ 为特征值，$u_i,v_i$ 为特征向量。

关键性质：

• 最大特征值：${\lambda }_1=1$（行随机矩阵的 Perron-Frobenius 性质）
• 特征值衰减：${\lambda }_1\gg {\lambda }_2\gg ...\gg {\lambda }_n$
• 谱间隙（spectral gap）：${\lambda }_1-{\lambda }_2$ 反映信息集中度

Rank Collapse 现象（Nait Saada et al. PMLR 2025）²：
当上下文长度 $n\to \infty$ 时，注意力矩阵的特征值集中在 $\frac{1}{n}$ 附近，导致所有 token 趋向均匀混合，丢失信息。

---

4. 多头注意力的张量视角

4.1 多头作为块对角变换

将多头注意力的多个头看作块对角矩阵：

$$\text{Multi-Head Attn}(X)=\text{Concat}(A_1{V}_1,A_2{V}_2,...,A_h{V}_h)W_O$$

其中 $A_i$ 是第 $i$ 个头的注意力矩阵。

有效维度：

• 单头：$d_k$
• 多头：$d_{\text{model}}=h\cdot d_k$
• 总注意力矩阵：${\mathbb{R}}^{n\times n\times h}$（按头）

4.2 表达能力

关键定理（Yau et al. NeurIPS 2025³）：线性 Transformer 的 PAC 可学习性界

定理（简化）：存在算法可在 $\text{poly}(n,d,1/ϵ)$ 时间内 PAC 学习单层线性 Transformer，错误率 $ϵ$。

直观意义：多头并行学习使得 Transformer 能同时建模多种关系（如主谓、修饰等）。

4.3 多头代码实现

class MultiHeadAttentionBlock(nn.Module):
    """完整的多头注意力Block"""
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 自注意力 + 残差
        attn_out, _ = self.attn(x, x, x, mask)
        x = self.ln1(x + self.dropout(attn_out))
        
        # FFN + 残差
        ffn_out = self.ffn(x)
        x = self.ln2(x + self.dropout(ffn_out))
        
        return x

---

5. 前馈网络（FFN）的矩阵分解

5.1 标准FFN

$$\text{FFN}(x)=\sigma (x{W}_1+b_1)W_2+b_2$$

其中 $W_1\in {\mathbb{R}}^{d\times d_{ff}}$，$W_2\in {\mathbb{R}}^{d_{ff}\times d}$，通常 $d_{ff}=4d$。

5.2 矩阵分解视角

将 FFN 视为秩 $d_{ff}$ 的矩阵分解：

$${\text{FFN}}_{\text{matrix}}=W_2\cdot \sigma (W_1\cdot x)$$

与注意力的对比：

组件	信息交互模式	矩阵结构
自注意力	Token之间	$n\times n$ 矩阵
FFN	通道之间	$d\times d$ 矩阵

现代洞见（Llama 系）：

• FFN占据了Transformer总参数的2/3以上
• 它是事实知识的存储位置⁴

5.3 SwiGLU 激活

现代主流（Llama, Mistral, Qwen）：

$$\text{SwiGLU}(x)=\text{SiLU}(x{W}_1)\odot (x{W}_2)$$

与标准FFN对比：

• 标准 FFN: 2 个矩阵
• SwiGLU: 3 个矩阵（门控 + 上投影 + 下投影）
• 参数：$3\cdot d\cdot d_{ff}$ 而非 $2\cdot d\cdot d_{ff}$

经验性能：相同计算量下 SwiGLU 一致优于 GELU/SiLU。

5.4 FFN 实现

class SwiGLU(nn.Module):
    """SwiGLU 激活的前馈网络"""
    
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        # SwiGLU 需要三个矩阵
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        gate = F.silu(self.w_gate(x))
        up = self.w_up(x)
        return self.dropout(self.w_down(gate * up))
 
 
class StandardFFN(nn.Module):
    """标准 FFN (GELU)"""
    
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        return self.dropout(self.linear2(F.gelu(self.linear1(x))))

---

6. 层归一化的线性代数

6.1 LayerNorm

$$\text{LayerNorm}(x)=\gamma \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$

其中 $\mu ,{\sigma }^2$ 沿特征维度计算。

作为仿射变换：可视为对角矩阵 + 平移。

6.2 RMSNorm

简化版（现代主流）：

$$\text{RMSNorm}(x)=\gamma \odot \frac{x}{\sqrt{\frac{1}{d}{\sum }_{i=1}^d{x}_i^2+ϵ}}$$

优势：

• 无需计算均值（节省约 10-15% 计算量）
• 无需偏置 $\beta$
• 实践中性能相当甚至更好

6.3 LayerNorm 的几何解释

关键观察：LayerNorm 使每层的输入位于单位超球面附近。

import torch
import torch.nn as nn
 
 
class RMSNorm(nn.Module):
    """RMSNorm 实现"""
    
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        # x: (B, n, d)
        norm = x.norm(dim=-1, keepdim=True) * (x.shape[-1] ** -0.5)
        return x / (norm + self.eps) * self.weight
 
 
# 测试
x = torch.randn(2, 5, 8)
norm = RMSNorm(8)
out = norm(x)
print(f"输入范数: {x.norm(dim=-1).mean():.4f}")
print(f"输出范数: {out.norm(dim=-1).mean():.4f}")

---

7. 残差连接的矩阵视角

7.1 残差作为隐式低秩正则

残差连接 $x+F(x)$ 的本质：

$$x_{l+1}=x_l+F_l(x_l)=x_0+\sum _{i=0}^l{F}_i(x_i)$$

解释：每层都是对残差路径的累加。

深度网络的有效路径：

# 从输入到输出的路径有 2^l 条（每层可选"通过"或"绕过"残差）
# 残差连接提供了类似 ensemble 的效果

7.2 Pre-norm vs Post-norm

Post-norm（原始 Transformer）：

$$x_{l+1}=\text{LN}(x_l+F_l(x_l))$$

Pre-norm（现代主流）：

$$x_{l+1}=x_l+F_l(\text{LN}(x_l))$$

关键差异：

• Pre-norm：残差路径是”干净”的恒等映射，深度网络更稳定
• Post-norm：每层都重新归一化，约束更强，但深度增加时易训练不稳定

为什么 Pre-norm 主导现代 LLM：

1. 训练稳定性：深层（>12层）Pre-norm 更稳定
2. 学习率宽容度：Pre-norm 对学习率不敏感
3. 可恢复性：某些层训练失败时，其他层仍可用

7.3 残差的谱解释

定理（简化）：残差连接 $x+F(x)$ 在谱上等价于”乘以 $(I+\text{谱变换})$“。

# 谱解释示例
def residual_spectral_analysis(F_func, x, num_iter=10):
    """演示残差连接的谱效应"""
    x_current = x.clone()
    singular_values = []
    
    for i in range(num_iter):
        F_x = F_func(x_current)
        x_next = x_current + F_x
        
        # 计算 Jacobian 的近似谱
        u, s, v = torch.svd(x_next - x_current, some=False)
        singular_values.append(s.mean().item())
        
        x_current = x_next
    
    return singular_values

---

8. 完整 Transformer Block 的实现

8.1 现代 Llama 风格 Block

class LlamaBlock(nn.Module):
    """现代 LLaMA 风格的 Transformer Block"""
    
    def __init__(
        self,
        d_model,
        num_heads,
        d_ff,
        max_seq_len,
        use_rope=True,
        dropout=0.0,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        
        # Pre-RMSNorm
        self.ln1 = RMSNorm(d_model)
        self.ln2 = RMSNorm(d_model)
        
        # 多头注意力 + RoPE
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        if use_rope:
            self.rope = RotaryPositionalEmbedding(d_model // num_heads)
        
        # SwiGLU FFN
        self.ffn = SwiGLU(d_model, d_ff, dropout)
        
        self.use_rope = use_rope
    
    def forward(self, x, mask=None, position_ids=None):
        # Pre-norm + 注意力 + 残差
        x_norm = self.ln1(x)
        if self.use_rope:
            x_norm = self.apply_rope(x_norm, position_ids)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm, mask)
        x = x + attn_out
        
        # Pre-norm + FFN + 残差
        x = x + self.ffn(self.ln2(x))
        
        return x
    
    def apply_rope(self, x, position_ids):
        B, n, _ = x.shape
        x = x.view(B, n, self.num_heads, -1).transpose(1, 2)  # (B, h, n, d_k)
        x = self.rope(x, position_ids)
        return x.transpose(1, 2).contiguous().view(B, n, -1)

8.2 完整的 Transformer 模型

class TransformerLM(nn.Module):
    """完整的 Transformer 语言模型"""
    
    def __init__(
        self,
        vocab_size,
        d_model=512,
        num_heads=8,
        num_layers=6,
        d_ff=2048,
        max_seq_len=2048,
        dropout=0.1,
    ):
        super().__init__()
        self.d_model = d_model
        self.max_seq_len = max_seq_len
        
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList([
            LlamaBlock(d_model, num_heads, d_ff, max_seq_len, dropout=dropout)
            for _ in range(num_layers)
        ])
        self.ln_f = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        
        # 权重绑定
        self.lm_head.weight = self.token_emb.weight
        
        # 初始化
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
    def forward(self, input_ids, mask=None):
        B, n = input_ids.shape
        position_ids = torch.arange(n, device=input_ids.device).unsqueeze(0)
        
        x = self.token_emb(input_ids)
        
        # 因果掩码
        if mask is None:
            mask = torch.triu(
                torch.ones(n, n, device=input_ids.device) * -1e9,
                diagonal=1
            ).unsqueeze(0).unsqueeze(0)
        
        for block in self.blocks:
            x = block(x, mask, position_ids)
        
        x = self.ln_f(x)
        logits = self.lm_head(x)
        
        return logits

---

9. 信息流几何学

9.1 Transformer 作为迭代映射

每一层的几何作用：

$$x_{l+1}=x_l+\text{Attn}(\text{LN}(x_l))+\text{FFN}(\text{LN}(x_l))$$

类比：每层是对输入空间的局部非线性流形变换。

9.2 Token 嵌入作为流形上的点

将序列视为流形上的离散点序列 $\{x_1,x_2,...,x_n\}\subset \mathcal{M}$。

注意力机制的作用：

• 局部信息：token 自身特征
• 全局信息：通过注意力加权聚合的邻居信息

形式化：

$$x_i^{l+1}=x_i^l+\sum _{j=1}^n{\alpha }_{ij}^l{W}_V^l{x}_j^l$$

9.3 通道混合 vs Token 混合

ConvNeXt 视角：

操作	混合维度	对应组件
Token 混合	$n$ 方向	自注意力
通道混合	$d$ 方向	FFN

重要结论：

• Transformer：强 token 混合 + 弱通道混合（注意力）
• CNN：强通道混合 + 弱 token 混合（卷积）
• MLP-Mixer：两者都通过 MLP 实现

# Token 混合 vs 通道混合的简单示例
class TokenMixer(nn.Module):
    """Token 混合（自注意力）"""
    def __init__(self, n, d):
        super().__init__()
        self.proj = nn.Parameter(torch.randn(d, d) * 0.02)
    
    def forward(self, x):
        # x: (B, n, d)
        # 计算 token 间关系
        return x @ self.proj  # 通道间
 
 
class ChannelMixer(nn.Module):
    """通道混合（FFN）"""
    def __init__(self, n, d):
        super().__init__()
        self.proj = nn.Parameter(torch.randn(n, n) * 0.02)
    
    def forward(self, x):
        # x: (B, n, d)
        # 计算 token 间关系
        return self.proj @ x  # token 间

---

10. Transformer 与卷积的对比

10.1 数学对比

性质	卷积	自注意力
权重	静态（与输入无关）	动态（依输入生成）
感受野	局部 + 堆叠扩展	全局（单层）
归纳偏置	平移等变性	无（数据驱动）
计算复杂度	$O(k^2\cdot n\cdot d)$	$O(n^2\cdot d)$
参数复杂度	$O(k^2\cdot d^2)$	$O(d^2)$

10.2 谱滤波视角

卷积：在傅里叶域是固定的乘法：

$$\mathcal{F}(\text{Conv}\ast x)(\omega )=\hat{K}(\omega )\cdot \hat{x}(\omega )$$

注意力：在某种意义上是输入相关的乘法：

$$\text{Attn}(X)=A(X)\cdot V(X)$$

关键差异：

• 卷积的滤波器 $\hat{K}$ 是固定的
• 注意力的权重 $A(X)$ 随输入 $X$ 变化（数据依赖）

10.3 全局 vs 局部归纳偏置

CNN 的归纳偏置：

• 平移等变性
• 局部性
• 尺度不变性（部分）

Transformer 的归纳偏置：

• 排列不变性（需加位置编码）
• 全局感受野
• 数据驱动的动态计算

---

11. 与核方法的桥梁

11.1 注意力作为核方法

关键观察（Tsai et al. 2019⁵）：注意力矩阵可以理解为核矩阵：

$$K(x_i,x_j)=\frac{\varphi (x_i{)}^T\psi (x_j)}{\sqrt{d}}$$

其中 $\varphi ,\psi$ 是通过 $W_Q,W_K$ 参数化的特征映射。

核视角的洞察：

1. 注意力学习的是输入空间的核函数
2. 不同头学习不同的核
3. 多头 = 多核学习

11.2 线性注意力

核心简化：去除 softmax

$$\text{LinearAttn}(Q,K,V)=\varphi (Q)(\varphi (K{)}^{T}V)$$

利用结合律：$(Q{K}^T)V=Q(K^{T}V)$

复杂度：$O(n\cdot d^2)$（线性于序列长度）

现代进展：

• Performers (Choromanski et al. 2021): 随机特征近似
• Linear Transformers (Katharopoulos et al. 2020): $\varphi =\text{elu}$
• cosFormer (Qin et al. 2022): 余弦注意力

11.3 核方法代码

class LinearAttention(nn.Module):
    """线性注意力（核方法）"""
    
    def __init__(self, d_model, num_heads, kernel_fn='elu'):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        # 特征映射
        if kernel_fn == 'elu':
            self.phi = lambda x: F.elu(x) + 1
        elif kernel_fn == 'relu':
            self.phi = F.relu
        else:
            self.phi = lambda x: x
    
    def forward(self, x):
        B, n, d = x.shape
        Q = self.phi(self.W_q(x).view(B, n, self.num_heads, self.d_k)).transpose(1, 2)
        K = self.phi(self.W_k(x).view(B, n, self.num_heads, self.d_k)).transpose(1, 2)
        V = self.W_v(x).view(B, n, self.num_heads, self.d_k).transpose(1, 2)
        
        # 线性注意力：利用结合律
        # (B, h, n, d_k) @ (B, h, d_k, n) -> (B, h, d_k, n) 然后 @ V
        KV = K.transpose(-2, -1) @ V  # (B, h, d_k, d_k)
        out = Q @ KV  # (B, h, n, d_k)
        
        # 归一化
        K_sum = K.sum(dim=-2, keepdim=True)  # (B, h, 1, d_k)
        Z = Q @ K_sum.transpose(-2, -1)  # (B, h, n, 1)
        out = out / (Z + 1e-6)
        
        out = out.transpose(1, 2).contiguous().view(B, n, d)
        return self.W_o(out)

---

12. 表达力界

12.1 通用逼近性

定理（简化版，Yun et al. 2020⁶）：单层多头注意力可以表示任何稀疏依赖的函数。

意义：从理论上保证 Transformer 能学习稀疏注意力模式。

12.2 TC⁰ 复杂性

重要结果（Merrill & Sabharwal 2023[^7]）：固定精度 Transformer 属于 ${\text{TC}}^0$（阈值电路的常数深度多项式规模）。

推论：

• 标准 Transformer 无法高效解决某些问题（如 Permutation）
• 但可以通过位置编码、CoT 等扩展

12.3 深度 vs 宽度

经验法则：

• 深度：对长距离依赖和层次化抽象至关重要
• 宽度：对并行特征和记忆容量至关重要

典型规模：

模型	深度	宽度	FFN
GPT-2 Small	12	768	3072
LLaMA-7B	32	4096	11008
LLaMA-70B	80	8192	28672

---

13. 训练动力学的线性代数

13.1 梯度流

设损失 $\mathcal{L}$，参数 $\theta$：

$$\dot{\theta }=-{\nabla }_{\theta }\mathcal{L}$$

Transformer 中梯度计算的关键矩阵：

• 注意力矩阵 $A$：影响 token 间梯度传播
• Jacobian $\frac{\partial x_{l+1}}{\partial x_l}$：控制信号传播

13.2 注意力矩阵与梯度消失

问题：深层 Transformer 中，梯度通过注意力矩阵的链式传播：

$$\frac{\partial \mathcal{L}}{\partial x_0}=\prod _{l=1}^L\frac{\partial x_l}{\partial x_{l-1}}$$

关键：

$$\frac{\partial x_l}{\partial x_{l-1}}=I+\frac{\partial F_l}{\partial x_{l-1}}$$

残差连接确保 $I$ 主导，避免梯度消失。

13.3 NTK 视角

NTK（Neural Tangent Kernel）：在无限宽度极限下，Transformer 训练等价于核回归。

关键矩阵：

$$\Theta (x,x^{\prime })={\mathbb{E}}_{\theta \sim \text{init}}\left[⟨{\nabla }_{\theta }f(x),{\nabla }_{\theta }f(x^{\prime })⟩\right]$$

Transformer 的 NTK：与RBF核相关，但有特殊结构。

---

14. 注意力模式的解析分析

14.1 常见注意力模式

import torch
import matplotlib.pyplot as plt
 
 
def visualize_attention_patterns(model, input_text):
    """可视化不同层的注意力模式"""
    attentions = []
    hooks = []
    
    def hook(module, input, output):
        if hasattr(output, '__len__') and len(output) > 1:
            attentions.append(output[1].detach())
    
    for block in model.blocks:
        block.attn.attention.register_forward_hook(hook)
    
    with torch.no_grad():
        model(input_text)
    
    for hook in hooks:
        hook.remove()
    
    return attentions
 
 
# 典型模式：
# 1. 起始 token sink（注意力集中在 BOS）
# 2. 局部 attention（相邻 token）
# 3. 列模式（特定 token 接收所有注意力）
# 4. 全局模式（均匀分布）

14.2 模式分类

模式	含义	典型位置
Sink	初始 token	多数头/层
Local	相邻 token	中间层
Column	特定位置	后层
Global	均匀分布	后层

---

15. 现代Transformer栈的关键设计选择

15.1 收敛栈（基于53个模型分析）

组件	主流选择	备选
Normalization	RMSNorm / Pre-norm	LayerNorm / Post-norm
位置编码	RoPE	ALiBi, Sinusoidal
激活函数	SwiGLU	GELU, ReLU²
KV Cache	GQA / MQA	MHA
偏置	无	有（GPT-2等）
注意力	全局	滑动窗口 + 全局

15.2 偏离收敛栈的著名模型

• Gemma 2：使用 Post-norm，但有 hybrid attention
• OLMo：使用 LayerNorm（非 RMS）
• Cohere Command R：使用 ALiBi 而非 RoPE

---

16. 实践中的工程考量

16.1 数值稳定性

常见陷阱：

1. 注意力 softmax 溢出：使用 -1e9 而非 -inf
2. LayerNorm 的零方差：加 eps
3. 混合精度：fp16/bf16 需要 loss scaling

16.2 内存优化

关键技术：

• Gradient Checkpointing：用时间换空间
• Flash Attention：避免显式 $n\times n$ 矩阵
• KV Cache：推理时避免重复计算

# Flash Attention 的核心思想（简化）
def flash_attention(Q, K, V, block_size=64):
    """分块计算注意力，避免完整矩阵"""
    B, h, n, d = Q.shape
    output = torch.zeros_like(Q)
    log_normalizer = torch.zeros(B, h, n, 1)
    
    for i in range(0, n, block_size):
        for j in range(0, n, block_size):
            Q_block = Q[:, :, i:i+block_size]
            K_block = K[:, :, j:j+block_size]
            V_block = V[:, :, j:j+block_size]
            
            scores = Q_block @ K_block.transpose(-2, -1) / (d ** 0.5)
            
            # 增量式 softmax
            new_max = torch.maximum(log_normalizer[i:i+block_size], scores.max(dim=-1, keepdim=True).values)
            ...
    
    return output

16.3 实现陷阱

常见错误：

1. 位置编码泄漏：在 attention 之前 vs 之后加
2. 注意力掩码方向：causal mask 的方向
3. 残差流的初始化：scaled init 避免激活爆炸

---

17. 进阶话题导航

17.1 与本专题其他文档的连接

• 注意力机制现代理论 — 共识、最优传输、知识容量
• RoPE位置编码理论 — 频率熵、相位调制
• 架构收敛模式 — 53个LLM分析
• Attention Sink分析 — 结构起源
• Transformer谱分析 — Rank Collapse
• Transformer-SSM混合 — Jamba等

17.2 与相关领域的连接

• 注意力机制的线性代数
• 谱分析与深度学习
• Transformer与注意力机制
• 反向传播

---

18. 关键论文推荐

必读

1. Vaswani et al. (2017) — Attention Is All You Need
2. Devlin et al. (2018) — BERT
3. Radford et al. (2019) — GPT-2
4. Brown et al. (2020) — GPT-3 / Scaling Laws

进阶

5. Su et al. (2021) — RoPE
6. Shazeer (2020) — SwiGLU (Noam Shazeer, GLU variants)
7. Child et al. (2019) — Sparse Attention
8. Choromanski et al. (2021) — Performers (Linear Attention)

前沿（2025-2026）

9. Yau et al. (NeurIPS 2025) — Learning Linear Attention in Polynomial Time
10. Abella et al. (PMLR 267) — Consensus Is All You Get
11. Nait Saada et al. (PMLR 267) — Spectral Analysis of Rank Collapse
12. Barbero et al. (ICLR 2025) — What Makes Rotary Useful?

---

19. 总结

Transformer 的本质：

视角	理解
代数	自注意力是 $Q{K}^T$ 的 softmax 加权，输出是 $V$ 的凸组合
几何	每层是流形上的局部非线性映射，残差提供”测地线”路径
核	注意力矩阵是输入依赖的核函数
谱	注意力矩阵是行随机矩阵，主特征值为1
动力学	残差使深度网络的 Jacobian 谱接近1

现代 LLM 的栈：RMSNorm + RoPE + SwiGLU + GQA + Bias-free + Pre-norm

未来方向：

• 混合架构（与 SSM 结合）
• 注意力替代（线性注意力、状态空间）
• 推理时扩展（CoT、Test-time Compute）

---

最后更新：2026-06-21
相关文档：transformer-architecture/ 专题目录

Footnotes

1. Vaswani et al. (2017). Attention Is All You Need. ↩

2. Nait Saada, Naderi, Tanner (2025). Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers. PMLR 267. ↩

3. Yau et al. (2025). Learning Linear Attention in Polynomial Time. NeurIPS 2025. ↩

4. Geva et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. ↩

5. Tsai et al. (2019). Transformer Dissection: A Unified Understanding of Transformer’s Attention via the Lens of Kernel. ↩

6. Yun et al. (2020). Are Transformers Universal Approximators of Sequence-to-Sequence Functions? ↩