Hugging Face 阅读：长程 Transformer

TL;DR

Hugging Face 分析了四种主要的架构方法——Longformer、Compressive Transformer、Linformer 和 Performer——以解决标准 Transformer 的二次方内存和时间复杂度瓶颈。这些方法能够处理远超传统 512 或 1024 token 限制的序列，这对于文档级 NLP、语音和蛋白质建模至关重要。

克服二次方瓶颈

标准 Transformer 的自注意力机制随序列长度 $n$ 呈二次方增长 ($O(n^2)$)，这为长文档带来了显著的内存和计算瓶颈。为了解决这个问题，研究人员开发了“高效 Transformer”（Efficient Transformers），旨在将这种复杂度降低到线性 ($O(n)$)。这些方法通常分为四类：自定义注意力模式、循环机制、低秩近似和核近似。

长程建模的架构方法

Longformer：自定义注意力模式

Longformer 使用窗口化（局部）和全局注意力的组合来取代标准自注意力，使其能够随序列长度线性扩展。

  • 机制：它在自回归语言建模中使用扩张窗口自注意力（dilated windowed self-attention），并在编码器预训练中使用局部窗口和全局双向注意力的混合模式。全局注意力应用于特定任务的 token（例如 [CLS] token 或问答任务中的问题 token），以允许信息在整个序列中流动。
  • 关键优势：自注意力层是一个即插即用的替代方案，这意味着预训练检查点可以适配长程输入，而无需进行昂贵的从头开始的预训练。
  • 权衡：滑动窗口注意力依赖于索引操作，在某些硬件（如 TPU）上可能会比较慢。

Compressive Transformer：循环机制

基于 Transformer-XL，Compressive Transformer 引入了压缩内存来存储原本会被丢弃的过去激活值。

  • 机制：它使用压缩函数（例如 max/mean pooling 或 1D convolution）将过去的激活值压缩 $c$ 倍。这使得模型能够同时关注到近期 token 的常规内存和更久远 token 的压缩内存。
  • 关键优势：它显著提高了在 enwik8 和 WikiText-103 等长程语言建模基准测试上的困惑度（perplexity），特别是对于在长距离内出现的罕见词。
  • 权衡：训练过程较为脆弱，需要专门的优化调度方案，即逐步增加有效 batch size。

Linformer：低秩近似

Linformer 通过将序列长度投影到较小的维度来降低复杂度，其依据是观察到自注意力矩阵是低秩的。

  • 机制：利用 Johnson-Lindenstrauss 引理，Linformer 学习注意力上下文矩阵的低秩分解。这确保了不会计算或存储任何 $n imes n$ 矩阵。
  • 关键优势：推理速度（时钟时间）不受序列长度增加的影响，且与标准 Transformer 相比，收敛速度保持稳定。
  • 权衡：分解是为训练时确定的固定上下文长度设计的，如果不进行适配，则无法推广到更长的序列。

Performer：核近似

Performer 使用 FAVOR+ (Fast Attention Via Orthogonal Random positive features) 算法来近似 softmax 注意力核，而不依赖于稀疏性或低秩先验。

  • 机制：它使用随机特征映射来近似 softmax 函数，从而允许在查询乘法之前执行矩阵乘法 $K imes V$。这有效地绕过了 $n imes n$ 注意力矩阵的计算。

  • 关键优势：由于它不对注意力矩阵的结构做任何假设，因此它在不同模态（包括语音和蛋白质序列）中都具有高度适用性。

  • 权衡：微小的近似误差可能会在多个 Transformer 层中传播，从而可能影响预训练网络的微调稳定性。

对比分析与讨论

归纳偏置与权衡

架构的选择取决于任务的具体要求和可用数据：

  • Longformer vs. Linformer：Longformer 使用固定的稀疏模式，而 Linformer 学习低秩分解。Longformer 通常比 Linformer 效率更低，但在序列长度方面更具灵活性（尽管 Linformer 受限于其训练上下文长度）。
  • Performer：与前两者不同，它通过近似核函数本身来实现，使其成为一种通用的即插即用替代方案，且不假设注意力矩阵是稀疏或低秩的。

位置编码

位置编码是长程效率的一个关键因素：

  • 相对位置编码：用于 Transformer-XL 和 Compressive Transformers；它们可以轻松扩展到未见过的序列长度，但计算成本较高。
  • 绝对位置编码：用于 Longformer 和 Linformer；它们在计算上更高效，但在处理比训练时见过的序列更长时，灵活性较低。
  • 位置注入注意力：由 Shortformer 引入的一种替代方案，它将位置信息添加到查询和键，而不是 token 嵌入。

训练策略

来自 Shortformer 和 Longformer 等模型的证据表明，在短序列上进行训练并逐渐增加长度，可以实现更快的训练速度和更强的下游任务性能，从而防止模型依赖于数据中的伪相关性。

Sources