為記憶體塞不下Transformer犯愁？OpenAI應用AI研究負責人寫了份指南

本文是一篇綜述性的部落格，探討總結當下常用的大型 transformer 效率最佳化方案。

大型 Transformer 模型如今已經成為主流，為各種任務創造了 SOTA 結果。誠然這些模型很強大，但訓練和使用起來代價非常昂貴。在時間和記憶體方面存在有極高的推理成本。概括來說，使用大型 Transformer 模型進行推理的難點，除了模型的規模不斷擴大外，還有兩個不可忽略的地方：

記憶體消耗大：推理時，需要把模型參數和中間狀態都保存到記憶體中。例如：KV 儲存機制下的快取中的內容在解碼期間需要儲存在記憶體中，舉例來說，對於 batch size 為 512，上下文長度為 2048 的設置來說，KV 快取裡需要的空間規模為 3TB，這是模型大小的 3 倍；注意力機制的推理成本和輸入序列的長度呈正相關；
低並行性：推理生成過程以自迴歸的方式執行，使解碼過程難以並行。

在這篇文章中，領導 OpenAI 應用研究的 Lilian Weng 寫了一篇部落格，文中介紹了幾種提高 transformer 推理效率的方法。一些是通用的網路壓縮方法，而另一些則應用於特定的 transformer 架構。

Lilian Weng 現為 OpenAI 應用人工智慧研究負責人，主要從事機器學習、深度學習等研究。她本科畢業於香港大學，碩士畢業於北京大學資訊系統與電腦科學系，之後前往印度安納大學布魯頓分校攻讀博士。

模型綜述

通常將以下內容視為模型推理最佳化的目標：

使用更少的 GPU 設備和更少的 GPU 記憶體，減少模型的記憶體佔用；
減少所需的 FLOP，降低計算複雜度；
減少推理延遲，運行得更快。

可以使用幾種方法來降低推理過程在記憶體中的成本，並且加快速度。

在多 GPU 上應用各種並行機制來實現對模型的擴展。模型元件和資料的智慧並行使得運行具有萬億級參數的大模型成為可能；
將暫時未使用的資料卸載到 CPU，並在以後需要時讀回。這樣做對記憶體比較友好，但會導致更高的延遲；
智慧批處理策略；例如 EffectiveTransformer 將連續的序列打包在一起，以刪除單個批次中的 padding；
神經網路壓縮技術，例如剪枝、量化、蒸餾。就參數數量或位寬而言，小尺寸的模型應該需要少量的記憶體，也就運行得更快；
特定於目標模型架構的改進。許多架構上的變化，尤其是注意力層的變化，有助於提高 transformer 的解碼速度。

本篇文章的重點是網路壓縮技術和 transformer 模型在特定體系結構下的改進。

量化策略

在深度神經網路上應用量化策略有兩種常見的方法：

訓練後量化（PTQ）：首先需要模型訓練至收斂，然後將其權重的精度降低。與訓練過程相比，量化操作起來往往代價小得多；
量化感知訓練 (QAT)：在預訓練或進一步微調期間應用量化。QAT 能夠獲得更好的性能，但需要額外的計算資源，還需要使用具有代表性的訓練資料。

值得注意的是，理論上的最優量化策略與實際在硬體核心上的表現存在著客觀的差距。由於 GPU 核心對某些類型的矩陣乘法（例如 INT4 x FP16）缺乏支持，並非下面所有的方法都會加速實際的推理過程。

Transformer 量化挑戰

許多關於 Transformer 模型量化的研究都有相同的觀察結果：訓練後將參數簡單地量化為低精度（例如 8 位）會導致性能顯著下降，這主要是由於普通的激活函數量化策略無法覆蓋全部的取值區間。

圖 1. 只將模型權重量化為 8 位，激活函數使用完整的精度的時候能取得較好的效果（W8A32）；激活函數量化為 8 位時，無論權重是否為低精度（W8A8 和 W32A8）效果都不如 W8A32。

Bondarenko 等人在一個小型 BERT 模型中觀察到，由於輸出張量中存在強異常值，FFN 的輸入和輸出具有非常不同的取值區間。因此，FFN 殘差和的逐個張量的量化可能會導致顯著的誤差。

隨著模型參數規模繼續增長到數十億的級別，高量級的離群特徵開始在所有 transformer 層中出現，導致簡單的低位量化效果不佳。Dettmers 等人觀察到大於 6.7B 參數的 OPT 模型就會出現這種現象。模型大了，有極端離群值的網路層也會變多，這些離群值特徵對模型的性能有很大的影響。在幾個維度上的激活函數異常值的規模就可以比其他大部分數值大 100 倍左右。

圖 2. 不同規模的 OPT 模型在四個語言任務（WinoGrande、HellaSwag、PIQA、LAMBADA）上的平均零樣本準確率。

混合精度量化

解決上述量化挑戰的最直接方法是以不同的精度對權重和激活函數進行量化。

GOBO 模型是首批將訓練後量化應用於 transformer 的模型之一（即小型 BERT 模型）。GOBO 假設每一層的模型權重服從高斯分佈，因此可以通過跟蹤每層的均值和標準差來檢測異常值。異常值特徵保持原始形式，而其他值被分到多個 bin 中，並且僅儲存相應的權重索引和質心值。

基於對 BERT 中只有某些激活層（例如 FFN 之後的殘差連接）導致性能大幅下降現象的觀察，Bondarenko 等人通過在有問題的激活函數上使用 16 位量化而在其他激活函數上使用 8 位來採用混合精度量化。

LLM.int8 () 中的混合精度量化是通過兩個混合精度分解實現的：

因為矩陣乘法包含一組行和列向量之間的獨立內積，所以可以對每個內積進行獨立量化。每一行和每一列都按最大值進行縮放，然後量化為 INT8；
異常值激活特徵（例如比其他維度大 20 倍）仍保留在 FP16 中，但它們只佔總權重的極小部分，不過需要經驗性地識別離群值。

圖 3.LLM.int8（）兩種混合精度分解方法。

細粒度量化

圖 4. 不同粒度量化對比。d 是模型大小 / 隱空間維度，h 是一個 MHSA（多頭自注意力）元件中的頭數。

簡單地量化一層中的整個權重矩陣（逐個張量或逐個層量化）是最容易實現的，但量化粒度往往不盡如人意。

Q-BERT 將分組量化應用於微調的 BERT 模型，將 MHSA（多頭自注意力）中每個頭的單個矩陣 W 視為一個組，然後應用基於 Hessian 矩陣的混合精度量化。

Per-embedding group (PEG) 激活函數量化的設計動機是觀察到離群值僅出現在少數幾個維度中。對每個嵌入層都量化的代價非常昂貴，相比之下，PEG 量化將激活張量沿嵌入維度分成幾個大小均勻的組，其中同一組中的元素共享量化參數。為確保所有異常值都分組在一起，PEG 應用了一種基於取值範圍的嵌入維度排列演算法，其中維度按其取值範圍排序。

ZeroQuant 與 Q-BERT 一樣都對權重使用分組量化，然後還對激活函數使用了 token-wise 量化策略。為了避免代價昂貴的量化和反量化計算，ZeroQuant 構建了獨特的核心來將量化操作與其之前的運算子融合。

使用二階資訊量化

Q-BERT 針對混合精度量化開發了 Hessian AWare 量化 (HAWQ)。其動機是，具有更高 Hessian 譜的參數對量化更敏感，因此需要更高的精度。這種方法本質上是一種識別異常值的方法。

從另一個角度來看，量化問題是一個最佳化問題。給定一個權重矩陣 W 和一個輸入矩陣 X ，想要找到一個量化的權重矩陣 W^ 來最小化如下所示的 MSE 損失：

GPTQ 將權重矩陣 W 視為行向量 w 的集合，並對每一行獨立量化。GPTQ 使用貪心策略來選擇需要量化的權重，並迭代地進行量化，來最小化量化誤差。更新被選定的權重會生成 Hessian 矩陣形式的閉合解。GPTQ 可以將 OPT-175B 中的權重位寬減少到 3 或 4 位，還不會造成太大的性能損失，但它僅適用於模型權重而不適用於激活函數。

異常值平滑

眾所周知，Transformer 模型中激活函數比權重更難量化。SmoothQuant 提出了一種智慧解決方案，通過數學等效變換將異常值特徵從激活函數平滑到權重，然後對權重和激活函數進行量化 (W8A8)。正因為如此，SmoothQuant 具有比混合精度量化更好的硬體效率。

圖 5. SmoothQuant 將尺度方差從激活函數遷移到離線權重，以降低激活函數量化的難度。由此產生的新權重和激活矩陣都易於量化。

基於每個通道的平滑因子 s，SmoothQuant 根據以下公式縮放權重：

根據平滑因子

可以很容易地在離線狀態下融合到前一層的參數中。超參數 α 控制從激活函數遷移到權重的程度。該研究發現 α=0.5 是實驗中許多 LLM 的最佳取值。對於激活異常值較大的模型，可以將 α 調大。

量化感知訓練 (QAT)

量化感知訓練將量化操作融合到預訓練或微調過程中。這種方法會直接學習低位表示的模型權重，並以額外的訓練時間和計算為代價獲得更好的性能。

最直接的方法是在與預訓練資料集相同或代表預訓練資料集的訓練資料集上量化後微調模型。訓練目標可以與預訓練目標相同（例如通用語言模型訓練中的 NLL/MLM）或特定於的下游任務（例如用於分類的交叉熵）。

另一種方法是將全精度模型視為教師模型，將低精度模型視為學生模型，然後使用蒸餾損失最佳化低精度模型。蒸餾通常不需要使用原始資料集。

剪枝

網路剪枝是在保留模型容量的情況下，通過修剪不重要的模型權重或連接來減小模型大小。剪枝可能需要也可能不需要重新訓練。剪枝可以是非結構化的也可以是結構化的。

非結構化剪枝允許丟棄任何權重或連接，因此它不保留原始網路架構。非結構化剪枝通常對硬體要求比較苛刻，並且不會加速實際的推理過程；
結構化剪枝不改變權重矩陣本身的稀疏程度，可能需要遵循某些模式限制才能使用硬體核心支持的內容。本文專注於那些能實現 transformer 模型的高稀疏性的結構化剪枝。

構建剪枝網路的常規工作流程包含三個步驟：

1. 訓練密集型的神經網路直到收斂；

2. 修剪網路以去除不需要的結構；

3. （可選擇）重新訓練網路，讓新權重保持之前的訓練效果。

通過剪枝在密集模型中發現稀疏結構，同時稀疏網路仍然可以保持相似性能的靈感是由彩票假設激發的：這是一個隨機初始化的密集前饋網路，它包含一個子網路池。其中只有一個子集（稀疏網路）是中獎彩票（winning tickets），這個中獎彩票在獨立訓練時可以達到最佳性能。

如何剪枝

Magnitude pruning 是最簡單但同時又非常有效的剪枝方法 – 只裁剪那些絕對值最小的權重。事實上，一些研究發現，簡單的量級剪枝方法可以獲得與複雜剪枝方法相當或更好的結果，例如變分 dropout 和 l_0 正則化。Magnitude pruning 很容易應用於大型模型，並在相當大的超參數範圍內實現相當一致的性能。

Zhu & Gupta 發現，大型稀疏模型能夠比小型但密集的模型獲得更好的性能。他們提出了 Gradual Magnitude Pruning (GMP) 演算法，該演算法在訓練過程中逐漸增加網路的稀疏性。在每個訓練步驟中，具有最小絕對值的權重被遮蔽為零以達到所需的稀疏度並且遮蔽的權重在反向傳播期間不會得到梯度更新。所需的稀疏度隨著訓練步驟的增加而增加。GMP 過程對學習率步長策略很敏感，學習率步長應高於密集網路訓練中所使用的，但不能太高以防止收斂。

迭代剪枝多次迭代上述三個步驟中的第 2 步（剪枝）和第 3 步（重新訓練），每次只有一小部分權重被剪枝，並且在每次迭代中重新訓練模型。不斷重複該過程，直到達到所需的稀疏度級別。

如何再訓練

再訓練可以通過使用相同的預訓練資料或其他特定於任務的資料集進行簡單的微調來實現。

Lottery Ticket Hypothesis 提出了一種權重 rewinding 再訓練方法：剪枝後，將未剪枝的權重重新初始化回訓練初期的原始值，然後以相同的學習率時間表進行再訓練。

學習率 rewinding 僅將學習率重置回其早期值，而保持未剪枝的權重自最後一個訓練階段結束以來不變。研究者觀察到 (1) 使用權重 rewinding 的再訓練結果優於通過跨網路和資料集進行微調的再訓練，以及 (2) 在所有測試場景中學習率 rewinding 與權重 rewinding 的效果持平甚至更優。

稀疏化

稀疏化是擴大模型容量同時保持模型推理計算效率的有效方法。本文考慮兩種類型的 transformer 稀疏性：

稀疏化的全連接層，包括自注意力層和 FFN 層；
稀疏模型架構，即 MoE 元件的合併操作。

通過剪枝實現的 N:M 稀疏化

N:M 稀疏化是一種結構化的稀疏化模式，適用於現代 GPU 硬體最佳化，其中每 M 個連續元素中的 N 個元素為零。例如，英偉達 A100 GPU 的稀疏張量核心支持 2:4 稀疏度以加快推理速度。

圖 6. 2:4 結構化稀疏矩陣及其壓縮表示。

為了使密集型神經網路的稀疏化遵循 N:M 結構化稀疏模式，英偉達建議使用三步操作來訓練剪枝後的網路：訓練 –> 剪枝以滿足 2:4 稀疏性 –> 重新訓練。

(1) 對矩陣中的列進行排列可以在剪枝過程中提供更多可能，以保持參數的數量或滿足特殊限制，如 N:M 稀疏性。只要兩個矩陣對應的軸按相同的順序排列，矩陣乘法的結果就不會改變。例如，(1) 在自注意力模組中，如果 query 的嵌入矩陣 Q 的軸 1 和 key 嵌入矩陣 K^⊤的軸 0 採用相同的排列順序，則 QK^⊤的矩陣乘法最終結果保持不變。

圖 7. Q（軸 1）和 K^⊤（軸 0）上相同排列，自注意力模組的結果不變。

(2) 在包含兩個 MLP 層和一個 ReLU 非線性層的 FFN 層內，可以將第一個線性權重矩陣 W_1 沿軸 1 排列，然後第二個線性權重矩陣 W_2 沿軸 0 按相同順序排列。

圖 8. W_1（軸 1）和 W_2（軸 0）上有著相同的排列，可以保持 FFN 層的輸出不變。為簡單起見，圖示省略了偏差項，但也應對它們應用相同的排列。

為了推動 N:M 結構稀疏化，需要將一個矩陣的列拆分為 M 列的多個 slide（也稱為 stripe），這樣可以很容易地觀察到每個 stripe 中的列順序和 stripe 的順序對 N:M 稀疏化產生的限制。

Pool 和 Yu 提出了一種迭代式的貪心演算法來尋找最優排列，使 N:M 稀疏化的權重幅度最大化。所有通道對都被推測性地交換，並且只採用幅度增加最大的交換，然後生成新的排列並結束單次迭代。貪心演算法可能只會找到局部極小值，因此他們引入了兩種技術來逃避局部極小值：

1. 有界迴歸：在實踐中，兩個隨機通道的最大交換次數是固定的。每次搜尋只有一個通道可以進行交換，以保持搜尋空間寬而淺；

2. 窄且深的搜尋：選擇多個 stripe 並同時最佳化它們。

圖 9. 貪心演算法實現迭代地尋找 N:M 稀疏化最佳排列的演算法。

與按默認通道順序對網路進行剪枝相比，如果在剪枝之前對網路進行置換，可以獲得更好的性能。

為了從頭開始訓練具有 N:M 稀疏化的模型，Zhou & Ma 擴展了常用於模型量化中的反向傳播更新的 STE，用於幅度剪枝和稀疏參數更新。

STE 計算剪枝後的網路

的密集參數的梯度

，並將其作為近似值應用於稠密網路 W：

STE 的擴展版本 SR-STE（稀疏精化 STE）通過以下方式更新稠密權重 W：

其中

是

的掩碼矩陣，⊙是元素對應位置相乘。SR-STE 通過（1）限制

中對權重的剪枝，以及（2）維持

中未被剪枝的權重，來防止二進位制掩碼劇烈變化。

圖 10. STE 和 SR-STE 的對比。⊙的比較是元素乘積；⊗是矩陣乘法。

與 STE 或 SR-STE 不同，Top-KAST 方法可以在前向和反向傳播的整個訓練過程中保持恆定的稀疏性，還不需要使用具有稠密參數或梯度的前向傳播。

在訓練到第 t 步時，Top-KAST 過程如下：

稀疏前向傳遞：選擇參數

的一個子集，包含每層按大小排列的前 K 個參數，限制為權重的前 D 比例。如果時間 t 的參數化 α^t 不在 A^t（活動權重）中，則參數化為零。

其中 TopK (θ,x) 是根據大小排序後從 θ 中的前 x 個權重。

稀疏向後傳遞：然後將梯度應用於更大的參數子集

, 其中 B 包含 (D+M)， A⊂B。擴大需要更新的權重比例可以更有效地探索不同的剪枝掩碼，從而更有可能將前 D% 的激活權重排列好。

訓練分為兩個階段，集合 B∖A 中的附加座標控制引入的探索量。探索量會在訓練過程中逐漸減少，最終掩碼會穩定下來。

圖 11. Top-KAST 的剪枝掩碼會隨時間穩定下來。

為了防止馬太效應，Top-KAST 通過 L2 正則化損失來懲罰激活權重，以鼓勵產生更多新的探索。在更新期間，B∖A 中的參數比 A 受到更多的懲罰以穩定掩碼。

稀疏 Transformer

稀疏 Transformer 將 Transformer 架構中的自注意力層和 FFN 層稀疏化，使單個樣本推理的速度提高了 37 倍。

圖 12. 當在不同網路層上應用稀疏化時，Transformer 模型解碼單個 token（非批量推理）的速度。

稀疏 FFN 層：每個 FFN 層包含 2 個 MLP 和中間的一個 ReLU。因為 ReLU 會引入很多零值，所以該方法在激活函數上設計了一個固定結構，來強制要求在一個包含 N 個元素的塊中只包含 1 個非零值。稀疏模式是動態的，每個 token 都不同。

其中 Y_(sparse ) 中的每個激活函數結果對應於 W_1 中的一列和 W_2 中的一行。控制器是一個低秩的 bottleneck 全連接層，其中

、

在訓練期間使用 argmax 進行推理以選擇哪些列應為非零和，以及 Gumbel-softmax 技巧。因為可以在載入 FFN 權重矩陣之前計算 Controller (x)，所以可以知道哪些列將被清零，因此選擇不將它們載入到記憶體中以加快推理速度。

圖 13. (a) 稀疏 FFN 層；紅色列未載入到記憶體中以進行更快的推理。(b) 1:4 稀疏度的稀疏 FFN 控制器。

稀疏注意力層：在注意力層中，維度 d_(model) 被劃分為 S 個模組，每個模組的大小為 M=d_(model)/S。為了確保每個細分都可以訪問嵌入的任何部分，Scaling Transformer 引入了一個乘法層（即，一個乘法層將來自多個神經網路層的輸入按元素相乘），它可以表示任意排列，但包含的參數少於全連接層。

給定輸入向量

，乘法層輸出

：

乘法層的輸出是一個大小為

的張量。然後由二維卷積層對其進行處理，其中 length 和 S 被視為圖像的高度和寬度。這樣的卷積層進一步減少了注意力層的參數數量和計算時間。

圖 14. (a) 引入乘法層以使分區能夠訪問嵌入的任何部分。(b) 乘法全連接層和二維卷積層的結合減少了注意力層的參數數量和計算時間。

為了更好地處理長序列資料，Scaling Transformer 進一步配備了來自 Reformer 的 LSH（局部敏感哈希）注意力和 FFN 塊循環，從而產生了 Terraformer 模型。

混合專家系統 MoE

專家混合系統 (MoE) 模型是一種專家網路的集合，每個樣本僅激活網路的一個子集來獲得預測結果。這個想法起源於上世紀九十年代並且與集成方法密切相關。有關如何將 MoE 模組合併到 Transformer 的詳細資訊，可以查看本文作者之前寫的關於大型模型訓練技術的帖子和 Fedus 等人關於 MoE 的論文。

使用 MoE 架構，在解碼時僅使用部分參數，因此節省了推理成本。每個專家的容量可以通過超參數容量因子 C 進行調整，專家容量定義為：

每個 token 需要選擇前 k 個專家。較大的 C 會擴大專家容量，提高性能，但這樣做計算成本更高。當 C>1 時，需要增加一個鬆弛容量；當 C<1 時，路由網路需要忽略一些 token。

路由策略改進

MoE 層有一個路由網路來為每個輸入 token 分配一個專家子集。原生 MoE 模型中的路由策略是將每個 token 以不同的方式路由到按自然順序出現的首選專家。如果路由到的專家已經沒有多餘的空間，token 將被標記為溢出並被跳過。

V-MoE 將 MoE 層添加到 ViT (Vision Transformer) 中。它與以前的 SOTA 模型的性能相匹配，但只需要一半的推理計算。V-MoE 可以擴展成一千五百萬個參數。有研究者在實驗中將 k=2、專家需要 32 位，每 2 位專家間放置一層 MoE。

由於每個專家的能力有限，如果某些重要且資訊豐富的 token 在預定義的序列順序（例如句子中的單詞順序或圖像中 patch 的順序）中出現得太晚，則可能不得不丟棄它們。為了避免原生路由方案中的這種缺陷，V-MoE 採用 BPR（批量優先路由）首先將專家分配給具有高優先級分數的 token。BPR 在專家分配之前計算每個 token 的優先級分數（前 k 名路由器得分的最大值或總和），並相應地更改 token 的順序。這保證了核心的 token 能優先使用專家容量的緩衝區。

圖 15. 當 C<1 時，根據優先級分數丟棄圖像 patch 的方式。

當 C≤0.5 時，BPR 比普通路由效果更好，此時模型開始丟棄大量 token。這使模型即使在非常低的容量下也能與稠密網路一較高低。

在研究如何解釋圖像的類別與專家之間的關係時，研究者觀察到早期的 MoE 層更通用，而後期的 MoE 層可以專門用於某類圖像。

任務級 MoE 將任務資訊考慮在內，並且將路由 token 在任務級的視角來處理。研究者以 MNMT（多語言神經機器翻譯）為例，根據目標語言或語言對進行翻譯任務分組。

Token 級路由是動態的，每個 token 的路由決策是不相交的。因此，在推理時，伺服器需要預載入所有專家。相比之下，任務級路由是靜態的，甚至是固定的任務，因此一個任務的推理伺服器只需要預載入 k 個專家（假設 top-k 才有路由）。根據研究者的實驗，與稠密模型的 baseline 相比，任務級 MoE 可以實現與 token MoE 類似的性能增益，峰值吞吐量高 2.6 倍，解碼器小 1.6%。

任務級 MoE 本質上是根據預定義的啟發式方法對任務分佈進行分類，並將此類人類知識納入路由器。當這種啟發式不存在時，任務級 MoE 就難以使用了。

PR MoE 讓每個 token 通過一個固定的 MLP 和一個選定的專家。由於靠後的 MoE 更有價值，PR MoE 在靠後的層上設計了更多的出口。DeepSpeed 庫實現了靈活的多專家、多資料並行，以支持使用不同數量的專家來訓練 PR MoE。

圖 16。PR MoE 架構與標準 MoE 的對比圖。

核心方面的改進措施

專家網路可以託管在不同的設備上。然而，當 GPU 數量增加時，每個 GPU 上的專家數量就會減少，專家之間的通訊成本變得更加昂貴。跨多個 GPU 的專家之間的多對多通訊依賴於 NCCL 的 P2P API，這個接口不能佔據高速鏈路所有的頻寬，這是因為使用的節點越多，單個 chunk 越小。現有的多對多演算法在大規模問題上性能較差，單個 GPU 的工作量不能提升。針對這種情況，有多種核心改進措施來實現更高效的 MoE 計算，例如使多對多通訊更便宜 / 更快。

DeepSpeed 庫和 TUTEL 都實現了基於樹的分層多對多演算法，該演算法在節點內使用多對多演算法處理，然後再在節點間實現多對多。這種演算法將通訊跳數從 O（G）減少到

，其中 G 是 GPU 節點的總數，G_(node) 是每個節點的 GPU 核心數。儘管在這樣的實現中通訊量增加了一倍，但當批大小較小時 1×1 卷積層存在延遲，因此可以更好地擴展 batch 的規模。

DynaMoE 使用動態再編譯使計算資源適應專家之間的動態工作負載。再編譯機制需要從頭開始編譯計算圖，並且只在需要時重新分配資源。它會琢磨分配給每個專家的樣本數量，並動態調整其容量因子 C，以減少運行時的記憶體和計算需求。這種方法基於在訓練早期對專家和樣本的分配關係的觀察，在模型收斂後引入樣本分配快取，然後使用再編譯演算法消除門控網路和專家之間的依賴性。

架構最佳化

論文《Efficient Transformers: A Survey》回顧了一系列新的 Transformer 架構，並針對提高計算和記憶體效率進行了一些改進，除此以外，大家還可以閱讀這篇文章《The Transformer Family》，以深入了解幾種類型的 Transformer 改進。

圖 17. 高效 transformer 模型的分類

自注意力機制的二次時間複雜度和記憶體複雜性問題是提高 transformer 解碼效率的主要瓶頸，因此所有高效 transformer 模型都對原本稠密的注意力層應用了某種形式的稀疏化措施。

1. 固定模式：使用預定義的固定模式限制注意力矩陣的感受野：

可以將輸入序列分成固定的塊；
圖像 transformer 使用了局部注意力；
稀疏 transformer 使用了跨線注意力模式；
Longformer 使用了 dilated 注意力窗口；
可以使用 strided 卷積壓縮注意力來減少序列長度。

2. 組合模式：對輸入的 token 進行排序 / 聚類以實現更最佳化的序列全局視圖，同時保持固定模式的效率優勢

稀疏 transformer 結合了跨步和局部注意力；
給定高維輸入張量，axial transformer 不會將輸入 flattened 後再使用注意力機制，而是使用多注意力機制，一個注意力對應著輸入張量的一個軸；
Big Bird 模型設計了一些關鍵元件，即（1）全局 token，（2）隨機注意力（query 向量隨機綁定 key 向量）和（3）固定模式（局部滑動窗口）。

3. 可學習模式：通過學習確定最佳注意力模式：

Reformer 使用局部敏感哈希將 token 聚類；
路由 transformer 用 k-means 將 token 聚類；
Sinkhorn 排序網路會對輸入序列塊的排序演算法進行學習。

4. 遞迴：通過遞迴連接多個 block/segment：

Transformer-XL 通過在 segment 之間重用隱藏狀態來獲取更長的上下文；
通用 transformer 將自注意力與 RNN 中的循環機制相結合；
Compressive transformer 是 Transformer-XL 的擴展，具有額外的記憶體，具有 n_m 個記憶體槽和 n_(cm) 個壓縮記憶體槽。每當有新的輸入段被輸入到模型當中時，主記憶體中最久未更新的前 n_s 個激活函數都會被轉移到壓縮記憶體中。

5.Side Memory：使用可以一次訪問多個 token 的 Side Memory 模組

Set Transformer 設計了一種受歸納點方法啟發的新注意力；
ETC（Extended transformer construction）是 Sparse Transformer 的變體，具有新的全局 – 局部注意力機制；
Longformer 也是 Sparse Transformer 的變體，使用 dilated 滑動窗口。隨著模型網路的深入，感受野也會逐漸增加。

6. 節省記憶體：更改架構以使用更少的記憶體：