O(n²)到O(n)的吸引力

标准注意力的O(n²)复杂度是长序列处理的核心障碍。线性注意力旨在将复杂度降至O(n)，使百万级token序列成为可能。但这个目标在保持注意力核心功能的同时实现极其困难——注意力之所以是O(n²)，正是因为它需要建模所有Query-Key对之间的关系。

线性注意力的核函数框架

基本公式

标准注意力：Attention(Q,K,V) = softmax(QK^T)V

如果我们将softmax替换为一个核函数 φ(·)，使得：

Attention(Q,K,V) = φ(Q)φ(K)^T V = φ(Q)(φ(K)^T V)

关键变换在于结合律：先计算 φ(K)^T V（与Q无关），得到一个 d×d 的矩阵，然后与 φ(Q) 相乘。这使得复杂度从 O(n²d) 降为 O(nd²)。当 n » d 时，这是显著的速度提升。

Performer的随机特征

Performer使用随机特征方法来近似softmax核：

def performer_attention(Q, K, V, n_features=256):
    """Performer的随机特征近似"""
    d = Q.shape[-1]
    
    # 生成随机投影矩阵
    omega = torch.randn(d, n_features) / math.sqrt(d)
    
    # 非负映射函数
    def phi(x):
        return torch.exp(x @ omega - 0.5 * (x ** 2).sum(-1, keepdim=True)) / math.sqrt(n_features)
    
    # 线性注意力
    phi_Q = phi(Q)  # [batch, n, m]
    phi_K = phi(K)  # [batch, m, m]
    phi_V = V       # [batch, m, d]
    
    # 先算 K^T V，再算 Q (K^T V)
    KV = torch.einsum('bmd,bme->bde', phi_K, phi_V)  # [batch, d, d]
    output = torch.einsum('bnd,bde->bne', phi_Q, KV)  # [batch, n, d]
    
    # 归一化
    normalizer = torch.einsum('bnd,bmd->bn', phi_Q, phi_K).unsqueeze(-1)
    output = output / (normalizer + 1e-6)
    
    return output

Performer的优势是无偏近似（随着特征数增大趋于精确），但需要较大的特征维度才能接近softmax注意力的性能。

Linear Transformer

更简单的方案是直接使用 elu(x)+1 作为核函数：

φ(x) = elu(x) + 1

这种方法无需随机特征，计算极其简单，但表达能力远不如softmax。在实践中，线性Transformer在短序列上表现尚可，但在长序列上与标准注意力的差距随序列长度增大而增大。

线性注意力的核心问题

表达能力损失

softmax函数的非线性是注意力表达力的关键来源。它提供了"赢家通吃"的尖锐注意力分布。线性核函数只能产生平滑的注意力分布，无法有效区分重要和不重要的Key。

数值稳定性

当序列很长时，φ(K)^T V 的累积可能溢出或下溢。需要仔细的归一化策略来维持数值稳定。

因果掩码的困难

标准因果注意力使用下三角掩码，这在线性注意力中难以高效实现——因为掩码破坏了结合律。RFA（Recurrent Flash Attention）通过维护一个运行状态来处理因果性，但状态更新引入了信息丢失。

状态空间模型：另一条路

线性注意力的困境促使研究者探索完全不同的架构。状态空间模型（SSM）从RNN和信号处理中汲取灵感，提供了另一种线性复杂度的序列建模方法。

S4与结构化SSM

S4模型使用一个结构化的状态矩阵来建模序列：

h(t) = Ah(t-1) + Bx(t)
y(t) = Ch(t)

通过HiPPO矩阵初始化A，S4能够在长序列上有效捕捉依赖关系。其核心创新是利用矩阵的结构（对角加低秩）将计算从O(n²L)降到O(nL)。

Mamba的选择性机制

Mamba在SSM的基础上引入了选择性——让状态更新依赖于输入：

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_state = d_state
        
        # 输入依赖的参数投影
        self.proj_B = nn.Linear(d_model, d_state, bias=False)
        self.proj_C = nn.Linear(d_model, d_state, bias=False)
        self.proj_delta = nn.Linear(d_model, 1, bias=True)
        
        # 可学习的A参数（固定）
        A = torch.randn(d_state)
        self.log_A = nn.Parameter(torch.log(-A))  # 确保A为负
        
    def forward(self, x):
        """简化的Mamba块前向传播"""
        batch_size, seq_len, _ = x.shape
        
        # 输入依赖的B, C, delta
        B = self.proj_B(x)  # [batch, seq, d_state]
        C = self.proj_C(x)  # [batch, seq, d_state]
        delta = F.softplus(self.proj_delta(x))  # [batch, seq, 1]
        
        A = -torch.exp(self.log_A)  # [d_state]
        
        # 离散化并递归计算
        h = torch.zeros(batch_size, self.d_state, device=x.device)
        outputs = []
        
        for t in range(seq_len):
            # 输入依赖的离散化
            dA = torch.exp(A * delta[:, t])  # [batch, d_state]
            h = dA * h + B[:, t] * x[:, t]   # 状态更新
            y = torch.einsum('bd,bd->b', C[:, t], h)  # 输出
            outputs.append(y)
        
        return torch.stack(outputs, dim=1).unsqueeze(-1)

Mamba的关键创新是让 Δ（步长）依赖于输入——重要token使用大步长（快速更新状态），不重要token使用小步长（缓慢更新）。这赋予了SSM类似于注意力"选择性聚焦"的能力。

Mamba-2与结构化对偶

Mamba-2建立了SSM与注意力之间的数学对偶关系。它证明了选择性SSM可以看作一种特殊的线性注意力，从而将SSM的优势（线性复杂度、选择性）与注意力的优势（表达力、并行训练）结合起来。

混合架构：取各家之长

2026年的趋势是将注意力、SSM和其他线性复杂度机制混合使用：

Jamba架构

AI21的Jamba模型交替使用注意力层和Mamba层：

  • 少数注意力层提供精确的长距离依赖建模
  • 多数Mamba层提供线性复杂度的序列处理
  • MoE层在FFN部分实现宽度稀疏

Zamba

Zamba在单个共享SSM块后添加一个注意力块，用注意力"修正"SSM的不足。这种设计减少了注意力的使用频率，同时保持了对复杂依赖的建模能力。

混合门控

在每个层动态选择使用注意力还是SSM：

class HybridLayer(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.attn = AttentionLayer(d_model)
        self.ssm = MambaBlock(d_model)
        self.gate = nn.Linear(d_model, 1)
    
    def forward(self, x):
        gate_score = torch.sigmoid(self.gate(x.mean(dim=1)))
        attn_out = self.attn(x)
        ssm_out = self.ssm(x)
        return gate_score * attn_out + (1 - gate_score) * ssm_out

线性注意力的评估基准

长距离依赖

在Long Range Arena（LRA）基准上，各类线性注意力方案的表现差异显著：

方法LRA平均分复杂度因果支持
Standard Transformer58.0O(n²)✅
Linear Transformer38.0O(n)✅
Performer46.0O(n)✅
Linformer52.0O(n)❌
S483.0O(n)✅
Mamba86.0O(n)✅

SSM类方法在长距离依赖上显著优于线性核注意力。

语言建模困惑度

在语言建模任务上，纯线性注意力模型的困惑度通常高于标准Transformer。混合架构（如Jamba）可以在接近标准Transformer性能的同时实现更低的推理成本。

2026年前沿方向

并行SSM训练

传统SSM的递归特性使得训练时无法并行。Mamba通过并行扫描算法解决了这个问题，但扫描的效率仍低于注意力的大规模矩阵乘法。最新的研究正在探索完全可并行的SSM形式。

多头SSM

借鉴多头注意力的思想，使用多个独立的SSM头，每个头学习不同的序列模式。这增加了模型的表达能力，但也增加了参数量。

SSM + 注意力的理论分析

为什么混合架构优于纯注意力或纯SSM？最新的理论研究正在揭示：注意力擅长精确的信息检索（如"第3个词是什么"），而SSM擅长连续的信息压缩（如"整体语义是什么"）。两者的互补性是混合架构成功的理论基础。

结语

线性注意力从最初的核函数近似发展到今天的状态空间模型和混合架构，已经成为长序列建模的重要方向。虽然纯线性注意力还无法完全替代标准注意力，但SSM和混合架构的进展让我们看到了在保持线性复杂度的同时实现接近注意力表达力的希望。在AGI需要处理超长上下文的未来，这条路线将越来越重要。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。