Reformer

Reformer 是一种节省内存的Transformer，由加州大学伯克利分校和 Google Research 的 Nikita Kitaev、Łukasz Kaiser 与 Anselm Levskaya 在 ICLR 2020 上发表的论文 Reformer：高效 Transformer 中提出。它针对限制序列长度的两项开销进行优化。局部敏感哈希（LSH）注意力取代完整的 n × n 注意力矩阵，长度为 L 的序列计算开销约为 L log L；可逆残差层则为整个网络只存储一次反向传播所需的激活值，而不是每层都存储一次。Google 报告称，Reformer 能在单个配备 16 GB 内存的加速器上处理最长 1 million 个词的上下文窗口，并将整本小说《罪与罚》作为一个训练样本进行处理。

LSH 注意力#

经过 softmax 后，每个查询的注意力主要集中在与其最相似的少数键上，因此 Reformer 只搜索这些键。它对查询和键使用同一个投影，并通过角度局部敏感哈希对每个向量进行哈希：与随机矩阵 R 相乘，并将 [xR, −xR] 中最大值的位置作为桶编号，因此指向相似方向的向量通常会落入同一个桶。随后，系统按桶对token排序，将排序后的序列切分为大小相同的块，并让每个 token 只关注自身所在块和前一个块。系统会并行运行多轮哈希，以捕捉被单轮哈希分开的相似向量对。

LSH 是近似最近邻向量搜索的经典方法：它只需精确搜索的一小部分开销，就能找到查询的大多数 k 近邻。分桶步骤足够简单，可以单独运行：

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

与查询的余弦相似度为 0.96 的近邻向量，在 1,000 轮中的 72% 落入同一桶；相似度为 0.04 的无关向量则有 9% 落入同一桶。因此，将注意力限制在同一个桶内仍能保留大多数重要的向量对。

可逆层与分块#

标准神经网络会为反向传播保存每一层的激活值，因此内存占用会随网络深度和序列长度一起增长。Reformer 改造了可逆残差网络（RevNets）：将激活值分为两半，并用另一半的函数更新每一半。

反向传播时，每一层都会根据输出恢复输入：x2 = y2 − FeedForward(y1)，然后 x1 = y1 − Attention(x2)，因此只需保存最终激活值，并对每一层重新计算一次。梯度检查点采用类似的权衡：在少数检查点保存激活值，并重新计算检查点之间的层。Reformer 还会按位置块运行前馈层，因为前馈层的隐藏层宽度是模型的数倍。激活值占用的内存决定了批次大小上限；在训练 Ultralytics YOLO26 时，使用 batch=-1，AutoBatch 会测量多个批次大小下的 GPU 内存占用，拟合一条直线，并选择预计占用可用内存 60% 的批次大小。

结果与实现#

Reformer 在 enwik8 的字符级文本和 ImageNet-64 的逐像素图像生成任务上达到了标准 Transformer 的水平，同时大幅减少内存占用，并在长序列上运行得更快。Hugging Face Transformers 将其实现为 ReformerModel，并发布了在 enwik8 和《罪与罚》上训练的检查点。

Reformer 与相关模型#

• Longformer： 预先固定注意力模式：每个 token 周围的滑动窗口加上少数全局 token。
• 稀疏注意力： 只为部分 token 对计算分数的一类方法。LSH 注意力就是其中一种基于内容的方法。
• 线性注意力： 用核特征映射取代 softmax，使注意力计算能够在线性时间内完成。
• Flash attention： 在 GPU 片上内存块中计算精确注意力，而不存储分数矩阵，因此内存占用线性增长，而分数数量仍呈二次增长。