模型壓縮95%，MIT韓松等人提出新型Lite Transformer

參與：小舟、魔王

Transformer 的高效能依賴於極高的算力，這讓移動端 NLP 嚴重受限。在不久之前的 ICLR 2020 論文中，MIT 與上海交大的研究人員提出了一種高效的移動端 NLP 架構 Lite Transformer，向在邊緣裝置上部署移動級 NLP 應用邁進了一大步。

雖然推出還不到 3 年，Transformer 已成為自然語言處理（NLP）領域裡不可或缺的一環。然而這樣流行的演算法卻需要極高的算力才能實現足夠的效能，這對於受到算力和電池嚴格限制的移動端來說有些力不從心。

在 MIT 最近的研究《Lite Transformer with Long-Short Range Attention》中，MIT 與上海交大的研究人員提出了一種高效的移動端 NLP 架構 Lite Transformer，向在邊緣裝置上部署移動級 NLP 應用邁進了一大步。該論文已被人工智慧頂會 ICLR 2020 收錄。

該研究是由 MIT 電氣工程和計算機科學系助理教授韓松領導的。韓松的研究廣泛涉足深度學習和計算機體系結構，他提出的 Deep Compression 模型壓縮技術曾獲得 ICLR2016 最佳論文，論文 ESE 稀疏神經網路推理引擎 2017 年曾獲得晶片領域頂級會議——FPGA 最佳論文獎，引領了世界深度學習加速研究，對業界影響深遠。

論文地址：
https://arxiv.org/abs/2004.11886v1

GitHub 地址：
https://github.com/mit-han-lab/lite-transformer

Transformer 在自然語言處理任務（如機器翻譯、問答）中應用廣泛，但它需要大量計算去實現高效能，而這不適合受限於硬體資源和電池嚴格限制的移動應用。

這項研究提出了一種高效的移動端 NLP 架構——Lite Transformer，它有助於在邊緣裝置上部署移動 NLP 應用。其核心是長短距離注意力（Long-Short Range Attention，LSRA），其中一組注意力頭（透過卷積）負責區域性上下文建模，而另一組則（依靠注意力）執行長距離關係建模。

這樣的專門化配置使得模型在三個語言任務上都比原版 transformer 有所提升，這三個任務分別是機器翻譯、文字摘要和語言建模。

在資源有限的情況下（500M/100M MACs），Lite Transformer 在 WMT’14 英法資料集上的 BLEU 值比分別比 transformer 高 1.2/1.7。Lite Transformer 比 transformer base 模型的計算量減少了 60%，而 BLEU 分數卻只降低了 0.3。結合剪枝和量化技術，研究者進一步將 Lite Transformer 模型的大小壓縮到原來的 5%。

對於語言建模任務，在大約 500M MACs 上，Lite Transformer 比 transformer 的困惑度低 1.8。

值得注意的是，對於移動 NLP 設定，Lite Transformer 的 BLEU 值比基於 AutoML 的 Evolved Transformer 高 0.5，而且它不需要使用成本高昂的架構搜尋。

從 Lite Transformer 與 Evolved Transformer、原版 transformer 的比較結果中可以看出，Lite Transformer 的效能更佳，搜尋成本相比 Evolved Transformer 大大減少。

那麼，Lite Transformer 為何能夠實現高效能和低成本呢？接下來我們來了解其核心思想。

長短距離注意力（LSRA）

NLP 領域的研究人員試圖理解被注意力捕捉到的上下文。Kovaleva 等人 (2019) 和 Clark 等人 (2020) 對 BERT 不同層的注意力權重進行了視覺化。

如下圖 3b 所示，權重 w 表示源句單詞與目標句單詞之間的關係（自注意力也是如此）。隨著權重 w_ij 的增加（顏色加深），源句中的第 i 個詞更加註意目標句中的第 j 個詞。注意力圖通常有很強的模式化特徵：稀疏和對角線。它們代表了一些特定單詞之間的關係：稀疏表示長距離資訊間的關係，對角線表示近距離資訊間的關係。研究者將前者稱為「全域性」關係，將後者稱為「區域性」關係。

圖 3：Lite Transformer 架構 (a) 和注意力權重的視覺化。傳統的注意力 (b) 過於強調區域性關係建模（參見對角線結構）。該研究使用卷積層專門處理區域性特徵提取工作，以高效建模區域性資訊，從而使注意力分支可以專門進行全域性特徵提取 (c)。

在翻譯任務中，注意力模組必須捕獲全域性和區域性上下文，這需要很大的容量。與專門化的設計相比，這並非最佳選擇。以硬體設計為例，CPU 等通用硬體的效率比 FPGA 等專用硬體低。研究者認為應該分別捕捉全域性和區域性上下文。模型容量較大時，可以容忍冗餘，甚至可以提供更好的效能。但是在移動應用上，由於計算和功率的限制，模型應該更加高效。因此，更需要專門化的上下文捕獲。

為了解決該問題，該研究提出一個更專門化的架構，即長短距離注意力（LSRA），而不是使用處理 “一般” 資訊的模組。該架構分別捕獲區域性和全域性上下文。

如圖 3a 所示，LSRA 模組遵循兩分支設計。左側注意力分支負責捕獲全域性上下文，右側卷積分支則建模區域性上下文。研究者沒有將整個輸入饋送到兩個分支，而是將其沿通道維度分為兩部分，然後由後面的 FFN 層進行混合。這種做法將整體計算量減少了 50%。

左側分支是正常的注意力模組（Vaswani et al. (2017)），不過通道維度減少了一半。至於處理區域性關係的右分支，一個自然的想法是對序列應用卷積。使用滑動視窗，模組可以輕鬆地覆蓋對角線組。為了進一步減少計算量，研究者將普通卷積替換為輕量級的版本，該版本由線性層和深度卷積組成。透過這種方式，研究者將注意力模組和卷積模組並排放置，引導它們對句子進行全域性和區域性的不同角度處理，從而使架構從這種專門化設定中受益，並實現更高的效率。

實驗設定

資料集和評估

研究者在機器翻譯、文字摘要和語言建模三個任務上進行了實驗和評估。

具體而言，機器翻譯任務使用了三個基準資料集：IWSLT’14 德語 - 英語 (De-En)、WMT 英語 - 德語 (En-De)、WMT 英語 - 法語(En-Fr)。

文字摘要任務使用的是 CNN-DailyMail 資料集。

語言建模任務則在 WIKITEXT-103 資料集上進行。

架構

模型架構是基於序列到序列學習的編碼器 - 解碼器。在機器翻譯任務中，針對 WMT 資料集，基線模型基於 Vaswani 等人提出的模型。對於 IWSLT 資料集，基線模型遵循 Wu 等人的設定。對於文字摘要任務，研究者採用了與 WMT 相同的模型。至於語言建模任務，模型與 Baevski & Auli (2019) 一致，但模型尺寸較小。

該研究提出的架構首先將 transformer base 模型中的 bottleneck 拉平，然後用 LSRA 替換自注意力。更具體地說，是使用兩個專門的模組，一個注意力分支和一個卷積分支。

實驗結果

機器翻譯

表 1 展示了 Lite Transformer 在 IWSLT’14 De-En 資料集上的定量結果，並與 transformer 基線方法和 LightConv 做了對比。在大約 100M Mult-Adds 時，Lite Transformer 模型的 BLEU 值比 transformer 高出 1.6。

表 1：IWSLT’14 De-En 上的結果。

表 2 中的定量結果表明，在 100M Mult-Adds 設定下，Lite Transformer 在 WMT En-De 資料集和 WMT En-Fr 資料集上的 BLEU 值分別比 Transformer 高出 1.2 和 1.7，在 300M Mult-Adds 設定下，也有 0.5 和 1.5 分的提升。

表 2：在 WMT’14 En-De 和 WMT’14 En-Fr 上的結果。

研究者還提供了模型在 WMT En-Fr 上的權衡曲線，如圖 4a 所示，Lite Transformer 一直優於原版 transformer。

圖 4：在 WMT En-Fr 資料集上的機器翻譯權衡曲線，以及在 WIKITEXT-103 資料集上的語言建模權衡曲線。兩個曲線都說明了在移動設定下，Lite Transformer 比 transformer 效能更佳（藍色區域）。

與自動化設計模型的對比

與基於 AutoML 的 Evolved Transformer（ET）相比，Lite Transformer 在移動設定中也有明顯的改進。此外，在 100M 和 300M 的 Mult-Adds 下，Lite Transformer 的 BLEU 值分別比 ET 高 0.5 和 0.2，詳見表 3。

表 3：不同 NMT 模型的效能和訓練成本。