[论文解读] Blockwise Self-Attention for Long Document Understanding

Jiezhong Qiu, Hao Ma+4 more|arXiv (Cornell University)|Nov 7, 2019

Topic Modeling参考文献 44被引用 16

一句话总结

BlockBERT 在 BERT 的自注意力机制中引入了稀疏块结构，通过将注意力限制在局部块而非完整序列中，降低了内存和计算成本。与 RoBERTa 相比，其内存使用量降低 18.7–36.1%，训练速度提升 12.0–25.1%，推理速度提升 27.8%，同时在长文档任务上的下游性能保持相当或有所提升。

ABSTRACT

We present BlockBERT, a lightweight and efficient BERT model for better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training/inference time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on language model pre-training and several benchmark question answering datasets with various paragraph lengths. BlockBERT uses 18.7-36.1% less memory and 12.0-25.1% less time to learn the model. During testing, BlockBERT saves 27.8% inference time, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

研究动机与目标

• 解决在长序列上训练和推理 BERT 时产生的高内存和计算成本问题，该问题限制了模型的可扩展性和部署能力。
• 识别点积自注意力是 BERT 中主要的内存瓶颈，因其与序列长度呈二次复杂度关系。
• 设计一种轻量级 BERT 变体，在显著降低内存和推理时间的同时，保持在长上下文任务上的强性能。
• 通过注意力矩阵中的结构化稀疏性实现对长距离依赖的有效建模，而不牺牲模型准确性。
• 提供一种简单且可实现的替代方案，用于替代依赖自定义内核或有限硬件支持的复杂稀疏注意力机制。

提出的方法

• 通过将序列划分为固定大小的块，并将自注意力限制在块内及选定的块间连接，引入块稀疏注意力。
• 用块稀疏矩阵结构替换 BERT 中原有的完整密集注意力矩阵，使浮点运算量（FLOPs）和内存使用量按稀疏模式成比例减少。
• 保留多头注意力机制，但将每头的注意力范围限制在预定义的块内或块之间，从而同时支持局部和长距离建模。
• 采用一种简单且可微分的稀疏模式，与标准深度学习框架兼容，避免使用自定义 CUDA 内核。
• 保留原始 BERT 的架构和预训练目标，仅通过修改注意力机制来提升效率。
• 在所有 Transformer 层中应用块稀疏注意力，确保与现有微调流水线的端到端兼容性。

实验结果

研究问题

• RQ1 在不降低长上下文任务性能的前提下，注意力矩阵中的结构化稀疏性是否能有效降低 BERT 的内存和计算成本？
• RQ2 与密集注意力相比，块稀疏注意力在下游任务准确率方面表现如何，特别是在长文档任务中？
• RQ3 块稀疏注意力在多大程度上能够减少训练和推理时间，同时保持对长距离依赖的建模能力？
• RQ4 块稀疏设计是否能有效使注意力头同时捕捉短距离和长距离依赖关系，如设计初衷所示？
• RQ5 与 RoBERTa 等最先进的 BERT 变体相比，BlockBERT 的效率提升在实际部署场景中表现如何？

主要发现

• 与 RoBERTa 相比，BlockBERT 在训练过程中内存消耗降低了 18.7–36.1%，主要归因于注意力矩阵的稀疏化。
• 由于浮点运算量减少和内存压力降低，从而支持更大的批量大小，训练时间减少了 12.0–25.1%。
• 与 RoBERTa 相比，推理时间减少了 27.8%，使其更适合实时或资源受限的应用场景。
• 在不同段落长度的多个问答基准测试中，BlockBERT 的性能与 RoBERTa 相当或更优。
• 块稀疏注意力机制成功实现了注意力头对局部和长距离依赖关系的有效建模。
• 该方法简单且与标准深度学习框架兼容，无需使用自定义内核或复杂的稀疏性管理。