Outrageously Large Neural Networks The Sparsely-Gated Mixture-of-Experts Layer
概
Mixture-of-Experts (MoE).
MoE
-
通过一 gating network 选择不同的 expert:
\[y = \sum_{i=1}^n G(x)_i E_i(x), \]若 \(G(x)_i = 0\), 则我们不需要计算 \(E_i(x)\).
-
\(E_i(x)\) 可以是任意的网络, 所以现在的问题主要是如何设计 \(G\). 倘若我们希望选择 \(k\) 给 experts, 可以:
\[G(x) = \text{Softmax}( \text{KeepTopK}(H(x), k), ) \\ H(x)_i = (x \cdot W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x \cdot W_{noise})_i), \\ \text{KeepTopK}(v, k)_i = \left \{ \begin{array}{ll} v_i & \text{if } v_i \text{ is in the top} k \text{ elements of } v. \\ -\infty & \text{otherwise}. \end{array} \right . \] -
特别的是, 这里加了高斯噪声, 并用 \(W_{noise}\) 去调节不同位置的噪声的比重, 从而可以实现负载平衡 (?).
训练
-
如果不对 \(G\) 加以额外的限制, 容易出现某些 experts 持续获得较大的权重, 所以本文引入了一个 soft constraint
\[L_{importance}(X) = w_{importance} \cdot CV(Importance (X))^2, \\ Importance(X) = \sum_{x \in X} G(x) \]CV 作者说是 variation, 是方差吗?
-
有了 soft constraint, 依然会出现每个 expert 接受的样本数量的差别很大 (有些 expert \(i\) 可能会接受很少的样本但是其上 \(G(x)_i\) 都很大, 有些 expert \(i\) 可能接受很多的样本, 但是其上 \(G(x)i\) 都很小). 所以作者额外添加了对于选择概率的约束.
-
对于样本 \(x\), expert \(i\) 被选择的概率为 (感觉这个定义应该是有问题的)
\[P(x, i) = Pr\bigg( (x \cdot W_g)_i + \text{StandardNormal}() \cdot \text{Softplus}((x \cdot W_{noise})_i) > kth_excluding (H(x), k, i) \bigg). \]其中 \(kth_excluding(v, k, i)\) 表示 \(v\) 中的 k-th 大的值 (排除 \(i\)).
-
所以,
\[P(x, i) = \Phi( \frac{(x \cdot W_g)_i - kth_excluding(H(x), k, i)}{ \text{Softplus}((x \cdot W_{noise})_i) } ). \] -
定义
\[Load(X)_i = \sum_{x \in X} P(x, i), \]则
\[L_{load}(X) = w_{load} \cdot \text{CV}(Load(X))^2. \]
