← 全部论文 / 2020

Longformer: The Long-Document Transformer

2020 稀疏化与内存高效扩展
下载 PDF

Longformer: The Long-Document Transformer - 中文验证版

英文原文卡片：longformer_2020.md

状态：已翻译。

元数据

计算设置

论文在工作的若干部分明确列出 GPU 硬件。对于自回归字符语言建模，附录说明实验在 48GB RTX8000 GPU 上运行，使用混合精度训练、梯度检查点，且注意力保持 fp32，因为 fp16 注意力在训练后期导致溢出和 NaN。小字符模型在 4 块 RTX8000 GPU 上运行 16 天，大模型在 8 块 RTX8000 GPU 上运行 13 天。阶段性字符 LM 调度从序列长度 2048 开始，在阶段 5 达到 23040，标记为 GPU 内存限制。

对于预训练 Longformer，论文从 RoBERTa 继续而非从头训练。它训练基础和大型模型 65K 梯度更新，序列长度 4096，batch size 64，最大学习率 3e-5，500 步预热。附录还指出某些任务模型在 32GB V100 GPU 上运行，小模型约需 4 块 GPU 1 天，大模型约需 8 块 GPU 1 天。WikiHop 微调列为单 RTX8000 GPU 训练，Longformer-base 约需 1 天完成 5 个 epoch。

瓶颈

瓶颈是那些证据长于 BERT 512-token 窗口的任务中的二次注意力矩阵。Longformer 直接记录数据集长度：WikiHop 平均 1535 个 wordpiece，第 95 百分位 3627；TriviaQA 平均 6589，第 95 百分位 17126；HotpotQA 平均 1316；OntoNotes 平均 506；Hyperpartisan 平均 705，第 95 百分位 1975。论文的动机是分块或截断会丢失跨块信息，并迫使使用复杂的任务特定拼接模型。

还存在实现瓶颈。滑动窗口注意力矩阵是带状的，但 PyTorch 和 TensorFlow 不直接提供仅计算 QK^T 选定对角线的带状矩阵乘法。朴素循环实现内存高效但慢得无法使用。矢量化分块实现快速，因为它使用单次 PyTorch 矩阵乘法，但它计算了一些零值并消耗约 2 倍于完美优化实现的内存。当最长序列和膨胀因子重要时，需要自定义 CUDA 路径。

方法适配

Longformer 围绕局部性加少量全局位置适配注意力。每个 token 关注一个固定大小的窗口，对于固定窗口 w 复杂度为 O(n x w)。堆叠层增长感受野，膨胀窗口可以达到数万 token 而不增加每层窗口大小。对于需要全文档聚合的任务输出，模型添加对称全局注意力：选定 token 关注所有 token，所有 token 关注它们。论文对分类使用 [CLS]，对 QA 使用问题 token，对候选答案式输入使用任务特定全局位置。

设备适配分三层。Longformer-loop 简单且内存高效，但对实验太慢。Longformer-chunks 是预训练/微调的主力，因为它将计算保持在大型密集 PyTorch 乘法中，以额外零值工作换取速度。Longformer-cuda 是 TVM 生成的自定义 CUDA 内核，用于带状注意力；它支持膨胀和自回归模式，是内存效率最高的实现，据报告速度与优化后的全自注意力相当。这种划分反映了硬件权衡：内存允许时使用密集库内核，当序列长度或膨胀推动内存预算时使用自定义内核。

编码器-解码器 LED 变体在更高层次上做出相同的计算选择。编码器使用 Longformer 局部加全局注意力，解码器对生成前缀和编码 token 保留全注意力。这是合理的，因为长文档摘要生成有非常长的输入和相对短的输出。

证据

语言模型证据展示了内存目标。Longformer 评估序列长度 32256 的字符 LM，并通过分阶段训练达到 23040 token，因为 GPU 内存限制。小模型在 text8 上报告 BPC 1.10，enwik8 上 1.00，论文称其在该设置下达到了当时的最新水平。大 enwik8 模型报告 BPC 0.99，与更大的长上下文模型具有竞争力。

预训练证据也是计算特定的。在继续预训练之前，仅将 RoBERTa 改为序列长度 4096 并复制位置嵌入的 Longformer 给出基础 BPC 1.957；2K 步更新后改进到 1.753，65K 步更新后到 1.705。大模型类似地在 65K 步更新后改进到 1.358。这表明模型可以重用在 4096-token 稀疏注意力范式中学习操作的 RoBERTa 权重。

下游结果支持长上下文预算。Longformer-base 在 WikiHop 75.0 vs 72.4、TriviaQA 75.2 vs 74.3、HotpotQA 64.4 vs 63.5、IMDB 95.7 vs 95.3、Hyperpartisan 94.8 vs 87.4 上均超越 RoBERTa-base。Longformer-large 在 WikiHop 上报告排行榜 F1 81.9，TriviaQA 上 77.3，领先于表格中当时的最新水平。WikiHop 消融尤其具有计算结构意义：同时移除独立全局投影和全局注意力会使准确率降至 65.5，而全 4096-token Longformer 根据训练长度达到 73.8 或 75.0。LED 在 arXiv 摘要生成上展示输入长度缩放，16384-token LED-large 达到 ROUGE-1/2/L 46.63/19.62/41.83。

历史影响

Longformer 帮助建立了长文档 Transformer 作为一个 GPU 工程问题，而不仅是架构草图。它将简单的稀疏模式（局部加全局注意力）与如何使带状注意力在真实 GPU 内存限制下运行的实际问题联系起来。它还表明稀疏长上下文注意力可以是预训练/微调工作流中 RoBERTa 注意力的直接替代，这对采用很重要。

局限

局限来自相同的设计选择。全局注意力是任务指定的，因此模型需要关于哪些 token 应为全局的人类选择。最快和内存效率最高的路径依赖专门 CUDA/TVM 实现，分块实现通过计算掩码零值付出额外内存代价。论文报告字符 LM 中 23040 字符是 GPU 内存限制，预训练 Longformer 主要是 4096-token 模型，因此更长输入仍需要仔细批处理、截断或专门 LED 式设置。

链接