bigbird長文本預訓練模型介紹

MLNLP(機器學習算法與自然語言處理)社區是國內外最大的自然語言處理社區之一，匯聚超過50w訂閱者，受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。
社區的願景是促進國內外自然語言處理，機器學習學術界、產業界和廣大愛好者之間的交流和進步。本博客翻譯自huggingface blog(https://huggingface.co/blog/big-bird)。前言基於Transformer的模型已經被證明了在許多NLP任務中的價值，但這類模型的時間複雜度、內存使用複雜度都是 n平方（n為序列長度），因此當序列長度超過常規的512時，模型對算力的要求將會大幅提高。最近的一些文章Longformer, Performer, Reformer, Clustered attention都試圖通過近似全主力機制改善該問題。例如這個帖子（https://huggingface.co/blog/long-range-transformers）就是介紹這些模型的。 BigBird論文是處理這類問題最新模型的其中之一，它使用block sparse attention替換了原類似Bert一樣的全注意力機制，在與BERT一樣的計算力情況下，可以處理的序列長度達到4096。它已經在很多長文本序列的任務上達到SOTA效果，例如長文本摘要、長文本問答。BigBird RoBERTa模型現在已經可以在Transformers倉庫中使用。這篇博客的目的是為了讓讀者深入理解big bird的運行機制，快速使用Transformers倉庫上手BigBird模型。但在開始之前，我們需要知道，BigBird的注意力機制是一個近似BERT的全注意力機制，因此它不是說比BERT的注意力機制效果更好，而是運行效率更高。BERT的注意力機制的存儲與序列長度是二次方關係，在長文本情況下的存儲需求就已經開始令人難以忍受，而BigBird的block sparse attention就是為了解決這個問題。也就是說，在無窮長度序列上，計算無窮次時，我們應該把BERT的全注意力機制換成block sparse attention。如果你想知道，為什麼在計算長序列時，我們需要更多的算力，這篇博客正好適合你。當使用標準BERT一樣的注意力機制時，主要的問題可能有：注意力應該考慮哪些詞？首先我們研究一個注意力工作的例子「BigBird is now available in HuggingFace for extractive question answering」。在全注意力機制中，每個詞都會在注意力中計算其他所有詞。用數學表達式表示即為： 

讓我們思考一下，如果通過掩碼，讓query-token只與最重要的key-token做注意力計算可能是一個更明智的方法，下面我們列舉一下對query-token來說，重要的key-token都有哪些。

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

# 現在query_token的例子為available
>>> query_token = 'available'

# 初始化一個空的key_token列表，下面我們逐漸添加
>>> key_tokens = []

相鄰的詞是非常重要的，當前詞的含義非常依賴於它鄰近的前幾個詞和鄰近的後幾個詞，這就是sliding attention

# 舉例窗口為3，即只考慮available詞的前一個詞和後一個詞
# 即左邊的詞」now「；右邊的詞」in「
>>> sliding_tokens = ["now", "available", "in"]

# 更新一下key_tokens集合
>>> key_tokens.append(sliding_tokens)

長程依賴（Long-range dependencies）：有的任務，需要捕獲距離較遠的詞之間的關係，例如問答模型需要比較原文中每個詞與整個問題，從而發現原文中哪些詞序列更適合做正確答案。如果大多數原文詞與其他原文詞做注意力計算，而不是與問題做注意力計算，這對模型來說將會更困難的篩出重要的原文詞。

BigBird有兩種長程注意力方式，可以讓計算變的更有效：

# 例如第一個詞和最後一個詞是全局的
>>> global_tokens = ["BigBird", "answering"]

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

# 現在用詞」is「做隨機詞
>>> random_tokens = ["is"]

>>> key_tokens.append(random_tokens)

>>> key_tokens
{'now', 'is', 'in', 'answering', 'available', 'BigBird'}

# 現在，詞」available「可以只與這些詞做注意力計算，而不是所有詞

通過這種方法，query token只與部分詞做注意力計算，同時與全注意力機制保持了相近的效果。這些操作會用於其他query_token，但是，所有這些操作都是為了在更節省的算力下，接近BERT的全注意力效果。在當今的硬體如GPU上，像BERT一樣的模型，簡單的將query token與所有的key token計算全注意力機制的方式，因為注意力時計算序列長度的平方矩陣乘法，所以非常消耗計算資源。然而如果通過結合sliding attention、global attention、random attention三種注意力機制，計算方式降為稀疏矩陣乘法，計算起來就更容易。BigBird就是通過block sparse注意力機制，結合了sliding attention、global attention、random attention三種方法。下面我們詳細介紹通過圖結構理解global, sliding, random keys首先，我們通過圖來理解為什麼三種注意力機制結合可以達到接近Bert的全注意力機制效果。

上圖給出了global、sliding、random三種連接，每個節點代表了詞，每條線代表了注意力，如果兩個詞之間沒有連接，意味著兩個詞中間沒有注意力計算。

BigBird block sparse attention結合了sliding、global、random的十條連接。正常注意力機制是有6個節點的所有15條連接。正常注意力與所有詞都為全局注意力是等價的。

Normal attention:模型可以在一層之內，把每個詞的信息都與其他任意詞直接交互。以上圖為例，如果模型需要兩個詞」going「與」now「互相聯繫，就需要在一層內，建立他們的直接連接。

Block sparse attention：由於節點不是全部互相連接的，那麼模型的節點（或詞）需要互相分享信息時，信息不得不通過最短路徑的其他節點。例如，當「going」需要連接至「now」，如果只考慮窗口為左右各一個詞的sliding attention，那麼「going」與「now」之間的信息，必須通過going -> am -> i -> now，才能被「now」接收。因此，我們需要多層「Block sparse」才能有效捕獲序列中所有的組合信息，而全連接的注意力機制只在一層內就可以捕獲序列的所有組合信息。在極端情況下，如兩個詞距離最遠，需要的層數，與輸入的詞數要一樣多。當加入global connection的連接時，信息可以通過going -> i -> now的方式傳遞。如果加入random connections，信息可以通過going -> am -> now的方式傳遞。所以在有global connections、random connections的參與下，信息可以在token之間傳遞更有效。如果當我們有很多global tokens時，就已經有很多足夠短的連接了，那麼就不再需要random connections。通過設置num_random_tokens = 0，就得到這樣變體的BigBird模型，也稱為ETC。BigBird block sparese attention通過以上討論可知，Block sparse attention是一個有效的運行方式。每個詞都是與其他global tokens、 sliding tokens、 random tokens進行計算，而不是與其他所有詞進行計算。作者已經在代碼中分別實現了多種注意力矩陣，因此實現了訓練和預測時間更短的小技巧。圖中可以看到2個向左平移和向右平移的額外句子。這是計算sliding attention時需要執行的。實現block_sparse注意力的代碼(https://github.com/vasudevgupta7/transformers/blob/5f2d6a0c93ca2017961199aa04a344b9b779d454/src/transformers/models/big_bird/modeling_big_bird.py#L513)看起來長，在這裡會逐步解釋。Global Attention對於global attention，注意力的每次query都是與輸入的所有其他詞進行計算。例如第一個詞[Vasudev]和最後一個詞[them]。如下代碼所示

# pseudo code

Q -> Query martix (seq_length, head_dim)
K -> Key matrix (seq_length, head_dim)

# 1st & last token attends all other tokens
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]]

# 1st & last token getting attended by all other tokens
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]]

Sliding Attention

計算注意力時序列中的key序列複製兩次，一個向左平移一個位置、一個向右平移一個位置，然後用query序列與三個key序列直接向量相乘，就實現了所有sliding tokens的計算。計算複雜度僅為O(3xn) = O(n)。在上圖中，橙色格子為sliding attention。在表格上方的三個句子代表三個key序列，即原key序列，向左平移的key序列、向右平移的key序列

# what we want to do
Q[i] x [K[i-1], K[i], K[i+1]] for i = 1:-1

# efficient implementation in code (assume dot product multiplication )
[Q[0], Q[1], Q[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]]

# Each sequence is getting multiplied by only 3 sequences to keep `window_size = 3`.
# Some computations might be missing; this is just a rough idea.

Random Attention

Random attention可以在計算注意力時，讓query詞與其他隨機的幾個key詞進行計算。在實際運行中，模型會隨機收到幾組注意力得分。

# r1, r2, r are some random indices; Note: r1, r2, r3 are different for each row
Q[1] x [Q[r1], Q[r2], ., Q[r]]
.
.
.
Q[n-2] x [Q[r1], Q[r2], ., Q[r]]

# leaving 0th & (n-1)th token since they are already global

模型推斷

在一般的Bert注意力中，序列為$X=x_1, x_2,...,x_n$,由稠密向量$Q,K,V$計算的注意力為$Z=Softmax(QK^T)$。在BigBird注意力計算中，同樣的算法就僅在個別的query向量和key向量之間計算。

下面我們來展開bigbird的block sparse attention的執行方式。首先我們假設b,r,s,g分別代表block_size, num_random_blocks, num_sliding_blocks, num_global_blocks。當b=4,r=1,g=2,s=3,d=5時:

注意力得分$q_1,q_2,q_3,q_{n-2},q_{n-1},q_n$分別由以下方式計算

注意力得分q_1由a_1表示，第一個詞與序列其他所有詞計算得來 

這裡 q_1 代表第一個塊， g_i 代表第i個塊，因此只是正常的注意力計算可得 q_1&g

在計算第二塊時，需要前三塊、最後一塊、第五塊，計算公式為 

圖中的 g,r,s 代表global,random,sliding。

當計算注意力得分 q_{3:n-2} 時，需要用global, sliding, random三種key，與 q_{3:n-2} 進行注意力計算。這裡的sliding key的計算由上面介紹的平移方法得到。

當計算倒數第二個詞（ q_{n-1} ）的得分時，需要第一個塊、最後三塊和第三塊。計算公式為 

當計算最後一個詞時，公式為 

 ，這裡的注意力計算為最後一個詞與其他所有詞的計算。

結合以上所有矩陣計算為：

在神經網絡前向傳播時，因為矩陣中白色的塊不會參與到注意力計算，所以計算整體會更輕量。
時間空間複雜度比較

BigBird time complexity = O(w x n + r x n + g x n)
BERT time complexity = O(n^2)

Assumptions:
w = 3 x 64
r = 3 x 64
g = 2 x 64

When seqlen = 512
=> time complexity in BERT = 512^2

When seqlen = 1024
=> time complexity in BERT = (2 x 512)^2
=> time complexity inBERT = 4 x 512^2

=> time complexity in BigBird = (8 x 64) x (2 x 512)
=> time complexity in BigBird = 2 x 512^2

When seqlen = 4096
=> time complexity in BERT = (8 x 512)^2
=> time complexity in BERT = 64 x 512^2

=> compute in BigBird = (8 x 64) x (8 x 512)
=> compute in BigBird = 8 x (512 x 512)
=> time complexity in BigBird = 8 x 512^2

ITC vs ETC

BigBird模型可以由兩種策略訓練：ITC & ETC。ITC(internal transformer construction)即為上面討論的計算方式。ETC(external transformer construction)為指定部分詞為全注意力計算方式。

ITC需要更少計算量，因為只有很少的詞是全局計算的，同時模型通過random attention保持全局信息。而ETC是設置了更多的全局計算詞，面向一些需要一些額外的全局信息任務，例如在閱讀理解任務中，問題需要與所有的答案進行注意力計算。

注意：在原論文中，ETC實驗中，不含random token夾帶私貨現已開源中文bigbird預訓練模型，從tiny至base共5個級別預訓練模型。可以從huggingface hub(https://huggingface.co/models?language=zh&sort=downloads&search=bigbird)直接下載使用，如果對您有用，歡迎star我的Github(https://github.com/LowinLi/chinese-bigbird)。

關於我們

MLNLP(機器學習算法與自然語言處理) 社區是由國內外自然語言處理學者聯合構建的民間學術社區，目前已經發展為國內外最大的自然語言處理社區之一，匯聚超過50w訂閱者，旗下包括萬人頂會交流群、AI臻選匯、AI英才匯以及AI學術匯等知名品牌，旨在促進機器學習，自然語言處理學術界、產業界和廣大愛好者之間的進步。社區可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。