MoE(Mixture-of-Experts, ICLR 2017)

Methods

Figure 1: A Mixture of Experts layer embedded within a recurrent language model.

 

conditional computation이라는 방법은 각 example에 따라 subnetwork가 active되어 model capacity를 늘리는 방법이다. model의 크기를 키우는 것에 비해서 더 많은 capacity를 늘릴 수 있다는 장점이 있는데, 여기에는 기술적인 문제가 몇 개 있다:

 

1) GPU는 branching보다 arithmetic에 유리하다.

2) dataset과 batch size가 아주 커져야 한다. 

3) network bandwidth가 bottleneck이 된다. 

 

여기서는 이를 해결하기 위한 방법인 Sparsely-Gated Mixture-of-Experts Layer(MoE)를 제안한다. 

 

Fig. 1에서 보이듯 각 layer는 $n$개의 expert network $E_1, \cdots E_n$와 gating network $G$로 구성된다. 이때 gating network의 output은 sparse $n$-dimensional vector이다.

 

이를 이용한 MoE module의 output은 다음과 같이 계산된다:

 

$$ y = \sum ^n _{i=1} G(x)_i E_i(x)$$

 

$G(x)_i=0$이면 계산을 생략한다. 그 결과 실제 experiment에서는 수천 개의 expert 중 몇 개만 activate된다. 

 

 

Gating Network

가장 simple한 형태의 gating은 softmax를 사용하는 것이다: 

 

$$ G_σ(x) = Softmax(x \cdot W_g)$$

 

그러나 이 방법은 sparse하지 않으므로 computationally efficient하지 않다. 따라서 Noisy Top-K gating을 사용해서 sparsity를 달성한다. tunable Gaussian noise를 더한 뒤 top k value를 제외한 모든 값을 -∞로 설정해서 sparse하도록 만든다. 이는 다음과 같이 표현된다:

 

$$ G(x) = Softmax(KeepTopK(H(x), k)) $$

$$ H(x)_i = (x\cdot W_g)_i + StandardNormal()\cdot Softplus ((x\cdot W_{noise})_i)$$

$$ KeepTopK(v,k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise}. \end{cases} $$

 

이 gating network는 다른 network와 함께 train된다.

 

 

Balancing Expert Utilization

그러나 이렇게 할 경우 실제로는 일부 expert에 large weight가 부과되어 빠르게 local minima로 converge하는 문제가 보고된 바 있다. 따라서 constraint를 부과한다. 

 

$$ Importance(X) = \sum _{x\in X} G(x)$$

$$L _{importance} (X) = w_{importance} \cdot CV(Importance(X))^2 $$

 

즉 위와 같이 importance term의 coefficient of variation을 loss term에 더해준다. 이때 importance는 batchwise로 더해지므로, cv square를 loss term으로 추가하는 것은, 특정 expert가 너무 많은 importance를 갖지 않도록 regularize한다.  

 

 

Experiments

Figure 2: Model comparison on 1-Billion-Word Language-Modeling Benchmark.

Fig. 2의 왼쪽을 보면 model의 크기를 키울 수 있고, 이 경우 perplexity가 감소한다. 우측에서는 같은 computational budget에서의 perplexity를 LSTM과 비교한다. 

 

---

References

 

Footnotes