深入理解BigBird的塊稀疏高效實(shí)現(xiàn)方案

引言

基于 transformer 的模型已被證明對很多 NLP 任務(wù)都非常有用。然而， 的時間和內(nèi)存復(fù)雜度 (其中 是序列長度) 使得在長序列 () 上應(yīng)用它們變得非常昂貴，因而大大限制了其應(yīng)用。最近的幾篇論文，如 Longformer 、Performer 、Reformer 、簇狀注意力 都試圖通過對完整注意力矩陣進(jìn)行近似來解決這個問題。如果你不熟悉這些模型，可以查看 之前的 博文。

BigBird (由 該論文 引入) 是解決這個問題的最新模型之一。 BigBird 依賴于 塊稀疏注意力 而不是普通注意力 ( 即 BERT 的注意力)，與 BERT 相比，這一新算法能以低得多的計算成本處理長達(dá) 4096 的序列。在涉及很長序列的各種任務(wù)上，該模型都實(shí)現(xiàn)了 SOTA，例如長文檔摘要、長上下文問答。

RoBERTa 架構(gòu)的 BigBird 模型現(xiàn)已集成入 transformers 中。本文的目的是讓讀者 深入 了解 BigBird 的實(shí)現(xiàn)，并讓讀者能在 transformers 中輕松使用 BigBird。但是，在更深入之前，一定記住 BigBird 注意力只是 BERT 完全注意力的一個近似，因此我們并不糾結(jié)于讓它比 BERT 完全注意力 更好，而是致力于讓它更有效率。有了它，transformer 模型就可以作用于更長的序列，因?yàn)?BERT 的二次方內(nèi)存需求很快會變得難以為繼。簡而言之，如果我們有 計算和 時間，那么用 BERT 注意力就好了，完全沒必要用本文討論的塊稀疏注意力。

如果你想知道為什么在處理較長序列時需要更多計算，那么本文正合你意！

在使用標(biāo)準(zhǔn)的 BERT 類注意力時可能會遇到以下幾個主要問題:

每個詞元真的都必須關(guān)注所有其他詞元嗎？

為什么不只計算重要詞元的注意力？

如何決定哪些詞元重要？

如何以高效的方式處理少量詞元？

本文，我們將嘗試回答這些問題。

應(yīng)該關(guān)注哪些詞元？

下面，我們將以句子 BigBird is now available in HuggingFace for extractive Question Answering 為例來說明注意力是如何工作的。在 BERT 這類的注意力機(jī)制中，每個詞元都簡單粗暴地關(guān)注所有其他詞元。從數(shù)學(xué)上來講，這意味著每個查詢的詞元,將關(guān)注每個鍵詞元。

我們考慮一下 每個查詢詞元應(yīng)如何明智地選擇它實(shí)際上應(yīng)該關(guān)注的鍵詞元 這個問題，下面我們通過編寫偽代碼的方式來整理思考過程。

假設(shè) available 是當(dāng)前查詢詞元，我們來構(gòu)建一個合理的、需要關(guān)注的鍵詞元列表。

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

#假設(shè)當(dāng)前需要計算'available'這個詞的表征
query_token='available'

#初始化一個空集合，用于放'available'這個詞的鍵詞元
key_tokens=[]#=>目前，'available'詞元不關(guān)注任何詞元

鄰近詞元當(dāng)然很重要，因?yàn)樵谝粋€句子 (單詞序列) 中，當(dāng)前詞高度依賴于前后的鄰近詞?；瑒幼⒁饬?即基于該直覺。

#考慮滑動窗大小為3,即將'available'的左邊一個詞和右邊一個詞納入考量
#左詞:'now';右詞:'in'
sliding_tokens=["now","available","in"]

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

長程依賴關(guān)系: 對某些任務(wù)而言，捕獲詞元間的長程關(guān)系至關(guān)重要。 例如 ，在問答類任務(wù)中，模型需要將上下文的每個詞元與整個問題進(jìn)行比較，以便能夠找出上下文的哪一部分對正確答案有用。如果大多數(shù)上下文詞元僅關(guān)注其他上下文詞元，而不關(guān)注問題，那么模型從不太重要的上下文詞元中過濾重要的上下文詞元就會變得更加困難。

BigBird 提出了兩種允許長程注意力依賴的方法，這兩種方法都能保證計算效率。

全局詞元: 引入一些詞元，這些詞元將關(guān)注每個詞元并且被每個詞元關(guān)注。例如，對 “HuggingFace is building nice libraries for easy NLP” ，現(xiàn)在假設(shè) 'building' 被定義為全局詞元，而對某些任務(wù)而言，模型需要知道 'NLP' 和 'HuggingFace' 之間的關(guān)系 (注意: 這 2 個詞元位于句子的兩端); 現(xiàn)在讓 'building' 在全局范圍內(nèi)關(guān)注所有其他詞元，會對模型將 'NLP' 與 'HuggingFace' 關(guān)聯(lián)起來有幫助。

#我們假設(shè)第一個和最后一個詞元是全局的，則有:
global_tokens=["BigBird","answering"]

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

隨機(jī)詞元: 隨機(jī)選擇一些詞元，這些詞元將通過關(guān)注其他詞元來傳輸信息，而那些詞元又可以傳輸信息到其他詞元。這可以降低直接從一個詞元到另一個詞元的信息傳輸成本。

#現(xiàn)在，我們可以從句子中隨機(jī)選擇`r`個詞元。這里，假設(shè)`r`為1，選擇了`is`這個詞元
>>>random_tokens=["is"]#注意:這個是完全隨機(jī)選擇的，因此可以是任意詞元。

#將隨機(jī)詞元加入到集合中
key_tokens.append(random_tokens)

#現(xiàn)在看下`key_tokens`集合中有哪些詞元
key_tokens
{'now','is','in','answering','available','BigBird'}

#至此，查詢詞'available'僅關(guān)注集合中的這些詞元，而不用關(guān)心全部

這樣，查詢詞元僅關(guān)注所有詞元的一個子集，該子集能夠產(chǎn)生完全注意力值的一個不錯的近似。相同的方法將用于所有其他查詢詞元。但請記住，這里的重點(diǎn)是盡可能有效地接近 BERT 的完全注意力。BERT 那種簡單地讓每個查詢詞元關(guān)注所有鍵詞元的做法可以建模為一系列矩陣乘法，從而在現(xiàn)代硬件 (如 GPU) 上進(jìn)行高效計算。然而，滑動、全局和隨機(jī)注意力的組合似乎意味著稀疏矩陣乘法，這在現(xiàn)代硬件上很難高效實(shí)現(xiàn)。BigBird 的主要貢獻(xiàn)之一是提出了 塊稀疏 注意力機(jī)制，該機(jī)制可以高效計算滑動、全局和隨機(jī)注意力。我們來看看吧！

圖解全局、滑動、隨機(jī)注意力的概念

首先，我們借助圖來幫助理解“全局”、“滑動”和“隨機(jī)”注意力，并嘗試?yán)斫膺@三種注意力機(jī)制的組合是如何較好地近似標(biāo)準(zhǔn) BERT 類注意力的。

BigBird 塊稀疏注意力 是滑動連接、全局連接和隨機(jī)連接 (總共 10 個連接) 的組合，如上圖左側(cè)動圖所示。而 完全注意力 圖 (右側(cè)) 則是有全部 15 個連接 (注意: 總共有 6 個節(jié)點(diǎn))。你可以簡單地將完全注意力視為所有詞元都是全局詞元 。

完全注意力: 模型可以直接在單個層中將信息從一個詞元傳輸?shù)搅硪粋€詞元，因?yàn)槊總€詞元都會對每個其他詞元進(jìn)行查詢，并且受到其他每個詞元的關(guān)注。我們考慮一個與上圖類似的例子，如果模型需要將 'going' 與 'now' 關(guān)聯(lián)起來，它可以簡單地在單層中執(zhí)行此操作，因?yàn)樗鼈儍蓚€是有直接連接的。

塊稀疏注意力: 如果模型需要在兩個節(jié)點(diǎn) (或詞元) 之間共享信息，則對于某些詞元，信息將必須經(jīng)過路徑中的各個其他節(jié)點(diǎn); 因?yàn)椴皇撬泄?jié)點(diǎn)都有直接連接的。例如 ，假設(shè)模型需要將 going 與 now 關(guān)聯(lián)起來，那么如果僅存在滑動注意力，則這兩個詞元之間的信息流由路徑 going -> am -> i -> now 來定義，也就是說它必須經(jīng)過 2 個其他詞元。因此，我們可能需要多個層來捕獲序列的全部信息，而正常的注意力可以在單層中捕捉到這一點(diǎn)。在極端情況下，這可能意味著需要與輸入詞元一樣多的層。然而，如果我們引入一些全局詞元，信息可以通過以下路徑傳播 going -> i -> now ，這可以幫助縮短路徑。如果我們再另外引入隨機(jī)連接，它就可以通過 going -> am -> now 傳播。借助隨機(jī)連接和全局連接，信息可以非?？焖俚?(只需幾層) 從一個詞元傳輸?shù)较乱粋€詞元。

如果我們有很多全局詞元，那么我們可能不需要隨機(jī)連接，因?yàn)樾畔⒖梢酝ㄟ^多個短路徑傳播。這就是在使用 BigBird 的變體 (稱為 ETC) 時設(shè)置 num_random_tokens = 0 的動機(jī) (稍后部分將會詳細(xì)介紹)。

在這些圖中，我們假設(shè)注意力矩陣是對稱的 即 因?yàn)樵趫D中如果某個詞元 A 關(guān)注 B，那么 B 也會關(guān)注 A。從下一節(jié)所示的注意力矩陣圖中可以看出，這個假設(shè)對于 BigBird 中的大多數(shù)詞元都成立。

原始完全注意力即 BERT 的注意力，而塊稀疏注意力則是 BigBird 的注意力。想知道 block_size 是什么？請繼續(xù)閱讀下文。_現(xiàn)在，為簡單起見，將其視為 1。_

BigBird 塊稀疏注意力

BigBird 塊稀疏注意力是我們上文討論的內(nèi)容的高效實(shí)現(xiàn)。每個詞元都關(guān)注某些 全局詞元 、 滑動詞元 和 隨機(jī)詞元，而不管其他 所有 詞元。作者分別實(shí)現(xiàn)了每類查詢注意力矩陣，并使用了一個很酷的技巧來加速 GPU 和 TPU 上的訓(xùn)練/推理。

BigBird 塊稀疏注意力

注意: 在上圖的頂部有 2 個額外的句子。正如你所注意到的，兩個句子中的每個詞元都只是交換了一個位置。這就是滑動注意力的實(shí)現(xiàn)方式。當(dāng) q[i] 與 k[i,0:3] 相乘時，我們會得到 q[i] 的滑動注意力分?jǐn)?shù) (其中i 是序列中元素的索引)。

你可以在 這兒 找到 block_sparse 注意力的具體實(shí)現(xiàn)?，F(xiàn)在看起來可能非?？膳拢@篇文章肯定會讓你輕松理解它。

全局注意力

對于全局注意力而言，每個查詢詞元關(guān)注序列中的所有其他詞元，并且被其他每個詞元關(guān)注。我們假設(shè) Vasudev (第一個詞元) 和 them (最后一個詞元) 是全局的 (如上圖所示)。你可以看到這些詞元直接連接到所有其他詞元 (藍(lán)色框)。

#偽代碼

Q->Querymartix(seq_length,head_dim)
K->Keymatrix(seq_length,head_dim)

#第一個和最后一個詞元關(guān)注所有其他詞元
Q[0]x[K[0],K[1],K[2],......,K[n-1]]
Q[n-1]x[K[0],K[1],K[2],......,K[n-1]]

#第一個和最后一個詞元也被其他所有詞元關(guān)注
K[0]x[Q[0],Q[1],Q[2],......,Q[n-1]]
K[n-1]x[Q[0],Q[1],Q[2],......,Q[n-1]]

滑動注意力

鍵詞元序列被復(fù)制兩次，其中一份每個詞元向右移動一步，另一份每個詞元向左移動一步?，F(xiàn)在，如果我們將查詢序列向量乘以這 3 個序列向量，我們將覆蓋所有滑動詞元。計算復(fù)雜度就是 O(3n) = O(n) 。參考上圖，橙色框代表滑動注意力。你可以在圖的頂部看到 3 個序列，其中 2 個序列各移動了一個詞元 (1 個向左，1 個向右)。

#我們想做的
Q[i]x[K[i-1],K[i],K[i+1]]fori=1:-1

#高效的代碼實(shí)現(xiàn)(乘法為點(diǎn)乘)
[Q[0],Q[1],Q[2],......,Q[n-2],Q[n-1]]x[K[1],K[2],K[3],......,K[n-1],K[0]]
[Q[0],Q[1],Q[2],......,Q[n-1]]x[K[n-1],K[0],K[1],......,K[n-2]]
[Q[0],Q[1],Q[2],......,Q[n-1]]x[K[0],K[1],K[2],......,K[n-1]]

#每個序列被乘3詞，即`window_size=3`。為示意，僅列出主要計算，省略了一些計算。

隨機(jī)注意力

隨機(jī)注意力確保每個查詢詞元也會關(guān)注一些隨機(jī)詞元。對實(shí)現(xiàn)而言，這意味著模型隨機(jī)選取一些詞元并計算它們的注意力分?jǐn)?shù)。

#r1,r2,r為隨機(jī)索引;注意r1,r2,r每行取值不同
Q[1]x[Q[r1],Q[r2],......,Q[r]]
.
.
.
Q[n-2]x[Q[r1],Q[r2],......,Q[r]]

#不用管第0個和第n-1個詞元，因?yàn)樗鼈円呀?jīng)是全局詞元了。

注意: 當(dāng)前的實(shí)現(xiàn)進(jìn)一步將序列劃分為塊，并且每個符號都依塊而定義而非依詞元而定義。我們在下一節(jié)中會更詳細(xì)地討論這個問題。

實(shí)現(xiàn)

回顧: 在常規(guī) BERT 注意力中，一系列詞元，即 通過線性層投影到 ，并基于它們計算注意力分?jǐn)?shù) ，公式為 。使用 BigBird 塊稀疏注意力時，我們使用相同的算法，但僅針對一些選定的查詢和鍵向量進(jìn)行計算。

我們來看看 BigBird 塊稀疏注意力是如何實(shí)現(xiàn)的。首先，我們用 分別代表 block_size 、num_random_blocks 、num_sliding_blocks 、num_global_blocks 。我們以 為例來說明 BigBird 塊稀疏注意力的機(jī)制部分，如下所示:

的注意力分?jǐn)?shù)分別計算如下:

的注意力分?jǐn)?shù)由 表示，其中 ，即為第一塊中的所有詞元與序列中的所有其他詞元之間的注意力分?jǐn)?shù)。

BigBird 塊稀疏注意力

表示第 1 塊， 表示第 塊。我們僅在 和 (即所有鍵) 之間執(zhí)行正常的注意力操作。

為了計算第二塊中詞元的注意力分?jǐn)?shù)，我們收集前三塊、最后一塊和第五塊。然后我們可以計算。

BigBird 塊稀疏注意力

這里，我用 表示詞元只是為了明確地表示它們的性質(zhì) (即是全局、隨機(jī)還是滑動詞元)，只用 無法表示他們各自的性質(zhì)。

為了計算 的注意力分?jǐn)?shù)，我們先收集相應(yīng)的全局、滑動、隨機(jī)鍵向量，并基于它們正常計算 上的注意力。請注意，正如前面滑動注意力部分所討論的，滑動鍵是使用特殊的移位技巧來收集的。

BigBird 塊稀疏注意力

為了計算倒數(shù)第二塊 (即 ) 中詞元的注意力分?jǐn)?shù)，我們收集第一塊、最后三塊和第三塊的鍵向量。然后我們用公式進(jìn)行計算。這和計算 非常相似。

BigBird 塊稀疏注意力

最后一塊 的注意力分?jǐn)?shù)由 表示，其中，只不過是最后一塊中的所有詞元與序列中的所有其他詞元之間的注意力分?jǐn)?shù)。這與我們對 所做的非常相似。

BigBird 塊稀疏注意力

我們將上面的矩陣組合起來得到最終的注意力矩陣。該注意力矩陣可用于獲取所有詞元的表征。

BigBird 塊稀疏注意力

上圖中 藍(lán)色 -> 全局塊 、紅色 -> 隨機(jī)塊 、橙色 -> 滑動塊 。在前向傳播過程中，我們不存儲“白色”塊，而是直接為每個單獨(dú)的部分計算加權(quán)值矩陣 (即每個詞元的表示)，如上所述。

現(xiàn)在，我們已經(jīng)介紹了塊稀疏注意力最難的部分，即它的實(shí)現(xiàn)。希望對你更好地理解實(shí)際代碼有幫助?，F(xiàn)在你可以深入研究代碼了，在此過程中你可以將代碼的每個部分與上面的某個部分聯(lián)系起來以助于理解。

時間和內(nèi)存復(fù)雜度

BERT 注意力和 BigBird 塊稀疏注意力的時間和空間復(fù)雜度之比較。

展開以了解復(fù)雜度的計算過程。

BigBird時間復(fù)雜度=O(wxn+rxn+gxn)
BERT時間復(fù)雜度=O(n^2)

假設(shè):
w=3x64
r=3x64
g=2x64

當(dāng)序列長度為512時
=>**BERT時間復(fù)雜度=512^2**

當(dāng)序列長度為1024時
=>BERT時間復(fù)雜度=(2x512)^2
=>**BERT時間復(fù)雜度=4x512^2**

=>BigBird時間復(fù)雜度=(8x64)x(2x512)
=>**BigBird時間復(fù)雜度=2x512^2**

當(dāng)序列長度為4096時
=>BERT時間復(fù)雜度=(8x512)^2
=>**BERT時間復(fù)雜度=64x512^2**

=>BigBird時間復(fù)雜度=(8x64)x(8x512)
=>BigBird時間復(fù)雜度=8x(512x512)
=>**BigBird時間復(fù)雜度=8x512^2**

ITC 與 ETC

BigBird 模型可以使用 2 種不同的策略進(jìn)行訓(xùn)練: ITC 和 ETC。 ITC (internal transformer construction，內(nèi)部 transformer 構(gòu)建) 就是我們上面討論的。在 ETC (extended transformer construction，擴(kuò)展 transformer 構(gòu)建) 中，會有更多的全局詞元，以便它們關(guān)注所有詞元或者被所有詞元關(guān)注。

ITC 需要的計算量較小，因?yàn)楹苌儆性~元是全局的，同時模型可以捕獲足夠的全局信息 (也可以借助隨機(jī)注意力)。而 ETC 對于需要大量全局詞元的任務(wù)非常有幫助，例如對 問答 類任務(wù)而言，整個問題應(yīng)該被所有上下文關(guān)注，以便能夠?qū)⑸舷挛恼_地與問題相關(guān)聯(lián)。

注意: BigBird 論文顯示，在很多 ETC 實(shí)驗(yàn)中，隨機(jī)塊的數(shù)量設(shè)置為 0?？紤]到我們上文圖解部分的討論，這是合理的。

下表總結(jié)了 ITC 和 ETC：

在 Transformers 中使用 BigBird

你可以像使用任何其他 模型一樣使用 BigBirdModel 。我們看一下代碼:

fromtransformersimportBigBirdModel

#從預(yù)訓(xùn)練checkpoint中加載bigbird模型
model=BigBirdModel.from_pretrained("google/bigbird-roberta-base")
#使用默認(rèn)配置初始化模型，如attention_type="block_sparse"，num_random_blocks=3，block_size=64
#你也可以按照自己的需要改變這些參數(shù)。這3個參數(shù)只改變每個查詢詞元關(guān)注的詞元數(shù)。
model=BigBirdModel.from_pretrained("google/bigbird-roberta-base",num_random_blocks=2,block_size=16)

#通過把a(bǔ)ttention_type設(shè)成`original_full`，BigBird就會用復(fù)雜度為n^2的完全注意力。此時，BigBird與BERT相似度為99.9%。
model=BigBirdModel.from_pretrained("google/bigbird-roberta-base",attention_type="original_full")

截至現(xiàn)在， Hub 中總共有 3 個 BigBird checkpoint: bigbird-roberta-base，bigbird-roberta-large 以及 bigbird-base-trivia-itc。前兩個檢查點(diǎn)是使用 masked_lm 損失 預(yù)訓(xùn)練 BigBirdForPretraining 而得; 而最后一個是在 trivia-qa 數(shù)據(jù)集上微調(diào) BigBirdForQuestionAnswering 而得。

讓我們看一下如果用你自己喜歡的 PyTorch 訓(xùn)練器，最少需要多少代碼就可以使用 的 BigBird 模型來微調(diào)你自己的任務(wù)。

#以問答任務(wù)為例
fromtransformersimportBigBirdForQuestionAnswering,BigBirdTokenizer
importtorch

device=torch.device("cpu")
iftorch.cuda.is_available():
device=torch.device("cuda")

#我們用預(yù)訓(xùn)練權(quán)重初始化bigbird模型，并隨機(jī)初始化其頭分類器
model=BigBirdForQuestionAnswering.from_pretrained("google/bigbird-roberta-base",block_size=64,num_random_blocks=3)
tokenizer=BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")
model.to(device)

dataset="torch.utils.data.DataLoaderobject"
optimizer="torch.optimobject"
epochs=...

#最簡訓(xùn)練循環(huán)
foreinrange(epochs):
forbatchindataset:
model.train()
batch={k:batch[k].to(device)forkinbatch}

#前向
output=model(**batch)

#后向
output["loss"].backward()
optimizer.step()
optimizer.zero_grad()

#將最終權(quán)重存至本地目錄
model.save_pretrained("")

#將權(quán)重推到Hub中
fromhuggingface_hubimportModelHubMixin
ModelHubMixin.push_to_hub("",model_id="")

#使用微調(diào)后的模型，以用于推理
question=["Howareyoudoing?","Howislifegoing?"]
context=["",""]
batch=tokenizer(question,context,return_tensors="pt")
batch={k:batch[k].to(device)forkinbatch}

model=BigBirdForQuestionAnswering.from_pretrained("")
model.to(device)
withtorch.no_grad():
start_logits,end_logits=model(**batch).to_tuple()
#這里，你可以使用自己的策略對start_logits，end_logits進(jìn)行解碼

#注意:
#該代碼段僅用于展示即使你想用自己的PyTorch訓(xùn)練器微調(diào)BigBrid，這也是相當(dāng)容易的。
#我會建議使用Trainer，它更簡單，功能也更多。

使用 BigBird 時，需要記住以下幾點(diǎn):

序列長度必須是塊大小的倍數(shù)，即 seqlen % block_size = 0 。你不必?fù)?dān)心，因?yàn)槿绻?batch 的序列長度不是 block_size 的倍數(shù)， transformers 會自動填充至最近的整數(shù)倍。

目前，Hugging Face 的實(shí)現(xiàn) 尚不支持 ETC，因此只有第一個和最后一個塊是全局的。

當(dāng)前實(shí)現(xiàn)不支持 num_random_blocks = 0 。

論文作者建議當(dāng)序列長度 < 1024 時設(shè)置 attention_type = "original_full" 。

必須滿足: seq_length > global_token + random_tokens + moving_tokens + buffer_tokens ，其中 global_tokens = 2 x block_size 、 sliding_tokens = 3 x block_size 、 random_tokens = num_random_blocks x block_size 且 buffer_tokens = num_random_blocks x block_size 。如果你不能滿足這一點(diǎn)， transformers 會自動將 attention_type 切換為 original_full 并告警。

當(dāng)使用 BigBird 作為解碼器 (或使用 BigBirdForCasualLM ) 時， attention_type 應(yīng)該是 original_full 。但你不用擔(dān)心， transformers 會自動將 attention_type 切換為 original_full ，以防你忘記這樣做。

下一步

@patrickvonplaten 建了一個非?？岬?筆記本，以展示如何在 trivia-qa 數(shù)據(jù)集上評估 BigBirdForQuestionAnswering 。你可以隨意用這個筆記本來玩玩 BigBird。

審核編輯:黃飛