Feedforward network를 하위 집합인 expert로 나누어 소수의 적합한 expert만이 계산에 참여함
[Github]
[arXiv](Current version v1)
Abstract
Feedfoward의 하위 네트워크인 Sparsely-Gated Mixture-of-Experts layer (MoE)를 통해 약간의 계산 효율성 손실만으로 모델 용량을 1000배 이상 향상하고 성능을 크게 향상함.
The Structure of The Mixture-of-Experts Layer
MoE는 n개의 전문가 네트워크 세트와 희소 벡터를 출력하는 gating network G로 구성된다.
i번째 expert의 출력을 Ei(x), gating output을 G(x)라 하면 MoE의 출력은 다음과 같다.
G(x)의 대부분의 출력이 0이므로 소수의 expert만 계산하면 되고 expert가 매우 많다면 2단계의 계층적 MoE를 사용할 수 있다.
Gating network
Softmax Gating
훈련 가능한 가중치 행렬을 곱한 다음 softmax 적용.
Noisy Top-K Gating
Softmax gating에 노이즈와 top-k 추가.
Training the Gating Network
나머지 모델과 함께 간단한 역전파를 통해 gating network를 훈련한다.
Addressing Performance Challenges
The shrinking batch problem
k개의 expert만 선택함으로써 훈련 효율성이 저하되어 배치 크기를 늘려야 하는데 활성화를 저장하는데 필요한 메모리 때문에 배치를 무한정 늘릴 수는 없다.
Mixing Data Parallelism and Model Parallelism
데이터 병렬에 대해 일반 계층과 gating network는 기존의 방법에 따라 분배하지만, 각 expert의 복사본은 하나로 유지하고 각 expert는 모든 데이터 병렬 입력 배치에서 관련 예제로 구성된 결합 배치를 받는다. (라고 하는데, d개의 장치에 대해 일반 계층은 d개로 복사하고, MoE에 대해서는 각 expert의 하위 집합을 지네들끼리 나눠 가진 다음, gating network의 분류에 따라 다시 취합해서 해당 expert를 가진 장치로 보내준다는 뜻인 듯?)
Increasing Batch Size for a Recurrent MoE
메모리를 절약하기 위해 expert의 hidden layer의 활성화를 저장하지 않고 역전파 시 다시 계산하는 활성화 재계산을 활용한다.
Network bandwidth
분산 컴퓨팅의 주요 문제는 네트워크 대역폭이다. 통신 시간에 비해 계산 시간이 적으면 통신 오버헤드가 발생한다. Expert에서 더 많은 hidden unit을 사용하여 쉽게 계산 효율을 높일 수 있다.
Balancing Expert Utilization
소수의 expert가 gating network에 의해 훨씬 더 많이 선택될 수 있다.
배치 당 gating network의 출력의 합을 중요도로 정의하고
중요도의 제곱에 대한 변동 계수(Coefficient of Variation)에 직접 지정한 가중치를 곱한 손실을 통해 모든 expert가 동등한 중요도를 갖도록 보장한다.
하지만 어떤 expert는 큰 가중치를 가진 예제를 많이 받고, 다른 expert는 작은 가중치를 가진 예제를 많이 받을 수 있으며, 이 또한 성능 문제를 일으킬 수 있다. 따라서 두 번째 손실 함수 Lload 추가.
Load Balancing Loss
위와 비슷하지만 G 대신 입력 x와 전문가 i에 대해 gating network가 0이 아닌 값을 출력할 확률 P를 사용한다.
