引言本研究展示了一種新型Transformer的語言模型：Mixture-of-Depths Transformer，該模型能夠動態地分配計算資源到輸入序列的特定位置，而不是像傳統模型那樣均勻地分配計算資源。通過動態計算分配方式，可以在保持性能的同時顯著提高模型速度，可比isoFLOP最優基線模型快66%！ 背景介紹生活中，并非所有的問題都需要相同的時間來解決。同樣在語言模型中也是，并非所有Token和序列都需要相同的算力來進行預測。然而，Transformer模型在前向傳播中為每個Token花費了相同的計算量，針對這個問題，我們能否讓Transformer省去這些不必要的計算呢？ 條件計算（Conditional computation）是一種減少總計算量的技術，它只會在需要時才會進行計算。何時需要計算、需要多少的計算量，目前已經有了多種解決方案。然而，這些算法在現有硬件下并不一定適用，因為它們傾向于引入動態計算圖，而現有硬件更傾向于使用的靜態計算圖。 為了克服這一挑戰，本文作者考慮在靜態計算預算下進行語言建模，并且該靜態預算可以小于普通Transformer所需的計算預算。其中，對于每一層的Token決策，神經網絡必須要學會如何動態分配計算資源。在實現過程中，總計算量是由用戶定義的，并且在訓練前是不變的，而不是網絡動態決策的函數。因此，硬件使用效率的提升可以根據占用內存的減少、每次前向傳播FLOP的減少進行提前預測。 本文使用了一個類似專家混合（MoE）Transformer的方法，其中動態Token級別路由決策是在整個網絡深度上做出的。不同于MoE，本文選擇對Token應用計算（類似標準Transformer）、或者通過殘差連接進行前向傳播。此外，本文還將這種路由方法同時應用于前向多層感知機（MLPs）和多頭注意力機制。因此，這也影響了對keys和query的處理，路由不僅決定更新哪些Token，還決定了哪些Token用于注意力機制。本文將這種策略稱為深度混合（Mixture-of-Depths, MoD）。 MoD實現MoD Transformers方法是通過設置一個靜態的計算預算，這個預置計算量比傳統Transformer模型要小。這個計算預算通過限制在給定層中可以參與自注Token意力和多層感知機（MLP）計算的Token數量來實現。這種方法使用了一個每層的路由器來決定哪些Token應該參與計算，哪些應該通過殘差連接繞過計算，從而節省計算資源。  具體實現方式： 「1.定義計算預算」 即通過限制序列中可以參與計算的Token數量來強制執行總體計算預算。為了在Transformer模型中有效控制計算資源，作者通過“容量”概念來限制每次計算的輸入Token數量。傳統Transformer的自注意力和MLP使用全部Token，而MoE Transformer則為每個專家分配較少的Token，以平衡計算負載。 計算資源的分配取決于Token容量，即使在條件計算中也是如此。通過減少計算容量，可以降低每次前向傳播的計算需求，但如果處理不當，可能會影響模型性能。作者認為，不是所有Token都需要同等程度的處理，因此模型可以通過學習來識別哪些Token更重要。這樣，網絡可以在保持性能的同時，更高效地使用計算資源。 「2.圍繞Transformer塊的路由」 路由機制可以讓模型決定哪些數據需要進行密集計算，哪些可以跳過。這是通過為序列中的每個數據項分配一個權重來實現的，權重高的數據項會參與完整的計算過程，而權重低的則通過一個簡單的跳過步驟，以節省計算資源。這種動態選擇的方法使得模型在保持處理質量的同時，能夠更高效地運行。通過調整這個機制，模型可以在速度和性能之間找到最佳平衡。 「3.路由方案」 用來決定哪些數據項參與復雜計算，哪些可以簡化處理的策略。主要有兩種方案：1）基于Token的路由：每個數據項根據偏好選擇參與計算的路徑，但可能導致處理不均衡。2）基于專家的路由：每個計算路徑選擇一定數量的數據項，保證處理均衡，但可能使某些數據項被過度或不足處理。 「4.采樣」 在MoD Transformer模型的自回歸采樣階段，面臨著如何在不依賴未來Token信息的情況下進行有效路由決策的挑戰。為了應對這一挑戰，文中提出了兩種策略。第一種是引入輔助損失，通過二元交叉熵損失函數調整路由器輸出，使得模型能夠基于當前和過去的Token信息做出因果路由決策。 第二種策略是使用一個輔助預測器，它作為一個小型的輔助網絡，預測每個Token是否應該參與計算，從而在采樣過程中提供必要的路由信息。這兩種方法都避免了對未來Token的依賴，確保了模型在序列生成時的高性能和效率。 「5.模型訓練」所有模型都使用相同的基本超參數配置（例如，128batch、2048 序列長度）。 實驗結果「速度提升」 下圖展示了MoD超參數微調結果，其中包括不同模型變體的性能比較，以及學習曲線，說明了模型在保持相同性能的同時，速度比isoFLOP最優基線模型快66%。 「isoFLOP分析」 如下圖所示，存在一些MoD變體在步驟速度上比isoFLOP最優基線模型更快，同時實現更低的訓練損失。這些結果表明MoD模型在保持性能的同時，能夠實現更高的計算效率 AINLPer的星球投稿或尋求報道聯系：ainlperbot 點擊下方鏈接??關注我們 |
|
|
本文最終選擇了基于專家的路由方案，因為它可以更有效地平衡計算資源，并且簡化了實施過程，如上圖所示。通過這種方法，模型能夠在保持性能的同時減少計算量，提高運行效率。
