作者?| kaiyuan? ?整理 |?NewBeeNLP

一篇來自ICLR2020，facebook人工智能實驗室的論文，關于改進Transformer推理速度的。論文地址：https://arxiv.org/abs/1910.10073

寫在前面

大模型在困難任務上表現非常好，而小模型也可以在比較簡單的任務上表現出色。但是目前大模型在應用簡單任務時的運算量同復雜任務一樣，這就導致了不必要的資源浪費，此外還會導致推理速度變慢。對此，提出了一種depth-adaptive transformer，可以在推理階段針對不同的token自動選擇在合適的層進行停止并輸出，在不損失模型效果的同時大大提高推理速度。

Anytime Prediction

傳統的transformer decoder，所有的token都需要經過N個block，且最后的分類器是唯一的。但是一些簡單的任務并不需要推理如此多步，可能在第 層模型已經得出結果了，再往上走并沒有意義，于是可以選擇在第 層就輸出，這稱為「動態計算(Dynamic Computation)」，

并且，每一層的分類器 可以選擇不同或者相同。接下去就是怎么去實現上面的想法，即對于每個輸入token，如何自動確定其在哪一層退出？文中提出了兩種策略：

Aligned training

對齊訓練會同時優化所有分類器，并假設當前狀態之前的所有狀態 均可以被看見，然后計算 項loss，如下圖(a)

其中， 是第 層所有 損失和， 是 的加權平均(論文附錄A實驗說明平均加權的效果最好)。

Mixed training

但是上述對齊訓練是不好實現的，因為不是所有token都在同一層輸出，那么當有的token在較早的層輸出后，在后續層計算時就無法看見它。解決這個問題一個直觀的想法就是將已經輸出的單元直接copy到上面一層即可，如下圖(b)。對輸入序列采樣 個退出序列 ，并對每個退出序列計算損失，

Adaptive Depth Estimation

自適應深度估計。上文還有一個問題待解決：如何采樣退出序列？提出了兩種策略：「sequence-specific」：對應于aligned training，所有token都在同一層退出；

「token-specific」：對應于mixed training，每個token可能在不同的層退出；

表示第 步對應的退出層參數分布，其中 表示 在第 層退出并得到預測 的概率， 表示真實分布，使用下式優化：

總體的損失函數為：

Sequence-Specific Depth

在該種模式下，所有token都在同一層退出，所以我們不考慮time-step 。通過將encoder輸出的平均值 送到多項式分類器，我們得到分布 ：

接下來怎么得到真實分布 呢？本文提出了兩種策略：「Likelihood-based」：基于每個塊之后整個序列的似然性，并且以Dirac delta為中心對退出進行了優化，從而使序列似然性最高，

并且還可以加入正則項鼓勵更低的退出層「Correctness-based」：基于似然的方式無法判斷最大的似然是不是正確的序列，因此基于正確性的方式選擇正確字符最大分數的最低的層。具體來說，對每個層，我們計算在該層上所有正確預測的字符數，然后選擇最大正確數的一層，并同樣加上正則項：

Token-Specific Depth

基于token的方式對每個token自動學習在哪一層退出，提出兩種exit distribution 的計算方式：「Multinomial」：多項分布， 表示第一層decoder的輸出「Poisson Binomial」：泊松二項分布， 表示停止概率(halting probability)，其對應有一個閾值 ，當 超過這個閾值時即退出否則繼續往上走直至結束

接下去也是對真實分布的計算，提出了三種策略：「Likelihood-based」：基于似然的方式在每一步選擇似然最大的層退出，

但是顯然，這種方法是貪心做法，我們還想要考慮未來的字符，所以本文考慮一個光滑化(smoothed)似然，它考慮附近的幾個字符：「Correctness-based」：基于正確性的方式，「Confidence thresholding」：基于自信度閾值，定義了一個自信度閾值 ，當當前輸出 的最大得分超過閾值時，即選擇退出。其中 是通過驗證集確定的。

上述幾種方法的示意圖：

一些結論Aligned training效果要優于Mixed training；

用對齊訓練方法得到的模型，在 的時候，差不多就能超過baseline模型(傳統Transformer)，而Possion方法大體上比多項分布好；

更多更詳細的實驗分析推薦閱讀reference里的作者PPT分享；

整體感覺和之前介紹過的Universal Transformer很像，都是采用dynamic computation的思想。但是很多細節還是不同，比如沒有大型共享的Transformer層，并且本文主要關注動態停止策略設計等等：There are a number of differences to universal transformer (UT): UT repeatedly applies the same layer for a number of steps, while as our approach applies different layers at every step. The dynamic computation in UT considers a single type of mechanism to estimate the number of steps (=network depth) while as we consider a number of mechanisms and supervision regimes. Moreover, dynamic computation was not actually used for the machine translation experiments in the UT paper (it was used for other tasks though), the authors used a fixed number of steps for every input instead.Open Review[1]DAT作者PPT分享[2]

本文參考資料[1]

Open Review: https://openreview.net/forum?id=SJg7KhVKPH[2]

DAT作者PPT分享: http://elbayadm.github.io/assets/talks/anytime/decore_04102019.pdf