摘要
基於Transformer的模型由於自注意力操作不能處理長序列,自注意力操作是序列長度的二次方。為了定位這個限制,我們提出一種Longformer的方法,它使用的注意力機制能夠隨着序列長度線性增長,使得能夠很容易的處理具有數千或者更長token的文檔。Longformer的注意力機制可以使用任務驅動的全局注意力來隨時替換典型的自注意力機制。
我們使用字符級語言模型來評估Longformer,並在text8 和 enwik8上獲得SOTA結果。同RoBERTa相比,我們的Longformer模型在長文本上的效果始終優於它。
最后,我們提出了Longformer-Encoder-Decoder(LED),一個Longformer的變形,用來支持長文本生成seq2seq任務,並在arXiv的摘要數據集上證明它的有效性。
Longformer
傳統的自注意力機制,其時間復雜度為O(n ^ 2),空間復雜度為O(n),這里提出了一種新的注意力模式,其時間復雜度同輸入序列長度成線性比例,使其對於長序列更有效。具體方法是,使用“注意模式”來稀疏化完整的自注意矩陣,該模式制定了一對相互關注的輸入位置。(個人理解:事實上時間復雜度應該是O(m * n^2),m是向量維度,如768)
注意模式
1、滑窗:用一個固定大小的滑窗來關注每一個token附近的信息。使用多層堆疊的窗口注意力會導致大的接受范圍,但是頂層可以利用所有輸入的位置信息,並能構建包含所有信息的表征,類似CNN。
給定一個固定的窗口大小w,一個token兩側分別關注w/2長度的token,則計算復雜度變為O(n * w)。如果有L層的話,則頂層的接受范圍大小是L * w(假設w大小對於所有的層是固定的)。我們可以在每一層使用不同的w來平衡效率和模型表征能力。
2、伸縮滑窗(dilated sliding window):想擴充接受范圍,但又不想增加計算量,可以使用伸縮滑窗。假設所有的層都有一個固定的d和w,則可接受的范圍是L * d * w,這樣即使一個很小的d值,都能接受到數以千計的tokens。(d的值到底是啥?)
在multi-head注意力里,每個注意力都計算一個不同的注意力分數。我們發現every head中不同的dilation config的設置會提升性能,通過允許一部分注意力不使用dilation來集中注意力到局部上下文,而其他的注意力使用dilation來關注更長的上下文。
3、全局注意(Global Attention):在MLM中,模型使用局部上下文信息來預測被遮掩的詞,對於分類,模型聚合整個句子的表征到一個特殊的token([CLS])中。對於QA,問題和文本會連接在一起讓模型通過自注意力來比較問題和文本。
在我們的情況下,滑窗和伸縮滑窗不夠靈活,不能處理特定任務表征。因此,我們在一些預先選擇好的輸入位置上增加“全局注意”。更重要的是,我們使得這些注意操作是對稱的:一個有全局注意的token,關注序列中的所有token,同時,序列中的所有token也都會關注它。
全局注意的應用舉例如下:在分類中,全局注意用於[CLS]這個token,而在QA中,全局注意用於所有的問題token中。因為全局注意標記的數量相對於序列長度n來說很小且獨立於n,因此局部注意和全局注意的復雜度仍然是O(n)。
全局注意的線性映射:原始的注意力計算公式為
$$
Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V
$$
這里我們使用兩組映射,$Q_s,K_s,V_s$計算滑窗注意的注意力分數,$Q_g,K_g,V_g$來計算全局注意的注意力分數。新增的映射提供了為不同注意力建模的靈活性,這對於下游任務的最佳性能至關重要。其中,兩組映射的Q,K,V使用相同的初始值。
自回歸語言模型
1、注意模式:在較低的層使用小的窗口size,而在較高的層上增加窗口size。這樣做的結果是,最頂層的layers會學習到整個序列的更高層級的表征,而較低的層則抓取局部信息。同時,這樣做也會平衡效率和性能。但是在較低的層上我們不會使用伸縮滑窗,這樣做可以最大化其學習和利用當前局部信息的能力。而對於更高的層,我們只在2個頭上(on 2 heads)上使用少量的遞增伸縮值。這給了模型在不犧牲局部上下文信息的情況下直接處理懲處token的能力。
2、實驗設置
-
訓練:字符集語言模型(character-level LM)
理想情況下,我們想用當前GPU內存能訓練的最大窗口size和序列長度來訓練模型。然而,我們發現模型在學會利用更長的信息之前先要進行大量的梯度更新來學習局部信息。考慮到這點,我們使用一個暫存的訓練步驟,這里我們可以增加注意窗口size和序列長度來跨越多個訓練階段。尤其,在第一階段我們用一個短序列長度和小的窗口size,然后在之后的每個階段,我們翻倍加大序列長度和窗口size,然后讓學習率減半。這會加快訓練,而保證慢的部分(最長的序列和窗口size)能持續運行到最后。
我們訓練了5個階段,序列初始長度是2048,最后一個階段的序列最終長度是23040。
-
評估:我們使用序列長度是32256來評估。我們將數據集拆分為32256、步長為512的重疊序列,並報告序列上最后512個token的性能。
預訓練和微調
在文本集上訓練Longformer,並針對6個任務微調它,包括分類,QA和指代消解。最終的模型可以處理長度為4096的序列(是BERT的8倍)。
我們使用MLM來預訓練Longformer,目標是發現序列中隨機遮掩的tokens。由於MLM的預訓練過程昂貴,我們基於RoBERTa來預訓練,這樣可以最小的改變來支持Longformer的注意力機制。注意到我們的注意力模式可以應用在任何預訓練的transformer模型中,卻不需要更改模型結構。
- 注意力模式:使用滑窗注意力,窗口size設為512,因此同RoBERTa一樣的計算量。
- 位置編碼:RoBERTa使用學習的最大位置為512的絕對位置編碼。為了支持更長的文本,我們增加另外的位置編碼來支持4096長度的位置。在長度為4096的序列中,我們多次復制RoBERTa的長度為512的序列的位置編碼進行位置編碼的初始化。使用該方法進行初始化能夠保存除了分區邊界以外的局部結構。
- 連續MLM預訓練:我們使用fairseq在長文本集上預訓練Longformer。我們訓練兩個size的模型,一個基本模型和一個大模型。每個模型的訓練參數都是:長度為4096的序列、65K的梯度更新,batch size 64,最大學習率是3e-5,線性warmup為500步。
- 凍結RoBERTa權重:預訓練Longformer時我們會凍結所有的RoBERTa權重,之訓練新的位置編碼。這樣設置的東西是為了能夠完美保留RoBERTa在短文本上的性能,這個設置獲得1.850的BPC值,但是如果所有的權重都參與訓練的話,會獲得更高的值1.705。
任務
baseline是一個基於RoBERTa的模型,它將上下文分解為最長的片段,將每個片段單獨傳遞給RoBERTa,並連接已激活部分進行進一步處理。
-
QA(問答)
使用兩階段:先選擇最大程度相關段落,然后把他們傳進第二階段進行答案提取。每個階段都連接問題和文本為一個序列,把它傳進Longformer中,然后使用特定任務的預測層。我們用多任務的方法訓練模型,這樣就能聯合預測relevant paragraphs,evidence sentences,answer spans and question types(yes/no/span)。
-
指代消解
在這個任務中沒有使用全局注意力。
Longformer-Encoder-Decoder(LED)
為了促進seq2seq學習對長序列建模,我們提出一個Longformer變體,它同時使用編碼解碼的Transformer,但是卻用有效的Longformer的局部+全局注意力模式來代替編碼器中的完整的自注意力機制。解碼器使用完整的自注意力機制來解碼完整的已編碼的tokens和之前已經解碼的位置信息。我們叫這個模型為LED。我們用BART來初始化LED參數,用層數和隱藏層大小來獲取准確的BART架構。唯一的不同是LED能處理更長的輸入序列,我們將位置編碼增加到16K(BART只有1K的tokens),而且我們重復復制BART的位置編碼(1K)16次,來獲取LED的初始的位置編碼。
我們使用arXiv摘要數據集來評估LED在摘要任務上的能力。LED的編碼器閱讀文本,解碼器生成輸出的摘要。編碼器使用窗口大小為1024的局部注意力,並在第一個token上使用全局注意力。解碼器使用完全的原始的注意力機制來解碼全部的編碼和之前的解碼位置。如同經典的seq2seq模型,LED使用teacher forcing在標注的訓練摘要上,並在推理上使用beam search。
LED在摘要上略優於BigBird(注意,BigBird摘要模型雖然支持序列長度為4000,但是它始於並基於Pegasus模型進行預訓練,Pegasus是一個專門為文本摘要設計和預訓練的模型)。在沒有預訓練或者特定任務的初始化參數時,LED卻能略優於BigBird,足以證明LED的優越性。
結論
Longformer提出一個注意力模式,能夠結合線性比例與序列長度的局部和全局信息。預訓練的Longformer在長文本上由於RoBERTa,並且在WikiHop和TriviaQA數據集上取得SOTA。同時,基於Longformer提出的LED,在arXiv長文本摘要任務上取得SOTA。未來,我們會研究更多的預訓練模型,尤其是LED,增加其序列長度,並探索其他可能從我們模型上獲益的任務。