Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

[arXiv](2024/03/12 version v1)

 

Abstract

LLM을 각 도메인에 대해 개별 훈련한 후 단일 MoE 모델로 통합

 

 

Branch-Train-MiX

Branch-Train-Merge와 Mixture-of-Experts를 결합하였다.

 

Branch & Train: Embarrassingly Parallel Expert Training

각 LLM을 각 데이터셋에 대해 훈련한다.

Branch-Train-Merge에서는 이렇게 훈련된 LLM 중 사용할 LLM들을 선택하고 단순히 평균을 계산하였다.

 

MiX: Combining Separate Experts to be a Mixture-of-Experts

Feedforward layer에서는 MoE와 같은 구조를 사용하며

 

self-attention이나 embedding 등 나머지 레이어는 FF layer 보다 전문화 수준이 낮다는 가정 하에 단순히 평균을 계산한다.

 

Mixed model은 당연히 최적이 아니지만 소량의 fine-tuning 만으로도 성능이 크게 개선된다고 한다.

 

Load balancing

배치 B에서 각 전문가가 처리한 토큰의 비율 u와 라우팅 확률 p의 내적으로 p가 uniform일 때 최소가 된다.

 

 

 

Experiments

Math, Code, Wikipedia 도메인에 대해 훈련하며 시드 모델인 LLaMA-2와 함께 4개의 전문가 사용.

 

레딧 평가가 좋길래 읽어봤는데 영...

결과로 제시하는 것들도 별로 의미 없어 보이고 애초에 4개의 전문가로 무슨 비교를 해보겠다는 건지...

Meta 논문인데 니네 돈 없는 것도 아니잖아...?

근데 여기서 MoE 모델은 Sparse upcycling 밖에 없고... 흠... 전문 데이터셋을 보지 못한 모델들과 평균을 비교하는 게 무슨 의미가 있는지...

 

그리고 사실 본문에도 이것저것 내용이 더 있었지만 뭐 성능 향상은 없었고...