大鸟：适用于更长序列的Transformer

非 Google 官方产品。

什么是 BigBird？

BigBird 是一种基于稀疏注意力机制的 Transformer 模型，它将 BERT 等基于 Transformer 的模型扩展到处理更长的序列。此外，BigBird 还提供了对完整 Transformer 能力的理论理解，说明稀疏模型能够处理的能力范围。

由于能够处理更长的上下文，BigBird 在问答和摘要等多种自然语言处理任务上显著提升了性能。

更多详细信息和对比结果，请参阅我们的 演示文稿。

引用

如果您觉得本项目有用，请引用我们的 NeurIPS 2020 论文：

@article{zaheer2020bigbird,
  title={Big bird: Transformers for longer sequences},
  author={Zaheer, Manzil and Guruganesh, Guru and Dubey, Kumar Avinava and Ainslie, Joshua and Alberti, Chris and Ontanon, Santiago and Pham, Philip and Ravula, Anirudh and Wang, Qifan and Yang, Li and others},
  journal={Advances in Neural Information Processing Systems},
  volume={33},
  year={2020}
}

代码

最重要的目录是 core。core 目录下有三个主要文件：

• attention.py： 包含 BigBird 的线性注意力机制
• encoder.py： 包含主要的长序列编码器堆栈
• modeling.py： 包含封装了 BigBird 注意力机制的 BERT 和 seq2seq Transformer 模型

Colab/IPython 笔记本

在 imdb.ipynb 中提供了一个用于文本分类的快速微调示例。

创建 GCP 实例

请先创建一个项目，并在具有以下配额的区域中创建实例：

gcloud compute instances create \
  bigbird \
  --zone=europe-west4-a \
  --machine-type=n1-standard-16 \
  --boot-disk-size=50GB \
  --image-project=ml-images \
  --image-family=tf-2-3-1 \
  --maintenance-policy TERMINATE \
  --restart-on-failure \
  --scopes=cloud-platform

gcloud compute tpus create \
  bigbird \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=2.3.1

gcloud compute ssh --zone "europe-west4-a" "bigbird"

为了便于说明，我们使用了实例名称 bigbird 和区域 europe-west4-a，但您可以根据需要进行更改。有关创建 Google Cloud TPU 的更多详细信息，请参阅 在线文档。

安装与检查点

git clone https://github.com/google-research/bigbird.git
cd bigbird
pip3 install -e .

您可以在我们的 Google Cloud Storage 存储桶中找到预训练和微调后的检查点。

您也可以使用 gsutil 下载它们：

mkdir -p bigbird/ckpt
gsutil cp -r gs://bigbird-transformer/ bigbird/ckpt/

该存储桶包含：

• 预训练的 BERT 模型，分为 base（bigbr_base）和 large（bigbr_large）两种尺寸。这些模型仅对应于类似 BERT/RoBERTa 的编码器部分。遵循原始 BERT 和 RoBERTa 的实现，它们采用后归一化结构，即层归一化发生在注意力层之后。然而，根据 Rothe 等人的研究，我们可以通过耦合编码器和解码器参数，将其部分地以编码器-解码器的方式使用，如 bigbird/summarization/roberta_base.sh 启动脚本所示。
• 预训练的 Pegasus 编码器-解码器 Transformer，large 尺寸（bigbp_large）。同样遵循 Pegasus 的原始实现，这些模型采用前归一化结构，拥有完整的独立编码器和解码器权重。此外，针对长文档摘要数据集，我们还为每个数据集转换了 Pegasus 检查点（model.ckpt-0），并提供了适用于更长文档的微调检查点（model.ckpt-300000）。
• 微调后的 tf.SavedModel 用于长文档摘要，可以直接用于预测和评估，如 colab 笔记本 所示。

运行分类任务

要快速开始使用 BigBird，可以从 classifier 目录中的分类实验代码入手。运行代码只需执行以下命令：

export GCP_PROJECT_NAME=bigbird-project  # 替换为您自己的项目名称
export GCP_EXP_BUCKET=gs://bigbird-transformer-training/  # 替换
sh -x bigbird/classifier/base_size.sh

使用 BigBird 编码器代替 BERT/RoBERTa

要直接使用编码器而不是例如 BERT 模型，可以使用以下代码：

from bigbird.core import modeling

bigb_encoder = modeling.BertModel(...)

它可以轻松替代 BERT 的编码器。

或者，您也可以尝试操作 BigBird 编码器的各个层：

from bigbird.core import encoder

only_layers = encoder.EncoderStack(...)

标志与配置说明

所有标志和配置都在 core/flags.py 中进行了说明。以下是一些重要的配置参数：

attention_type 用于选择使用的注意力类型。将其设置为 block_sparse 可以运行 BigBird 的注意力模块。

flags.DEFINE_enum(
    "attention_type", "block_sparse",
    ["original_full", "simulated_sparse", "block_sparse"],
    "选择注意力实现方式。"
    "'original_full'：原始 BERT 的全注意力机制。"
    "'simulated_sparse'：模拟的稀疏注意力机制。"
    "'block_sparse'：稀疏注意力的分块实现。")

block_size 用于定义块的大小，而 num_rand_blocks 则用于设置随机块的数量。当前代码使用 3 个窗口块和 2 个全局块。目前代码仅支持静态张量。

需要注意的要点：

• 隐藏层维度必须能被头数整除。
• 当前代码仅处理静态形状的张量，因为它主要是为 TPU 设计的，而 TPU 只能处理静态形状的张量。
• 对于序列长度小于 1024 的情况，建议使用 original_full，因为此时使用稀疏的 BigBird 注意力并无优势。

对比

最近，Long Range Arena 提供了一个包含六个需要较长上下文的任务的基准测试，并对现有的所有长距离 Transformer 模型进行了实验比较。结果如下所示。与其他模型相比，BigBird 模型在不牺牲性能的情况下显著降低了内存消耗。

Big Bird 快速上手指南

Big Bird 是一种基于稀疏注意力机制的 Transformer 模型，旨在将 BERT 等模型扩展至处理更长的序列（Longer Sequences）。它在问答、摘要等需要长上下文的 NLP 任务中表现卓越。

环境准备

• 系统要求：推荐在 Linux 环境下运行。
• 硬件加速：官方代码主要针对 Google Cloud TPU 优化（需静态形状张量），但也支持 GPU/CPU 运行（需注意显存限制）。
• 前置依赖：
  • Python 3.x
  • TensorFlow 2.3.1+ (官方示例基于此版本)
  • git, pip
  • (可选) Google Cloud SDK (gcloud, gsutil)：如需使用官方提供的预训练权重或 TPU 实例。

注意：由于该工具由 Google Research 发布且强依赖 GCP 生态，国内用户若无 GCP 环境，建议在本地安装 TensorFlow 后尝试运行分类示例，但下载预训练模型可能需要配置网络代理或使用镜像加速。

安装步骤

1. 克隆代码库

git clone https://github.com/google-research/bigbird.git
cd bigbird

2. 安装依赖包

使用可编辑模式安装，以便直接修改核心代码：

pip3 install -e .

3. 获取预训练模型（可选）

官方权重存储在 Google Cloud Storage。如果您已配置好 gsutil 且有访问权限，可执行以下命令下载：

mkdir -p bigbird/ckpt
gsutil cp -r gs://bigbird-transformer/ bigbird/ckpt/

包含内容：

• bigbr_base / bigbr_large: 类 BERT/RoBERTa 的编码器模型。
• bigbp_large: 类 Pegasus 的 Encoder-Decoder 模型（适用于长文档摘要）。
• 针对特定数据集微调后的检查点。

基本使用

方式一：运行文本分类示例（推荐起步）

仓库提供了一个基于 IMDB 数据集的微调演示，位于 classifier 目录。

1. 设置环境变量（请替换为您自己的 GCP 项目名称和存储桶路径，若仅在本地测试可暂时忽略或设为本地路径）：

   export GCP_PROJECT_NAME=bigbird-project
   export GCP_EXP_BUCKET=gs://bigbird-transformer-training/
   

2. 执行基础尺寸模型的训练脚本：

   sh -x bigbird/classifier/base_size.sh
   

方式二：在代码中直接调用 Encoder

您可以直接将 Big Bird Encoder 作为 BERT 的替代品嵌入到您的项目中。

替换 BERT 编码器：

from bigbird.core import modeling

# 初始化 Big Bird BERT 模型
# 参数配置需根据具体任务调整，此处省略具体参数
bigb_encoder = modeling.BertModel(...)

仅使用编码器堆栈（底层调用）：

from bigbird.core import encoder

# 直接使用 EncoderStack
only_layers = encoder.EncoderStack(...)

关键配置说明

在运行模型时，需通过 flags 配置注意力机制类型。核心参数如下：

• attention_type: 选择注意力实现方式。

  • "block_sparse": Big Bird 核心模式，启用块稀疏注意力。
  • "original_full": 原始 BERT 的全量注意力（序列长度 < 1024 时建议使用此项，稀疏模式无优势）。
  • "simulated_sparse": 模拟稀疏注意力。

• block_size: 定义块的大小。

• num_rand_blocks: 设置随机块的数量。

示例配置逻辑（参考 core/flags.py）：

# 伪代码示例，实际运行通常通过命令行参数传递
# --attention_type=block_sparse

提示：当前代码主要设计用于 TPU，要求张量形状静态固定。隐藏维度（Hidden dimension）必须能被注意力头数（Number of heads）整除。