这次我们来看两个能显著降低注意力机制计算复杂度的关键工作：Linformer 和 Performer。对于任何在本地部署或微调大语言模型（LLM）的开发者来说，注意力机制的 O(n²) 复杂度都是一个绕不开的瓶颈，它直接限制了模型能处理的序列长度，并推高了显存和计算成本。Linformer 和 Performer 分别通过“低秩投影”和“核化+结合律”这两条不同的技术路径，将复杂度从 O(n²) 降至 O(n)，让长文本处理在有限硬件上成为可能。

如果你关心如何在消费级显卡上跑更长的上下文、降低推理延迟，或者想深入理解如何优化 Transformer 架构的核心模块，这篇文章会直接切入核心。我们将重点拆解这两个方法的核心思想、实现差异、硬件门槛，并通过一个概念性的代码示例，展示它们如何被集成到现有的注意力计算流程中。本文不会停留在理论公式，而是聚焦于它们“能不能用”、“怎么用”以及“用了之后效果如何”的工程视角。

1. 核心能力速览

2. 适用场景与使用边界

适合谁用？

• 模型研究者与算法工程师：需要深入理解并尝试改进 Transformer 效率，为自家模型集成更高效的注意力模块。
• 本地部署 LLM 的开发者：受限于 8G、12G 等消费级显卡显存，希望在不升级硬件的情况下处理更长的输入文本（如整个 PDF 文档、长对话历史）。
• 需要处理长序列任务的应用方：例如法律文档分析、长文本摘要、代码仓库理解、高分辨率图像分块处理等。

能解决什么问题？

1. 突破序列长度瓶颈：将传统注意力无法处理的超长序列（如 10k+ tokens）变为可能。
2. 大幅降低显存占用：避免存储巨大的 n×n 注意力矩阵，这是 OOM（内存溢出）的常见原因。
3. 降低计算延迟：线性复杂度意味着处理长序列时，计算时间增长更平缓，提升推理速度。

不适合什么场景？

• 极短序列：当序列长度 n 很小时（如 < 512），传统注意力的开销本身不大，引入近似可能带来不必要的精度损失和实现复杂度。
• 对注意力权重有精确解释性要求的场景：近似方法无法提供精确的、可逐点解释的注意力分布图。
• 某些特定的预训练模型微调：如果下游任务极度依赖预训练阶段学到的精确注意力模式，直接替换为近似注意力可能需要谨慎的再训练或适配。

使用边界与合规性：

• 本文讨论的线性注意力是通用的模型架构优化方法，不涉及特定数据、模型或应用。
• 在实际应用中，若使用基于线性注意力改进的模型处理用户数据，需遵守数据隐私与安全规范。
• 使用相关开源实现时，请遵循其对应的许可证（如 MIT、Apache 2.0）。

3. 环境准备与前置条件

要理解或实验 Linformer 和 Performer，你需要一个能够运行 PyTorch 或 JAX 的深度学习环境。以下是一个通用的环境检查清单：

1. 操作系统: Linux (Ubuntu 20.04+ 推荐), Windows (WSL2), macOS。Linux 环境对深度学习支持最友好。
2. Python: 3.8 或 3.9 版本。建议使用 conda 或 venv 创建虚拟环境。
3. 深度学习框架:
   • PyTorch: 1.9+ 版本。这是大多数研究和工程实现的首选。
   • JAX(可选): 如果你要深入研究 Performer 的官方实现或相关变体。
4. CUDA 与显卡驱动(GPU 环境):
   • 确保安装与 PyTorch 版本匹配的 CUDA Toolkit (如 CUDA 11.7, 11.8)。
   • 更新显卡驱动至最新稳定版。
5. 硬件建议:
   • GPU: 至少 8GB 显存，用于体验长序列（>2048）与标准注意力的显存差异。拥有更多显存（12G/24G）可以测试更极端的序列长度。
   • CPU/RAM: 作为备选，可以在 CPU 上运行小规模实验，但需要足够的内存（32GB+）来加载模型和中间变量。
6. 代码与库:
   • 安装基础科学计算库：pip install numpy
   • 为了后续可能的代码实验，可以安装transformers库和一些工具：pip install transformers datasets tqdm

4. 原理精讲与代码概念演示

本章节将深入两者的核心机制，并用高度简化的代码说明其如何改变计算流程。

4.1 Linformer：低秩投影的直觉

Linformer 的核心假设是：在 Transformer 的自注意力中，经过 softmax 后的 n×n 注意力矩阵是低秩的。这意味着，我们可以用两个更小的矩阵来近似它。

具体操作：

1. 对于长度为n的序列，我们有两个投影矩阵E_i,F_i∈ R^{k×n}，其中k是一个远小于n的常数（如 256）。
2. 将原始的 Key (K) 和 Value (V) 矩阵（形状为n×d）分别与这两个投影矩阵相乘：
   • K_compressed = E_i · K(形状: k×d)
   • V_compressed = F_i · V(形状: k×d)
3. 注意力计算变为：Attention(Q, K, V) = softmax(Q·K_compressed^T / sqrt(d)) · V_compressed
4. 计算流程从Q(n×d) @ K^T(d×n) -> (n×n) @ V(n×d)变为Q(n×d) @ K_compressed^T(d×k) -> (n×k) @ V_compressed(k×d)。复杂度从 O(n²d) 降为 O(nkd)。当 k 固定时，即为 O(n)。

概念代码 (PyTorch):

import torch import torch.nn as nn import torch.nn.functional as F class LinformerAttention(nn.Module): def __init__(self, d_model, n_heads, seq_len, k=256): super().__init__() self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads self.seq_len = seq_len self.k = k # 定义投影矩阵 E 和 F self.E = nn.Parameter(torch.randn(n_heads, k, seq_len)) self.F = nn.Parameter(torch.randn(n_heads, k, seq_len)) # 标准的 Q, K, V 投影 self.q_proj = nn.Linear(d_model, d_model) self.k_proj = nn.Linear(d_model, d_model) self.v_proj = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x): # x: (batch_size, seq_len, d_model) batch_size, seq_len, _ = x.shape Q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) K = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) V = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) # Q, K, V: (batch_size, n_heads, seq_len, head_dim) # Linformer 压缩步骤 # 将 E 和 F 应用到 K 和 V 的序列维度上 # 这里为了清晰，我们循环处理每个头。实际高效实现会使用张量运算。 K_compressed = torch.zeros(batch_size, self.n_heads, self.k, self.head_dim, device=x.device) V_compressed = torch.zeros(batch_size, self.n_heads, self.k, self.head_dim, device=x.device) for h in range(self.n_heads): # 使用第 h 个头的投影矩阵 E_h = self.E[h] # (k, seq_len) F_h = self.F[h] # (k, seq_len) K_h = K[:, h, :, :] # (batch_size, seq_len, head_dim) V_h = V[:, h, :, :] # (batch_size, seq_len, head_dim) # 压缩: (batch_size, seq_len, head_dim) -> (batch_size, k, head_dim) K_compressed[:, h, :, :] = torch.matmul(E_h, K_h) V_compressed[:, h, :, :] = torch.matmul(F_h, V_h) # 线性注意力计算 attn_scores = torch.matmul(Q, K_compressed.transpose(-2, -1)) / (self.head_dim ** 0.5) attn_weights = F.softmax(attn_scores, dim=-1) # (batch_size, n_heads, seq_len, k) attn_output = torch.matmul(attn_weights, V_compressed) # (batch_size, n_heads, seq_len, head_dim) # 恢复形状并输出投影 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.out_proj(attn_output)

关键点：E和F是可学习参数，将seq_len维度从n映射到k。计算的核心变成了(n×d) @ (d×k) -> (n×k) @ (k×d)。

4.2 Performer：核化与结合律的魔法

Performer 采用了更数学化的方法。它使用一个核函数φ来近似原始的 softmax 注意力，使得softmax(QK^T)可以写成φ(Q) · φ(K)^T的形式。然后利用矩阵乘法的结合律(QK^T)V = Q(K^TV)，但这里用φ替换后，计算顺序变为(φ(Q) φ(K)^T) V = φ(Q) (φ(K)^T V)。

具体操作：

1. 核函数选择：例如，使用随机特征映射来近似 exp(q·k)（即 softmax 的分子部分）。常用的是基于随机傅里叶特征的方法。
2. 特征映射：设计一个函数φ(x)，将每个查询向量q_i和键向量k_j映射到一个高维随机特征空间，使得φ(q_i)·φ(k_j) ≈ exp(q_i·k_j)。
3. 改变计算顺序：
   • 传统：A = softmax(QK^T/√d),O = A V。计算 A 需要 O(n²)。
   • Performer:O‘ ≈ φ(Q) · [ φ(K)^T · V ]。
   • 先计算φ(K)^T · V，这是一个(m×d)的矩阵（m 是随机特征维度），复杂度 O(ndm)。
   • 再计算φ(Q) · (上一个结果)，复杂度 O(ndm)。
   • 因为 m 是固定常数，总复杂度为 O(n)。

概念代码 (PyTorch, 使用 FAVOR+ 机制):

import torch import torch.nn as nn import torch.nn.functional as F from math import log, pi def orthogonal_random_matrix(num_rows, num_cols): """生成正交随机矩阵，用于随机特征映射""" q, _ = torch.linalg.qr(torch.randn(num_cols, num_rows)) return q.T # (num_rows, num_cols) class PerformerAttention(nn.Module): def __init__(self, d_model, n_heads, m=256): # m: 随机特征维度 super().__init__() self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads self.m = m # 标准的 Q, K, V 投影 self.q_proj = nn.Linear(d_model, d_model) self.k_proj = nn.Linear(d_model, d_model) self.v_proj = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) # 为每个注意力头生成/定义随机特征映射的参数（这里简化为共享） # 在实际 Performer 中，使用 FAVOR+ 机制，包含随机矩阵和可选的确定性映射 self.w = orthogonal_random_matrix(self.m, self.head_dim) # (m, head_dim) self.b = torch.rand(self.m) * 2 * pi # 相位偏移 def random_feature_map(self, x): """随机傅里叶特征映射 φ(x) 的近似实现""" # x: (..., head_dim) # self.w: (m, head_dim), self.b: (m) proj = torch.matmul(x, self.w.T.to(x.device)) + self.b.to(x.device) # (..., m) # 使用 cos 和 sin 并缩放 return torch.cat([torch.cos(proj), torch.sin(proj)], dim=-1) / (self.m ** 0.5) # (..., 2m) def forward(self, x): batch_size, seq_len, _ = x.shape Q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) K = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) V = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) # 应用随机特征映射到 Q 和 K Q_prime = self.random_feature_map(Q) # (batch_size, n_heads, seq_len, 2m) K_prime = self.random_feature_map(K) # (batch_size, n_heads, seq_len, 2m) # Performer 线性注意力计算: φ(Q) * [φ(K)^T * V] # 1. 先计算 K^T V (利用结合律，但这里用特征映射后的 K') KV = torch.matmul(K_prime.transpose(-2, -1), V) # (batch_size, n_heads, 2m, head_dim) # 2. 再计算 Q' * KV attn_output = torch.matmul(Q_prime, KV) # (batch_size, n_heads, seq_len, head_dim) # 恢复形状并输出投影 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.out_proj(attn_output)

关键点：random_feature_map函数是关键，它将head_dim维的向量映射到2m维。计算顺序的改变避免了构造n×n矩阵。

5. 功能测试与效果验证思路

由于 Linformer 和 Performer 是底层架构组件，其“功能测试”更接近于在具体任务（如语言建模、长文本分类）上的性能评估和效率对比。以下是一个通用的验证流程：

5.1 验证目标

1. 正确性：在短序列上，近似注意力模块的输出应与标准注意力模块的输出大致相同（允许微小误差）。
2. 效率提升：随着序列长度n增加，线性注意力模块的内存占用增长应远慢于标准注意力（O(n) vs O(n²)）。
3. 下游任务性能：在保持模型其他部分不变的情况下，将标准注意力替换为线性注意力后，在验证集上的性能（如准确率、困惑度）下降应在可接受范围内。

5.2 测试步骤（概念性）

环境：准备一个标准的 Transformer 编码器或解码器层。对照组：使用标准的多头自注意力。实验组：使用 LinformerAttention 或 PerformerAttention 模块。

步骤 1：初始化与数据准备

import torch from transformers import AutoTokenizer, AutoModelForCausalLM # 假设我们有一个简单的测试 seq_lengths = [128, 256, 512, 1024, 2048, 4096] d_model = 768 n_heads = 12 batch_size = 2 # 生成随机数据模拟输入 for n in seq_lengths: dummy_input = torch.randn(batch_size, n, d_model) print(f"\n--- 测试序列长度: {n} ---")

步骤 2：内存占用对比

# 测试标准注意力 (这里需要实现或调用一个标准模块) standard_attn = StandardAttention(d_model, n_heads) torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None out_std = standard_attn(dummy_input) mem_std = torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0 print(f"标准注意力峰值内存: {mem_std / 1024**2:.2f} MB") # 测试线性注意力 (以 Linformer 为例) linformer_attn = LinformerAttention(d_model, n_heads, seq_len=n, k=256) torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None out_lin = linformer_attn(dummy_input) mem_lin = torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0 print(f"Linformer峰值内存: {mem_lin / 1024**2:.2f} MB") print(f"内存节省比例: {(mem_std - mem_lin) / mem_std * 100:.1f}%")

步骤 3：输出相似度对比（短序列）

# 在短序列上（如 n=128），检查输出是否相似 if n <= 512: # 使用余弦相似度或 MSE cos_sim = F.cosine_similarity(out_std.flatten(), out_lin.flatten(), dim=0) mse_loss = F.mse_loss(out_std, out_lin) print(f"输出余弦相似度: {cos_sim.item():.4f}") print(f"输出 MSE: {mse_loss.item():.6f}")

步骤 4：速度基准测试（可选）

import time num_iterations = 100 start = time.time() for _ in range(num_iterations): _ = linformer_attn(dummy_input) torch.cuda.synchronize() if torch.cuda.is_available() else None linformer_time = time.time() - start print(f"Linformer 平均迭代时间: {linformer_time/num_iterations*1000:.2f} ms")

预期结果：

• 当n较小时，内存节省可能不明显，甚至因额外投影而略高，但输出应基本相似。
• 当n增大（如 >1024），标准注意力的内存占用会急剧上升，而线性注意力的增长平缓，内存节省效果显著。
• 速度上，线性注意力在长序列上应有明显优势。

6. 接口 API 与批量任务集成

线性注意力模块本身不直接提供 HTTP API，但它可以作为核心组件被集成到模型服务中。例如，你可以使用 FastAPI 部署一个集成了 Performer 的文本生成模型。

假设场景：部署一个用于长文本摘要的模型，该模型使用了 Linformer 编码器。

服务端代码框架 (FastAPI):

from fastapi import FastAPI, HTTPException from pydantic import BaseModel import torch from your_model_arch import LongDocSummarizer # 你的模型，内部使用 Linformer app = FastAPI() model = None tokenizer = None device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class SummarizeRequest(BaseModel): text: str max_length: int = 150 min_length: int = 30 @app.on_event("startup") async def load_model(): global model, tokenizer print("加载模型和分词器...") # 初始化你的自定义模型和分词器 # model = LongDocSummarizer.from_pretrained(...) # tokenizer = AutoTokenizer.from_pretrained(...) model.to(device) model.eval() print("模型加载完毕。") @app.post("/summarize") async def summarize(request: SummarizeRequest): try: # 1. 文本预处理与分词 inputs = tokenizer(request.text, truncation=True, padding=True, return_tensors="pt", max_length=8192) # 支持长文本 inputs = {k: v.to(device) for k, v in inputs.items()} # 2. 模型推理 with torch.no_grad(): # 模型内部使用 Linformer 处理长序列 summary_ids = model.generate( **inputs, max_length=request.max_length, min_length=request.min_length, num_beams=4, early_stopping=True ) # 3. 解码输出 summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True) return {"summary": summary, "status": "success"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) # 批量任务处理（伪代码） @app.post("/summarize_batch") async def summarize_batch(request: List[SummarizeRequest]): results = [] for req in request: # 这里可以引入任务队列（如 Celery）进行异步处理 result = await summarize(req) # 注意：这里需要异步处理 results.append(result) return {"results": results}

客户端调用示例 (Python):

import requests import json url = "http://localhost:8000/summarize" headers = {"Content-Type": "application/json"} # 模拟一个长文档 long_document = "..." # 非常长的文本内容 data = { "text": long_document, "max_length": 200, "min_length": 50 } response = requests.post(url, headers=headers, data=json.dumps(data)) if response.status_code == 200: result = response.json() print(f"摘要: {result['summary']}") else: print(f"请求失败: {response.status_code}, {response.text}")

关键点：API 服务封装了模型细节。用户只需发送文本，服务端利用集成的线性注意力模型高效处理长输入，并返回结果。批量任务可以通过循环或消息队列实现。

7. 资源占用与性能观察

理解线性注意力如何影响资源占用至关重要。

1. 显存占用分析

• 标准注意力 (Softmax)：主要开销在于存储QK^T矩阵，大小为[batch_size, num_heads, seq_len, seq_len]。显存占用与seq_len²成正比。例如，seq_len=4096,num_heads=12,batch_size=1，仅该矩阵就需要约1 * 12 * 4096 * 4096 * 4 bytes ≈ 805 MB（float32）。这还不包括Q,K,V等。
• Linformer：存储压缩后的K_compressed和V_compressed，大小为[batch_size, num_heads, k, head_dim]。显存占用与k * seq_len成正比（因为Q仍是n×d）。若k=256，则上述例子的关键中间变量大小约为1 * 12 * 256 * 64 * 4 bytes ≈ 0.8 MB，加上Q的1 * 12 * 4096 * 64 * 4 bytes ≈ 12.6 MB，总量远小于标准注意力。
• Performer：存储映射后的Q_prime和K_prime，大小为[batch_size, num_heads, seq_len, 2m]，以及中间结果KV([batch_size, num_heads, 2m, head_dim])。显存占用与m * seq_len成正比。m通常也在几百量级，因此也是线性增长。

观察方法：

• 在 PyTorch 中，使用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()来测量特定操作前后的显存变化。
• 使用nvtop(Linux) 或nvidia-smi命令实时监控 GPU 利用率。

2. 计算速度分析

• 标准注意力的计算量随n²增长，在n很大时，矩阵乘法(n×d)@(d×n)和(n×n)@(n×d)都非常耗时。
• 线性注意力将计算量转化为O(ndk)或O(ndm)。当n很大时，k或m是常数，因此速度优势明显。
• 注意：线性注意力引入了额外的投影或特征映射操作，在n很小时，这些开销可能使其速度不如标准注意力。优势区间通常在 n > 512 或 1024 之后。

性能测试建议：

• 编写基准测试脚本，循环测试不同seq_len下的前向传播时间。
• 使用 PyTorch 的torch.cuda.Event进行精确的 GPU 时间测量。
• 对比相同硬件下，标准注意力与线性注意力模块的耗时-序列长度曲线。

8. 常见问题与排查方法

9. 最佳实践与使用建议

1. 从小开始，逐步验证：

   • 不要一开始就在完整模型和全量数据上替换注意力。先在一个简单的任务（如字符级语言建模）或一个小型 Transformer 模块上测试 Linformer/Performer，验证其正确性和效率增益。

2. 参数选择：

   • Linformer 的k：通常设置为 256 或 512。可以通过在验证集上做小网格搜索来确定。k越大，近似越精确，但计算成本也越高。
   • Performer 的m(随机特征数)：类似地，128, 256, 512 是常见起点。更多的特征通常意味着更好的近似，但计算量增加。

3. 训练策略：

   • 微调而非从头训练：如果有一个预训练好的标准 Transformer 模型，想为其增加长文本处理能力，建议采用“微调”策略。即用线性注意力替换原有注意力，然后在长文本下游任务数据上微调整个模型，而不是从头训练。
   • 学习率调整：引入新的可学习参数（如 Linformer 的投影矩阵）后，可能需要调整学习率或使用分层学习率。

4. 模型架构调整：

   • 混合注意力：对于某些任务，模型底层（靠近输入）可能需要更精细的局部注意力，而高层（靠近输出）可以进行更强的压缩。可以设计模型，前几层使用标准注意力，后几层使用线性注意力。
   • 因果注意力：对于自回归生成模型（如 GPT），需要确保线性注意力实现是因果的（即当前位置不能关注未来位置）。Performer 和 Linformer 都有对应的因果掩码实现方式，需仔细检查。

5. 工程化部署：

   • 内核融合：高效的线性注意力实现往往需要自定义 CUDA 内核来融合操作（如投影与注意力计算），以最大化性能。生产环境应考虑使用优化好的库，如xformers库中提供的memory_efficient_attention。
   • 量化与加速：部署时，可以考虑对线性注意力模型进行量化（INT8），以进一步减少内存占用和加速推理。

10. 总结与下一步

Linformer 和 Performer 为我们提供了打破 Transformer 序列长度瓶颈的实用工具箱。Linformer 通过低秩投影直接压缩 Key/Value，思路直观，易于实现和集成；Performer 则基于坚实的数学推导，通过核化与结合律实现线性复杂度，具有更好的理论保证和灵活性。

最值得尝试的点：如果你正在被长文本任务的显存溢出（OOM）所困扰，或者希望你的模型能处理超过 2048 甚至 8192 个 token 的上下文，那么将模型中的标准注意力替换为线性注意力变体，是当前最直接有效的解决方案之一。

最先应该验证的功能：在你的开发环境中，用一个简单的脚本，对比标准注意力与线性注意力模块在不同序列长度下的显存占用和前向传播时间。这个直观的对比能立刻让你感受到线性复杂度的优势。

最容易踩的坑：

1. 参数设置不当：k或m太小，导致信息损失严重，模型性能下降。
2. 训练不充分：替换注意力后，没有在足够的数据上进行微调，直接评估导致效果差。
3. 忽略因果性：在生成任务中，使用了非因果的线性注意力实现，导致模型泄露未来信息。

后续扩展方向：

1. 探索其他线性注意力变体：如Linear Transformer(Katharopoulos et al.),Fast Attention Via Positive Orthogonal Random Features(FAVOR++)，它们各有特点和优化。
2. 集成到流行框架：学习如何将线性注意力模块无缝集成到 Hugging Facetransformers、Fairseq 等库中，方便调用和微调现有大模型。
3. 硬件感知优化：研究针对特定硬件（如 NVIDIA GPU, Apple Silicon）的线性注意力内核优化，追求极致的推理速度。

建议将本文提及的核心代码片段和测试方法保存下来，作为你探索高效 Transformer 架构的起点。在实际项目中，结合具体任务和数据，耐心进行调试和验证，线性注意力很可能成为你解决长序列问题的关键利器。