深入理解 BigBird 的块稀疏注意力

引言

基于 transformer 的模型已被证明对很多 NLP 任务都非常有用。然而， 的时间和内存复杂度 (其中 是序列长度) 使得在长序列 () 上应用它们变得非常昂贵，因而大大限制了其应用。最近的几篇论文，如 Longformer 、Performer 、Reformer 、簇状注意力 都试图通过对完整注意力矩阵进行近似来解决这个问题。如果你不熟悉这些模型，可以查看 🤗 之前的 博文。

BigBird (由 该论文 引入) 是解决这个问题的最新模型之一。 BigBird 依赖于 块稀疏注意力 而不是普通注意力 ( 即 BERT 的注意力)，与 BERT 相比，这一新算法能以低得多的计算成本处理长达 4096 的序列。在涉及很长序列的各种任务上，该模型都实现了 SOTA，例如长文档摘要、长上下文问答。

RoBERTa 架构的 BigBird 模型现已集成入 🤗 transformers 中。本文的目的是让读者 深入 了解 BigBird 的实现，并让读者能在 🤗 transformers 中轻松使用 BigBird。但是，在更深入之前，一定记住 BigBird 注意力只是 BERT 完全注意力的一个近似，因此我们并不纠结于让它比 BERT 完全注意力 更好，而是致力于让它更有效率。有了它，transformer 模型就可以作用于更长的序列，因为 BERT 的二次方内存需求很快会变得难以为继。简而言之，如果我们有 计算和 时间，那么用 BERT 注意力就好了，完全没必要用本文讨论的块稀疏注意力。

如果你想知道为什么在处理较长序列时需要更多计算，那么本文正合你意！

---

在使用标准的 BERT 类注意力时可能会遇到以下几个主要问题:

• 每个词元真的都必须关注所有其他词元吗？

• 为什么不只计算重要词元的注意力？

• 如何决定哪些词元重要？

• 如何以高效的方式处理少量词元？

---

本文，我们将尝试回答这些问题。

应该关注哪些词元？

下面，我们将以句子 BigBird is now available in HuggingFace for extractive Question Answering 为例来说明注意力是如何工作的。在 BERT 这类的注意力机制中，每个词元都简单粗暴地关注所有其他词元。从数学上来讲，这意味着每个查询的词元 , 将关注每个键词元 。

我们考虑一下 每个查询词元应如何明智地选择它实际上应该关注的键词元 这个问题，下面我们通过编写伪代码的方式来整理思考过程。

假设 available 是当前查询词元，我们来构建一个合理的、需要关注的键词元列表。

# 以下面的句子为例
example = ['BigBird', 'is', 'now', 'available', 'in', 'HuggingFace', 'for', 'extractive', 'question', 'answering']

# 假设当前需要计算 'available' 这个词的表征
query_token = 'available'

# 初始化一个空集合，用于放 'available' 这个词的键词元
key_tokens = [] # => 目前，'available' 词元不关注任何词元

邻近词元当然很重要，因为在一个句子 (单词序列) 中，当前词高度依赖于前后的邻近词。滑动注意力 即基于该直觉。

# 考虑滑动窗大小为 3, 即将 'available' 的左边一个词和右边一个词纳入考量
# 左词: 'now'; 右词: 'in'
sliding_tokens = ["now", "available", "in"]

# 用以上词元更新集合
key_tokens.append(sliding_tokens)

长程依赖关系: 对某些任务而言，捕获词元间的长程关系至关重要。 例如 ，在问答类任务中，模型需要将上下文的每个词元与整个问题进行比较，以便能够找出上下文的哪一部分对正确答案有用。如果大多数上下文词元仅关注其他上下文词元，而不关注问题，那么模型从不太重要的上下文词元中过滤重要的上下文词元就会变得更加困难。

BigBird 提出了两种允许长程注意力依赖的方法，这两种方法都能保证计算效率。

• 全局词元: 引入一些词元，这些词元将关注每个词元并且被每个词元关注。例如，对 “HuggingFace is building nice libraries for easy NLP” ，现在假设 'building' 被定义为全局词元，而对某些任务而言，模型需要知道 'NLP' 和 'HuggingFace' 之间的关系 (注意: 这 2 个词元位于句子的两端); 现在让 'building' 在全局范围内关注所有其他词元，会对模型将 'NLP' 与 'HuggingFace' 关联起来有帮助。

# 我们假设第一个和最后一个词元是全局的，则有:
global_tokens = ["BigBird", "answering"]

# 将全局词元加入到集合中
key_tokens.append(global_tokens)

• 随机词元: 随机选择一些词元，这些词元将通过关注其他词元来传输信息，而那些词元又可以传输信息到其他词元。这可以降低直接从一个词元到另一个词元的信息传输成本。

# 现在，我们可以从句子中随机选择 `r` 个词元。这里，假设 `r` 为 1， 选择了 `is` 这个词元
>>> random_tokens = ["is"] # 注意: 这个是完全随机选择的，因此可以是任意词元。

# 将随机词元加入到集合中
key_tokens.append(random_tokens)

# 现在看下 `key_tokens` 集合中有哪些词元
key_tokens
{'now', 'is', 'in', 'answering', 'availa