背景

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

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

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

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

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

CogLTX 的工作流程

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

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

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

MemRecall 關(guān)鍵信息抽取

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

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

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

訓(xùn)練

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

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

實(shí)驗(yàn)

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

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

小結(jié)

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