← 全部论文 / 2020

Big Bird: Transformers for Longer Sequences

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

Big Bird: Transformers for Longer Sequences - 中文验证版

英文原文卡片：bigbird_2020.md

状态：已翻译。

元数据

计算设置

BigBird 在加速器包络方面异常明确。在主实验中，作者指出模型在"16GB memory/chip"上训练，batch size 为 32-64，并将可行性归因于分块加上稀疏注意力结构。附录 E 给出了确切的 TPU 切片。基础 MLM 预训练用于 BigBird-ITC 和 BigBird-ETC，使用最大序列长度 4096、batch size 256、12 层、12 个 head、隐藏维度 768，计算资源为 8 x 8 TPUv3。大版本保持 4096 token，扩展到 24 层、16 个 head、隐藏维度 1024，并使用更大的预训练 batch size 2048。QA 微调表格根据任务和模型大小列出 4 x 2、4 x 4 或 4 x 8 TPUv3 切片；摘要生成对基础 BigBird-RoBERTa 模型使用 4 x 4 TPUv3，对大 BigBird-Pegasus 模型使用 4 x 8 TPUv3。基因组学预训练和下游运行也列出 8 x 8 TPUv3。

瓶颈

论文将全自注意力视为一道计算和内存墙，而非仅仅建模上的不便。BERT 式密集注意力将实际上下文大致限制在 512 token，而长文档工作负载有更长的证据。QA 附录列出了 Natural Questions 的中位数和最大实例长度分别为 3258 和 77962 token，TriviaQA 为 4900 和 32755，WikiHop 为 1541 和 20337。在摘要生成中，作者指出输出较短，中位数约 200 token，而输入的中位数长度通常超过 3000。瓶颈因此集中在编码器端：阅读证据或源文档占主导地位。

第二个瓶颈是设备结构。数学上稀疏的邻接矩阵在 GPU 或 TPU 上并不会自动快速运行。实现附录指出硬件加速器"在合并内存操作上表现优异"，而来自滑动窗口或随机元素查询的小规模零星查找效率低下。论文还指出，在此设置下稀疏矩阵乘法在 GPU 上并未高效实现，因为成千上万的并行核心需要规整的内存访问和密集张量操作来保持占用率。

方法适配

BigBird 围绕渐近稀疏性和加速器友好布局两方面适配 Transformer 注意力。注意力模式结合了一小组全局 token、局部窗口注意力和随机注意力。全局 token 保持全序列通信路由，并在理论部分用于保留通用近似和图灵完备性属性。局部窗口保持近处上下文廉价。随机边提供图连通性，而无需物化全部成对 token 交互。这使注意力从密集全对工作量缩减为在固定分块、窗口、随机和全局计数下与序列长度线性关系的被关注块数量。

硬件适配是更重要的计算设备要点。BigBird 没有使用任意稀疏操作，而是对查询和键进行"分块化"。查询和张量重塑为块张量，局部注意力通过滚动复制的键块张量计算，全局块拼接在一起，只有少量随机组件需要 gather。最终打包的键张量形状与查询块数量乘以小的被关注块预算成比例，因此注意力分数通过密集张量乘法产生。这正是 TPU/GPU 矩阵单元高效处理的操作类型。换句话说，BigBird 并非仅从注意力矩阵中移除条目；它重新打包剩余条目，使稀疏性可以批处理。

编码器-解码器设计也遵循计算剖面。对于摘要生成，论文仅在编码器上使用稀疏 BigBird 注意力，解码器保留全注意力，因为输出序列远短于输入序列。该选择将稀疏注意力工程投入到最节省内存和带宽的地方。

证据

规模证据是具体的。MLM 预训练使用拆分或打包到 4096 token 的文档。留存 BPC 从序列长度 512 的 RoBERTa（基础 1.846、大 1.496）改善到序列长度 4096 的 BigBird-ETC（基础 1.611、大 1.274）；论文报告同等长度的 Longformer 为基础 1.705、大 1.358。这些数字不仅是精度声明：它们表明，在列出的 16 GB/chip TPUv3 设置上可以训练 8 倍更长的上下文。

下游任务展示了在更长输入上花费计算的价值。在文档分类中，BigBird 报告 Arxiv F1 92.31，对比 RoBERTa 87.42 和此前最新水平 87.96；同一表格在 IMDb 上显示较小收益，因为只有少部分文档超过 512 token。在长文档摘要生成中，BigBird-Pegasus 在 Arxiv 上达到 ROUGE-1/2/L 46.63/19.02/41.77，PubMed 上 46.32/20.65/42.33，BigPatent 上 60.64/42.46/50.01。基因组学实验在另一领域阐述了相同的计算观点：DNA MLM BPC 从 512 token 的 BERT（1.23）改善到 4096 token 的 BigBird（1.12）；启动子预测达到 F1 99.9，染色质组蛋白标记 AUC 从 DeepSea 的 85.6 改善到 88.7。

历史影响

BigBird 使稀疏 Transformer 注意力看起来像一条实用的 TPU 时代缩放路径，而非纯粹的理论技巧。其持久的计算教训是，长上下文既需要图模式也需要内存布局。随机/局部/全局模式解决了连通性和表达能力；分块实现解决了实际的 TPU/GPU 执行模型。论文还帮助规范了以下观念：长上下文模型应在输入真正超过 512 token 的工作负载上评估，包括 QA、长文档分类、摘要生成和基因组学。

局限

主要局限是 BigBird 并非对密集注意力的普遍免费午餐。稀疏模式比普通注意力更复杂，实现和调优难度更大，且任务特定的全局 token 选择很重要。若干最佳结果使用了精心结构的全局 token 或 ETC 变体，而非简单的通用随机稀疏掩码。实验大多限制在 4096 token，因此有数万 token 的文档仍需要截断、拆分或滑动窗口处理。最后，计算表格依赖于可观的 TPUv3 切片；该方法减少了内存增长，但报告的训练和微调设置仍然是工业级加速器设置。

链接