Big Bird:适用于更长序列的Transformer模型
摘要
基于Transformer的模型,如BERT,已成为自然语言处理(NLP)中最成功的深度学习模型之一。然而,它们的一个核心限制是由于其全注意力机制,对序列长度的二次依赖(主要是在内存方面)。为了解决这个问题,我们提出了BIGBIRD,一种稀疏注意力机制,将这种二次依赖减少到线性。我们证明BIGBIRD是序列函数的通用逼近器,并且是图灵完备的,从而保留了二次全注意力模型的这些特性。在此过程中,我们的理论分析揭示了拥有O(1)全局token(如CLS)的一些好处,这些token作为稀疏注意力机制的一部分,关注整个序列。所提出的稀疏注意力可以处理长度达到之前使用类似硬件可能处理的8倍的序列。由于能够处理更长的上下文,BIGBIRD在各种NLP任务(如问答和摘要)上大幅提高了性能。我们还提出了对基因组数据的新应用。
1. 引言
基于Transformer的模型[91],如BERT[22, 63],在各种自然语言处理(NLP)任务中取得了巨大成功,因此成为现代NLP研究的主流。它们的多功能性和鲁棒性是Transformer广泛采用的主要原因。该模型很容易适应各种基于序列的任务——作为翻译[91]、摘要[66]、生成[15]等的序列到序列模型,或作为情感分析[83]、词性标注[65]、机器阅读理解[93]等的独立编码器——并且已知其性能远超之前的序列模型,如LSTM[37]。Transformer的关键创新是引入了自注意力机制,该机制可以并行评估输入序列的每个token,消除了像LSTM这样的循环神经网络中的顺序依赖。这种并行性使Transformer能够充分利用现代SIMD硬件加速器(如GPU/TPU)的全部能力,从而促进了在空前规模的数据集上训练NLP模型。这种在大规模数据上训练的能力导致了像BERT[22]和T5[75]这样的模型的出现,它们在大型通用语料库上预训练Transformer,并将知识转移到下游任务。预训练在低数据量的下游任务[51]以及有足够数据的任务[101]上带来了显著改进,因此成为Transformer在当代NLP中无处不在的主要推动力。
自注意力机制通过允许输入序列中的每个token独立关注序列中的每个其他token,克服了RNN的限制(即RNN的顺序性)。这一设计选择有几个有趣的影响。特别是,全自注意力的计算和内存需求是序列长度的二次方。我们注意到,虽然语料库可以很大,但在许多应用中提供上下文的序列长度非常有限。使用当前常见的硬件和模型大小,这一需求大致转化为能够处理长度为512个token的输入序列。这减少了其在需要更大上下文的任务(如问答[60]、文档分类等)中的直接适用性。然而,虽然我们知道自注意力和Transformer是有用的,但我们的理论理解还很初级。自注意力模型的哪些方面对其性能是必要的?我们能对Transformer和类似模型的表达能力说些什么?从设计上看,甚至不清楚所提出的自注意力机制是否像RNN一样有效。例如,自注意力甚至不遵守序列顺序,因为它是排列等变的。这一担忧已部分解决,因为Yun等人[104]表明,Transformer足够表达,可以捕获具有紧凑域的所有连续序列到序列函数。同时,Pérez等人[72]表明,完整的Transformer是图灵完备的(即可以模拟完整的图灵机)。两个自然问题出现了:我们能否使用更少的内积实现完全二次自注意力方案的经验优势?这些稀疏注意力机制是否保留了原始网络的表达能力和灵活性?
在本文中,我们解决了上述两个问题,并提出了一种稀疏注意力机制,该机制在需要长上下文的多项任务上提高了性能。我们系统地开发了BIGBIRD,一种注意力机制,其复杂度在token数量上是线性的(第2节)。我们从图稀疏化方法中汲取灵感,并理解当全注意力放松以形成所提出的注意力模式时,Transformer表达能力证明的崩溃点。这一理解帮助我们开发了BIGBIRD,它在理论上同样具有表达力,并且在实际应用中也很有用。特别是,我们的BIGBIRD由三个主要部分组成:
• 一组g个全局token,关注序列的所有部分。
• 所有token关注一组w个局部邻近token。
• 所有token关注一组r个随机token。
这导致了一种高性能的注意力机制,可以扩展到更长的序列长度(8倍)。总结一下,我们的主要贡献是:
-
BIGBIRD满足全Transformer的所有已知理论特性(第3节)。特别是,我们展示了添加额外token可以仅用O(n)内积表达所有连续序列到序列函数。此外,我们展示了在关于精度的标准假设下,BIGBIRD是图灵完备的。
-
我们通过实验证明,BIGBIRD建模的扩展上下文有益于各种NLP任务。我们在多个不同数据集上实现了问答和文档摘要的最先进结果。这些结果的总结见第4节。
-
最后,我们介绍了一种新的基于注意力的模型应用,其中长上下文有益:提取基因组序列(如DNA)的上下文表示。通过更长的掩码LM预训练,BIGBIRD在下游任务(如启动子区域和染色质谱预测)上提高了性能(第5节)。
1.1 相关工作
有许多有趣的尝试旨在缓解Transformer的二次依赖,这些尝试可以大致分为两个方向。第一条工作路线接受长度限制并围绕其开发方法。这一类中最简单的方法仅采用滑动窗口[93],但一般来说,大多数工作都符合以下一般范式:使用其他机制选择相关上下文的较小子集输入Transformer,并可选地迭代,即每次使用不同的上下文多次调用Transformer块。最显著的是,SpanBERT[42]、ORQA[54]、REALM[34]、RAG[57]在不同任务中取得了强劲的性能。然而,值得注意的是,这些方法通常需要大量的工程努力(如通过大规模最近邻搜索进行反向传播)并且难以训练。第二条工作路线质疑全注意力是否必要,并尝试提出不需要全注意力的方法,从而减少内存和计算需求。显著的是,Dai等人[21]、Sukhbaatar等人[82]、Rae等人[74]提出了自回归模型,这些模型在从左到右的语言建模中表现良好,但在需要双向上下文的任务中表现不佳。Child等人[16]提出了一个稀疏模型,将复杂度降低到O( n n n\sqrt{n} nn),Kitaev等人[49]通过使用LSH计算最近邻进一步将复杂度降低到O( n log ( n ) n\log(n) nlog(n))。
Ye等人[103]提出了数据的二进制分区,而Qiu等人[73]通过使用块稀疏性降低了复杂度。最近,Longformer[8]引入了一种基于局部滑动窗口的掩码,带有少量全局掩码以减少计算,并将BERT扩展到基于更长序列的任务。最后,我们的工作与扩展Transformer构造[4]的工作密切相关并建立在其基础上。该工作旨在为Transformer编码文本结构。他们广泛使用了全局token的概念来实现其目标。我们的理论工作也可以看作是为这些模型的成功提供了理论依据。需要注意的是,大多数上述方法都是基于启发式的,并且在经验上不如原始Transformer那样多功能和鲁棒,即相同的架构在多个标准基准测试中未能达到最先进的水平。(Longformer是一个例外,我们将其包含在所有比较中,详见附录E.3以获取更详细的比较)。此外,这些近似方法没有理论保证。
2 BIGBIRD架构
在本节中,我们使用广义注意力机制来描述BIGBIRD模型,该机制用于Transformer的每一层,处理输入序列 X = ( x 1 , . . . , x n ) ∈ R n × d X=(x_1,...,x_n)\in\mathbb{R}^{n\times d} X=(x1,...,xn)∈Rn×d。广义注意力机制通过一个有向图 D D D来描述,其顶点集为 [ n ] = { 1 , … , n } [n]=\{1,\ldots,n\} [n]={1,…,n}。弧集(有向边)表示注意力机制将考虑的内积集合。设 N ( i ) N(i) N(i)表示节点 i i i在 D D D中的出邻居集,则广义注意力机制的第 i i i个输出向量定义为:
A T T N D ( X ) i = x i + ∑ h = 1 H σ ( Q h ( x i ) K h ( X N ( i ) ) T ) ⋅ V h ( X N ( i ) ) \begin{aligned}\mathrm{ATTN}_{D}(X)_{i}=x_{i}+\sum_{h=1}^{H}\sigma\left(Q_{h}(x_{i})K_{h}(X_{N(i)})^{T}\right)\cdot V_{h}(X_{N(i)})\end{aligned} ATTND(X)i=xi+h=1∑Hσ(Qh(xi)Kh(XN(i))T)⋅Vh(XN(i))
其中 Q h , K h : R d → R m Q_h, K_h: \mathbb{R} ^d\to \mathbb{R} ^m Qh,Kh:Rd→Rm分别是查询和键函数, V h : R d → R d V_h: \mathbb{R} ^d\to \mathbb{R} ^d Vh:Rd→Rd是值函数, σ \sigma σ是评分函数(例如softmax或hardmax), H H H表示头数。还要注意 X N ( i ) X_N(i) XN(i)对应于仅由 { x j : j ∈ N ( i ) } \{\boldsymbol{x}_j:j\in N(i)\} {xj:j∈N(i)}堆叠而成的矩阵,而不是所有输入。如果 D D D是完全有向图,我们恢复了Vaswani等人[91]的完全二次注意力机制。为了简化我们的描述,我们将操作图 D D D的邻接矩阵 A A A,即使底层图可能是稀疏的。详细来说, A ˇ ∈ [ 0 , 1 ˇ ] n × n \check{A}\in[0,\check{1}]^{n\times n} Aˇ∈[0,1ˇ]n×n,如果查询 i i i关注键 j j j,则 A ( i , j ) = 1 A(i,j)=1 A(i,j)=1,否则为零。例如,当 A A A是全1矩阵(如BERT中)时,它会导致二次复杂度,因为所有token都关注每个其他token。这种将自注意力视为完全连接图的观点使我们能够利用现有的图论来帮助降低其复杂度。降低自注意力二次复杂度的问题现在可以看作是一个图稀疏化问题。众所周知,随机图是扩展图,并且可以在许多不同上下文中近似完全图,包括其谱特性[80,38]。我们相信用于注意力机制的稀疏随机图应具有两个需求:节点之间的短平均路径长度和局部性概念,我们将在下面讨论这些内容。让我们考虑最简单的随机图构造,称为Erdős-Rényi模型,其中每条边以固定概率独立选择。在这样一个只有 Θ ~ ( n ) \tilde{\Theta}(n) Θ~(n)边的随机图中,任何两个节点之间的最短路径是节点数的对数[17,43]。因此,这样的随机图在谱上近似完全图,其邻接矩阵的第二特征值与第一特征值相距甚远[9,10.6]。这一特性导致图中随机游走的快速混合时间,非正式地表明信息可以在任何一对节点之间快速流动。因此,我们提出了一种稀疏注意力机制,其中每个查询关注 r r r个随机数量的键,即 A ( i , ⋅ ) = 1 A(i,\cdot)=1 A(i,⋅)=1对于 r r r个随机选择的键(见图[1a])。
第二个启发BIGBIRD创建的观点是,NLP和计算生物学中的大多数上下文数据都表现出很强的局部性。在这种现象中,关于某个token的大量信息可以从其邻近的token中推导出来。最重要的是,Clark等人[19]研究了NLP任务中的自注意力模型,并得出结论:邻近的内积极为重要。局部性(即语言结构中token的邻近性)的概念也构成了各种语言学理论(如转换生成语法)的基础。在图论术语中,聚类系数是连接局部性的度量,当图中包含许多团或近团(几乎完全互连的子图)时,聚类系数较高。简单的Erdős-Rényi随机图不具有高聚类系数[84],但一类称为小世界图的随机图表现出高聚类系数[94]。Watts和Strogatz[94]提出的一个特定模型对我们非常相关,因为它在平均最短路径和局部性概念之间实现了良好的平衡。其模型的生成过程如下:构建一个规则环格,这是一个具有 n n n个节点的图,每个节点连接到 w w w个邻居,每侧各 w / 2 w/2 w/2个。
换句话说,我们从一个节点上的滑动窗口开始。然后,将所有连接中的随机子集(k%)替换为随机连接。其他(100 - k)%的局部连接被保留。然而,在现代硬件上删除这样的随机边可能效率低下,因此我们保留它,这不会影响其属性。总之,为了捕捉上下文中的这些局部结构,在BIGBIRD中,我们定义了一个滑动窗口注意力,使得在宽度为 w w w的自注意力中,位置 i i i的查询从 i − w / 2 i-w/2 i−w/2到 i + w / 2 i+w/2 i+w/2的键进行关注。在我们的符号中, A ( i , i − w / 2 : i + w / 2 ) = 1 A(i,i-w/2:i+w/2)=1 A(i,i−w/2:i+w/2)=1(见图[1]b)。作为初步的合理性检查,我们进行了基本实验,以测试这些直觉是否足以获得接近BERT模型的性能,同时保持注意力在token数量上是线性的。我们发现,随机块和局部窗口不足以捕捉与BERT性能竞争所需的所有上下文。BIGBIRD的最后一部分受到我们理论分析(第B节)的启发,这对实证性能至关重要。更具体地说,我们的理论利用了“全局token”(关注序列中所有token并被所有token关注的token)的重要性(见图|To)。这些全局token可以通过两种方式定义:
-
BIGBIRD-ITC:在内部Transformer构造(ITC)中,我们使一些现有token“全局化”,这些token关注整个序列。具体来说,我们选择一个索引子集 G G G( g : = ∣ G ∣ g:=|G| g:=∣G∣),使得对于所有 i ∈ G i\in G i∈G, A ( i , : ) = 1 A(i,:)=1 A(i,:)=1且 A ( : , i ) = 1 A(:,i)=1 A(:,i)=1。
-
BIGBIRD-ETC:在扩展Transformer构造(ETC)中,我们包含额外的“全局”token,如CLS。具体来说,我们添加 g g g个全局token,这些token关注所有现有token。在我们的符号中,这对应于通过向矩阵 A A A添加 g g g行来创建一个新矩阵 B ∈ [ 0 , 1 ] ( N + g ) × ( N + g ) B\in[0,1]^{(N+g)\times(N+g)} B∈[0,1](N+g)×(N+g),使得对于所有 i ∈ { 1 , 2 , … g } i\in\{1,2,\ldots g\} i∈{1,2,…g}, B ( i , : ) = 1 B(i,:)=1 B(i,:)=1且 B ( : , i ) = 1 B(:,i)=1 B(:,i)=1,并且 B ( g + i , g + j ) = A ( i , j ) ∀ i , j ∈ { 1 , … , N } B(g+i,g+j)=A(i,j)\forall i,j\in\{1,\ldots,N\} B(g+i,g+j)=A(i,j)∀i,j∈{1,…,N}。这增加了额外的位置来存储上下文,正如我们将在实验中看到的,这提高了性能。
BIGBIRD的最终注意力机制(图[1d])具有以下三个特性:
- 查询关注 r r r个随机键:每个查询会随机关注 r r r个键。
- 局部滑动窗口:每个查询会关注其位置左侧的 w / 2 w/2 w/2个token和右侧的 w / 2 w/2 w/2个token。
- 全局token:包含 g g g个全局token(这些全局token可以是现有token,也可以是额外添加的token)。
我们在附录[D]中提供了具体的实现细节。
3 关于稀疏注意力机制的理论结果
在本节中,我们将从两个方面证明稀疏注意力机制与全注意力机制具有相同的强大能力和表达能力。首先,我们证明当稀疏注意力机制用于独立编码器(如BERT)时,它们是序列到序列函数的通用逼近器,类似于Yun等人[104]的风格。我们注意到,这一特性在Yun等人[105]的同期工作中也得到了理论上的探讨。其次,与[105]不同,我们进一步证明稀疏编码器-解码器Transformer是图灵完备的(假设与[72]中定义的相同条件)。作为上述积极结果的补充,我们还表明,转向稀疏注意力机制会带来一定的代价,即没有免费的午餐。在第3.4节中,我们通过展示一个自然任务来证明下界,其中任何足够稀疏的机制都需要多项式级别的更多层。3.1 符号
完整的Transformer编码器堆栈不过是对单层编码器(具有独立参数)的重复应用。我们表示使用广义编码器(第[2]节)定义的此类Transformer编码器堆栈的类别为 T D H , m , q \mathcal{T}_D^{H,m,q} TDH,m,q,其中包含 H H H个头,头大小为 m m m, q q q是输出网络的隐藏层大小,注意力层由有向图 D D D定义。我们提出的注意力机制与Vaswani等人[91]和Yun等人[104]的关键区别在于,我们在每个序列的开头添加了一个特殊token,并为其分配了一个特殊向量。我们将此称为 x 0 x_0 x0。因此,我们的图 D D D将具有顶点集 { 0 } ∪ [ n ] = { 0 , 1 , 2 , … , n } \{0\}\cup[n]=\{0,1,2,\ldots,n\} {0}∪[n]={0,1,2,…,n}。我们假设这个额外的节点及其相应的向量将在Transformer的最终输出层被丢弃。为了避免繁琐的符号,我们仍然将Transformer视为将序列 X ∈ R n × d X\in\mathbb{R}^{n\times d} X∈Rn×d映射到 R n × d \mathbb{R}^n\times d Rn×d。我们还允许Transformer在输入层将位置嵌入 E ∈ R d × n E\in\mathbb{R}^d\times n E∈Rd×n附加到矩阵 X X X上。
最后,我们需要定义函数类和距离度量,以证明通用逼近性质。设 F C D \mathcal{F}_CD FCD表示连续函数 f : [ 0 , 1 ] n × d → R n × d f:[0,1]^{n\times d}\to\mathbb{R}^{n\times d} f:[0,1]n×d→Rn×d的集合,这些函数在由 ℓ p \ell_p ℓp范数定义的拓扑下是连续的。回想一下,对于任何 p ≥ 1 p\geq1 p≥1, ℓ p \ell_p ℓp距离为 d p ( f 1 , f 2 ) = d_p(f_1,f_2)= dp(f1,f2)= ( ∫ ∥ f 1 ( X ) − f 2 ( X ) ∥ p p d X ) 1 / p . \left(\int\|f_{1}(X)-f_{2}(X)\|_{p}^{p}dX\right)^{1/p}. (∫∥f1(X)−f2(X)∥ppdX)1/p.
3.2 通用逼近器
定义1。以0为中心的星图 S S S是定义在 { 0 , … , n } \{0, \ldots, n\} {0,…,n}上的图。所有顶点 i i i的邻居为 N ( i ) = { 0 , i } N(i) = \{0, i\} N(i)={0,i},其中 i ∈ { 1 … n } i\in \{1\ldots n\} i∈{1…n},且 N ^ ( 0 ) = { 1 , … n } \hat{N}(0) = \{1, \ldots n\} N^(0)={1,…n}。
我们的主要定理是,由任何包含 S S S的图定义的稀疏注意力机制是一个通用逼近器:
定理1。给定 1 < p < ∞ 1<p<\infty 1<p<∞和 ϵ > 0 \epsilon>0 ϵ>0,对于任何 f ∈ F C D f\in\mathcal{F}_{CD} f∈FCD,存在一个具有稀疏注意力的Transformer, g ∈ T D H , m , q g\in \mathcal{T}_{D}^{H, m, q} g∈TDH,m,q,使得 d p ( f , g ) ≤ ϵ d_{p}(f, g) \leq \epsilon dp(f,g)≤ϵ,其中 D D D是任何包含星图 S S S的图。为了证明该定理,我们将遵循[104]中概述的标准证明结构。
步骤1:用分段常数函数逼近 F C D \mathcal{F}_{CD} FCD。由于 f f f是一个定义在有界域 [ 0 , 1 ) n × d [0,1)^{n\times d} [0,1)n×d上的连续函数,我们将用一个合适的分段常数函数来逼近它。这是通过将区域 [ 0 , 1 ) [0,1) [0,1)划分为粒度 δ \delta δ的网格来实现的,从而得到一个离散集 G δ G_{\delta} Gδ。因此,我们可以假设我们正在处理一个函数 f ˉ : G δ → R n × d \bar{f}:\mathbb{G}_{\delta}\to\mathbb{R}^{n\times d} fˉ:Gδ→Rn×d,其中 d p ( f , f ‾ ) ≤ ϵ 3 d_p(f,\overline{f})\leq\frac\epsilon3 dp(f,f)≤3ϵ。
步骤2:用改进的Transformer逼近分段常数函数。这是证明的关键步骤,其中自注意力机制用于生成输入的上下文映射。非正式地说,上下文映射是由矩阵 ( X , x i ) (X, x_i) (X,xi)和一列组成的对的唯一编码。其唯一性使得前馈层可以使用每个编码将其映射到唯一的输出列。
主要的技术挑战是仅使用稀疏注意力机制计算上下文映射。这在[04]中使用“选择性”移位运算符完成,该运算符将特定区间内的条目上移。他们证明的关键在于,移位范围恰好是从最大条目到最小条目的范围。
使用稀疏注意力机制创建上下文映射是一个相当大的挑战。特别是,由于每个查询仅关注少数键,因此尚不清楚是否可以收集足够的信息来生成整个矩阵的上下文嵌入。为了解决这个问题,我们开发了一个稀疏移位运算符,如果矩阵的条目位于某个范围内,则对其进行移位。移位的具体量由有向稀疏注意力图 D D D控制。第二个关键要素是使用额外的全局token。通过将运算符仔细应用于一组选定的范围,我们将展示每列将包含完整映射的唯一映射。因此,我们可以通过使用多层和辅助全局token来弥补自注意力机制中内积的损失。
步骤3:用原始Transformer逼近改进的Transformer**:最后一步是通过使用ReLU和softmax的原始Transformer来逼近改进的Transformer。我们在附录A中提供了完整的细节。
3.3 图灵完备性
Transformer是一个非常通用的类别。在Vaswani等人[91]的原始论文中,它们被同时用于编码器和解码器。虽然上一节概述了仅编码器的强大能力,但另一个自然的问题是,编码器和解码器的结合会带来什么额外的能力?Pérez等人[72]证明了基于二次注意力机制的完整Transformer是图灵完备的。这一结果做了一个不现实的假设,即模型在任意精度模型上工作。当然,这是必要的,否则Transformer是有限状态机,不可能是图灵完备的。
很自然地会问,全注意力机制是否是必要的。或者稀疏注意力机制是否也可以用来模拟任何图灵机?我们证明情况确实如此:我们可以使用稀疏编码器和稀疏解码器来模拟任何图灵机。
为了在Transformer架构中使用稀疏注意力机制,我们需要定义一个合适的修改,其中每个token仅对前面的token作出反应。与BERT的情况不同,BERT中整个注意力机制应用一次,而在完整的Transformer中,解码器侧的稀疏注意力机制是逐个token使用的。其次,Pérez等人[72]的工作使用每个token作为磁带历史的表示,并使用全注意力来移动和检索正确的磁带符号。Pérez等人[72]的大部分构造适用于稀疏注意力,除了它们用于指向历史中的寻址方案([72]中的引理B.4)。我们展示了如何使用稀疏注意力机制来模拟这一点,并将细节推迟到附录B中。
3.4 局限性
我们展示了一个自然任务,该任务可以通过全注意力机制在 O ( 1 ) O(1) O(1)层内解决。然而,在标准复杂性理论假设下,对于任何具有 O ~ ( n ) \tilde{O}(n) O~(n)条边的稀疏注意力层(不仅仅是BIGBIRD),这个问题需要 Ω ~ ( n ) \tilde{\Omega}(n) Ω~(n)层来解决。(这里 O ~ \tilde{O} O~隐藏了多对数因子)。考虑一个简单的问题:在给定的长度为 n n n的序列中,为每个向量找到对应的最远向量。形式化地定义如下:
任务1。给定 n n n个单位向量 { u 1 , … , u n } \{u_1,\ldots,u_n\} {u1,…,un},找到 f ( u 1 , … , u n ) → ( u 1 • , … , u n • ) f(u_1,\ldots,u_n)\to(u_1•,\ldots,u_n•) f(u1,…,un)→(u1•,…,un•),其中对于固定的 j ∈ [ n ] j\in[n] j∈[n],我们定义 j ∗ = arg max k ∥ u k − u j ∥ 2 2 j^*=\arg\max_k\|u_k-u_j\|_2^2 j∗=argmaxk∥uk−uj∥22。
对于单位向量来说,找到最远的向量可以归结为最小化内积搜索。对于具有适当查询和键的全注意力机制,这个任务非常简单,因为我们可以评估所有成对的内积。
稀疏注意力机制的不可能性源于正交向量猜想(OVC)[II]ZZZ 96]的硬度结果。OVC是细粒度复杂性中广泛使用的假设。非正式地说,它表明无法在次二次时间内确定 n n n个布尔向量的最小内积是否为0。在附录[C]中,我们展示了使用 OVC ^ \hat{\text{OVC}} OVC^的归约,证明如果对于任何稀疏有向图 D D D,注意力Transformer g ∈ T D H = 1 ^ , m = 2 d , q = 0 g\in\mathcal{T}_{D}^{H=\hat{1},m=2d,q=0} g∈TDH=1^,m=2d,q=0可以评估任务1,那么它可以解决正交向量问题。
命题1。存在一个单层全自注意力 g ∈ T H = 1 , m = 2 d , q = 0 g\in\mathcal{T}^{H=1,m=2d,q=0} g∈TH=1,m=2d,q=0,可以评估任务1,即 g ( u 1 , . . . , u n ) = [ u 1 • , . . . , u n • ] g(u_{1}, . . . , u_{n}) = [u_{1}• , . . . , u_{n}•] g(u1,...,un)=[u1•,...,un•],但对于任何具有 O ~ ( n ) \tilde{O}(n) O~(n)条边(即内积评估)的稀疏注意力图 D D D,将需要 Ω ~ ( n 1 − o ( 1 ) ) \tilde{\Omega}(n^{1-o(1)}) Ω~(n1−o(1))层。
我们在附录C中给出了这一事实的正式证明。
4 实验:自然语言处理
在本节中,我们的目标是展示为NLP任务建模更长输入序列的好处,为此我们选择了三个代表性任务。我们从基本的掩码语言建模(MLM;Devlin等人,22)开始,以检查是否可以通过利用更长的连续序列学习更好的上下文表示。接下来,我们考虑带有支持证据的问答(QA),处理更长序列的能力将使我们能够使用像TF-IDF/BM25这样的简单系统检索更多证据。
最后,我们处理长文档分类任务,其中区分信息可能不在前512个token中。以下我们总结了使用序列长度为4096的BIGBIRD的结果,而将所有其他设置细节(包括计算资源、批量大小、步长等)推迟到附录E中。
预训练和掩码语言建模(MLM)
我们遵循[22, 63]创建BIGBIRD的基础版和大型版,并使用MLM目标进行预训练。该任务涉及预测被掩码的随机子集token。我们使用四个标准数据集进行预训练(列在附录E.1,表9中),并从公开的RoBERTa检查点2进行热启动。我们按照[8]的方法,以每字符比特数(bits per character)来比较预测掩码token的性能。如附录E.1表10所示,BIGBIRD和Longformer的表现均优于有限长度的RoBERTa,其中BIGBIRD-ETC表现最佳。我们注意到,我们的模型是在16GB内存/芯片上训练的,批量大小为32-64。我们的内存效率归功于第2节中描述的稀疏注意力机制的高效分块和稀疏结构。
问答(QA)
我们考虑了以下四个具有挑战性的数据集:
- Natural Questions[52]:对于给定的问题,从给定的证据中找到简短的答案(SA),并突出显示包含正确答案信息的段落(LA)。
- HotpotQA-distractor[100]:类似于Natural Questions,它需要从给定的证据中找到答案(Ans)以及支持多跳推理的支持事实(Sup)。
- TriviaQA-wiki[41]:我们需要使用提供的Wikipedia证据为给定问题提供答案,但答案可能不在给定的证据中。在一个较小的已验证问题子集中,给定的证据保证包含答案。尽管如此,我们在此情况下也将答案建模为跨度选择问题。
- WikiHop[95]:通过聚合分布在多个文档中的信息,从多项选择题(MCQ)中选择正确答案。
由于这些任务竞争激烈,已经为每个数据集设计了多个高度工程化的系统,以符合各自的输出格式。为了公平比较,我们不得不使用一些额外的正则化来训练BIGBIRD,详细信息与架构描述一起在附录E.2中提供。我们使用基础大小的模型进行实验,并为每个数据集选择开发集上的最佳配置(如表2所示)。我们可以看到,BIGBIRD-ETC通过扩展的全局token始终优于所有其他模型。因此,我们选择此配置来训练大型模型,用于在隐藏测试集上进行评估。在表3中,我们将BIGBIRD-ETC模型与排行榜上除BIGBIRD外的前三名进行比较。可以清楚地看到,使用更长上下文的重要性,因为Longformer和BIGBIRD的表现均优于使用较短上下文的模型。此外,值得注意的是,BIGBIRD的提交是单一模型,而Natural Questions的其他前三名是集成模型,这可能解释了在精确答案短语选择上略低的准确性。
分类
我们在不同长度和内容的数据集上进行了实验,特别是各种文档分类和GLUE任务。遵循BERT,我们在第一个[CLS]token上使用了一层交叉熵损失。我们发现,当文档较长且训练样本较少时,使用BIGBIRD的收益更为显著。例如,使用基础大小的模型,BIGBIRD将Arxiv数据集的最新技术水平提高了约5个百分点。在Patents数据集上,与使用简单的BERT/RoBERTa相比有所改进,但由于训练数据规模较大,与最先进技术(非基于BERT)的改进并不显著。需要注意的是,对于较小的IMDb数据集,这种性能提升并不明显。我们在附录E.4中提供了详细的实验结果和实验设置细节,展示了具有竞争力的性能。
4.1 编码器-解码器任务
在编码器-解码器设置中,可以很容易地看出,由于全自注意力的存在,两者都存在二次复杂度的问题。我们专注于仅在编码器侧引入BIGBIRD的稀疏注意力机制。这是因为在实际的生成应用中,输出序列的长度通常比输入序列小得多。例如,在文本摘要任务中,我们观察到在实际场景中(参见附录E.5表18),输出序列的中位长度约为200,而输入序列的中位长度超过3000。对于此类应用,在编码器中使用稀疏注意力机制,而在解码器中使用全自注意力机制更为高效。
摘要生成
文档摘要生成是一项为文本文档创建简短而准确的摘要的任务。我们使用了三个长文档数据集来测试我们的模型,详细信息见表18。在本文中,我们专注于长文档的抽象摘要生成,使用更长的上下文编码器应该会提高性能。原因有两个:首先,重要内容可能均匀分布在长文档中,而不仅仅集中在前512个token中,这是BigPatents数据集[78]的设计特点。其次,长文档表现出更丰富的语篇结构,摘要也更加抽象,因此观察更多上下文有助于生成更好的摘要。
正如最近指出的[76, 107],预训练对生成任务有帮助,我们从基础大小模型的通用MLM预训练开始,并在大型模型上利用Pegasus[107]的最先进摘要生成特定预训练。在这些长文档数据集上训练BIGBIRD稀疏编码器和完整解码器的结果见表4。我们可以清楚地看到,建模更长的上下文带来了显著的改进。除了超参数外,我们还在附录E.5中展示了在较短但更广泛的数据集上的结果,这些结果表明使用稀疏注意力机制也不会影响性能。
5 实验:基因组学
近年来,使用深度学习处理基因组数据的趋势迅速增长[86, 106, 13],这已经在多个具有生物学意义任务中提高了性能,例如启动子位点预测[71]、甲基化分析[55]、预测非编码变体的功能效应[109]等。这些方法以DNA序列片段作为输入,因此我们相信BIGBIRD处理更长输入序列的能力将非常有益,因为DNA中的许多功能效应具有高度非局部性[12]。此外,受NLP的启发,我们通过MLM预训练利用丰富的未标记数据(例如人类参考基因组、酵母基因组数据库)学习DNA片段的强大上下文表示。接下来,我们展示了我们的长输入BIGBIRD以及所提出的预训练在两个下游任务中显著提高了性能。这两个任务的详细实验设置在附录F中提供。
预训练和掩码语言建模(MLM)
正如Liang [58]所探讨的,我们建议首先将DNA分割为token,而不是直接操作碱基对,以进一步增加上下文长度(附录F,图7)。具体来说,我们为DNA序列构建了一个大小为32K的字节对编码(BPE)表[50],每个token平均代表8.78个碱基对。我们使用MLM目标在人类参考基因组(GRCh37)3上学习这些token的上下文表示。然后,我们在一个保留集上报告每字符比特数(BPC),如表5所示。我们发现,基于注意力的DNA上下文表示确实改善了BPC,而使用更长的上下文进一步提高了这一性能。
启动子区域预测
启动子是通常位于基因上游的DNA区域,是转录起始的位点。由于理解基因调控的重要第一步是识别启动子区域,因此已经提出了多种方法来识别给定DNA序列中的启动子区域[99, 59, 11, 98, 71]。相应的机器学习任务是将给定的DNA片段分类为启动子或非启动子序列。我们使用了由Oubounyt等人[71]编译的数据集,该数据集基于真核启动子数据库(EPDnew)[24]4构建。我们使用训练数据对上述预训练的BIGBIRD模型进行微调,并在测试数据集上报告F1分数。我们将我们的结果与之前报道的最佳方法进行比较,如表6所示。我们发现,BIGBIRD几乎达到了完美的准确率,比之前报道的最佳准确率提高了5%。
染色质谱预测
DNA的非编码区域不编码蛋白质。大多数疾病和其他与性状相关的单核苷酸多态性(SNP)与非编码基因组变异相关[109, 46]。因此,理解DNA非编码区域的功能效应是一项非常重要的任务。正如Zhou和Troyanskaya[109]所定义的,这一过程中的一个重要步骤是从非编码基因组序列中预测大规模的染色质谱。为此,DeepSea[109]从DNA元素百科全书(ENCODE)5和Roadmap表观基因组学项目6中编译了240万个非编码变体的919个染色质谱。相应的机器学习任务是预测给定DNA非编码区域的这919个染色质谱,包括160种不同转录因子(TF)的690个TF结合谱、125个DNase I敏感性(DHS)谱和104个组蛋白标记(HM)谱。我们联合训练了919个二分类器,从DNA片段序列中预测这些功能效应。在保留的染色体上,我们将AUC与基线进行比较,如表7所示,并发现我们在更困难的任务HM上显著提高了性能,已知HM比其他任务具有更长的相关性[27]。
6 结论
我们提出了BIGBIRD:一种在token数量上呈线性复杂度的稀疏注意力机制。BIGBIRD满足多项理论结果:它是序列到序列函数的通用逼近器,并且是图灵完备的。从理论上讲,我们利用额外的全局token的能力来保持模型的表达能力。我们通过展示转向稀疏注意力机制确实会带来一定的代价来补充这些结果。从实证上看,BIGBIRD在多项NLP任务(如问答和长文档分类)中提供了最先进的性能。我们进一步引入了基于注意力的DNA上下文语言模型,并针对下游任务(如启动子区域预测和非编码变体效应预测)进行了微调。