摘要：近來，增大模型規模成為了提升模型性能的主要手段。特别是NLP領域的自監督預訓練語言模型，規模越來越大，從GPT3的1750億參數，到Switch Transformer的16000億參數，又是一個數量級的增加。
前言
近來，增大模型規模成為了提升模型性能的主要手段。特别是NLP領域的自監督預訓練語言模型，規模越來越大，從GPT3的1750億參數，到Switch Transformer的16000億參數，又是一個數量級的增加。
模型規模的數量級的增大，雖然取得了一定程度的性能提升，甚至産生了某些意想不到的“神奇”效果（如GPT3），但其背後的計算開銷成了最大的問題，比如GPT3訓練使用了萬級的GPU和數周的訓練時間。如何既能利用超大規模的參數來提升模型表達和性能，又能控制計算量的較小的增加，成為了最主要的挑戰之一。以MoE為代表的動态神經網絡技術被重點引入。大腦是典型的低能耗高效率的計算模式，稀疏激活是最重要的特性。除了巨型模型在訓練推理特别是訓練時的計算效能挑戰外，目前巨型模型的訓練優化算法另一個更大的挑戰是（不在此處讨論），BP算法是目前最為可用的深度網絡優化，但更理想的優化算法需要高并行、優化過程非對稱、并能夠在時空次元通過局部持續優化完成整體優化。
1. 傳統的神經網絡模型，前饋的時候，輸入的batch中，每一個樣本的處理，都将激活網絡中的每一個參數參與計算。
2. 條件計算最寬松的定義，指僅激活網絡中某些部分的一類算法。Conditional Computation refers to a class of algorithms that activate only some of the different parts in a network. 在具體某類條件計算實作中，條件選擇模式，可能按照輸入的batch中每sample獨立激活網絡不同部分，可能按照輸入資料空間上不同的部分（比如image不同區域或者channel），可能按照輸入資料時間上不同的部分（比如time series的不同slide window或者video的不同的frame。），可能按照目标任務的不同每task獨立的，可能按照非可學習的固定的随機配置設定不同的子網獨立計算。
3. 對不同的輸入（原始或者前層），按照一定條件，選擇性的執行後續部分網絡的計算，這個技術下，有一些近似或相關的技術，如：dynamic neural network(s), conditional computing, conditional activation, sparse activating, selective execution, mixture of experts (MoE), dynamic routing, …；強相關的一些模型比如 Switch Transformer等。
條件計算的分類（廣義）
1. 按照routing是否可學習可以分為：learnable routing conditional computation和 unlearnable routing conditional computation.
2. 按照activation是否不執行non-activation計算，可以分為：hard conditional computation和soft conditional computation。對于hard-mode的條件計算，通過tensor挑選切分等操作，無論何種條件選擇模式，不需要激活的資料将完全不參與不激活的網絡部分的計算；soft-mode的條件計算，可能僅采取将相關資料置零等方式來避免産生計算效果，但還是和不需要激活網路部分實際執行計算過程。
條件計算的主要優勢
1. 計算有效，降低能耗：通過部分激活部分計算，以每樣本條件激活的條件計算為例，單個樣本隻需要經過整個SuperNet的一部分參與計算。
2. 更大網絡，表達更強：由于一處到多處的Route，各處（層）的Input被路由到不同的子網獨立計算，不同的輸入的互相在各層的表達相對獨立沒有影響，表達能力更強，網絡可以更大，但表達效率降低了。
條件計算的網絡和計算形式
條件計算的網絡和計算形式比較靈活，部分建構形式如：
1. 按照CV等task的特點，用多個獨立的CNN作為expert網絡，按照task來獨立路由，尾部組合後給一個大網絡。
2. 使用更複雜的cascading等形式組合不同層級的不同的expert網絡。
3. 通過決策樹等方法做資料變換實作路由。
4. 通過可學習的網絡來選擇路由。其中政策學習的損失有多種建構形式：直接使用分類等任務的主損失，對不同專家的重要性和負載建構損失作為輔助損失等等。
條件計算的路由政策
1. non-learnable/hard-mode，通過某種确定性政策，如LSH等方式計算路由。
2. learnable-mode，通過可學習網絡計算路由。網絡規模可大可小，簡單的可學習路由為單層權重：G(x) = P(X*W)，G(x)為路由Gate函數，X為輸入， W為通損失函數來度量的可學習路由權重，P為某種挑選函數（如topk, sort等），在實際實作中，X*W的輸入與權重計算結果可能作為後續網絡的輸入資訊的一部分，不僅僅利用G(x)來選擇路由，則需要對X*W的結果做歸一化，更典型的形式則為：G(x)=P(N(X*W))，其中N為表達Normalization函數，如Softmax。
條件計算的備援政策
條件計算的備援政策，可分為無備援條件計算和備援條件計算：
1. 無備援條件計算可通過P(.)函數的實作如topk(k=1,…)來實作；
2. 備援條件計算，可以多種實作形式，可以通過P(.)函數的實作如topk(k=n,…)，n>=2來實作，也可以通過硬備援模式，整個網絡中支援輸入的複制和多路計算實作。
條件計算的挑戰
1. 路由算法對模型品質的影響無論輸入和路由權重作用的資訊（X*W），是僅作為路由選擇并作為後續網絡單元的輸入，還是直接作為後續網絡單元的輸入的一部分，路由算法決定了輸入資訊的處理流向，對模型的整體品質都有很大影響。2. 路由(routing)/門(gate)的穩定性随機初始化的路由/門的權重，權重自身在不斷被訓練調整；在前後層的網絡持續訓練變化，同一樣本在訓練的不同階段會被分派到不同的後續網絡單元中，這種動态變化過于劇烈，将嚴重影響整個網絡訓練過程的穩定性和收斂速度。3、路由的專家樣本重要性和負載的平衡性
訓練階段，每專家和樣本批次中樣本的關聯度重要性，和每批次中樣本被均衡分派到不同專家的負載平衡性，這兩個名額既相關又沖突。需要分别建構損失函數作為輔助損失，來優化這兩個名額。在arxiv:1701.06538《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》做了相關讨論。
關于條件計算/動态神經網絡
關于條件計算/動态神經網絡，更多的資訊在《Dynamic Neural Networks: A Survey》arxiv:2102.04906 一文中，作者對廣義的動态神經網絡，将各種動态網絡相關的技術按照執行個體級、時間級、空間級做了分類。
- 1. Instance-wise Dynamic NN：逐執行個體動态，每樣本獨立激活不同的網絡和參數（MoE為這個方向）。Dynamic Architecture：Dynamic Depth、Dynamic Width、Dynamic Routing/MoE；Dynamic Parameter：Parameter Adjustment、Parameter Prediction、Dynamic Feature(s)
- 2. Spatial-wise Dynamic NN：空間級動态：圖像等不同空間位置激活後續不同網絡和參數。(CNN等)：Pixel Level、Region Level、Resolution Level
- 3. Temporal-wise Dynamic NN：時間級動态：時序資料按時序維切分激活後續不同網絡和參數。（video-frames, text-sequence, time-series, stream, ...)Text-SequenceVideo-Frames
上述為該綜述論文對Dynamic NN的總體分類。
從超大規模網絡動态網絡技術支撐角度，高表達能力，低計算代價為主的來考慮分類，從兩個次元對動态網絡技術分類:
1. 按照在前饋計算時是否部分激活：
Hard-Dynamic：在前饋的時候，部分網絡絕對不激活參與計算
Soft-Dynamic：在前饋的時候，部分網絡經過softmax等gate/route後，通過張量元素置零等方式，失去表達能力，但會參與計算。
2. 按照動态激活判定算法的輸入：
- 逐樣本級：（在輸入層）按照每樣本的執行個體來決定動态網絡的後續激活。
- 亞樣本級：（在輸入層）樣本内時間/空間級激活不同的後續網絡單元。一般深度網絡，不僅在輸入層會被選擇性激活執行，在中間層也類似。
其中，智能平台支援Hard-Dynamic逐樣本級的動态神經網絡，能比較自然的獲得網絡結構大顆粒的稀疏激活，在超大模型中能實作訓練和推理的高能效。
動态神經網絡相比與靜态結構的神經網絡，在相關研究中，從效能，表達，泛化、魯棒，可解釋等方面做了大量對比研究。從智能平台通過計算成本盡量低的支援超大規模網絡來提升模型性能的角度看，Efficiency和Representation最為重要：
1、Efficiency：靜态網絡“牽一發而動全身”，每一個樣本輸入整個網絡/所有參數都要響應，這對超大網絡來取得領先效果的模型能耗挑戰太大。
2、Representation: 參數量更大，表達容量更大；但MoE等結構在深度網絡的各層特征的表達上，複用降低，每參數的表達效率更低。
實作政策
實作各種模型的帶有動态路由稀疏激活的超大規模參數版本，需要分模型研究和實作。
以Switch Transformer為例，其參數擴充到部分在Transformer的FFN部分。其MoE化擴充，如下圖：
(圖檔來源：Switch Transformer論文)
可見，MoE化主要變化在需要Expert子網絡前後增加MoE相關的邏輯。本文主要介紹平台上的實作。動态路由條件計算，主要包括四個步驟：路由計算、資料分派、獨立計算，結果合并。
1. 路由計算-Gate：根據輸入（可以為整個網絡的輸入，或者前面網絡單元/層的輸出），在路由單元完成計算，在以batch内sample-wise的路由中，計算出每個樣本要分派的後續網絡路由（Mixture-of-Experts/MoE中的專家）。
2. 資料分派-Dispatch：從輸入的整體的Tensor中，按照路由計算的樣本-專家關系，收集合并出每個專家需要處理的Tensor。如果在固定expert-batch的設計中，要平衡每批訓練中，分派到每個專家的樣本數和專家每輪訓練最大容量，由于樣本輸入的随機性，很難保證較為均勻的分派，對于低于最大容量的批次，對固定batch-size的要做pad，對于高于最大容量的樣本，可以采用延後重采樣等方式。為了維護正确的輸入輸出關系（Input/X – Label/Y）和訓練是反向傳播的求導關系，實作中需要維護原始batch到每專家的sub-batch的index關系，在後來求導和結合合并時使用。
3. 獨立計算-Expert：并發（邏輯上可以先後）調用各個專家處理對應的sub-batch。這也是智能平台要支援的并發API之一。
4. 結果合并-Combine：合并每專家的結果tensor到整個batch的tensor，并按照資料分派索引，交換到原始輸入的順序。
在主流的深度學習智能平台中，可以采用兩類主要的實作政策：
張量置零：對需要分派到不同的後續網絡單元（專家網絡子網等），對需要分派的專家拷貝若幹份tensor，對于不應輸入目前專家處理的資料次元置零。該方式在保證置零計算邏輯正确的情況下，實作簡單，全張量操作，對平台無特殊要求，适用于算法研究，僅展現條件計算前序資料被動态路由到不同的後續網絡單元，分析算法的效果。如果通過置零方式，該方法每個專家處理的tensor在batch次元大小是全batch，不能節省計算量和記憶體使用量。
張量整理：對需要分派到不同的後續網絡單元（專家網絡子網等），對需要分派的專家拷貝若幹份tensor，對于不應輸入目前專家處理的資料次元不保留。并維護好sample級的index在變換前後的對應關系。在分布式友好的實作中，如果專家子網為機關被劃分到不同的計算節點，那麼專家網絡的實作最好從子網級的平台對象繼承後實作，比如：MindSpore中的mindspore.nn.Cell。詳細實作細節參見後續技術實作章節。
核心代碼
核心代碼：路由計算、資料分派、獨立計算，結果合并
參考代碼采用MindSpore示意實作。(注：import mindspore as ms)
Mixture of Experts的核心邏輯，對輸入I，經過routing_network(最簡單*W即可)，然後topk(若變種算法需要gate權重則需要softmax，否則可不)，然後用tensor的操作(可按照batch)選擇出每個subnetwork/expert的張量。
為友善調試，采用了規模極小的非随機的确定數值構造輸入和路由權重，路由網絡采用簡單的X*W。
1、路由計算
當上述輸入5行（僅3類，希望分派給3個專家）樣本，和Gate權重做矩陣乘後，可以明确算出每個樣本要分派的專家。可以用matmul，也可以類似gates_weighted = einsum('bd,de->be', [data_inputs, gate_weights])第一輪矩陣乘的結果為：
輸入和權重乘法，在python中可以采用@，也可以采用matmul，也可以采用愛因斯坦求和簡記憶法函數einsum。當是簡單的矩陣乘的時候，采用einsum在計算圖編譯的時候實際會拆分成多個算法，性能并不好；但當輸入和權重超過2維，需要以batch維固定做路由計算的時候，使用einsum可以程式設計實作很簡單。
2、分派
條件計算的分派，主要邏輯是根據路由網絡的輸出，為每個樣本計算出top-k的專家。其實作可以通過topk函數實作。由于top選擇score可作為後續網絡單元的輸入資訊（含路由的資訊），是以一般要對路由輸出做softmax做歸一化。
按需計算1：all-N專家之間的歸一化權重 (please refer to #2) ，gates_weighted一樣，按照dim=-1做了歸一化而已其輸出為：
為batch中每個sample選擇Top-K個專家 這裡為batch中每個的專家權重，可以從softmax-ed來top-k，也可以直接從gates_weighted來top-k；由于這裡可能不做softmax或者延後，是以可gates_weighted，這裡為batch中每個的專家序号
其輸出為：
接着：
按需計算2: top-n專家之間的歸一化權重
如何根據分派索引，從原始的輸入中，為每個專家提取出屬于該專家處理的tensor，在目前的主流智能平台，都沒有專門的算子，可以通過其他算子的組合來實作類似的效果。在MindSpore中，可以通過底層的C++實作算子，也可以通過Python中繼承Cell并實作bprob， 然後将原始 gate tensor中按照index組織到目标輸出中。這裡我們實作一個Dispatch類
3、獨立計算
直接并行調用後續的專家網絡。并行部分可以通過平台來支援。可以通過特殊的函數或者annotation等辨別，也可以由平台編譯時優化為并行執行。（在非動态路由條件計算的網絡模型中，一般不存在類似的優化。）
4、合并
合并的邏輯相對簡單，先通過cat按照batch次元做拼接，然後構造正确的zeros tensor用index_add按照索引将各個專家網絡的結果在保持input序合并到一起，做為該MoE子產品的輸出。
上述完成了整個MoE的完整計算過程。
代碼架構
我們按照上述基本動态路由條件計算的張量操作為主的邏輯，擴充到一個完整的訓練代碼架構中：
- class Dispatch(ms.nn.Cell): 實作路由中的分派邏輯
- class Combine(ms.nn.Cell): 實作路由中的組裝邏輯
- class Route(ms.nn.Cell): 完成整個動态路由邏輯，可以實作為相對通用的類
- class Expert(ms.nn.Cell): 平台使用者自定義的專家網絡
- class Network(ms.nn.Cell): 平台使用者自定義的大網絡
- class MSELoss(ms.nn.Cell)：實作MSE損失，實作輔助損失的邏輯
- class OutputLossGraph(ms.nn.Cell)：輸出infer和loss，PyNative模式單步
- class Dataset: 資料集類，僅滿足輸入shape和X-Y合理對應關系，僅僅示例def train( …): 訓練入口
條件計算實作技術點
1、動态路由
- 不可學習路由
如使用LSH (locality sensitive hashing)做路由：在整個可學習網絡的前端，使用LSH來分派樣本，這樣可以避免LSH部分求導問題；如果在網絡中間增加LSH子產品，需要通過梯度估計完成确定性算法部分梯度傳遞。
- 可學習路由
簡單的做法，定義gate_weights為可學習Parameter，對于二維的張量，通過python@或者matmul等完成權重路由計算；如果是更高次元的張量，且需固定batch維，einsum('bd*,*de->b*e')的形式完成計算。
2、topk和softmax的前後關系
在G_1(x)=softmax(topk(X*W)))和G_2(x)=topk(softmax(X*W)))兩類Gate實作中，
将softmax置于Topk前後，對top-k的選擇不變；當需要将G_*作為後序網絡輸入的一部分，即将路由權重資訊作為後續網絡輸入資訊，則需要考慮：需要all-N專家之間的歸一化權重，則softmax置于top-k之前；否則softmax置于top-k之後，來計算top-N專家之間的歸一化權重。
3、如何每專家在批次進行中平衡
按照每樣本的路由權重求和，即對batch單個樣本被配置設定的1+個export的重要性和權重求和，計算出importance；按照每樣本的路由權重中非0的求和，計算出有負載的專家來求得load。将coefficient_of_variation(importance) + coefficient_of_variation(load)作為auxiliary_loss參與優化，來平衡importance和load。變異系數(Coefficient of Variation)是用于無量綱度量資料的離散程度，越離散在此處表示均衡性越差，需要向更小優化。
在Transformer等多層（多處）MoE的模型中，将多組auxiliary_loss聯合作為auxiliary_loss, 在加dominated_loss之後即可。