기존 sparse MoE의 문제

• Load Imbalance
  • 과사용되는 전문가가 없도록 토큰을 균등하게 분배하는 것이 중요하다.
    • 하지만 균등하게 분배되지 않는 경우가 종종 발생함
      • 일부 전문가에 많은 토큰이 할당돼서, 20~40% 정도의 토큰이 드랍되는 것을 관찰함
• 전문가의 전문성이 떨어짐
  • 균등하게 분배하도록 auxiliary loss를 사용하는데, auxiliary loss의 weight에 따라 전문성이 떨어질 수 있다.
    • 이유: weight를 너무 키우면 균등하게 분배하는 것에 초점이 맞춰져서 전문가 네트워크가 전문성을 갖기 어려워진다. 전문성을 갖기 위해, "비슷한" 토큰이 특정 전문가에 라우팅 돼야 하는데, 이는 불균등 분배로 이어질 수 있기 때문
• 토큰의 종류와 상관없이 동일한 양의 계산(compute)으로 토큰을 처리함
  • 각 토큰을 동일한 크기와 구조의 k 개 전문가 네트워크로 처리함

아키텍처

• GLaM과 큰 차이 없음
• 라우터가 토큰을 할당하는 방법
  • $topK(transpose(softmax(XW_g)), k)$
    • $topK$는 입력 행렬의 각 행벡터에서 제일 큰 K 개를 고르는 연산
    • $X$는 $n*d$. $n=$ 배치의 토큰 수, $d=$ 모델 차원
    • $W_g$는 $d * e$. $e=$ 전문가 네트워크 수
    • 이는 GLaM이나 swtich transformer와 같은 "token choice" 모델의 $topK(softmax(XW_g),k)$와 유사함. 다만 "token choice" 모델의 라우터는 각 토큰을 k 개의 전문가 네트워크에 할당하고, 논문에서 제시한 "expert choice" 모델의 라우터는 $transpose$를 사용함으로써 각 전문가 네트워크에 k 개의 토큰을 할당함. 이 방식을 통해, 모든 전문가 네트워크에 완벽히 동일한 개수의 토큰을 할당할 수 있음

한계

• 과거와 미래의 토큰 모두 고려하여 라우팅하기 때문에, autoregressive 모델에 바로 사용하기 어렵다.
• 추론 때 작은 크기의 배치를 갖는다면 이 방법은 적합하지 않다.
  • 나의 추측 1: $tranpose(softmax(XW_g))$의 행벡터 차원이 k보다 작을 수 있음. 이 경우 k 개의 토큰이 없기에 각 전문가 네트워크에 k 개의 토큰을 할당하는 건 불가능
  • 나의 추측 2: 행벡터 크기가 k보다 '유의미하게' 크지 않다면 dense 네트워크와 큰 차이 없음. k 개의 토큰을 처리한다고 가정해보자. 그러면 각 전문가 네트워크에 정확히 k 개의 토큰이 들어가기에 dense 네트워크보다 계산량이 적어지지 않는다.
  • 가능한 방법: 추론 때 "token choice"를 사용하는 방법을 제시함.
    • $W_g$가 token/expert choice 둘 다 $d*e$ 차원을 가지기에 가능함
    • 다만 성능이 어떨지 실험하지 않음
• 균등하게 토큰을 분배하지만, 토큰이 생략(token dropping)되는 문제는 해결하지 못함