背景

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

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

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

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

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

CogLTX 的工作流程

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

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

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

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

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

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

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

訓(xùn)練

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

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

實(shí)驗(yàn)

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

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

小結(jié)

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