논문 리뷰) Mixtral of Experts (Mixtral 8x7B)

Mistral AI는 파리에 본사를 두고 Google의 DeepMind와 Meta 출신이 공동 창립한 스타트업이다.

Mixtral 8x7B 모델로 우리에게 많이 알려졌는데,

Mistral 7B 모델의 기본 구조를 그대로 가져가면서, Mixture of Expert를 적용한 것이 큰 특징이다.

(Mistral 7B 모델의 기본 구조에 대한 내용은 이전 글에서 소개)

 

 

[통합 정리]
- 대부분 Mistral 7B와 같은 구조이지만 Mixture of Expert를 도입한 것이 큰 특징

두 모델의 Parameter 비교

- FFN(Feed Forward Network)이 8개의 Feedforward blocks(이를 experts라 부름)와 Gate의 구성으로 대체된 차이점 존재.

- Gate(router)의 역할은 각 토큰에 대해 두 전문가(experts)를 선택하고, 그 출력을 결합하는 것.

Mixture of Experts Layer

  
- 각 토큰은 오직 두 전문가들만 선택하지만, 각 timestamp마다 선택된 전문가들은 다를 수 있다. 
- 결과적으로 각 토큰은 47B parameters에 접근할수 있으면서도, 

   실제 inference 동안엔 오직 13B parameters만을 사용하는셈. 
- 13B 파라미터만 활성화되므로, Computation cost 이득. 
  메모리 같은 경우에는 47B(전체 모델 사이즈) 전부 올라가겠지만, 

  그래도 라마2 70B보단 적음. (메모리 이득 x, 연산 이득 o)
- Mixtral은 32k 토큰의 context size로 훈련
- 평가된 모든 벤치마크에서 LLaMA 2 70B와 GPT-3.5를 동등/상회

 

---

본 글은 Mixture of Expert라는 개념에 집중하여 리뷰를 진행하였습니다.

 

< Dense layer(ex. FFN) 방식과 Sparse Expert (Mixture of Expert) 방식 >

출처 : https://4n3mone.tistory.com/10

- Sparse Expert의 목적 : 더 적은 연산량과 특화된 expert의 파라미터로써, '자원'과 '성능' 측면에서 이점을 얻길 기대
- Sparse Expert의 단점 :

  Sparse expert models may introduce training difficulties over a vanilla Transformer. 

  (Sparse expert 모델은 vanilla transformer보다 학습에 어려움이 있다)
  Instability can result because of the hard-switching (routing) decisions at each of these layers. 

  (각 Layer에서 hard-switching 결정하므로 불안정성이 발생 가능) 

  위 내용은 'Switch transformer' 논문 중에서 발췌한 내용으로 

  Switch transformer에서는 Selective precision with large sparse models 방식으로 안정성을 도모했다고 함.

 

 

< MoE : Mixture of Expert >

- MoE가 최초 제안된 논문은 Geoffrey Hintton의 Adaptive Mixtures of Local Experts (1991)
- LSTM을 활용한 MoE도 존재 (Outrageously Large Neural Networks: The Sparsely-gated Mixture-of-Experts Layer(‘17)
- GShard 논문에서 'Transformer' Encoder의 FFN 레이어를 MoE로 대체하는 방식을 제안.

LSTM MoE

GShard

 

 

< Mixtral 8x 7B와 GShard >

 Mixtral 8x 7B의 MoE 방식은 GShard의 아키텍처와 유사하나,  
 Mixtral은 모든 FFN을 MoE 구조로 대체하였고, GShard는 1개만 대체하였다. 

출처 : https://jun048098.tistory.com/8

 

 

< Mixtral 8x 7B에서 MoE 함수와 출력 >

따라서, Mixtral 8x7B에서 MoE 함수와 출력은 아래와 같이 표현될수 있다.

 

< Expert에서 적용된 SwiGLU Activation Function > 

- GLU Variants Improve Transformer 논문에서 소개.
- SwiGLU : Swish Activation Function (Swi) + Gated Linear Unit (GLU)

SwiGLU Activation Function

출처 : https://blog.paperspace.com/swish-activation-function/

- 많은 계산량의 단점에도 불구, 높은 성능으로 LLM에도 많이 쓰이고 있음 (LLama2, Mixtral 8 x 7B 등)
- GLU Variants Improve Transformer 논문에서 성능이 높게 나왔다고, 실험 결과를 공개했으나

  왜 성능이 높은지 명시하지 않았음(혹은 못했음)

출처 : GLU Variants Improve Transformer

 

※ Swish Activation Function (Swi)
    - Beta 값에 따라서, ReLu와 Linear를 유사하게 표현 가능. 

※ Gated Linear Unit (GLU)

 

 

 

[3. Result - 벤치마크 결과]

결과는 평가된 모든 벤치마크에서 LLaMA 2 70B와 GPT-3.5를 동등/상회

Table 3

 

[4. Instruction]

SFT + DPO 조합으로 fine-tuning하였음.

 

[5. Routing Analysis]

expert 별로 특정 도메인에 특화되는 지를 확인해봤다.
그러나, expert 별로 별다른 패턴이 관찰되지 않았다.