筆記

第四節的主題是在過去一兩年已經成為 LLM 顯學的架構 — Mixture of Experts (MoE)

MoE 簡介

什麼是 MoE?

image.png

上圖左側是原始的 transformer block，在前面完成 attention 以及 add & normalize 的計算後，每個 token 的 hidden state 會輸入至 feed forward network (FFN) 做非線性轉換。

簡單來說，MoE 就是將單一個參數量龐大的 FFN，替換為多個較小的 FFN（即 Experts），並加上一個 router (或 selector)。透過 router 來決定每個 token 要送交給哪幾個 expert 進行計算。這樣做的好處在於，可以大幅增加模型中的專家數量（即總參數量），但完全不會增加整體模型在做 generation 時所需的 FLOPs 。

要注意的是， "Mixture of Experts" 其實是一個很容易讓人誤解的命名。聽到 “experts”，人們常誤以為模型中有專門處理特定領域的模組（例如：寫程式的專家、英文專家等），但實際上並非如此。MoE 純粹是一種架構設計，指的是擁有多個「稀疏啟動 (sparsely activated)」子元件的網路架構。

因此，在討論 MoE 時，我們看重的往往不是模型的「總參數量」，而是「活躍參數量 (Activated Parameters)」。也就是每次計算時真正被啟動拿來用的參數量，這才是決定訓練與推論時 FLOPs 消耗的核心指標。

為什麼 MoE 在過去兩年突然變得大受歡迎

  1. 維持相同 FLOPs，但參數更多、效能更好：下圖可以發現，在維持相同 FLOPs 的狀況下，增加 MoE 的 experts 數量可以有效降低模型的 training Loss，得到更好的效能表現

image.png

[Fedus et al 2022]

  1. 模型收斂速度更快：從下圖可見，使用 MoE 的模型在訓練時，收斂速度比原本的 Dense model 更快（即達到相同效能時，MoE 模型所需的 training steps 更少）

image.png

近年發表的各種 MoE 模型效能，也已遠超過去同等參數量級的 Dense model

image.png

  1. 有利於多設備上的平行計算：MoE 非常適合目前訓練與部署大型模型的常見硬體配置。它可以將不同的 expert 部署在不同的 GPU 設備上，這種做法被稱為 expert parallelism，能完美發揮多 GPU 平行計算的優勢

image.png