1. 项目概述与核心需求解析1.1 这个综述到底在讲什么2020年谷歌放出的这篇《Efficient Transformers: A Survey》在圈子里引起的讨论热度一直没降过。我最早是在一个技术群里看到有人转链接当时第一反应是终于有人把散落在各个角落的优化方案系统性地梳理了一遍。Transformer架构从2017年那篇开山之作开始几乎统治了自然语言处理的各个任务但它的核心问题也一直被人诟病——自注意力机制的计算复杂度和内存占用都是序列长度的平方级。序列短的时候还能忍一旦长度上千甚至上万显存直接爆炸训练时间也成倍往上翻。这篇综述做的事情就是把当时能解决这个问题的各路方案做了一个全景式的分类和对比。它把Efficient Transformer拆成了几个大的方向固定模式、组合模式、可学习模式、低秩方法、基于核的方法、以及缓存和分块复用等。每个方向下面又细分了具体的实现思路比如固定模式里有局部窗口注意力、块状注意力、膨胀滑动窗口等低秩方法里有Linformer那种把注意力矩阵做低秩近似的思路基于核的方法里有Performer那种用随机特征映射来逼近softmax核的技巧。我读完最大的感受是它不是在推销某一个具体方案而是给你一张地图让你知道自己遇到的问题该往哪个方向去找答案。适合谁看如果你正在做长文本建模、语音识别、图像生成这些序列长度动不动就上千的场景或者你只是单纯想搞清楚为什么标准Transformer在长序列上这么吃力这篇综述都值得花时间啃一啃。1.2 为什么标准Transformer在长序列上会“卡脖子”要理解这篇综述的价值得先回到问题的源头。标准自注意力的计算过程是这样的给定一个长度为n的序列每个token会生成Query、Key、Value三个向量然后计算Q和K的点积得到注意力分数矩阵这个矩阵的大小是n×n。接着做softmax归一化再和V相乘得到输出。整个过程的计算复杂度是O(n²·d)内存复杂度也是O(n²)其中d是特征维度。我拿具体数字算一下你就知道有多夸张了。假设序列长度n4096特征维度d512那么注意力矩阵就有4096×4096≈1677万个元素。如果用32位浮点数存储光这一个矩阵就要占64MB左右的显存。这还只是一层、一个注意力头的情况。实际模型里通常有12层、12个头那显存占用直接奔着几十GB去了。序列再长一点比如n8192矩阵元素数量翻四倍显存需求也跟着翻四倍。这就是为什么很多长文本任务在实际操作中不得不做截断或者用滑动窗口的方式硬拆。但截断会丢失长距离依赖信息滑动窗口又可能把跨窗口的关联切断。所以Efficient Transformer要解决的核心矛盾就是如何在尽量不损失模型表达能力的前提下把O(n²)的复杂度降下来。注意复杂度降低从来不是免费的午餐。每一种高效注意力方案本质上都是在做某种形式的近似或约束关键是要搞清楚你的任务对哪种信息最敏感然后选择对应的方案。1.3 综述的分类框架与我的理解这篇综述把Efficient Transformer分成了五大类我按自己的理解重新梳理一下固定模式注意力范围是预先定义好的不随输入变化。比如只关注局部窗口、只关注固定间隔的token、或者按块做注意力。代表方案有Local Attention、Blockwise Attention、Longformer的滑动窗口等。组合模式把多种固定模式组合起来用比如局部加全局、局部加随机。BigBird就是典型代表它同时用了局部窗口、全局token和随机连接三种模式。可学习模式注意力模式不是固定的而是通过某种方式学习出来的。比如Reformer用的LSH局部敏感哈希来聚类相似的Query和Key然后只在同一个桶内做注意力。低秩方法假设注意力矩阵是低秩的用低秩分解来近似。Linformer把Key和Value投影到低维空间把复杂度从O(n²)降到O(n)。基于核的方法用核函数来重新表达注意力计算把softmax注意力拆解成两个独立的映射从而避免显式计算n×n矩阵。Performer是这方面的代表。这个分类框架的好处是它不纠结于具体实现细节而是从“你动了注意力的哪个部分”这个角度来划分。你拿到一个新方案先判断它属于哪一类就能快速理解它的核心思路和潜在局限。2. 核心细节解析与实操要点2.1 固定模式注意力的设计逻辑与参数选择固定模式注意力是最容易理解的一类因为它的注意力范围是人为设定的不依赖输入内容。最常见的做法是局部窗口注意力每个token只关注它前后各w个token窗口大小w是超参数。这样注意力矩阵就变成了带状矩阵计算复杂度从O(n²)降到O(n·w)。我实际用过的一个配置是w256序列长度n4096。这样每个token最多关注512个token前后各256计算量大约是原来的1/8。但这里有个坑窗口大小不能设得太小否则模型学不到长距离依赖。我试过w64结果在需要跨段落推理的任务上表现明显下降。后来调到w256效果就稳了很多。另一个常见的固定模式是块状注意力把序列切成固定大小的块每个token只关注自己所在块内的token。这种做法的好处是可以并行计算但缺点是块与块之间的信息完全隔离。为了解决这个问题通常会加一些全局token或者跨块的连接。实操心得固定模式注意力的窗口大小和块大小需要根据任务的平均序列长度来定。如果大部分样本长度在512左右窗口设128到256比较合适如果长度普遍在2048以上窗口可以适当放大到512但要注意显存占用。2.2 组合模式如何用多种模式互补短板组合模式的核心思想是“不要把所有鸡蛋放在一个篮子里”。单一固定模式总有覆盖不到的地方那就把几种模式拼起来用。BigBird的做法很有代表性它同时使用了三种注意力模式——局部窗口、全局token和随机连接。局部窗口负责捕捉邻近token之间的细粒度关系全局token负责汇聚整个序列的概要信息随机连接则提供了一种“远距离跳跃”的能力让信息可以在不同区域之间流动。这三种模式组合起来理论上可以逼近完整注意力的表达能力但计算复杂度只有O(n)。我在一个长文档摘要任务上对比过BigBird和标准Transformer。标准Transformer在长度超过1024之后显存就不够了只能截断。BigBird可以处理4096的长度虽然单步训练时间比截断版的标准Transformer慢一些但因为它能看到完整文档摘要质量明显更好尤其是对文档后半部分内容的覆盖更完整。组合模式的关键参数是全局token的数量和随机连接的比例。全局token太少信息汇聚不够太多计算量又上去了。我的经验是全局token数量设为序列长度的1%到5%比较合适。随机连接的比例一般设在10%到20%之间太高会引入噪声太低又起不到远距离连接的作用。2.3 可学习模式Reformer的LSH聚类实操细节Reformer是这类方案里我踩坑最多的一个。它的核心是用局部敏感哈希把相似的Query和Key分到同一个桶里然后只在桶内做注意力计算。理想情况下每个桶的大小是均匀的计算复杂度可以降到O(n·log n)。但实际操作中LSH的桶大小很难控制得均匀。我遇到过好几次桶大小严重倾斜的情况有的桶里塞了几百个token有的桶里只有几个。这样一来计算量并没有降下来反而因为哈希计算本身的开销整体速度还不如标准注意力。后来我调整了哈希函数的数量和桶的数量情况才好转。具体来说哈希函数的数量一般设为4到8个桶的数量设为序列长度的平方根左右。比如n4096桶数量设在64到128之间比较合适。另外Reformer还用了可逆残差连接来减少内存占用这个技巧在显存紧张的时候非常有用但它要求每一层的输入输出维度完全一致所以网络结构设计上会受一些限制。注意Reformer的LSH方案对序列长度比较敏感。如果序列太短比如小于256LSH的聚类效果不明显反而增加了额外开销。这种情况下直接用标准注意力可能更划算。2.4 低秩方法与核方法数学上的取舍低秩方法的代表是Linformer它的核心假设是注意力矩阵是低秩的可以用两个低秩矩阵的乘积来近似。具体做法是把Key和Value从n×d投影到k×d其中k远小于n。这样注意力计算就变成了n×k的矩阵乘法复杂度从O(n²)降到O(n·k)。这个方法的参数选择很关键。k太小近似误差大k太大又失去了加速的意义。论文里建议k取序列长度的1/4到1/2但我实际测试下来k取128到256之间在大多数任务上都能取得不错的平衡。另外投影矩阵是共享的还是每层独立的也会影响效果。共享投影矩阵参数更少但表达能力弱一些每层独立投影效果更好但参数量增加。基于核的方法里Performer是我用得比较多的。它的思路是用随机特征映射来逼近softmax核函数从而把注意力计算拆成两个独立的线性变换。这样复杂度也是O(n)而且不需要像Linformer那样对序列长度做固定假设理论上可以处理任意长度的序列。但Performer的随机特征数量是个需要调的超参数。特征太少逼近误差大特征太多计算量又上去了。我的经验是特征数量设为特征维度的2到4倍比较合适。比如d512特征数量设在1024到2048之间。另外Performer对初始化比较敏感我一般会用比标准Transformer更小的学习率来做预热。3. 实操过程与核心环节实现3.1 环境准备与依赖安装我搭建实验环境的时候用的是PyTorch作为主要框架。下面是我实际用的一套依赖配置你可以直接参考pip install torch1.7.0 pip install transformers4.0.0 pip install datasets1.1.0 pip install tensorboard2.4.0这里说明一下版本选择的理由。PyTorch 1.7.0是我测试下来和大多数Efficient Transformer实现兼容性最好的版本再新的版本有些自定义CUDA算子会编译报错。Transformers库用4.0.0是因为它已经内置了一些高效注意力的实现比如Longformer和Reformer的接口省得自己从头写。如果你要用Performer还需要额外安装一个专门的包pip install performer-pytorch这个包里包含了Performer的核心实现包括随机特征映射和快速注意力计算。安装完之后你可以直接用它提供的PerformerAttention模块替换标准注意力层。提示不同版本的PyTorch对CUDA算子的支持不一样。如果你用的是较新的显卡可能需要把PyTorch版本往上调。我建议先在CPU上跑通小规模实验确认逻辑没问题之后再上GPU。3.2 数据准备与序列长度分布分析在选型之前我强烈建议你先做一件事统计你的任务里序列长度的分布。这个信息直接决定了你该选哪种高效注意力方案。我拿一个长文档分类任务举例。原始数据的序列长度分布是这样的长度区间样本占比累计占比0-25615%15%256-51225%40%512-102430%70%1024-204820%90%2048以上10%100%从这个分布可以看出70%的样本长度在1024以内标准Transformer勉强能处理。但剩下30%的样本长度超过1024如果直接截断这部分样本的信息损失会很大。这种情况下选择一种能处理2048以上长度的方案就比较有意义。如果大部分样本都在512以内那其实没必要上太复杂的方案标准Transformer加个梯度累积就能跑。如果大部分样本在2048以上那固定模式的局部窗口可能就不够了需要考虑组合模式或者基于核的方法。3.3 模型搭建与关键参数配置我以Longformer为例展示一下具体的模型搭建过程。Longformer属于组合模式它用了滑动窗口加全局token的注意力结构。from transformers import LongformerModel, LongformerConfig config LongformerConfig( vocab_size30522, hidden_size512, num_hidden_layers6, num_attention_heads8, intermediate_size2048, max_position_embeddings4096, attention_window256, # 滑动窗口大小 num_global_tokens1, # 全局token数量 ) model LongformerModel(config)这里有几个参数需要重点说明。attention_window设的是256意味着每个token关注前后各256个token。这个值是我根据序列长度分布调的因为大部分样本在1024以内256的窗口可以覆盖四分之一的序列再加上全局token的信息汇聚基本够用。num_global_tokens设的是1也就是在序列开头加一个全局token让它关注所有其他token同时也被所有token关注。这个全局token相当于一个信息中转站弥补了局部窗口无法直接传递远距离信息的缺陷。如果你用的是Reformer配置会不太一样from reformer_pytorch import ReformerLM model ReformerLM( num_tokens30522, dim512, depth6, max_seq_len4096, heads8, lsh_dropout0.1, n_hashes8, # 哈希函数数量 bucket_size64, # 桶大小 causalFalse, )Reformer的n_hashes和bucket_size是两个关键参数。n_hashes设8是我测试下来比较稳的值太少聚类效果差太多计算开销大。bucket_size设64是配合序列长度4096来的桶数量大约是4096/6464个每个桶平均64个token分布相对均匀。3.4 训练过程与显存监控训练的时候我习惯用torch.cuda.max_memory_allocated()来监控显存峰值。下面是我在同一个任务上对比三种方案的显存占用和训练速度方案序列长度显存峰值单步训练时间验证集准确率标准Transformer5128.2GB0.35s89.2%标准Transformer截断102414.5GB0.68s87.5%Longformer40969.8GB0.82s91.3%Reformer40967.6GB0.95s90.1%Performer40968.1GB0.71s90.8%从这张表可以看出来Longformer在显存占用和准确率之间取得了比较好的平衡。Reformer显存最省但训练时间最长因为LSH的哈希计算有额外开销。Performer的训练速度最快准确率也不错但它的随机特征映射需要更多的预热步数才能稳定。实操心得显存监控一定要在训练开始后的前100步就做不要等到跑了几千步才发现显存不够。我一般会在第一个epoch里每隔50步打印一次显存峰值确认稳定后再关掉监控。4. 常见问题与排查技巧实录4.1 训练不收敛或loss震荡怎么办这是我在用Efficient Transformer时遇到最多的问题。标准Transformer上跑得好好的学习率换到高效注意力方案上可能就不适用了。原因在于高效注意力本质上是对标准注意力的近似近似误差会引入额外的噪声导致优化过程更不稳定。我的排查顺序是这样的先检查学习率。高效注意力方案通常需要更小的学习率或者更长的预热步数。我一般会把标准Transformer的学习率乘以0.5到0.7预热步数从总步数的5%提高到10%。检查注意力掩码。固定模式注意力需要正确的掩码来限制注意力范围。如果掩码写错了模型可能会关注到不该关注的位置导致训练崩溃。检查初始化。有些高效注意力方案对初始化比较敏感比如Performer的随机特征映射。我一般会用更小的初始化标准差比如0.02而不是默认的0.1。下面这张表是我整理的一些常见现象和对应的排查方向现象可能原因排查方法loss震荡不下降学习率太大降低学习率增加预热步数loss下降但验证集不涨近似误差太大增大窗口/桶/特征数量训练初期loss为NaN初始化或数值不稳定检查初始化加梯度裁剪显存溢出序列长度或批次太大减小批次用梯度累积训练速度比标准Transformer还慢额外开销太大检查哈希/投影的计算量4.2 长序列推理时的显存优化技巧训练的时候可以用梯度累积、混合精度这些技巧来省显存但推理的时候这些方法就不太适用了。我在推理长序列时主要用以下几个策略分块推理是最直接的方法。把长序列切成固定大小的块逐块推理然后把结果拼接起来。但这里有个坑块与块之间的注意力信息会丢失。如果任务对跨块依赖比较敏感分块推理的效果会打折扣。我的做法是在块之间保留一定的重叠区域比如块大小1024重叠128这样可以在一定程度上缓解信息断裂的问题。KV缓存是另一个常用技巧尤其适合自回归生成任务。标准Transformer在生成每个token时都要重新计算所有历史token的Key和Value非常浪费。KV缓存把这些中间结果存下来生成新token时直接复用。但KV缓存本身也占显存序列越长占用越大。我一般会设置一个缓存上限超过上限就丢弃最早的缓存。混合精度推理也能省不少显存。把模型权重和中间激活值从32位浮点降到16位浮点显存占用直接减半。但要注意有些高效注意力方案对数值精度比较敏感比如Reformer的LSH在16位下可能会出现哈希冲突增多的问题。我一般会先用32位跑一遍确认结果正常再切到16位对比差异。4.3 不同方案之间的迁移与替换经验在实际项目中我经常需要把标准Transformer替换成某种高效方案或者在不同高效方案之间切换。这个过程有几个需要注意的地方参数命名和结构差异。不同库对高效注意力的实现方式不一样参数命名和模块结构都有差异。比如HuggingFace的Longformer和Reformer的接口就完全不同。替换的时候不能直接复制权重需要重新映射参数名。我一般会写一个小的映射脚本把旧模型的权重按名称对应到新模型上。位置编码的兼容性。标准Transformer用的是绝对位置编码而有些高效方案用的是相对位置编码或者旋转位置编码。替换的时候位置编码需要重新初始化不能直接沿用。我试过直接沿用绝对位置编码结果模型完全学不到位置信息效果一塌糊涂。注意力掩码的适配。标准注意力的掩码通常是上三角矩阵用于因果注意力但高效注意力的掩码结构更复杂。比如Longformer的滑动窗口掩码是一个带状矩阵Reformer的LSH掩码是动态生成的。替换的时候需要根据新方案的注意力模式重新生成掩码。注意迁移的时候一定要做小规模验证。我一般会先用一个很小的数据集比如100条样本跑几个epoch确认loss能正常下降再上全量数据。这样可以快速发现参数映射或掩码适配的问题避免浪费大量时间。4.4 效果评估与方案选型速查表最后我整理了一张方案选型的速查表方便你根据任务特点快速定位任务特点推荐方案关键参数注意事项序列长度512以内标准Transformer无需特殊配置直接用没必要上高效方案序列长度1024-2048显存紧张Longformer窗口256全局token 1-3窗口别太小否则长距离依赖丢失序列长度2048以上追求速度Performer特征数1024-2048需要较长预热学习率调小序列长度2048以上追求显存最省Reformer哈希数8桶大小64桶分布要监控太倾斜就调参需要处理超长序列8192以上组合模式BigBird类全局token 1%-5%随机连接10%-20%全局token别太多否则计算量回升自回归生成任务带KV缓存的标准注意力缓存上限根据显存定缓存别无限增长定期清理这张表是我踩了无数坑之后总结出来的但也不是万能的。每个任务的数据分布和计算资源都不一样最好的办法还是拿你的实际数据做一个小规模对比实验用数据说话。我个人在实际操作中的体会是Efficient Transformer这个领域发展很快新的方案层出不穷但这篇2020年的综述提供的分类框架到现在依然适用。你不需要记住每一个具体方案的细节但一定要理解每一类方案的核心取舍——是牺牲表达能力换速度还是牺牲速度换显存还是用近似换效率。搞清楚这个你就能在面对新方案时快速判断它适不适合你的场景。