Longformer:超越RoBERTa,为长文档而生的预训练模型
星標/置頂小屋,帶你解鎖
最萌最前沿的NLP、搜索與推薦技術
文 |?ChrisCao, 小軼
前言
今天要與大家分享的是AllenAI今年發表的最新工作,Longformer——一種可高效處理長文本的升級版Transformer。作者團隊提供了開源代碼,大家可快速復現,直接用于自己的任務。
傳統Tranformer-based模型在處理長文本時有著天然的劣勢。因為傳統模型采用的是“全連接”型的attention機制,即每一個token都要與其他所有token進行交互。其attention復雜度高達。此前的解決辦法是將長文切分為若干個較短的text span,然后逐個處理。這就導致不同的text span之間無法進行交互,因而必然存在大量information loss。當然,我們也可以通過添加一些其他機制來加強這種text span之間的交互。但這種新增機制實現起來通常比較復雜,而且往往是task-specific的,通用性不強。
本文提出的Longformer,改進了Transformer的傳統attention機制:對于每一個token,只對固定窗口大小的附近token計算local attention,并結合具體任務,計算少量的global attention。該方法的優點包括:
復雜度低,將attention機制的復雜度降至
通用性強,可用于各類文檔級任務
部署容易,作者在cuda內核上直接實現了Longformer的attention pattern,并提供了開源代碼。
Longformer在兩個字符級語言建模任務上都取得了SOTA的效果。并且,作者用Longformer的attention方法繼續預訓練RoBERTa。訓練得到的語言模型在多個長文檔任務上進行fine-tune后,性能全面超越Roberta。該預訓練模型也已開源,大家可以很方便地直接應用于自己的任務。
原文鏈接:
https://arxiv.org/pdf/2004.05150.pdf
Github:
https://github.com/allenai/longformer
Arxiv訪問慢的小伙伴也可以在訂閱號 「夕小瑤的賣萌屋」 后臺回復關鍵詞 【0716】下載論文PDF~
模型
作者共提出了三種新的attention pattern,來降低傳統self-attention的復雜度,分別是滑窗機制、膨脹滑窗機制、融合全局信息的滑窗機制。下圖展示了傳統attention與這三種attention pattern的示意圖。接下來將為大家分別講解。
1. 滑窗機制(Sliding window): 對于每一個token,只對其附近的w個token計算attention計算復雜度與文本序列長度成線性關系,為。作者認為,根據應用任務的不同可以對Transformer每一層施以不同的窗口大小,對模型表示能力可能有潛在幫助。
讀到這里的,大家可能和我一樣,誤認為這個窗口應該比較小,估計在16~64這個量級。但看到實驗部分會發現,作者在具體實現的時候,設置的窗口大小為512,和Bert的Input限制完全一樣。所以,大家不要存有“Longformer比Bert還要更輕量”的錯覺。
2. 膨脹滑窗機制(Dilated sliding window): 在對每一個進行token編碼時,普通滑窗機制只能考慮到長度為的上下文。作者進一步提出膨脹滑窗機制,在不增加計算負荷的前提下,拓寬模型“視場”。其做法借鑒了空洞卷積的思想[1]。如下圖所示,在滑動窗口中,被attend到的兩個相鄰token之間會存在大小為d的間隙。當transformer的層數為l時,則視場范圍可達到。實驗表明,由于考慮了更加全面的上下文信息,膨脹滑窗機制比普通的滑窗機制表現更佳。
3. 融合全局信息的滑窗機制(Global+sliding window): 我們知道Bert一類的語言模型在應用于具體任務時,實現方式略有不同。比如,對于文本分類任務,我們會在文本序列前添加[CLS]這一特殊token;而對于QA類任務,則會將問題與文本進行拼接后輸入。在Longformer中,作者也希望能夠根據具體任務的不同,在local attention的基礎上添加少量的global attention。比如,在分類任務上就會在[CLS]處添加一個global attention,而在QA任務上會對question中的所有token添加global attention。如下圖所示,對于添加了global attention的token,我們對其編碼時要對整個序列做attention。并且,編碼其他所有token時,也都要attend到它。
實驗
自定義CUDA內核
由于現有的深度學習庫中并沒有能直接實現膨脹滑窗機制的接口,為此作者直接自定義了CUDA內核操作,用于實現Longformer的attention pattern[2]。如下圖所示,Longformer的內存消耗與文本長度成線性關系(紅線)。用自定義CUDA來實現Longformer,相比于用Naive Pytorch來實現(藍線),運行速度加快了六倍。
Longformer在字符級別任務上的表現
作者在text8和enwik8兩個字符級任務上對Longformer進行實驗。實驗中,模型每一層采用了不同的窗口大小:底層使用較小的滑窗,以建模局部信息;在高層使用較大的滑窗,以擴大感受野。訓練時,理想狀況下當然是希望使用GPU所能承受的最大的window size和sequence len。但為了加快訓練速度,作者采用的是一種階段式的訓練方式:在學習更長的上下文之前,先學好局部的上下文。在第一階段,先設置較短的序列長度和窗口大小。在后續階段,window size和sequence length增加一倍,學習率減半 實驗結果如下圖所示,Longformer在這兩個數據集上皆達到了SOTA效果(注:測試指標為BPC,bits-per-character;BPC越小,性能越優)。
作者通過實驗,對滑窗機制的設置進行了進一步的討論。如下表所示:
表中第一組實驗(前三行)討論的是:如果transformer的不同層采用不同窗口大小,是否可以提高性能?實驗結果表明,由底層至高層遞增窗口大小,可提升性能;遞減則反而性能降低。
第二組實驗(后兩行)是對膨脹滑窗機制的消融實驗,證明了增加間隙后的滑窗機制,性能可以有小幅度提升
Longformer用于預訓練
MLM Pretraining
作者采用Longformer的方法在以下四個文檔級語料上進行預訓練,從而得到適于文檔級NLP任務的語言模型。作者并沒有完全從頭預訓練一個隨機初始化的模型,而是以RoBERTa為基礎,采用MLM(masked language modeling)的方法繼續預訓練。預訓練時,每一層都采用固定的大小為512的滑動窗口,暫不添加global attention。為支持長文本,論文作者把position embedding擴展到了4096個。
預訓練結束后,在多個文檔級任務上再進一步對預訓練模型做fine-tuning。fine-tuning時會根據任務增加global attention:共設置兩套映射矩陣,一套用于局部自注意力,另一套用于全局注意力。實驗表明,Longformer全面超越了RoBERTa的表現。
消融實驗
為了證明Longformer更優異的性能并不是因為對RoBERTa額外的預訓練帶來的,作者做了一組消融實驗。采用了與RoBERTa完全相同的序列長度和attention機制,在繼續預訓練后并沒有在文檔級任務上取得更優的性能。
賣萌屋成立了自然語言處理、搜索技術、推薦系統、算法崗求職等方向大群和知識圖譜與信息抽取、智能問答、對話系統等垂類討論群。不僅可以認識眾多志同道合的優秀小伙伴,而且還有若干賣萌屋美麗小姐姐(劃掉)、頂會審稿人、大廠研究院、知乎大V等你來撩哦。后臺回復關鍵詞【入群】即可獲取入口。
記得掃描下方二維碼關注并星標置頂,我才能來到你面前哦。
參考文獻
[1] Dilated CNN:CV中常用的圖像編碼方式,可拓寬模型編碼時的感受野。paper原文鏈接:https://arxiv.org/abs/1511.07122。另可參考知乎討論“如何理解空洞卷積(dilated convolution)?”,鏈接:https://www.zhihu.com/question/54149221。
[2] TVM: 關于如何自定義CUDA內核:這里作者使用了TVM (tensor virtual machine)(tvm.apache.org),2018年由華盛頓大學的SAMPL組貢獻的開源項目。TVM為不同的深度學習框架和硬件平臺實現了統一的編譯棧,從而實現將不同框架的深度學習模型到硬件平臺的快速部署。
總結
以上是生活随笔為你收集整理的Longformer:超越RoBERTa,为长文档而生的预训练模型的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 从我开发过的Tensorflow、飞桨、
- 下一篇: 深度学习如何均衡精度、内存、计算和通信开