MoE環遊記：從幾何意義出發

©PaperWeekly 原創 · 作者 | 蘇劍林

單位 | 科學空間

研究方向 | NLP、神經網路

前兩年福至心靈之下，開了一個“Transformer 升級之路”系列，陸續分享了主流 Transformer 架構的一些改進工作和個人思考，得到了部份讀者的認可。這篇文章開始，我們沿著同樣的風格，介紹當前另一個主流架構 MoE（Mixture of Experts）。

MoE 的流行自不必多說，近來火出圈的 DeepSeek-V3 [1] 便是 MoE 架構，傳言 GPT-4 也是 MoE 架構，國內最近出的一些模型也有不少用上了 MoE。

然而，雖然 MoE 的研究由來已久，但其應用長時間內都不慍不火，大致上是從去年初的《Mixtral of Experts》[2] 開始，MoE 才逐漸吸引大家的注意力，其顯著優點是引數量大，但訓練和推理成本都顯著低。

但同時 MoE 也有一些難題，如訓練不穩定、負載不均衡、效果不夠好等，這也是它早年沒有流行起來的主要原因。不過隨著這兩年關注度的提升，這些問題在很大程度上已經得到解決，我們在接下來的介紹中會逐一談到這些內容。

問題定義

首先要指出的是，這裡會用筆者自己的一種理解思路來介紹 MoE，在必要的地方會附上相應的參考文獻，但不會對 MoE 架構進行系統的追根溯源，還請讀者見諒。

我們知道，Transformer 模型由 Attention 層和 MLP 層組成，MoE 替換的是模型中 MLP 層。MLP 層又分 FFN（FeedForward Network）和 GLU（Gated Linear Unit）兩種，主流的是 GLU，但簡單起見我們還是以 FFN 為例：

其中是輸入向量（行向量），是兩個引數矩陣，f 是 Element-wise 的啟用函式。設 n 是一個能整除 D 的整數，那麼上述可以等價地用分塊矩陣寫成：

其中，這裡的切片按照 Python 規則來。由此可見，FFN 可以等價表示成 n 個向量之和，每個向量代表了一個小模型的輸出，每個小模型計算量相同，這些小模型就是 MoE 中的“Expert”。

MoE 提出的問題是：

能否只挑 k 個向量的和來逼近 n 個向量的和呢？這樣就可以將計算量降低到 k/n 了。

模長排序

這個問題其實我們在《低秩近似之路（三）：CR》已經探究過，寫成數學公式是：

記，那麼它又可以寫成：

這個問題的精確求解是比較困難的，但有一個簡單的近似解：當兩兩正交時，我們有：

上式最優解顯然就是讓模長最小的 n-k 個等於 1，這又等價於說挑出模長最大的 k 個向量來逼近 n 個向量之和。當不滿足兩兩正交的條件時，我們依然用它來作為一個近似解。它的幾何意義也很直觀，模長越大的向量，在求和過程中越不容易被抵消，從而作用越突出。

此外，在《低秩近似之路（三）：CR》中我們還討論了一種依機率取樣的逼近過程，在方差最小的假設下得到的最優取樣概率同樣有正比於模長的特點，所以總的來說按向量模長排序是一個簡單但不失有效的策略。

MoE初現

現在策略已經有了——“挑模長最大的 k 個向量”——可是細想之下我們會發現它並不實用：要挑模長最大的 k 個向量，就得把所有向量的模長都算出來，這又意味著要把所有的先算出來，可我們的原本目的卻是減少的計算量！

為了解決這個矛盾，我們需要重新設計每個 Expert 模型，使得它的模長可以低成本地計算出來。什麼意思呢？首先我們將歸一化得到，這樣每個的模長都相同了。接著我們定義：

其中是引數矩陣，是一個的啟用函式，說白了這就是一個 d 維到 n 維的線性變換加啟用函式，所以計算量是比較小的，這部分模型在 MoE 中被稱為“Router”。

的作用是什麼呢？預測每個 Expert 的模長！換言之，我們將作為第 i 個 Expert 的模長，才是完整的 Expert，它被分解為兩部分：計算量比較小的模長以及計算量比較大的方向。

為了減少計算量，我們先計算出，挑出最大的k個後才去計算相應的，最後乘上並求和：

這便是 MoE 模型的基本公式。由於計算中只保留了 Top-k 部分，所以它本質上屬於一種 Sparse 模型，而原本的 FFN 或者 k=n 時的模型，通常稱為對應的 Dense 模型。

思路概括

不管是熟悉 MoE 還是不熟悉 MoE 的讀者，可能都會對上述過程有點陌生，因為這是筆者自己閉門造車的一種 MoE 理解路線，但因為其幾何意義更明確，所以本質上應該是更好理解的。

我們再來整理一下整個思路：

1. 一個常規的 Dense 模型 FFN，可以等價改寫為 n 個 Expert 向量之和；

2. 為了節省計算量，我們試圖挑出 k 個向量求和來逼近原本的 n 個向量之和；

3. 轉化為數學問題求解後，我們發現挑選規則是模長最大的 k 個向量；

4. 直接去算 n 個 Expert 的模長然後選 k 個實際上是不省計算量的，所以要重新設計 Expert；

5. 將歸一化得到，然後用另外的小模型（Router）預測模長，最終的 Expert 為；

6. 此時，我們就可以先算全體，挑出 k 個後才去計算，達到節省計算量的目的。

為何如此

可能有些讀者疑問，為什麼要做這個看似複雜的過程？原本的 MoE 不是挺好理解的嗎？一般的 MoE 形式為：

也就是求和前少了對的歸一化，此時也沒有模長的意義，它純粹是一個用來對 Expert 排序的打分模型（即 Router）。可為什麼將乘到 Expert 上去就能讓 Router 學會正確排序 Expert 呢？筆者發現只有《Sparse Backpropagation for MoE Training》[3] 對此給出了一個解釋，但還是稍欠直觀。

而在本文的幾何視角下，我們會發現很多問題就“豁然開朗”了。我們將 Expert 重新引數化為後，Dense 模型對應於全體求和，而 MoE 對應於選 Top-k 後求和，這是 Dense 模型的一個有理論保證的逼近。我們沒有去考慮 Router 如何選擇 Expert，只是每一步都儘可能逼近 Dense 模型，這可以說是既要大引數、又要小計算量的最佳選擇。

現在的幾何意義是模長而不是機率，所以啟用函式就沒有歸一化的要求了，除了 Softmax 外，像 Sigmoid、ReLU 都可以考慮使用，也可以考慮我們在《Softmax後傳：尋找Top-K的光滑近似》介紹的 Top-k 光滑近似。Router 使用非歸一化的啟用函式，有助於避免 k > 1 時 Expert 之間的惡性競爭，有時候能取得更好的效果。

最後補充一點，我們前面定義，目的是讓所有模長相同，實際操作中不是一定要 L2 Normalize，也可以是其他等價操作，比如 gamma 引數恆等於 1 的 RMS Norm，它更符合我們的輸出習慣。