longformer
滑动窗口注意力机制,让 Transformer 突破 512 token 限制,支持最长 16K t
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
滑动窗口注意力机制,让 Transformer 突破 512 token 限制,支持最长 16K t
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
你有没有过这样的经历:想用 AI 分析一份几百页的法律合同、一整本学术论文,或者一整年的客服记录,结果模型只读了个开头就卡住了?这种尴尬,恰恰就是 Longformer 要解决的问题。 2020 年 3 月,艾伦人工智能研究所(Allen Institute for AI,简称 AI2)发布了一篇轰动 NLP 圈的技术论文,标题简洁有力——《Longformer: The Long-Document Transformer》。作者团队包括 Iz Beltagy、Matthew E. Peters 和 Arman Cohan,他们提出了一个核心理念:标准 Transformer 的 O(n²) 注意力机制是长文档的噩梦,必须改掉。
传统 Transformer(如 BERT、RoBERTa)在处理文本时,需要计算每个 token 与所有其他 token 的「注意力分数」。当输入是 4096 个 token 时,单层注意力就要做 4096 × 4096 = 1600 万次计算;如果是 16K tokens,这个数字会爆炸到 2.56 亿次。更糟糕的是,GPU 显存消耗随注意力矩阵的平方增长——大多数模型的上下文窗口被锁死在 512 tokens,背后就是这个硬约束。 一些研究者选择截断输入(head-truncation),损失信息;另一些用滑动窗口近似注意力,但精度受损。Longformer 的思路是用「局部注意力 + 全局注意力」的组合拳,在精度和效率之间找到最佳平衡点。
Longformer 的注意力模式分为三层: 局部滑动窗口注意力:每个 token 只关注周围固定窗口大小(如 512)的邻居,配合多层堆叠(layer stacking)实现全局感受野。想象你在看一张地图,每页只显示相邻城市,但翻几页后就能推断出全国布局。 全局注意力:针对特定关键 token(如分类任务的 [CLS] token、问答任务的问句 token),允许其与所有位置交互。全局注意力是模型「大局观」的来源,确保重要信息不被局部视野稀释。 Dilated Attention(扩张注意力):在局部窗口中跳过部分 token,类似空洞卷积,在不增加计算量的情况下扩大感受野。 这套组合让 Longformer 在 4096 tokens 上训练,推理时支持长达 16K tokens(LED 版本 16384),同时将计算复杂度从 O(n²) 降到接近 O(n)。
Longformer 采用了渐进式转换策略:不是在随机初始化的模型上从头训练,而是从已经预训练好的 RoBERTa checkpoint 出发,逐步将位置编码扩展到 4096,将注意力模式切换为滑动窗口,然后继续在长文档语料上微调。这种方式大幅节省了算力——从头训练一个大语言模型需要数万美元,而 Longformer 的转换成本要低得多。 预训练的长文档语料包括 BookCorpus、Stories、Wikipedia、Stories 和 RealNews 的长版本。模型在 TriviaQA、WikiHop、PopQA 等长文档问答任务上刷新了 SOTA。
Longformer 提供了两套注意力实现,各有优劣: 自定义 TVM CUDA 内核:AI2 团队用 TVM(端到端深度学习编译器)编写了专用 GPU 内核,将滑动窗口注意力编译为高度优化的 CUDA 代码,计算效率极高,但仅支持 Linux + GPU + FP32,不支持 CPU、TP、FP16。 PyTorch 纯实现(sliding_chunks):不依赖任何自定义编译,一条 PyTorch 语句走天下,支持 CPU、GPU、TPU 和 FP16,但显存消耗约是 TVM 版的 2 倍。对于大多数下游任务微调场景,这个权衡是值得的。
Longformer 在 2020 年 6 月正式合入 HuggingFace Transformers(v2.11.0),使得用户无需任何自定义安装,直接几行代码就能调用:
from transformers import LongformerModel
model = LongformerModel.from_pretrained('allenai/longformer-base-4096')
HuggingFace 集成自动处理了全局注意力的 mask 设置(attention_mask 中的数值 2 表示全局注意力),并提供了 LongformerForQuestionAnswering、LongformerForSequenceClassification 等下游任务专用模型。
适用场景:长文档问答(如法律文书分析)、多文档摘要、学术论文理解、长对话理解、科学文献挖掘。 局限性:
Longformer 的意义不只是论文本身,而在于它开启了长上下文模型的研究热潮。它的滑动窗口注意力设计,直接启发了 BigBird(Google)、LED 架构、以及后续所有长上下文 Transformer 的设计思路。 从增长曲线看,Longformer 在 GitHub 收获 2196 stars、286 forks,被引用超过 5000 次(Google Scholar),是 NLP 长文档处理领域最具影响力的工作之一。对于需要处理长文档但又不希望部署几百亿参数大模型的开发者来说,Longformer 依然是性价比最高的选择之一。 图1:Longformer 与标准 Transformer 的注意力模式对比
详见 Longformer 官方论文 Figure 1