DeepMind升級Transformer，前向通過FLOPs最多可降一半

機器之心報道

編輯：Panda W

Transformer 的重要性無需多言，目前也有很多研究團隊致力於改進這種變革性技術，其中一個重要的改進方向是提升 Transformer 的效率，比如讓其具備自適應計算能力，從而可以節省下不必要的計算。

正如不久前 Transformer 架構的提出之一、NEAR Protocol 聯合創始人 Illiya Polosukhin 在與黃仁勳的對話中說到的那樣：「自適應計算是接下來必須出現的。我們要關注，在特定問題上具體要花費多少計算資源。」

其實人類就天生具備自適應計算的能力 —— 人在解決各種不同的問題時，會自然地分配不同的時間和精力。

語言建模也應如此，爲了得到準確的預測結果，並不需要爲所有 token 和序列都投入同樣的時間或資源。但是，Transformer 模型在一次前向傳播中卻會爲每個 token 花費同等的計算量。這不禁讓人哀嘆：大部分計算都被浪費了！

理想情況下，如果可以不執行非必要的計算，就可以降低 Transformer 的計算預算。

條件式計算這種技術可在需要執行計算時才執行計算，由此可以減少總計算量。之前許多研究者已經提出了多種可以評估何時執行計算以及使用多少計算量的算法。

但是，對於這個頗具挑戰性的問題，普遍使用的解決形式可能無法很好地應對現有的硬件限制，因爲它們往往會引入動態計算圖。最有潛力的條件式計算方法反而可能是那些能協調使用當前硬件棧的方法，其會優先使用靜態計算圖和已知的張量大小（基於對硬件的最大利用而選取這個張量大小）。

近日，Google DeepMind 研究了這個問題，他們希望使用更低的計算預算來縮減 Transformer 使用的計算量。

他們設想：在每一層中，網絡必須學會爲每個 token 做決策，從而動態地分配可用計算預算。在他們的具體實現中，總計算量由用戶在訓練之前設定並且不再更改，而非網絡工作時執行決策的函數。這樣一來，便可以提前預知並利用硬件效率收益（比如內存足跡減少量或每次前向傳播的 FLOPs 減少量）。該團隊的實驗表明：可以在不損害網絡整體性能的前提下獲得這些收益。

DeepMind 的這個團隊採用了類似於混合專家（MoE） Transformer 的方法，其中會在整個網絡深度上執行動態 token 層面的路由決策。

而與 MoE 不同的是，這裡他們的選擇是：要麼是將計算應用於 token（和標準 Transformer 一樣），要麼就是通過一個殘差連接繞過它（保持不變，節省計算）。另一個與 MoE 的不同之處是：這裡是將這種路由機制同時用在 MLP 和多頭注意力上。因此，這也會影響網絡處理的鍵值和查詢，因此該路由不僅要決定更新哪些 token，還要決定哪些 token 可供關注。

DeepMind 將這一策略命名爲 Mixture-of-Depths（MoD），以突顯這一事實：各個 token 在 Transformer 深度上通過不同數量的層或模塊。我們這裡將其翻譯成「混合深度」，見圖 1。

MoD 支持使用者權衡考量性能與速度。一方面，使用者可以使用與常規 Transformer 同等的訓練 FLOPs 來訓練 MoD Transformer，這可爲最終的對數概率訓練目標帶來多達 1.5% 的提升。另一方面，MoD Transformer 使用更少的計算量就能達到與常規 Transformer 同樣的訓練損失 —— 每一次前向傳播的 FLOPs 可少最多 50%。

這些結果表明，MoD Transformer 可以學習智能地路由（即跳過不必要的計算）。

實現混合深度（MoD）Transformer

概況來說，其策略如下：

路由方案

該團隊考慮了兩種學習到的路由方案（見圖 2）：token 選擇型和專家選擇型。

在 token 選擇型路由方案中，路由算法會跨計算路徑（比如跨 MoE Transformer 中的專家身份）生成針對每個 token 的概率分佈。然後 token 會被傳送到它們偏好的路徑（即概率最高的路徑），而輔助損失可以確保所有 token 不會收斂到同一路徑。token 選擇型路由可能會有負載平衡問題，因爲不能確保 token 在可能的路徑之間劃分適當。

專家選擇型路由則是將上述方案反過來：不是讓 token 選擇它們偏好的路徑，而是讓每條路徑基於 token 偏好選擇前 k 個 token（top-k）。這能確保負載完美平衡，因爲每條路徑總是保證 k 個 token。但是，這也可能導致某些 token 被過處理或欠處理，因爲某些 token 可能是多條路徑的前 k 名，另一些 token 則可能沒有相應路徑。

DeepMind 的選擇是使用專家選擇型路由，原因有三。

第一，它無需輔助性的平衡損失。

第二，由於選取前 k 名這一操作取決於路由權重的幅度，因此該路由方案允許使用相對路由權重，這有助於確定當前模塊計算最需要哪些 token；路由算法可以通過適當地設定權重來盡力確保最關鍵的 token 是在前 k 名之中 —— 這是 token 選擇型路由方案無法做到的。在具體的用例中，有一條計算路徑本質上是 null 操作，因此應該避免將重要 token 路由到 null。

第三，由於路由只會經由兩條路徑，因此單次 top-k 操作就能高效地將 token 分成兩個互斥的集合（每條計算路徑一個集合），這能應對上面提及的過處理或欠處理問題。

該路由方案的具體實現請參看原論文。

採樣

儘管專家選擇型路由有很多優點，但它也有一個很明顯的問題：top-k 操作是非因果式的。也就是說，一個給定 token 的路由權重是否在前 k 名取決於其之後的路由權重的值，但在執行自迴歸採樣時，我們無法獲得這些權重。

爲了解決這個問題，該團隊測試了兩種方法。

第一種是引入一個簡單的輔助損失；實踐證明，其對語言建模主目標的影響程度爲 0.2%− 0.3%，但卻能夠讓模型自迴歸地採樣。他們使用了一個二元交叉熵損失，其中路由算法的輸出提供 logit，通過選取這些 logit 中的 top-k，就能提供目標（即，如果一個 token 在 top-k 中，就爲 1，否則爲 0）。

第二種方法是引入一個小的輔助 MLP 預測器（就像是又一個路由算法），其輸入與路由算法的一樣（具有 stop gradient），但其輸出是一個預測結果：token 是否在序列的 top-k 中。該方法不會影響語言建模目標，實驗表明也不會顯著影響該步驟的速度。

有了這些新方法，就可以通過選擇路由到的 token 來執行自迴歸採樣，也可以根據路由算法的輸出繞過一個模塊，這無需依賴任何未來 token 的信息。實驗結果表明，這是一種相對簡單輔助任務，可以很快實現 99% 的準確度。

結果

訓練，isoFLOP 比較

首先，該團隊訓練了一些 FLOP 預算相對較小（6e18）的模型，以確定最優的超參數（見下圖 3）。

總體而言，可以看到 MoD Transformer 會將基準 isoFLOP 曲線向右下方拖動。也就是說，最優的 MoD Transformer 的損失比最優的基準模型更低，同時參數也更多。這種效應帶來了一個幸運的結果：存在一些和最優基準模型表現一樣好甚至更好的 MoD 模型（同時步驟速度更快），儘管它們本身在其超參數設置下並不是 isoFLOP 最優的。舉個例子，一個 220M 參數量的 MoD 變體（圖 3 中的 3 號模型）稍優於 isoFLOP 最優基準模型（參數量也是 220M，圖 3 中的 1 號模型），但這個 MoD 變體在訓練期間的步驟速度快了 60% 以上。

下圖 4 給出了總 FLOPs 爲 6e18、2e19 和 1e20 時的 isoFLOP 分析。可以看到，當 FLOP 預算更大時，趨勢依然繼續。

下圖 5 給出了一個使用交織的路由模塊訓練的 MoD Transformer 的路由決策。儘管其中存在大量繞過模塊的情況，但這個 MoD Transformer 依然能實現優於常規 Transformer 的性能。

自迴歸評估

他們也評估了 MoD 變體的自迴歸採樣表現，結果見下圖 6。這些結果表明 MoD Transformer 所帶來的計算節省不僅僅侷限於訓練設置。

混合深度與專家（MoDE）

MoD 技術可以自然地與 MoE 模型整合起來，組成所謂的 MoDE 模型。下圖 7 展示了 MoDE 及其帶來的提升。

MoDE 有兩種變體：分階段 MoDE 和集成式 MoDE。

其中分階段 MoDE 是在自注意力步驟之前進行路由繞過或到達 token 的操作；而集成式 MoDE 則是通過在常規 MLP 專家之間集成「無操作」專家來實現 MoD 路由。前者的優勢是允許 token 跳過自注意力步驟，而後者的好處在於其路由機制很簡單。

該團隊注意到，以集成方式實現 MoDE 明顯優於直接降低專家的能力、依靠丟棄 token 來實現殘差路由的設計。