AI Transformer 变体解析：从 Linformer 到 Mamba 的注意力效率演进路径
一、标准注意力的计算瓶颈：O(n²) 为什么不可接受
Transformer 的核心组件是自注意力机制，其计算复杂度为 O(n²d)，其中 n 是序列长度，d 是隐藏维度。当序列长度从 512 增长到 8192 时，注意力矩阵的内存占用从 2MB 增长到 512MB（FP32），计算量增长 256 倍。这意味着标准 Transformer 处理长文档（如法律合同、学术论文）时，内存和计算成本急剧上升。
更深层的问题是注意力矩阵的"信息冗余"。大量研究表明，学习到的注意力模式往往是局部性的——大部分注意力权重集中在少数 Token 上，许多 Token 对之间的注意力权重接近零。这意味着 O(n²) 的计算中很大比例是在处理"无关紧要"的交互。如果能在计算前识别并跳过这些低价值交互，就能在不损失精度的前提下大幅降低复杂度。
二、效率优化的三条路径：稀疏、低秩与状态空间
Transformer 变体的效率优化可以归纳为三条路径：稀疏注意力（只计算部分 Token 对的交互）、低秩近似（用低秩矩阵近似完整注意力矩阵）和状态空间模型（用递归结构替代注意力机制）。
flowchart TB A[标准注意力 O n² d] --> B{优化路径} B -->|稀疏化| C[局部窗口注意力] B -->|稀疏化| D[Sparse Transformer] B -->|低秩近似| E[Linformer: K,V 投影] B -->|低秩近似| F[Performer: 随机特征] B -->|状态空间| G[Mamba: SSM + 选择性扫描] C --> C1[复杂度: O n w d, w 为窗口大小] D --> D1[复杂度: O n sqrt n d] E --> E1[复杂度: O n k d, k 为投影维度] F --> F1[复杂度: O n r d, r 为随机特征数] G --> G1[复杂度: O n d, 线性复杂度] C1 & D1 & E1 & F1 & G1 --> H[精度-效率权衡]Mamba 的创新在于完全抛弃了注意力机制，用选择性状态空间模型（Selective SSM）实现线性复杂度的序列建模。其核心思想是：不是所有历史信息都需要被显式存储和访问，选择性机制可以根据输入动态决定哪些信息需要被保留、哪些可以被遗忘。
三、关键变体的代码实现与对比
3.1 Linformer：低秩近似注意力
""" Linformer: 将 K、V 投影到低维空间 将 n×d 的 K/V 矩阵投影为 k×d（k << n） 复杂度从 O(n²d) 降低到 O(nkd) """ import torch import torch.nn as nn import math class LinformerAttention(nn.Module): """Linformer 注意力模块""" def __init__(self, dim: int, seq_len: int, k: int = 256, num_heads: int = 8): super().__init__() self.dim = dim self.seq_len = seq_len self.k = k # 投影维度，通常 64-256 self.num_heads = num_heads self.head_dim = dim // num_heads # Q、K、V 投影 self.q_proj = nn.Linear(dim, dim) self.k_proj = nn.Linear(dim, dim) self.v_proj = nn.Linear(dim, dim) # Linformer 核心：K 和 V 的降维投影矩阵 # E: n → k, F: n → k self.E = nn.Parameter(torch.randn(seq_len, k)) self.F = nn.Parameter(torch.randn(seq_len, k)) self.out_proj = nn.Linear(dim, dim) def forward(self, x: torch.Tensor) -> torch.Tensor: B, N, D = x.shape assert N <= self.seq_len, \ f"序列长度 {N} 超过预设 {self.seq_len}" # 计算 Q、K、V Q = self.q_proj(x).reshape( B, N, self.num_heads, self.head_dim).transpose(1, 2) K = self.k_proj(x).reshape( B, N, self.num_heads, self.head_dim).transpose(1, 2) V = self.v_proj(x).reshape( B, N, self.num_heads, self.head_dim).transpose(1, 2) # Linformer 核心：K 和 V 乘以投影矩阵 # K: [B, H, N, d] × E: [N, k] → [B, H, k, d] K_proj = torch.einsum("bhid,nk->bhkd", K, self.E[:N, :]) V_proj = torch.einsum("bhid,nk->bhkd", V, self.F[:N, :]) # 注意力计算：Q × K_proj^T → [B, H, N, k] scale = self.head_dim ** -0.5 attn = torch.matmul(Q, K_proj.transpose(-2, -1)) * scale attn = torch.softmax(attn, dim=-1) # 注意力加权：[B, H, N, k] × [B, H, k, d] → [B, H, N, d] out = torch.matmul(attn, V_proj) out = out.transpose(1, 2).reshape(B, N, D) return self.out_proj(out)3.2 Mamba：选择性状态空间模型
""" Mamba: 选择性状态空间模型 核心创新：参数化的 SSM，根据输入动态调整状态转移和输出 实现线性复杂度的序列建模 """ import torch import torch.nn as nn import torch.nn.functional as F class SelectiveSSM(nn.Module): """选择性状态空间模块（Mamba 核心）""" def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4, expand: int = 2): super().__init__() self.d_model = d_model self.d_state = d_state # SSM 状态维度 self.d_conv = d_conv # 局部卷积核大小 self.d_inner = d_model * expand # 内部扩展维度 # 输入投影 self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False) # 局部卷积（捕获短程依赖） self.conv1d = nn.Conv1d( in_channels=self.d_inner, out_channels=self.d_inner, kernel_size=d_conv, padding=d_conv - 1, groups=self.d_inner ) # SSM 参数投影（选择性机制的核心） # A、B、C、Δ 都是输入相关的，而非固定参数 self.x_proj = nn.Linear(self.d_inner, d_state * 2 + 1, bias=False) self.dt_proj = nn.Linear(1, self.d_inner, bias=True) # SSM 的 A 参数（对角矩阵，用 log 形式存储保证正定性） self.A_log = nn.Parameter( torch.log(torch.arange(1, d_state + 1).float() ).unsqueeze(0).expand(self.d_inner, -1)) self.D = nn.Parameter(torch.ones(self.d_inner)) # 跳跃连接 # 输出投影 self.out_proj = nn.Linear(self.d_inner, d_model, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: B, L, D = x.shape # 输入投影并分为两路 xz = self.in_proj(x) x_branch, z = xz.chunk(2, dim=-1) # 局部卷积 x_conv = self.conv1d( x_branch.transpose(1, 2))[:, :, :L].transpose(1, 2) x_conv = F.silu(x_conv) # 计算 SSM 参数（输入相关） ssm_params = self.x_proj(x_conv) B_param = ssm_params[:, :, :self.d_state] C_param = ssm_params[:, :, self.d_state:self.d_state * 2] dt = F.softplus( self.dt_proj(ssm_params[:, :, -1:])) # 步长参数 # SSM 扫描（简化实现，实际使用 CUDA 核函数加速） A = -torch.exp(self.A_log) # 负数保证稳定性 y = self._ssm_scan(x_conv, A, B_param, C_param, dt) # 跳跃连接 + 门控 y = y + self.D * x_conv y = y * F.silu(z) return self.out_proj(y) def _ssm_scan(self, x, A, B, C, dt): """ SSM 递归扫描 h_t = exp(A * dt) * h_{t-1} + B * x_t y_t = C * h_t """ B_batch, L, D_inner = x.shape N = self.d_state # 离散化 A dA = torch.exp(dt.unsqueeze(-1) * A.unsqueeze(1)) # 递归计算 h = torch.zeros(B_batch, D_inner, N, device=x.device) ys = [] for t in range(L): h = dA[:, t] * h + torch.einsum( "bd,bn->bdn", x[:, t], B[:, t]) y_t = torch.einsum("bdn,bn->bd", h, C[:, t]) ys.append(y_t) return torch.stack(ys, dim=1)四、效率优化的精度代价与适用边界
Linformer 的序列长度限制：Linformer 的投影矩阵 E、F 是针对固定序列长度训练的。如果推理时的序列长度超过训练时的seq_len，需要对 E、F 进行插值，但插值会引入近似误差，导致精度下降。建议在训练时使用可能遇到的最大序列长度，或使用可学习的插值策略。
Mamba 的长程依赖局限：Mamba 的 SSM 递归结构天然适合捕获局部依赖，但对长程依赖的建模能力弱于注意力机制。在需要跨段落推理的任务（如文档问答、长文本摘要）上，Mamba 的精度可能低于 Transformer。混合架构（Mamba + 局部注意力）是当前的折中方案。
稀疏注意力的模式设计：Sparse Transformer 需要人工设计稀疏模式（如 strided、fixed），不同任务的最优模式不同。模式设计不当可能导致关键交互被遗漏，精度显著下降。建议从局部窗口模式开始，逐步扩大感受野，通过验证集精度确定最优模式。
Mamba 的 CUDA 依赖：Mamba 的高效实现依赖自定义 CUDA 核函数，在 CPU 或非 NVIDIA GPU 上无法获得理论加速比。纯 Python 实现的递归扫描速度远慢于 CUDA 版本，不适合生产部署。如果目标平台不支持 CUDA，建议使用 Linformer 或 Performer 等纯 PyTorch 实现的方案。
五、总结
Transformer 效率优化有三条路径：稀疏化（局部窗口、Sparse Transformer）、低秩近似（Linformer、Performer）和状态空间模型（Mamba）。选型核心在于"序列长度-长程依赖-部署平台"三角权衡：短序列（< 4K）用标准 Transformer 即可；中等序列（4K-32K）用 Linformer 或局部窗口注意力；超长序列（> 32K）考虑 Mamba。如果任务需要强长程依赖，优先选择低秩近似方案而非 Mamba。Mamba 部署需确认目标平台支持 CUDA 核函数，否则退回到 Linformer 方案。
