Longformer

Longformer 和 LongformerEncoderDecoder (LED) 是针对长文档的预训练 Transformer 模型。

***** 2020年12月1日新功能：LongformerEncoderDecoder *****

现在提供了一个 LongformerEncoderDecoder (LED) 模型。它支持长输入的序列到序列任务。借助梯度检查点、fp16 和 48GB GPU，输入长度最高可达 16K tokens。请查看更新后的论文以获取模型详情和评估结果。

• 预训练模型：1) led-base-16384, 2) led-large-16384

• 要求：请确保使用 requirements.txt 中指定的 huggingface/transformers 分支。该分支增加了对梯度检查点的支持，并允许输入和输出设置不同的最大序列长度。您也可以运行 pip install git+https://github.com/allenai/longformer.git

• 请查看脚本 scripts/summarization.py 以了解如何使用该模型的示例。

***** 2020年7月23日新功能：速度下降 *****

最近在 hugginface/transformers 中发现并修复了一个显著的速度下降问题（详情请查看 此 PR）。为避免此问题，您可以使用旧版本 v2.11.0，但该版本不支持梯度检查点；或者使用 master 分支。下一个 hugginface/transformers 发布版应会修复此问题。

***** 2020年6月29日新功能：更易使用梯度检查点 *****

梯度检查点已随 huggingface/transformers v3.0.0 版本发布。梯度检查点可将内存减少 5 倍，从而在较小的 GPU 上处理更长的序列成为可能。要使用，可以尝试如下代码：

from transformers import LongformerModel
model = LongformerModel.from_pretrained('allenai/longformer-base-4096', gradient_checkpointing=True)

******* 2020年6月2日新功能：与 Huggingface 集成 + 训练自己的长模型 + 梯度检查点 *****

1. Longformer 现已集成到 huggingface/transformers v2.11.0 版本中。现在您可以这样操作：

from transformers import LongformerModel
model = LongformerModel.from_pretrained("allenai/longformer-base-4096")

该版本还包含 LongformerForQA 和其他 LongformerForTaskName，并自动设置全局注意力。

2. 我们添加了一个 notebook 来展示如何将现有预训练模型转换为“长”版本。

3. 梯度检查点已合并到 HF master 分支（查看 PR）。梯度检查点可显著减少内存占用（对于 longformer-base-4096 可减少 5 倍），从而在较小的 GPU 上处理更长的序列成为可能。

***** 2020年4月27日新功能：滑动窗口注意力的 PyTorch 实现 *****

我们添加了滑动窗口注意力的 PyTorch 实现，无需自定义 CUDA 内核。它的功能有限，但在下游任务微调时使用起来更方便。

优势：支持 CPU、TPU 和 fp16，而自定义 CUDA 内核不支持这些。

局限性：内存占用增加一倍（但 fp16 可部分抵消），且不支持扩张和自回归注意力（微调时不需要）。

因此，它适合下游任务微调，但不适合作为语言建模的首选。下面的代码片段和 TriviaQA 脚本已更新为使用这一新实现。

***** 新信息结束 *****

使用方法

1. 下载预训练模型

• longformer-base-4096
• longformer-large-4096

2. 安装环境和代码

   conda create --name longformer python=3.7
   conda activate longformer
   conda install cudatoolkit=10.0
   pip install git+https://github.com/allenai/longformer.git
   

3. 运行模型

   import torch
   from longformer.longformer import Longformer, LongformerConfig
   from longformer.sliding_chunks import pad_to_window_size
   from transformers import RobertaTokenizer
   
   config = LongformerConfig.from_pretrained('longformer-base-4096/') 
   # 选择注意力模式 'n2'、'tvm' 或 'sliding_chunks'
   # 'n2': 用于常规 n2 注意力
   # 'tvm': 我们的滑动窗口注意力的自定义 CUDA 内核实现
   # 'sliding_chunks': 我们的滑动窗口注意力的 PyTorch 实现
   config.attention_mode = 'sliding_chunks'
   
   model = Longformer.from_pretrained('longformer-base-4096/', config=config)
   tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
   tokenizer.model_max_length = model.config.max_position_embeddings
   
   SAMPLE_TEXT = ' '.join(['Hello world! '] * 1000)  # 长输入文档
   
   input_ids = torch.tensor(tokenizer.encode(SAMPLE_TEXT)).unsqueeze(0)  # 批次大小为 1
   
   # TVM 代码无法在 CPU 上运行。如果 `config.attention_mode = 'tvm'`，请取消注释
   # model = model.cuda(); input_ids = input_ids.cuda()
   
   # 注意力掩码值 —— 0: 不关注，1: 局部关注，2: 全局关注
   attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device) # 初始化为局部关注
   attention_mask[:, [1, 4, 21,]] =  2  # 根据任务设置全局关注。例如，
                                        # 分类：<s> 标记
                                        # QA：问题标记
   
   # 将序列长度填充到最接近 512 的倍数。这是‘sliding_chunks’注意力所必需的
   input_ids, attention_mask = pad_to_window_size(
           input_ids, attention_mask, config.attention_window[0], tokenizer.pad_token_id)
   
   output = model(input_ids, attention_mask=attention_mask)[0]
   

模型预训练

此 notebook 展示了我们从 RoBERTa 检查点开始训练 Longformer 的流程。同样的流程也可用于获得其他现有预训练模型的长版本。

TriviaQA

• 训练脚本：scripts/triviaqa.py
• 预训练大模型：这里（复现排行榜结果）
• 使用说明：scripts/cheatsheet.txt

CUDA内核

我们的自定义CUDA内核是在TVM中实现的。目前，该内核仅支持GPU和Linux系统。我们已在Ubuntu、Python 3.7、CUDA10以及PyTorch >= 1.2.0环境下进行了测试。如果它在您的环境中无法正常运行，请提交一个新的问题。

编译内核：我们已内置了CUDA内核的编译二进制文件，因此大多数用户无需再进行编译。不过，如果您有兴趣，可查看scripts/cheatsheet.txt以获取相关说明。

已知问题

请查看仓库的issues，了解我们近期计划解决的已知问题列表。如果您的问题未被提及，请创建一个新的issue。

引用

如果您在研究中使用了Longformer，请引用Longformer: The Long-Document Transformer。

@article{Beltagy2020Longformer,
  title={Longformer: The Long-Document Transformer},
  author={Iz Beltagy and Matthew E. Peters and Arman Cohan},
  journal={arXiv:2004.05150},
  year={2020},
}

Longformer是由艾伦人工智能研究所（AI2）开发的一个开源项目。AI2是一家非营利性研究机构，其使命是通过具有重大影响力的AI研究与工程，为全人类做出贡献。

Longformer 快速上手指南

环境准备

• 系统要求：Linux 或 macOS（CUDA 内核仅支持 Linux + GPU）
• Python 版本：3.7+
• 硬件建议：至少 16GB 显存（使用梯度检查点可降低至 8GB）
• 前置依赖：PyTorch ≥ 1.2.0，CUDA 10.0+（如使用 TVM 内核）

推荐使用国内镜像加速：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

conda create --name longformer python=3.7
conda activate longformer
conda install cudatoolkit=10.0
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple git+https://github.com/allenai/longformer.git

若仅需基础功能（推荐新手），可直接使用 Hugging Face Transformers：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple transformers>=3.0.0

基本使用

1. 加载预训练模型（推荐方式）

from transformers import LongformerModel, LongformerTokenizer

model = LongformerModel.from_pretrained('allenai/longformer-base-4096', gradient_checkpointing=True)
tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
tokenizer.model_max_length = 4096

2. 处理长文本示例

import torch

SAMPLE_TEXT = 'Hello world! ' * 1000  # 构造长文本
input_ids = tokenizer.encode(SAMPLE_TEXT, return_tensors='pt')

# 设置注意力掩码：1=局部注意力，2=全局注意力（如 [CLS] 位置）
attention_mask = torch.ones(input_ids.shape, dtype=torch.long)
attention_mask[0, 0] = 2  # 对 [CLS] token 设置全局注意力

# 模型推理
outputs = model(input_ids, attention_mask=attention_mask)
last_hidden_state = outputs[0]  # 获取输出

✅ 推荐使用 gradient_checkpointing=True 以节省显存，适用于消费级 GPU
✅ 支持 Hugging Face 生态，可直接用于 LongformerForSequenceClassification、LongformerForQuestionAnswering 等任务