本文是一篇綜述性的博客，探討總結(jié)當(dāng)下常用的大型 transformer 效率優(yōu)化方案。

大型 Transformer 模型如今已經(jīng)成為主流，為各種任務(wù)創(chuàng)造了 SOTA 結(jié)果。誠然這些模型很強(qiáng)大，但訓(xùn)練和使用起來代價(jià)非常昂貴。在時(shí)間和內(nèi)存方面存在有極高的推理成本。概括來說，使用大型 Transformer 模型進(jìn)行推理的難點(diǎn)，除了模型的規(guī)模不斷擴(kuò)大外，還有兩個(gè)不可忽略的地方：

• 內(nèi)存消耗大：推理時(shí)，需要把模型參數(shù)和中間狀態(tài)都保存到內(nèi)存中。例如：KV 存儲(chǔ)機(jī)制下的緩存中的內(nèi)容在解碼期間需要存儲(chǔ)在內(nèi)存中，舉例來說，對于 batch size 為 512，上下文長度為 2048 的設(shè)置來說，KV 緩存里需要的空間規(guī)模為 3TB，這是模型大小的 3 倍；注意力機(jī)制的推理成本和輸入序列的長度呈正相關(guān)；

• 低并行性：推理生成過程以自回歸的方式執(zhí)行，使解碼過程難以并行。

在這篇文章中，領(lǐng)導(dǎo) OpenAI 應(yīng)用研究的 Lilian Weng 寫了一篇博客，文中介紹了幾種提高 transformer 推理效率的方法。一些是通用的網(wǎng)絡(luò)壓縮方法，而另一些則應(yīng)用于特定的 transformer 架構(gòu)。

Lilian Weng 現(xiàn)為 OpenAI 應(yīng)用人工智能研究負(fù)責(zé)人，主要從事機(jī)器學(xué)習(xí)、深度學(xué)習(xí)等研究 。她本科畢業(yè)于香港大學(xué)，碩士畢業(yè)于北京大學(xué)信息系統(tǒng)與計(jì)算機(jī)科學(xué)系，之后前往印度安納大學(xué)布魯頓分校攻讀博士。

模型綜述

通常將以下內(nèi)容視為模型推理優(yōu)化的目標(biāo)：

• 使用更少的 GPU 設(shè)備和更少的 GPU 內(nèi)存，減少模型的內(nèi)存占用；

• 減少所需的 FLOP，降低計(jì)算復(fù)雜度；

• 減少推理延遲，運(yùn)行得更快。

可以使用幾種方法來降低推理過程在內(nèi)存中的成本，并且加快速度。

• 在多 GPU 上應(yīng)用各種并行機(jī)制來實(shí)現(xiàn)對模型的擴(kuò)展。模型組件和數(shù)據(jù)的智能并行使得運(yùn)行具有萬億級參數(shù)的大模型成為可能；

• 將暫時(shí)未使用的數(shù)據(jù)卸載到 CPU，并在以后需要時(shí)讀回。這樣做對內(nèi)存比較友好，但會(huì)導(dǎo)致更高的延遲；

• 智能批處理策略；例如 EffectiveTransformer 將連續(xù)的序列打包在一起，以刪除單個(gè)批次中的 padding；

• 神經(jīng)網(wǎng)絡(luò)壓縮技術(shù)，例如剪枝、量化、蒸餾。就參數(shù)數(shù)量或位寬而言，小尺寸的模型應(yīng)該需要少量的內(nèi)存，也就運(yùn)行得更快；

• 特定于目標(biāo)模型架構(gòu)的改進(jìn)。許多架構(gòu)上的變化，尤其是注意力層的變化，有助于提高 transformer 的解碼速度。

本篇文章的重點(diǎn)是網(wǎng)絡(luò)壓縮技術(shù)和 transformer 模型在特定體系結(jié)構(gòu)下的改進(jìn)。

量化策略

在深度神經(jīng)網(wǎng)絡(luò)上應(yīng)用量化策略有兩種常見的方法：

• 訓(xùn)練后量化（PTQ）：首先需要模型訓(xùn)練至收斂，然后將其權(quán)重的精度降低。與訓(xùn)練過程相比，量化操作起來往往代價(jià)小得多；

• 量化感知訓(xùn)練 (QAT)：在預(yù)訓(xùn)練或進(jìn)一步微調(diào)期間應(yīng)用量化。QAT 能夠獲得更好的性能，但需要額外的計(jì)算資源，還需要使用具有代表性的訓(xùn)練數(shù)據(jù)。

值得注意的是，理論上的最優(yōu)量化策略與實(shí)際在硬件內(nèi)核上的表現(xiàn)存在著客觀的差距。由于 GPU 內(nèi)核對某些類型的矩陣乘法（例如 INT4 x FP16）缺乏支持，并非下面所有的方法都會(huì)加速實(shí)際的推理過程。

Transformer 量化挑戰(zhàn)

許多關(guān)于 Transformer 模型量化的研究都有相同的觀察結(jié)果：訓(xùn)練后將參數(shù)簡單地量化為低精度（例如 8 位）會(huì)導(dǎo)致性能顯著下降，這主要是由于普通的激活函數(shù)量化策略無法覆蓋全部的取值區(qū)間。

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

Bondarenko 等人在一個(gè)小型 BERT 模型中觀察到，由于輸出張量中存在強(qiáng)異常值，F(xiàn)FN 的輸入和輸出具有非常不同的取值區(qū)間。因此，F(xiàn)FN 殘差和的逐個(gè)張量的量化可能會(huì)導(dǎo)致顯著的誤差。

隨著模型參數(shù)規(guī)模繼續(xù)增長到數(shù)十億的級別，高量級的離群特征開始在所有 transformer 層中出現(xiàn)，導(dǎo)致簡單的低位量化效果不佳。Dettmers 等人觀察到大于 6.7B 參數(shù)的 OPT 模型就會(huì)出現(xiàn)這種現(xiàn)象。模型大了，有極端離群值的網(wǎng)絡(luò)層也會(huì)變多，這些離群值特征對模型的性能有很大的影響。在幾個(gè)維度上的激活函數(shù)異常值的規(guī)模就可以比其他大部分?jǐn)?shù)值大 100 倍左右。

圖 2. 不同規(guī)模的 OPT 模型在四個(gè)語言任務(wù)（WinoGrande、HellaSwag、PIQA、LAMBADA）上的平均零樣本準(zhǔn)確率。

混合精度量化

解決上述量化挑戰(zhàn)的最直接方法是以不同的精度對權(quán)重和激活函數(shù)進(jìn)行量化。

GOBO 模型是首批將訓(xùn)練后量化應(yīng)用于 transformer 的模型之一（即小型 BERT 模型）。GOBO 假設(shè)每一層的模型權(quán)重服從高斯分布，因此可以通過跟蹤每層的均值和標(biāo)準(zhǔn)差來檢測異常值。異常值特征保持原始形式，而其他值被分到多個(gè) bin 中，并且僅存儲(chǔ)相應(yīng)的權(quán)重索引和質(zhì)心值。

基于對 BERT 中只有某些激活層（例如 FFN 之后的殘差連接）導(dǎo)致性能大幅下降現(xiàn)象的觀察，Bondarenko 等人通過在有問題的激活函數(shù)上使用 16 位量化而在其他激活函數(shù)上使用 8 位來采用混合精度量化。

LLM.int8 () 中的混合精度量化是通過兩個(gè)混合精度分解實(shí)現(xiàn)的：

• 因?yàn)榫仃嚦朔ò唤M行和列向量之間的獨(dú)立內(nèi)積，所以可以對每個(gè)內(nèi)積進(jìn)行獨(dú)立量化。每一行和每一列都按最大值進(jìn)行縮放，然后量化為 INT8；

• 異常值激活特征（例如比其他維度大 20 倍）仍保留在 FP16 中，但它們只占總權(quán)重的極小部分，不過需要經(jīng)驗(yàn)性地識別離群值。

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

細(xì)粒度量化

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

簡單地量化一層中的整個(gè)權(quán)重矩陣（逐個(gè)張量或逐個(gè)層量化）是最容易實(shí)現(xiàn)的，但量化粒度往往不盡如人意。

Q-BERT 將分組量化應(yīng)用于微調(diào)的 BERT 模型，將 MHSA（多頭自注意力）中每個(gè)頭的單個(gè)矩陣 W 視為一個(gè)組，然后應(yīng)用基于 Hessian 矩陣的混合精度量化。

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

ZeroQuant 與 Q-BERT 一樣都對權(quán)重使用分組量化，然后還對激活函數(shù)使用了 token-wise 量化策略。為了避免代價(jià)昂貴的量化和反量化計(jì)算，ZeroQuant 構(gòu)建了獨(dú)特的內(nèi)核來將量化操作與其之前的運(yùn)算符融合。

使用二階信息量化

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

從另一個(gè)角度來看，量化問題是一個(gè)優(yōu)化問題。給定一個(gè)權(quán)重矩陣 W 和一個(gè)輸入矩陣 X ，想要找到一個(gè)量化的權(quán)重矩陣  W^ 來最小化如下所示的 MSE 損失：

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

異常值平滑

眾所周知，Transformer 模型中激活函數(shù)比權(quán)重更難量化。SmoothQuant 提出了一種智能解決方案，通過數(shù)學(xué)等效變換將異常值特征從激活函數(shù)平滑到權(quán)重，然后對權(quán)重和激活函數(shù)進(jìn)行量化 (W8A8)。正因?yàn)槿绱?，SmoothQuant 具有比混合精度量化更好的硬件效率。

圖 5. SmoothQuant 將尺度方差從激活函數(shù)遷移到離線權(quán)重，以降低激活函數(shù)量化的難度。由此產(chǎn)生的新權(quán)重和激活矩陣都易于量化。

基于每個(gè)通道的平滑因子 s，SmoothQuant 根據(jù)以下公式縮放權(quán)重：

根據(jù)平滑因子可以很容易地在離線狀態(tài)下融合到前一層的參數(shù)中。超參數(shù) α 控制從激活函數(shù)遷移到權(quán)重的程度。該研究發(fā)現(xiàn) α=0.5 是實(shí)驗(yàn)中許多 LLM 的最佳取值。對于激活異常值較大的模型，可以將 α 調(diào)大。

量化感知訓(xùn)練 (QAT)

量化感知訓(xùn)練將量化操作融合到預(yù)訓(xùn)練或微調(diào)過程中。這種方法會(huì)直接學(xué)習(xí)低位表示的模型權(quán)重，并以額外的訓(xùn)練時(shí)間和計(jì)算為代價(jià)獲得更好的性能。

最直接的方法是在與預(yù)訓(xùn)練數(shù)據(jù)集相同或代表預(yù)訓(xùn)練數(shù)據(jù)集的訓(xùn)練數(shù)據(jù)集上量化后微調(diào)模型。訓(xùn)練目標(biāo)可以與預(yù)訓(xùn)練目標(biāo)相同（例如通用語言模型訓(xùn)練中的 NLL/MLM）或特定于的下游任務(wù)（例如用于分類的交叉熵）。

另一種方法是將全精度模型視為教師模型，將低精度模型視為學(xué)生模型，然后使用蒸餾損失優(yōu)化低精度模型。蒸餾通常不需要使用原始數(shù)據(jù)集。

剪枝

網(wǎng)絡(luò)剪枝是在保留模型容量的情況下，通過修剪不重要的模型權(quán)重或連接來減小模型大小。剪枝可能需要也可能不需要重新訓(xùn)練。剪枝可以是非結(jié)構(gòu)化的也可以是結(jié)構(gòu)化的。

• 非結(jié)構(gòu)化剪枝允許丟棄任何權(quán)重或連接，因此它不保留原始網(wǎng)絡(luò)架構(gòu)。非結(jié)構(gòu)化剪枝通常對硬件要求比較苛刻，并且不會(huì)加速實(shí)際的推理過程；

• 結(jié)構(gòu)化剪枝不改變權(quán)重矩陣本身的稀疏程度，可能需要遵循某些模式限制才能使用硬件內(nèi)核支持的內(nèi)容。本文專注于那些能實(shí)現(xiàn) transformer 模型的高稀疏性的結(jié)構(gòu)化剪枝。

構(gòu)建剪枝網(wǎng)絡(luò)的常規(guī)工作流程包含三個(gè)步驟：

1. 訓(xùn)練密集型的神經(jīng)網(wǎng)絡(luò)直到收斂；

2. 修剪網(wǎng)絡(luò)以去除不需要的結(jié)構(gòu)；

3. （可選擇）重新訓(xùn)練網(wǎng)絡(luò)，讓新權(quán)重保持之前的訓(xùn)練效果。

通過剪枝在密集模型中發(fā)現(xiàn)稀疏結(jié)構(gòu)，同時(shí)稀疏網(wǎng)絡(luò)仍然可以保持相似性能的靈感是由彩票假設(shè)激發(fā)的：這是一個(gè)隨機(jī)初始化的密集前饋網(wǎng)絡(luò)，它包含一個(gè)子網(wǎng)絡(luò)池。其中只有一個(gè)子集（稀疏網(wǎng)絡(luò)）是中獎(jiǎng)彩票（winning tickets），這個(gè)中獎(jiǎng)彩票在獨(dú)立訓(xùn)練時(shí)可以達(dá)到最佳性能。

如何剪枝

Magnitude pruning 是最簡單但同時(shí)又非常有效的剪枝方法 - 只裁剪那些絕對值最小的權(quán)重。事實(shí)上，一些研究發(fā)現(xiàn)，簡單的量級剪枝方法可以獲得與復(fù)雜剪枝方法相當(dāng)或更好的結(jié)果，例如變分 dropout 和 l_0 正則化。Magnitude pruning 很容易應(yīng)用于大型模型，并在相當(dāng)大的超參數(shù)范圍內(nèi)實(shí)現(xiàn)相當(dāng)一致的性能。

Zhu & Gupta 發(fā)現(xiàn)，大型稀疏模型能夠比小型但密集的模型獲得更好的性能。他們提出了 Gradual Magnitude Pruning (GMP) 算法，該算法在訓(xùn)練過程中逐漸增加網(wǎng)絡(luò)的稀疏性。在每個(gè)訓(xùn)練步驟中，具有最小絕對值的權(quán)重被屏蔽為零以達(dá)到所需的稀疏度并且屏蔽的權(quán)重在反向傳播期間不會(huì)得到梯度更新。所需的稀疏度隨著訓(xùn)練步驟的增加而增加。GMP 過程對學(xué)習(xí)率步長策略很敏感，學(xué)習(xí)率步長應(yīng)高于密集網(wǎng)絡(luò)訓(xùn)練中所使用的，但不能太高以防止收斂。

迭代剪枝多次迭代上述三個(gè)步驟中的第 2 步（剪枝）和第 3 步（重新訓(xùn)練），每次只有一小部分權(quán)重被剪枝，并且在每次迭代中重新訓(xùn)練模型。不斷重復(fù)該過程，直到達(dá)到所需的稀疏度級別。

如何再訓(xùn)練

再訓(xùn)練可以通過使用相同的預(yù)訓(xùn)練數(shù)據(jù)或其他特定于任務(wù)的數(shù)據(jù)集進(jìn)行簡單的微調(diào)來實(shí)現(xiàn)。

Lottery Ticket Hypothesis 提出了一種權(quán)重 rewinding 再訓(xùn)練方法：剪枝后，將未剪枝的權(quán)重重新初始化回訓(xùn)練初期的原始值，然后以相同的學(xué)習(xí)率時(shí)間表進(jìn)行再訓(xùn)練。

學(xué)習(xí)率 rewinding 僅將學(xué)習(xí)率重置回其早期值，而保持未剪枝的權(quán)重自最后一個(gè)訓(xùn)練階段結(jié)束以來不變。研究者觀察到 (1) 使用權(quán)重 rewinding 的再訓(xùn)練結(jié)果優(yōu)于通過跨網(wǎng)絡(luò)和數(shù)據(jù)集進(jìn)行微調(diào)的再訓(xùn)練，以及 (2) 在所有測試場景中學(xué)習(xí)率 rewinding 與權(quán)重 rewinding 的效果持平甚至更優(yōu)。

稀疏化

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

• 稀疏化的全連接層，包括自注意力層和 FFN 層；

• 稀疏模型架構(gòu)，即 MoE 組件的合并操作。

通過剪枝實(shí)現(xiàn)的 N:M 稀疏化

N:M 稀疏化是一種結(jié)構(gòu)化的稀疏化模式，適用于現(xiàn)代 GPU 硬件優(yōu)化，其中每 M 個(gè)連續(xù)元素中的 N 個(gè)元素為零。例如，英偉達(dá) A100 GPU 的稀疏張量核心支持 2:4 稀疏度以加快推理速度。

圖 6. 2:4 結(jié)構(gòu)化稀疏矩陣及其壓縮表示。

為了使密集型神經(jīng)網(wǎng)絡(luò)的稀疏化遵循 N:M 結(jié)構(gòu)化稀疏模式，英偉達(dá)建議使用三步操作來訓(xùn)練剪枝后的網(wǎng)絡(luò)：訓(xùn)練 –> 剪枝以滿足 2:4 稀疏性 –> 重新訓(xùn)練。

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

圖 7. Q（軸 1）和 K^?（軸 0）上相同排列，自注意力模塊的結(jié)果不變。

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

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

為了推動(dòng) N:M 結(jié)構(gòu)稀疏化，需要將一個(gè)矩陣的列拆分為 M 列的多個(gè) slide（也稱為 stripe），這樣可以很容易地觀察到每個(gè) stripe 中的列順序和 stripe 的順序?qū)?N:M 稀疏化產(chǎn)生的限制。

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

1. 有界回歸：在實(shí)踐中，兩個(gè)隨機(jī)通道的最大交換次數(shù)是固定的。每次搜索只有一個(gè)通道可以進(jìn)行交換，以保持搜索空間寬而淺；

2. 窄且深的搜索：選擇多個(gè) stripe 并同時(shí)優(yōu)化它們。

圖 9. 貪心算法實(shí)現(xiàn)迭代地尋找 N:M 稀疏化最佳排列的算法。

與按默認(rèn)通道順序?qū)W(wǎng)絡(luò)進(jìn)行剪枝相比，如果在剪枝之前對網(wǎng)絡(luò)進(jìn)行置換，可以獲得更好的性能。

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

STE 計(jì)算剪枝后的網(wǎng)絡(luò)的密集參數(shù)的梯度，并將其作為近似值應(yīng)用于稠密網(wǎng)絡(luò) W：

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

其中是的掩碼矩陣，⊙是元素對應(yīng)位置相乘。SR-STE 通過（1）限制中對權(quán)重的剪枝，以及（2）維持中未被剪枝的權(quán)重，來防止二進(jìn)制掩碼劇烈變化。

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

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

在訓(xùn)練到第 t 步時(shí)，Top-KAST 過程如下：

稀疏前向傳遞：選擇參數(shù)的一個(gè)子集，包含每層按大小排列的前 K 個(gè)參數(shù)，限制為權(quán)重的前 D 比例。如果時(shí)間 t 的參數(shù)化 α^t 不在 A^t（活動(dòng)權(quán)重）中，則參數(shù)化為零。

其中 TopK (θ,x) 是根據(jù)大小排序后從 θ 中的前 x 個(gè)權(quán)重。

稀疏向后傳遞：然后將梯度應(yīng)用于更大的參數(shù)子集, 其中 B 包含 (D+M)， A?B。擴(kuò)大需要更新的權(quán)重比例可以更有效地探索不同的剪枝掩碼，從而更有可能將前 D% 的激活權(quán)重排列好。

訓(xùn)練分為兩個(gè)階段，集合 B?A 中的附加坐標(biāo)控制引入的探索量。探索量會(huì)在訓(xùn)練過程中逐漸減少，最終掩碼會(huì)穩(wěn)定下來。

圖 11. Top-KAST 的剪枝掩碼會(huì)隨時(shí)間穩(wěn)定下來。

為了防止馬太效應(yīng)，Top-KAST 通過 L2 正則化損失來懲罰激活權(quán)重，以鼓勵(lì)產(chǎn)生更多新的探索。在更新期間，B?A 中的參數(shù)比 A 受到更多的懲罰以穩(wěn)定掩碼。

稀疏 Transformer

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

圖 12. 當(dāng)在不同網(wǎng)絡(luò)層上應(yīng)用稀疏化時(shí)，Transformer 模型解碼單個(gè) token（非批量推理）的速度。

稀疏 FFN 層：每個(gè) FFN 層包含 2 個(gè) MLP 和中間的一個(gè) ReLU。因?yàn)?ReLU 會(huì)引入很多零值，所以該方法在激活函數(shù)上設(shè)計(jì)了一個(gè)固定結(jié)構(gòu)，來強(qiáng)制要求在一個(gè)包含 N 個(gè)元素的塊中只包含 1 個(gè)非零值。稀疏模式是動(dòng)態(tài)的，每個(gè) token 都不同。

其中 Y_(sparse ) 中的每個(gè)激活函數(shù)結(jié)果對應(yīng)于 W_1 中的一列和 W_2 中的一行。控制器是一個(gè)低秩的 bottleneck 全連接層，其中、在訓(xùn)練期間使用 argmax 進(jìn)行推理以選擇哪些列應(yīng)為非零和，以及 Gumbel-softmax 技巧 。因?yàn)榭梢栽诩虞d FFN 權(quán)重矩陣之前計(jì)算 Controller (x)，所以可以知道哪些列將被清零，因此選擇不將它們加載到內(nèi)存中以加快推理速度。

圖 13. (a) 稀疏 FFN 層；紅色列未加載到內(nèi)存中以進(jìn)行更快的推理。(b) 1:4 稀疏度的稀疏 FFN 控制器。

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

給定輸入向量 ，乘法層輸出 ：

乘法層的輸出是一個(gè)大小為 的張量。然后由二維卷積層對其進(jìn)行處理，其中 length 和 S 被視為圖像的高度和寬度。這樣的卷積層進(jìn)一步減少了注意力層的參數(shù)數(shù)量和計(jì)算時(shí)間。

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

為了更好地處理長序列數(shù)據(jù)，Scaling Transformer 進(jìn)一步配備了來自 Reformer 的 LSH（局部敏感哈希）注意力和 FFN 塊循環(huán)，從而產(chǎn)生了 Terraformer 模型。

混合專家系統(tǒng) MoE

專家混合系統(tǒng) (MoE) 模型是一種專家網(wǎng)絡(luò)的集合，每個(gè)樣本僅激活網(wǎng)絡(luò)的一個(gè)子集來獲得預(yù)測結(jié)果。這個(gè)想法起源于上世紀(jì)九十年代并且與集成方法密切相關(guān)。有關(guān)如何將 MoE 模塊合并到 Transformer 的詳細(xì)信息，可以查看本文作者之前寫的關(guān)于大型模型訓(xùn)練技術(shù)的帖子和 Fedus 等人關(guān)于 MoE 的論文。

使用 MoE 架構(gòu)，在解碼時(shí)僅使用部分參數(shù)，因此節(jié)省了推理成本。每個(gè)專家的容量可以通過超參數(shù)容量因子 C 進(jìn)行調(diào)整，專家容量定義為：

每個(gè) token 需要選擇前 k 個(gè)專家。較大的 C 會(huì)擴(kuò)大專家容量，提高性能，但這樣做計(jì)算成本更高。當(dāng) C>1 時(shí)，需要增加一個(gè)松弛容量；當(dāng) C<1 時(shí)，路由網(wǎng)絡(luò)需要忽略一些 token。

路由策略改進(jìn)

MoE 層有一個(gè)路由網(wǎng)絡(luò)來為每個(gè)輸入 token 分配一個(gè)專家子集。原生 MoE 模型中的路由策略是將每個(gè) token 以不同的方式路由到按自然順序出現(xiàn)的首選專家。如果路由到的專家已經(jīng)沒有多余的空間，token 將被標(biāo)記為溢出并被跳過。

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

由于每個(gè)專家的能力有限，如果某些重要且信息豐富的 token 在預(yù)定義的序列順序（例如句子中的單詞順序或圖像中 patch 的順序）中出現(xiàn)得太晚，則可能不得不丟棄它們。為了避免原生路由方案中的這種缺陷，V-MoE 采用 BPR（批量優(yōu)先路由）首先將專家分配給具有高優(yōu)先級分?jǐn)?shù)的 token。BPR 在專家分配之前計(jì)算每個(gè) token 的優(yōu)先級分?jǐn)?shù)（前 k 名路由器得分的最大值或總和），并相應(yīng)地更改 token 的順序。這保證了核心的 token 能優(yōu)先使用專家容量的緩沖區(qū)。

圖 15. 當(dāng) C<1 時(shí)，根據(jù)優(yōu)先級分?jǐn)?shù)丟棄圖像 patch 的方式。

當(dāng) C≤0.5 時(shí)，BPR 比普通路由效果更好，此時(shí)模型開始丟棄大量 token。這使模型即使在非常低的容量下也能與稠密網(wǎng)絡(luò)一較高低。

在研究如何解釋圖像的類別與專家之間的關(guān)系時(shí)，研究者觀察到早期的 MoE 層更通用，而后期的 MoE 層可以專門用于某類圖像。

任務(wù)級 MoE 將任務(wù)信息考慮在內(nèi)，并且將路由 token 在任務(wù)級的視角來處理。研究者以 MNMT（多語言神經(jīng)機(jī)器翻譯）為例，根據(jù)目標(biāo)語言或語言對進(jìn)行翻譯任務(wù)分組。

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

任務(wù)級 MoE 本質(zhì)上是根據(jù)預(yù)定義的啟發(fā)式方法對任務(wù)分布進(jìn)行分類，并將此類人類知識納入路由器。當(dāng)這種啟發(fā)式不存在時(shí)，任務(wù)級 MoE 就難以使用了。

PR MoE 讓每個(gè) token 通過一個(gè)固定的 MLP 和一個(gè)選定的專家。由于靠后的 MoE 更有價(jià)值，PR MoE 在靠后的層上設(shè)計(jì)了更多的出口。DeepSpeed 庫實(shí)現(xiàn)了靈活的多專家、多數(shù)據(jù)并行，以支持使用不同數(shù)量的專家來訓(xùn)練 PR MoE。

圖 16。PR MoE 架構(gòu)與標(biāo)準(zhǔn) MoE 的對比圖。

內(nèi)核方面的改進(jìn)措施

專家網(wǎng)絡(luò)可以托管在不同的設(shè)備上。然而，當(dāng) GPU 數(shù)量增加時(shí)，每個(gè) GPU 上的專家數(shù)量就會(huì)減少，專家之間的通信成本變得更加昂貴?？缍鄠€(gè) GPU 的專家之間的多對多通信依賴于 NCCL 的 P2P API，這個(gè)接口不能占據(jù)高速鏈路所有的帶寬，這是因?yàn)槭褂玫墓?jié)點(diǎn)越多，單個(gè) chunk 越小?，F(xiàn)有的多對多算法在大規(guī)模問題上性能較差，單個(gè) GPU 的工作量不能提升。針對這種情況，有多種內(nèi)核改進(jìn)措施來實(shí)現(xiàn)更高效的 MoE 計(jì)算，例如使多對多通信更便宜 / 更快。

DeepSpeed 庫和 TUTEL 都實(shí)現(xiàn)了基于樹的分層多對多算法，該算法在節(jié)點(diǎn)內(nèi)使用多對多算法處理，然后再在節(jié)點(diǎn)間實(shí)現(xiàn)多對多。這種算法將通信跳數(shù)從 O（G）減少到，其中 G 是 GPU 節(jié)點(diǎn)的總數(shù)，G_(node) 是每個(gè)節(jié)點(diǎn)的 GPU 內(nèi)核數(shù)。盡管在這樣的實(shí)現(xiàn)中通信量增加了一倍，但當(dāng)批大小較小時(shí) 1×1 卷積層存在延遲，因此可以更好地?cái)U(kuò)展 batch 的規(guī)模。

DynaMoE 使用動(dòng)態(tài)再編譯使計(jì)算資源適應(yīng)專家之間的動(dòng)態(tài)工作負(fù)載。再編譯機(jī)制需要從頭開始編譯計(jì)算圖，并且只在需要時(shí)重新分配資源。它會(huì)琢磨分配給每個(gè)專家的樣本數(shù)量，并動(dòng)態(tài)調(diào)整其容量因子 C，以減少運(yùn)行時(shí)的內(nèi)存和計(jì)算需求。這種方法基于在訓(xùn)練早期對專家和樣本的分配關(guān)系的觀察，在模型收斂后引入樣本分配緩存，然后使用再編譯算法消除門控網(wǎng)絡(luò)和專家之間的依賴性。

架構(gòu)優(yōu)化

論文《Efficient Transformers: A Survey》回顧了一系列新的 Transformer 架構(gòu)，并針對提高計(jì)算和內(nèi)存效率進(jìn)行了一些改進(jìn)，除此以外，大家還可以閱讀這篇文章《The Transformer Family》，以深入了解幾種類型的 Transformer 改進(jìn)。

圖 17. 高效 transformer 模型的分類

自注意力機(jī)制的二次時(shí)間復(fù)雜度和內(nèi)存復(fù)雜性問題是提高 transformer 解碼效率的主要瓶頸，因此所有高效 transformer 模型都對原本稠密的注意力層應(yīng)用了某種形式的稀疏化措施。

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

• 可以將輸入序列分成固定的塊；

• 圖像 transformer 使用了局部注意力；

• 稀疏 transformer 使用了跨線注意力模式；

• Longformer 使用了 dilated 注意力窗口；

• 可以使用 strided 卷積壓縮注意力來減少序列長度。

2. 組合模式：對輸入的 token 進(jìn)行排序 / 聚類以實(shí)現(xiàn)更優(yōu)化的序列全局視圖，同時(shí)保持固定模式的效率優(yōu)勢

• 稀疏 transformer 結(jié)合了跨步和局部注意力；

• 給定高維輸入張量，axial transformer 不會(huì)將輸入 flattened 后再使用注意力機(jī)制，而是使用多注意力機(jī)制，一個(gè)注意力對應(yīng)著輸入張量的一個(gè)軸；

• Big Bird 模型設(shè)計(jì)了一些關(guān)鍵組件，即（1）全局 token，（2）隨機(jī)注意力（query 向量隨機(jī)綁定 key 向量）和（3）固定模式（局部滑動(dòng)窗口）。

3. 可學(xué)習(xí)模式：通過學(xué)習(xí)確定最佳注意力模式：

• Reformer 使用局部敏感哈希將 token 聚類；

• 路由 transformer 用 k-means 將 token 聚類；

• Sinkhorn 排序網(wǎng)絡(luò)會(huì)對輸入序列塊的排序算法進(jìn)行學(xué)習(xí)。

4. 遞歸：通過遞歸連接多個(gè) block/segment：

• Transformer-XL 通過在 segment 之間重用隱藏狀態(tài)來獲取更長的上下文；

• 通用 transformer 將自注意力與 RNN 中的循環(huán)機(jī)制相結(jié)合；

• Compressive transformer 是 Transformer-XL 的擴(kuò)展，具有額外的內(nèi)存，具有 n_m 個(gè)內(nèi)存槽和 n_(cm) 個(gè)壓縮內(nèi)存槽。每當(dāng)有新的輸入段被輸入到模型當(dāng)中時(shí)，主內(nèi)存中最久未更新的前 n_s 個(gè)激活函數(shù)都會(huì)被轉(zhuǎn)移到壓縮內(nèi)存中。

5.Side Memory：使用可以一次訪問多個(gè) token 的 Side Memory 模塊

• Set Transformer 設(shè)計(jì)了一種受歸納點(diǎn)方法啟發(fā)的新注意力；

• ETC（Extended transformer construction）是 Sparse Transformer 的變體，具有新的全局 - 局部注意力機(jī)制；

• Longformer 也是 Sparse Transformer 的變體，使用 dilated 滑動(dòng)窗口。隨著模型網(wǎng)絡(luò)的深入，感受野也會(huì)逐漸增加。

6. 節(jié)省內(nèi)存：更改架構(gòu)以使用更少的內(nèi)存：

• Linformer 將 key 和 value 的代表長度的維度投影到低維表示（N→k），因此內(nèi)存復(fù)雜度從 N×N 降低到 N×k；

• Shazeer 等人提出了多 query 注意力，在不同注意力頭之間共享 key 和 value，大大減少了這些張量的大小和內(nèi)存成本。

7. 使用內(nèi)核：使用內(nèi)核可以讓自注意力機(jī)制的公式書寫起來更簡單。需要注意的使，這里的內(nèi)核是指內(nèi)核方法中的內(nèi)核，而不是 GPU 操作程序。

8. 自適應(yīng)注意力：讓模型學(xué)習(xí)最佳注意力廣度，或決定何時(shí)按每個(gè) token 提前退出：

• 自適應(yīng)注意力廣度訓(xùn)練模型，通過 token 和其他 key 之間的 soft mask 機(jī)制，為每個(gè) token、每個(gè)注意力頭學(xué)習(xí)最佳的注意力廣度；

• 通用 transformer 結(jié)合了循環(huán)機(jī)制，并使用 ACT（自適應(yīng)計(jì)算時(shí)間）來動(dòng)態(tài)決定循環(huán)幾次；

• 深度自適應(yīng) transformer 和 CALM 使用一些置信度度量方法來學(xué)習(xí)何時(shí)提前退出每個(gè) token 的計(jì)算層，這樣可以在性能和效率之間找到一種平衡。

原文鏈接：https://lilianweng./posts/2023-01-10-inference-optimization/