Reformer: 使用内存高效的 Transformers 推动语言建模的极限

Overview

Reformer 模型解决了标准 Transformer 在处理长序列时的内存瓶颈。通过重新设计自注意力、前馈层、残差连接和位置编码，它能够在低于 8 GB RAM 的情况下训练多达 500 000 个标记，远超如 BERT 等模型的 512‑标记限制。

Reformer Self-Attention Layer

Reformer 将全局自注意力替换为两种内存高效的替代方案：局部自注意力和 LSH 自注意力。

Local Self-Attention

局部自注意力将输入分割为长度为 config.local_chunk_length 的块，并在每个块内应用全局自注意力。通过添加重叠（config.local_num_chunks_before 和 config.local_num_chunks_after），每个标记可以关注相邻块的有限上下文，将二次内存成本降低到 O(n × chunk_length)。仅此仍不足以满足需要长距离依赖的任务。

LSH Self-Attention

LSH 自注意力通过使用局部敏感哈希将查询（和键）向量哈希到桶中来近似全局注意力。哈希到同一桶的向量被视为相似，因此自注意力仅在每个桶内计算。然后将置换后的输入进行分块（带重叠）并进行注意力计算，使得远距离标记的信息能够混合，同时保持内存使用为 O(n × hashes × chunk_length)。多个哈希轮次（config.num_hashes）可以组合以提高准确性。

Benchmark

对 google/reformer-enwik8 模型的基准测试显示了内存节省：

  • 使用全局自注意力（块长度设置为 8192）时，模型在约 16 K 标记时超出 GPU 内存。
  • 使用默认的局部 + LSH 自注意力，相同模型在 11 GB GPU 上可运行至约 16 K 标记才出现内存不足，表明随着序列长度增加，内存消耗的增长被抑制。

Chunked Feed Forward Layers

宽型 Transformer 中的大型前馈中间矩阵占据了主要内存。Reformer 引入了分块前馈层，其中线性层以大小为 config.chunk_size_feed_forward 的小块进行处理。这避免了存储完整的中间张量，用额外的计算换取较低的内存使用。

Benchmark

当前馈尺寸被放大（例如，达到 16384）且注意力头数减少时，启用分块（chunk_size_feed_forward=1）可使序列长度为 4096、批量大小为 8 时的峰值内存从约 9 GB 降至约 6 GB，这证实了在前馈层为瓶颈的模型中分块的好处。

Reversible Residual Layers

训练标准 Transformer 时会存储所有中间激活，导致内存随深度线性增长。Reformer 使用可逆残差层，使得激活可以在反向传播过程中重新计算而不是被保存。只需保留最终层的输出，即可将每层的内存开销从几百兆字节降低到低于 100 MB。

Benchmark

在序列长度为 512、批量大小为 8 的情况下，随着层数增加（4、8、12），比较 BERT 和 Reformer 可得：

  • BERT‑12‑Layers：约 7.4 GB
  • Reformer‑12‑Layers：约 5.4 GB 因此，Reformer 每层增加的内存远少，使得在相同硬件限制下可以构建更深的模型。

Axial Positional Encodings

标准位置嵌入随最大序列长度线性增长，对于非常长的输入（例如，0.5 M 标记需要约 2 GB）会变得不可行。轴向位置编码将位置空间分解为两个较小的维度（config.axial_pos_shape），并拆分隐藏大小（config.axial_pos_embds_dim）。得到的嵌入表大小为 shape[0]×dim[0] + shape[1]×dim[1]，而不是 max_len×hidden_size。

Benchmark

对于 google/reformer-crime-and-punishment 模型（能够处理 0.5 M 标记，隐藏大小 256）：

  • 默认位置嵌入：524 288 × 256 参数（约 1.36 亿）。
  • 形状为 (512，1024)、维度为 (64，192) 的轴向位置嵌入：约 2.6 M 参数。 在批量大小为 8、序列长度为 512 时，推理阶段的内存从约 959 MB 降至约 447 MB，说明了显著的减少。

Implications

通过结合这四种机制，Reformer 使得在不需要专用硬件的情况下，训练非常长文档（如完整书籍或代码库）上的语言模型成为可能。这扩展了可以从全上下文理解中受益的 NLP 任务范围，包括长距离摘要、在大型语料库上的问答以及建模依赖跨越数十万标记的序列数据。

Sources