线性注意力实战：Linformer与Performer的工程落地指南

这类把注意力复杂度从 O(n²) 降到 O(n) 的方案，最值得先看的不是数学推导，而是它们到底能不能在真实任务里稳定跑起来，以及为了换速度牺牲了什么。Linformer 和 Performer 是两个典型的代表，一个用低秩投影，一个用核化加结合律，目标都是让长序列处理变得可行。如果你正在处理文本、代码或时间序列，感觉传统 Transformer 的显存和速度成了瓶颈，那这两个方法的工程实现和落地选择就值得仔细拆一遍。

我一般会先跑通标准注意力，再换线性注意力对比，这样对速度提升和效果变化才有体感。下面按实际落地顺序，从它们各自的核心思路、代码怎么改、参数怎么调、以及什么情况下会出问题，完整走一遍。

1. 先拆清楚 Linformer 和 Performer 各自解决了什么问题

很多人一上来就去看公式，但更容易迷糊。其实可以从一个具体场景切入：假设你有一个长度为 4096 的序列，标准自注意力需要计算一个 4096×4096 的矩阵，这不仅是显存杀手，计算量也巨大。Linformer 和 Performer 给出了两种不同的“近似”思路，但它们的适用场景和代价不一样。

1.1 Linformer：用低秩投影强行压缩 KV，适合内存敏感但对精度要求不极端的场景

Linformer 的核心假设是：注意力矩阵是低秩的。既然如此，就没必要老老实实算那个巨大的 n×n 矩阵。它的做法是把原始的 Key 和 Value 投影到一个低维空间（比如从 n 维投影到 k 维，k 远小于 n）。这样，计算复杂度就从 O(n²) 降到了 O(nk)。当 k 是一个固定常数时，复杂度就是 O(n)。

这里最关键的一个工程参数是投影维度k。它不是超参数，而是你根据任务和序列长度预设的一个值。论文里常用 256 或 512。设置时，你需要权衡：

• k 太小：压缩太狠，信息丢失严重，模型效果会明显下降。
• k 太大：压缩效果有限，速度提升不明显。 一个经验是，k可以设为序列长度n的对数级别，或者通过一小部分验证集搜索一个性价比最高的值。

Linformer 的实现通常很直接，就是在标准的注意力模块前，给 K 和 V 各加一个线性投影层。它的优势是改动小，易于插入现有的 Transformer 架构。但缺点也明显：这个低秩假设并非永远成立，对于某些需要精细 token-to-token 交互的任务（比如语法解析），性能损失可能会比较大。

1.2 Performer：用核函数+结合律重写注意力，适合需要严格理论保证和稳定训练的场景

Performer（FAVOR+）的思路更巧妙。它不直接近似注意力矩阵，而是利用数学变换，将标准的点积注意力（计算 exp(QKᵀ)）重写为一种可以通过“结合律”先聚合再计算的形式。

简单来说，它找到一个随机特征映射函数 φ(·)，使得 exp(q·k) ≈ φ(q)·φ(k)。这样一来，注意力计算就变成了：Attention = (Q' * (K'ᵀ * V))，其中 Q‘ = φ(Q)， K’ = φ(K)。 由于矩阵乘法满足结合律，你可以先计算(K'ᵀ * V)，这是一个 (k_dim × d_v) 的矩阵，然后再与 Q‘ 相乘。这样，无论序列多长，你都不需要显式地构造 n×n 矩阵。

Performer 最需要关注的参数是随机特征的数量m（即映射后的维度）。m越大，对原始点积的近似就越精确，但计算量也越大。实践中，m通常设置在 64 到 256 之间。Performer 提供了理论上的误差界，这是它的一大优点，意味着你可以预估近似带来的最大误差。

它的优势是提供了无偏的近似，并且支持双向、因果（单向）等多种注意力模式，通用性更强。但代价是，随机特征映射本身需要一些计算，并且对于非常短的序列，可能带来额外的开销。

1.3 一句话选择建议：要插拔简便选 Linformer，要理论稳健选 Performer

如果你的需求是快速将一个现有模型（比如 BERT）改造得能处理更长文本，并且可以接受一定的精度损失来做实验，那么 Linformer 的侵入性小，更容易集成。

如果你的项目从头开始，或者对近似的理论可靠性有要求，需要处理各种复杂的注意力掩码（比如因果语言建模），那么 Performer 是更稳健的选择。它的实现现在也很成熟，在主流深度学习框架里都有较好的支持。

2. 环境准备与最小化实现验证

在动手改任何代码之前，先把实验环境搭好，并准备好一个可以复现问题的基线。这能帮你快速判断线性注意力带来的到底是收益还是新坑。

2.1 基础环境与基线模型

你需要一个标准的 Transformer 实现作为基线。这里以 PyTorch 和 Hugging Facetransformers库为例，因为它最普遍。

# 基础环境 pip install torch torchvision torchaudio pip install transformers pip install datasets # 可选，用于性能监控 pip install nvitop

我们先定义一个简单的测试脚本，用标准注意力跑一个前向传播，记录时间和显存：

import torch import torch.nn as nn from transformers import AutoModel, AutoTokenizer import time device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # 1. 加载一个标准 Transformer 作为基线 model_name = "bert-base-uncased" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name).to(device) # 2. 构造一个长序列输入（这里模拟 1024 长度） batch_size = 2 seq_length = 1024 input_ids = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_length)).to(device) # 3. 预热并测量基线性能 model.eval() with torch.no_grad(): # 预热 for _ in range(5): _ = model(input_ids) torch.cuda.synchronize() if device.type == 'cuda' else None # 正式测速和显存 start_mem = torch.cuda.memory_allocated(device) if device.type == 'cuda' else 0 start_time = time.time() outputs = model(input_ids) torch.cuda.synchronize() if device.type == 'cuda' else None elapsed_time = time.time() - start_time end_mem = torch.cuda.memory_allocated(device) if device.type == 'cuda' else 0 print(f"[Baseline] SeqLen {seq_length}, Time: {elapsed_time:.4f}s, GPU Mem Delta: {(end_mem - start_mem) / 1024**2:.2f} MB") print(f"Output shape: {outputs.last_hidden_state.shape}")

这个基线能告诉你，在当前硬件上，处理 1024 长度的序列，标准注意力需要多少时间和显存。这是所有对比的起点。

2.2 集成 Linformer：替换注意力层

对于 Linformer，我们可以使用社区实现的库，比如linformer（注意，这可能不是官方唯一实现）。这里演示如何替换 BERT 中的自注意力层。

pip install linformer-pytorch

然后，修改模型创建部分：

from linformer import LinformerSelfAttention from transformers import BertConfig, BertModel import torch.nn as nn class LinformerBertModel(nn.Module): def __init__(self, config, k=256): super().__init__() self.config = config self.embeddings = ... # 复用 BERT 的 embeddings # 关键：用 LinformerSelfAttention 替换原来的 BertSelfAttention self.encoder = nn.ModuleList([ LinformerBertLayer(config, k) for _ in range(config.num_hidden_layers) ]) # ... 其他部分 class LinformerBertLayer(nn.Module): def __init__(self, config, k): super().__init__() # 创建 Linformer 注意力层 self.attention = LinformerSelfAttention( dim = config.hidden_size, # 输入维度 seq_len = config.max_position_embeddings, # 最大序列长度 heads = config.num_attention_heads, # 头数 k = k, # 投影维度，核心参数！ one_kv_head = True, # 共享 KV 投影，节省参数 share_kv = True ) self.output = ... # 前馈网络等 # 创建模型并测试 config = BertConfig.from_pretrained("bert-base-uncased") config.max_position_embeddings = 2048 # 可以尝试更大的长度 linformer_model = LinformerBertModel(config, k=256).to(device) # 使用同样的输入进行测试，对比时间和显存

注意：这里的关键是k参数。你需要用同样的输入跑一遍，观察速度提升和显存下降，同时也要验证输出是否“合理”（比如 shape 正确，没有 NaN）。

2.3 集成 Performer：使用现成库

Performer 的实现更复杂，但幸运的是，有非常成熟的库如performer-pytorch。

pip install performer-pytorch

集成方式类似：

from performer_pytorch import Performer from transformers import BertConfig import torch.nn as nn class PerformerBertModel(nn.Module): def __init__(self, config, dim_head=64, heads=12, causal=False): super().__init__() self.config = config self.embeddings = ... # 复用 BERT embeddings # 使用 Performer 作为编码器 self.encoder = Performer( dim = config.hidden_size, depth = config.num_hidden_layers, heads = heads, dim_head = dim_head, causal = causal, # 是否为因果（解码器）注意力 ff_mult = 4, # 前馈网络扩展倍数 feature_redraw_interval = 1000, # 重绘随机特征的间隔，重要！ generalized_attention = True, # 使用广义注意力（默认） ) # ... 其他部分 # 创建模型 performer_model = PerformerBertModel(config).to(device) # 进行测试

这里需要特别关注feature_redraw_interval参数。Performer 的随机特征在训练过程中需要定期重新采样以保持稳定性。这个参数控制了重绘的步数间隔。设置太小影响速度，设置太大可能影响近似质量。在推理（evaluation）模式下，可以关闭这个功能。

2.4 验证正确性：不只是看能不能跑通

模型能跑起来只是第一步。你需要验证：

1. 输出形状：确保last_hidden_state的形状是(batch, seq_len, hidden_dim)，和基线一致。
2. 前向传播稳定性：检查输出张量里有没有inf或nan。
3. 小数据拟合：用一个极小的、随机的数据集（比如 10 条数据），让模型过拟合。如果线性注意力版本也能快速拟合（损失降到接近 0），说明其基本的前向和反向传播功能是正常的。这是检查模型是否“学得进去”的快速方法。

# 简单的过拟合测试 test_input = torch.randn(4, 128, config.hidden_size).to(device) test_target = torch.randn(4, 128, config.hidden_size).to(device) optimizer = torch.optim.Adam(performer_model.parameters(), lr=1e-3) performer_model.train() for i in range(100): optimizer.zero_grad() output = performer_model(test_input) loss = nn.MSELoss()(output, test_target) loss.backward() optimizer.step() if i % 20 == 0: print(f"Step {i}, Loss: {loss.item():.6f}") # 观察损失是否持续下降

3. 核心参数调优与性能-效果权衡

模型跑起来后，接下来就是精细调整，在速度和效果之间找到最佳平衡点。这个阶段最容易犯的错是只盯着训练速度，忽略了下游任务的效果。

3.1 Linformer 的k：在内存、速度和精度间做选择

k（投影维度）是 Linformer 的命门。你可以设计一个简单的扫描实验：

seq_length = 2048 batch_size = 1 k_values = [64, 128, 256, 512, 1024] # 1024相当于接近标准注意力了 results = [] for k in k_values: model = LinformerBertModel(config, k=k).to(device) model.eval() with torch.no_grad(): input_ids = torch.randint(0, 1000, (batch_size, seq_length)).to(device) start = time.time() _ = model(input_ids) torch.cuda.synchronize() elapsed = time.time() - start mem = torch.cuda.max_memory_allocated(device) / 1024**2 results.append((k, elapsed, mem)) print(f"k={k}: Time={elapsed:.3f}s, Peak Mem={mem:.1f}MB")

把结果画成图，你会看到一条典型的曲线：k很小时，内存和时间开销很小，但模型容量低；随着k增大，开销向标准注意力逼近。你的任务是在这条曲线的“肘部”附近选一个值——即性能提升开始放缓的那个点。

对于文本分类、情感分析这类高层语义任务，k=128或256可能就够了。对于需要细粒度理解的任务（如问答、命名实体识别），可能需要k=512或更大。

3.2 Performer 的dim_head和随机特征数

Performer 的参数更复杂一些：

• dim_head：每个注意力头的维度。dim_head * heads ≈ hidden_size。较小的dim_head可以减少计算量，但可能会限制表达能力。通常保持和原模型一致即可。
• generalized_attention：是否使用广义注意力（默认的 FAVOR+ 机制）。建议保持True，除非你有特殊理由。
• feature_redraw_interval：这是训练时最重要的参数之一。它控制随机特征矩阵的重绘频率。如果训练不稳定（损失震荡），尝试减小这个值（如从 1000 调到 500）。如果追求更快的训练速度，可以适当增大（如 2000），但要监控验证集性能。推理时可以设为 None 或一个很大的数。

Performer 的随机特征数m通常由库内部根据dim_head自动决定，但有些实现允许你手动指定num_features。更多特征意味着更好的近似，但计算量也线性增加。除非你怀疑近似误差是性能瓶颈，否则先用默认值。

3.3 注意力头数 (heads) 的影响

无论是 Linformer 还是 Performer，注意力头数都是一个重要参数。原版 Transformer 使用多头注意力来并行学习不同子空间的信息。

• 减少头数：这是最直接的加速和节省显存的方法。例如，将 BERT-base 的 12 个头减少到 6 个。但这样会直接降低模型的表达能力，可能对效果有较大影响。
• Linformer 的one_kv_head：这是一个特殊的优化。当设置为True时，所有注意力头共享同一套 Key 和 Value 的投影。这能显著减少参数和计算量，但同样会限制灵活性。对于初步实验，可以打开这个选项以获得最大速度；如果效果下降太多，再关闭它。

一个实用的策略是：先用默认或较小的头数、开启one_kv_head（对于 Linformer）、使用默认feature_redraw_interval（对于 Performer）跑通实验流程。如果效果达标，再尝试增加头数或关闭共享来提升效果，同时观察资源消耗是否仍在可接受范围内。

4. 从单任务到批量处理：稳定性与生产化考量

在单条样本上测试成功，不代表能稳定处理批量数据或生产环境中的流式数据。这里有几个工程上必须检查的点。

4.1 动态序列长度与填充处理

训练数据通常被填充到统一长度。但实际应用中，序列长度是变化的。线性注意力层是否能正确处理填充 token (pad_token_id)？

• 标准做法：在计算注意力权重前，需要生成一个attention_mask，将填充位置的值设为极大的负数（如 -1e9），这样经过 softmax 后，这些位置的权重接近 0。
• Linformer：由于它对 K 和 V 进行了投影，你需要确保投影操作不会破坏 mask 的语义。通常，mask 应在投影之后、计算注意力权重之前应用。检查你使用的 Linformer 实现是否支持传入attention_mask。
• Performer：同样，需要确认其前向传播函数是否接收mask参数。performer-pytorch库通常支持mask参数。

验证方法：构造一个 batch，其中包含两个长度差异很大的序列（如长度 50 和 200）。用同一个模型（带 mask）处理，观察输出中填充位置的特征向量是否接近 0（或与有效位置有显著差异）。

4.2 长序列下的数值稳定性

当序列长度n非常大时（例如上万），即使复杂度是 O(n)，中间计算也可能出现数值问题。

• Softmax 溢出：标准注意力中，exp(QKᵀ)在QKᵀ值很大时容易溢出。Transformer 通常使用scaled注意力（除以sqrt(d_k)）来缓解。
• Linformer：低秩投影后，数值范围可能发生变化。确保你的实现包含了适当的缩放。
• Performer：核化方法本身就是为了避免计算大矩阵的exp，通常数值稳定性更好。但仍需关注随机特征映射是否引入了数值异常。

排查：在长序列输入下，监控模型内部张量的统计信息（均值、标准差、最大值、最小值），看是否有异常值。

4.3 训练技巧与学习率调整

线性注意力是对原始注意力机制的近似，这可能会改变模型的优化地貌。因此，直接沿用原模型的学习率可能不是最优的。

• 学习率预热 (Warmup)：对于从头开始训练，学习率预热变得更加重要。可以考虑使用更长的预热步数。
• 学习率缩放：由于模型结构变化，最优学习率可能不同。可以尝试以基线学习率为中心，进行小范围的网格搜索（例如[0.5x, 1x, 2x]）。
• 梯度裁剪：在训练初期，梯度可能不稳定，使用梯度裁剪可以防止梯度爆炸。

建议：在正式训练前，用一个小的开发集（dev set）跑几个 epoch，观察训练损失和验证损失曲线。如果损失震荡剧烈或下降缓慢，优先调整学习率和预热策略。

4.4 与现有预训练模型的结合：继续预训练 vs. 微调

你想在 BERT 或 GPT 上使用线性注意力，有两种策略：

1. 结构替换后直接微调：将预训练模型中的注意力层替换为线性注意力层，然后在你的下游任务数据上直接微调。这种方法快，但线性注意力层是随机初始化的，破坏了预训练权重的一致性，效果可能打折扣。
2. 结构替换后继续预训练：替换注意力层后，用大量文本数据对模型进行一段时间的继续预训练（可能只需要几个 epoch），让线性注意力层适应新的计算方式，然后再微调。这种方法效果通常更好，但成本更高。

如何选择：如果你的下游任务数据量足够大（十万级以上），可以尝试方法 1。如果数据量小，或者你对效果要求高，建议采用方法 2，即使只继续预训练一个 epoch，也能带来显著提升。

5. 效果评估与常见问题排查

模型训完了，速度也上去了，但效果怎么样？如果效果下降，问题出在哪里？

5.1 建立系统的评估基准

不要只看最终的准确率或 F1 分数。建立一个多维度的评估表：

这个表格能帮你一目了然地看到权衡。速度提升 50%，但精度只下降 0.7%，在很多场景下是完全可接受的。

5.2 效果下降的排查路径

如果效果下降超出预期（比如 >2%），按以下顺序排查：

1. 检查输入输出对齐：确保替换注意力层后，模型输入输出的维度和类型没有错误。特别是attention_mask是否被正确传递和处理。
2. 验证近似本身的质量：对于 Linformer，可以计算一下原始注意力矩阵和低秩近似后的注意力矩阵的差异（Frobenius 范数）。对于 Performer，可以检查随机特征映射的近似误差。这能帮你确认问题是否出在近似算法本身。
3. 调参：回到第 3 节。尝试增大 Linformer 的k，或调整 Performer 的feature_redraw_interval和dim_head。
4. 检查训练动态：学习率是否合适？损失曲线是否正常？梯度是否有爆炸或消失？可以尝试使用更小的学习率、更长的 warmup、或更强的梯度裁剪。
5. 任务是否对注意力矩阵敏感：有些任务极度依赖精确的 token-to-token 交互，例如句法分析、指代消解。对于这些任务，线性注意力的近似可能天生就不适合。此时，可以考虑混合使用，只在高层（语义）层使用线性注意力，在底层（语法）层保留标准注意力。

5.3 特定场景下的性能陷阱

• 超短序列：当序列长度n很小（比如小于 64）时，线性注意力带来的加速可能微乎其微，甚至因为额外的投影或映射操作而更慢。在这种情况下，可以考虑动态选择：当n小于某个阈值时，回退到标准注意力计算。
• 推理部署：在部署到生产环境时，需要考虑：
  • 算子融合：线性注意力中的连续线性层是否可以融合以提高效率？
  • 量化支持：你的线性注意力实现是否支持 INT8 量化？量化后的精度损失是否可控？
  • 框架兼容性：是否能在 TensorRT、ONNX Runtime 等推理引擎中高效运行？
• 分布式训练：线性注意力层是否支持模型并行或数据并行？在分布式环境下，通信开销是否会成为新的瓶颈？

6. 总结：如何将线性注意力技术落地

经过上面这些步骤，你应该对 Linformer 和 Performer 有了从理论到实操的完整认识。最后，抛开具体代码，我总结一下落地时的核心思路：

不要一上来就全盘替换。先从模型的一两个层开始替换，做对比实验。比如，只替换最后 4 层的注意力机制，看看效果和速度的变化。这能帮你以最小的代价评估这项技术对你的任务是否有效。

** profiling 是你的朋友**。在优化前和优化后，一定要用 profiling 工具（如 PyTorch Profiler、Nsight Systems）分析一下，时间到底省在哪里了，显存又省在哪里了。有时候你以为的瓶颈可能不是真正的瓶颈。

为“近似”付出代价是正常的。线性注意力用数学上的近似换取了计算效率，这意味着效果上轻微的妥协是符合预期的。关键是要量化这个妥协：速度提升了多少？显存节省了多少？精度损失了多少？然后根据你的业务需求判断这个交易是否划算。

关注社区动态。线性注意力是一个活跃的研究领域，除了 Linformer 和 Performer，还有像 Linear Transformer、Synthesizer、Longformer（虽然它不是严格的 O(n)）等多种变体。保持关注，新的实现可能在某些方面更有优势。

最终，是否采用线性注意力，不是一个纯技术问题，而是一个工程权衡。如果你的应用场景中，序列长度是核心瓶颈，且可以容忍微小的精度损失，那么 Linformer 或 Performer 就是一个非常有力的工具。反之，如果序列不长，或者任务对注意力精度极其敏感，那么坚持标准注意力可能是更稳妥的选择。我的建议是，建立一个快速的评估管道，用你的实际数据和任务去验证，让数据告诉你答案。