1 算法原理
Reformer 是 Google 于 2020 年提出的一种高效序列建模范式,旨在解决传统 Transformer 在处理超长序列时面临的二次方计算复杂度与显存瓶颈问题。其核心贡献在于引入局部敏感哈希(LSH)注意力机制与可逆残差网络(Reversible Layers),并辅以分段式前馈网络设计,使模型在保持较强表示能力的前提下,将计算与内存开销从 O(L²) 降至近似 O(L log L),其中 L 为序列长度。
1.1 局部敏感哈希注意力
局部敏感哈希注意力(LSH Attention)通过哈希函数将输入序列中的查询和键向量映射到多个桶中,仅在同一个桶内计算注意力,从而避免全局的二次方点积。该机制假设相似的向量(即注意力权重高的向量对)更容易被哈希到相同桶内,从而近似完整自注意力的结果。
1.1.1 哈希分桶机制
Reformer 采用随机投影哈希:对每个查询向量 q 和键向量 k,先应用随机旋转矩阵,然后按角度分区。具体而言,将向量投影到一个随机超平面上,根据投影值的符号或连续区间划分桶。为降低哈希冲突导致的近似误差,Reformer 使用多轮哈希(如 4 轮或 8 轮),每轮独立生成不同的哈希函数。在每一轮中,每个 token 被分配到一个桶,桶内 token 数量通过 padding 对齐到固定大小以适应批处理。
1.1.2 分块点积计算
在每个哈希桶内,只对该桶中的查询和键进行点积注意力计算。由于桶内 token 数量远小于序列总长度,且每个桶独立并行处理,整体计算量大幅下降。为处理桶间边界,Reformer 还引入“分块”技巧:在每个桶内,将查询和键按块划分,块内计算注意力后再聚合。分块大小通常取固定值(如 128 或 256),进一步减少显存占用。
1.2 可逆残差网络
传统的 Transformer 每层需缓存中间激活值用于反向传播,导致显存随层数线性增长。Reformer 采用可逆残差网络,通过特定设计的转换使得网络在反向传播时可以仅用当前层的输出和输入重构出前向计算的中间结果,从而无需存储激活值。
1.2.1 前向与反向重构策略
可逆残差网络的每个子层由两个可逆变换组成:第一部分通常为注意力子层,第二部分为前馈子层。输入 x 被拆分为两个通道 x₁ 和 x₂。前向传播时:
- y₁ = x₁ + Attn(x₂)
- y₂ = x₂ + FFN(y₁)
反向传播时,给定梯度 dy₁ 和 dy₂,可依次计算出 dx₂、dx₁,以及中间输入的梯度,而无需任何额外激活缓存。这一策略将每层的显存占用从 O(L×d) 降低到常数级别(仅需存储输出和部分中间变量)。
1.2.2 显存节省机制
结合 LSH 注意力与可逆层,Reformer 在训练 64K 长度序列时显存占用仅为传统 Transformer 的 1/4 以下。显存节省主要来自两方面:一是可逆层消除了对每层激活的存储需求;二是 LSH 注意力避免了完整注意力矩阵的存储。
1.3 分段式前馈网络(Chunked FFN)
传统 Transformer 的前馈网络(FFN)对序列中每个 token 独立处理,计算量与序列长度成正比。当序列极长时,一次性计算所有 token 的 FFN 仍会占用大量显存。Reformer 将序列沿长度维度切分为固定大小的块(如 1024 个 token),分别对每个块执行 FFN 计算。这一“分块”操作不改变数学结果,但允许模型在块间进行梯度累积,从而适应有限显存的硬件环境。
2 组件与机制
2.1 自注意力变体
Reformer 的自注意力机制在原始 Transformer 的自注意力之上进行了三项关键修改:首先,用 LSH 注意力替代全局注意力;其次,引入多轮哈希以保证近似精度;最后,在查询和键共享的设定下(即 Q=K),进一步简化计算——Reformer 默认使用共享的 Q 和 K 参数矩阵,这符合 LSH 的“相似向量映射到相同桶”的直觉。
2.2 位置编码替代方案
Reformer 默认不使用传统的正弦位置编码或可学习位置编码,而是采用“临时位置编码”方案:在每个桶内,为 token 分配其在桶内的相对位置索引(0 到桶大小-1),并将此索引作为编码输入。该设计旨在使位置信息不受全局序列长度变化的影响,且与 LSH 分桶机制兼容。
2.3 层归一化位置
Reformer 在层归一化位置的选择上提供了灵活性,支持两种主流策略:前置归一化(Pre-LN)和后置归一化(Post-LN),用户可根据任务需求调整。
2.3.1 前置归一化(Pre-LN)
前置归一化将层归一化置于子层(注意力或 FFN)之前,即顺序为:LN → 子层 → 残差连接。该策略在训练深层网络时更容易稳定梯度,是 Reformer 的默认选项。
2.3.2 后置归一化(Post-LN)
后置归一化将层归一化置于子层之后,即顺序为:子层 → 残差连接 → LN。此方案在原始 Transformer 中被采用,但深层训练时容易出现梯度爆炸或消失。Reformer 允许用户选择以兼容现有模型架构。
3 训练与优化
3.1 长序列训练的梯度稳定性
长序列训练中,梯度传播路径长度增加,容易出现梯度消失或爆炸。Reformer 通过可逆残差网络保持了梯度的流动稳定性:由于可逆层的前向与反向传播具有确定性重构,梯度不会因中间激活的舍弃而扩散或衰减。同时,LSH 注意力的局部性质也避免了远距离 token 间梯度的过度耦合。
3.2 批大小与哈希轮次选择
训练时需平衡批次大小与哈希轮次。较大的批次能更好地近似完整注意力,但会增加显存占用。实践中,Reformer 推荐每轮使用 4-8 轮哈希,并将每个桶的 token 数量限制在 128 或 256。对于序列长度超过 16K 的任务,可增加轮次至 8-16 轮以确保近似精度。批次大小通常设为 8 或 16,以适配 GPU 显存。
3.3 损失平滑技术
由于 LSH 注意力引入了随机采样误差,模型在训练初期收敛可能震荡。Reformer 采用标签平滑(label smoothing)技术缓解此问题:将目标分布的 one-hot 硬标签替换为均匀分布与硬标签的加权平均,平滑系数常设为 0.1。此外,在损失函数中增加对哈希桶内 padding 部分的掩码,避免无效 token 影响梯度。
4 性能与比较
4.1 推理速度与显存占用
在序列长度 L=64K、隐藏维度 d=512 的设定下,Reformer 的推理速度比原始 Transformer 快约 3-5 倍,显存占用减少 4-8 倍。具体而言,LSH 注意力使每次点积的计算量从 O(L² d) 降至 O(L·b·d),其中 b 为平均桶大小(约 128)。可逆层则使每层的显存占用从 O(Ld) 降至 O(2d)(仅存储两个通道的输出)。
4.2 与 Transformer、Longformer 的对比
与原始 Transformer 相比,Reformer 在序列长度超过 4K 时表现出明显的计算与内存优势,但在短序列上因哈希开销略慢于 Transformer。与 Longformer(采用滑动窗口注意力)相比,Reformer 在处理稀疏但非局部的长程依赖时表现更优(如整本书级别的文本),而 Longformer 在局部密集交互任务上更快。整体上,Reformer 在极长序列(>32K)任务中占优势。
4.3 在长文本分类与生成任务上的表现
在长文本分类测试(如 IMDB 评论、PubMed 摘要)中,Reformer 在 F1 分数上达到与 Transformer 相当的水平,但训练时间减少 40%。在生成长文档(如维基百科文章)任务中,Reformer 能生成连贯的 64K token 文本,而相同显存下的 Transformer 仅能处理 8K 长度。在逐段生成测试中,Reformer 的困惑度仅比完整注意力模型高 0.3-0.5 个点。
5 应用场景
5.1 长文档问答
Reformer 可直接处理整篇长文档(如 5000 页的技术报告)作为上下文,无需分块或滑动窗口。在自然问题(Natural Questions)等数据集上,Reformer 在长上下文检索准确率上比基于窗口的方法提升 12%。
5.2 基因组序列建模
基因组序列常包含数百万个碱基对,Reformer 的线性复杂度使其能直接对整条染色体进行建模。在预测基因表达调控区域任务中,Reformer 比传统 CNN 模型在长程依赖捕捉上显著更优,训练速度也快于基于 GPU 的 Transformer 变体。
5.3 高分辨率图像生成
将图像分割为极小的 patch(如 4×4像素),序列长度随图像分辨率呈平方增长。Reformer 可用于直接生成 1024×1024 分辨率图像,而传统 Transformer 因显存限制只能处理 256×256。在 ImageNet 生成任务中,Reformer 的 FID 分数与 BigGAN 接近,但训练资源更少。
5.4 另类解读:网络梗时代“改革者”的变体
在中文互联网社区中,“Reformer”因其字面含义“改革者”而衍生出非技术性的文化用法。程序员群体常将其作为幽默梗,用以自嘲或调侃团队内部的技术重构行为。
5.4.1 作为团队自嘲昵称
当某个团队频繁进行代码重构、架构调整或流程“优化”时,成员可能戏称自己为“Reformer”,暗示虽然改动很多但效果未必立竿见影。例如:“这周我们又当了一回 Reformer——需求改了五遍,最后改回第一版。”
5.4.2 与“重构者”(Refactorer)的混用趣事
由于“Reformer”与英文“Refactorer”(重构者)发音和拼写相近,非专业用户经常将二者混淆。在一些技术聊天群中,常出现“我们需要一个勇于 Reformer 的人来重构这个模块”的误用,成为圈内幽默素材。这种混用甚至催生了“Reformer → Refactorer”的谐音梗,用来指代那些以“改革”为名、实则只改代码格式的开发者。