背景

自從谷歌的BERT預訓練模型橫空出世,預訓練-下游任務微調(diào)的方式便成了自然語言處理任務的靈丹妙藥。然而,復雜度高、顯存消耗大等問題一直困擾著BERT等預訓練模型的優(yōu)化;由于BERT中Transformer(多層自注意力)關于輸入文本長度L有的O()的時間空間復雜度,長文本消耗顯存陡然增加。想象一下,一位工程師興致勃勃地將數(shù)據(jù)在設計好的下游任務上微調(diào),滿懷期待地盼望著結果的提升,卻因為其中的一些長文本使得顯存溢出或超過位置嵌入(position embedding)最大長度,該是一件多么沮喪的事情。

解決這個問題最直接的方法是滑動窗口(sliding window)對每個512(通常BERT位置嵌入的最大長度)字符的窗口分別預測,最終合并不同窗口的結果的方式隨著具體下游任務的不同略有差異,例如閱讀理解問答可以輸出各段中總評分最高的小段(span)作為答案。然而,如果問題需要長程注意力,也就是兩個關鍵的句子分布在段落中相距較遠位置的時候,這種方法的效果就會大打折扣,下圖就是一個例子。

解決這個問題的另一種思路是優(yōu)化Transformer結構,這一條思路的工作有很多,例如Longformer[1]、BlockBert、最近的BigBird等……但是這些工作通常只是將文本長度從512擴展幾倍(基于現(xiàn)有的硬件條件),讓BERT一次“看到”更多的文本;然而,人類并不需要如此強的瞬時閱讀能力——實際上人類同時在工作記憶里存儲的元素通常只有5-7個——也能閱讀并理解長文本,那么人類是如何做到的呢?

認知中的工作記憶和調(diào)度

“工作記憶的核心是一個中央處理機制,它協(xié)調(diào)來自于多種來源的信息”, 并且 “它發(fā)揮一個有限容量的注意力系統(tǒng)的作用,這個系統(tǒng)能選擇和操作控制過程和策略”, 這是工作記憶的提出者Baddeley [2] 在他1992年《Science》著作中的論斷。事實上,人腦正是通過回憶和注意力,協(xié)調(diào)長期記憶和短期記憶(工作記憶)的使用策略來完成對長文本的理解。下圖是分層注意力機制的圖解,工作記憶從當前的感知記憶空間或者長期記憶中抽取關鍵信息進行深層理解,然而這些信息如果沒有被不斷重演(rehearsal)就會在大約5-15s后逐漸忘掉,剩余的有用的信息來進行繼續(xù)的推理。

CogLTX 的工作流程

受到人的認知過程啟發(fā),我們用同樣的方法來處理長文本。如果將BERT的512輸入字符限制比作人的工作記憶,那么既然人思考問題時能夠找到關鍵的少量信息,并在工作記憶中推理出結果,BERT的512也應該遠遠足夠,關鍵是對于特定的問題,我們要最終用的真正關鍵的那部分信息。

因此,CogLTX遵循一種特別簡單直觀的范式,即 抽取關鍵的句子 => 通過BERT得到答案 這樣的兩步流程。常見的幾種任務都可以用這種范式來解決。比如下圖列舉了語段抽取、序列級別任務、字詞級別任務的處理方法。

在這里,我們將完成任務要訓練的BERT稱為推理機(reasoner),解決問題的關鍵語段記為z,CogLTX通過一個被稱為“MemRecall”的過程,如同人類調(diào)度工作記憶一樣的方式來抽取關鍵的語段z。

MemRecall 關鍵信息抽取

對于關鍵信息的認識本身也是智能的重要部分,這并非易事。最直觀的想法是通過信息檢索的辦法(例如BM25)來抽取關鍵句,但是仔細一想就會發(fā)現(xiàn)這其實是不可行的,因為下游任務的不確定性,無法建模成信息檢索的形式。例如,文本分類任務如果用BM25去檢索,則無法定義查詢(query)是什么。因此抽取的模型也要與任務息息相關。

其次就是直接檢索的方式過于粗糙,同時對于無法處理多跳推理的信息。而人在工作記憶中的思考是一個時序的過程,會不斷忘記被錯誤檢索的信息而將空間留給新信息。因此,我們模仿這樣的過程設計了MemRecall,其核心是一個對于每個字詞進行相關度打分的評分機(judge)模型,也用BERT實現(xiàn)。MemRecall的過程如下圖所示。

我們首先將長文本切塊,用現(xiàn)有的關鍵信息去連接每一塊,通過評分機獲取塊平均相似度得分,其中最高分的塊被放入“工作記憶”——但是這是暫時的,正如人腦的工作記憶一樣,我們?nèi)绻�只對重要的信息進行重演,其他信息就會很快忘掉——在MemRecall中我們將這些最高分的塊一起通過評分機,信息充分交換后再進行評分,并且“忘掉”那些得分不高的塊。新的關鍵信息將用來重復這一過程,完成多步推理。

訓練

在模型訓練時,我們考慮兩種情況:第一種是閱讀理解問答這樣的任務,由于信息句可以從答案所在句推斷出來,因此是監(jiān)督學習。此時評分機和推理機的訓練(finetuning)都比較簡單,只需將真正的關鍵句和一些負樣本信息句組合,然后像正常BERT那樣訓練即可;第二種是文本分類這種,數(shù)據(jù)集中往往不會提供關鍵句的標注,這就需要我們自己推斷。關鍵句的一個特性是,如果缺少關鍵句將不能推斷到正確答案,因此我們先用詞向量等方法初始化關鍵句標簽后,再訓練中調(diào)整關鍵句標簽,如果某個句子剔出后損失函數(shù)驟然增加那么就必然是關鍵句,如果可有可無則不是,根據(jù)這個方法在調(diào)整關鍵句標簽后可重新進行下一輪訓練,具體算法如下:

從隱變量的角度,如果認為關鍵句是隱變量,那么CogLTX的算法則可以看成是由于隱變量z取值空間離散且較大,選擇的一種點估計的近似。

實驗

文章在NewsQA、HotpotQA問答數(shù)據(jù)集,20NewsGroup文本分類和Alibaba淘外文本多標簽分類等幾個任務上進行試驗,結果均超過或類似于目前最好的模型效果,具體數(shù)據(jù)在論文中列舉。

同時,CogLTX犧牲了部分推理的時間,換取了與文本長度無關的訓練空間開銷。下圖展示了BERT-large在batch size為1的時候的時空開銷對比,可以看出CogLTX(橙色)消耗的空間是固定的。

小結

對于BERT處理長文本時遇到的困境,通常的做法都會考慮輕量化Transformer的思路,然而如果能從人類處理信息的方式得到啟發(fā),另辟蹊徑從下游任務微調(diào)的流程上考慮,更直接地解決這個問題。