Branch-Train-MiX: 전문가 LLM을 Mixture-of-Experts LLM으로 통합하기

초록

우리는 대규모 언어 모델(LLM)이 코딩, 수학적 추론, 세계 지식과 같은 여러 전문 분야에서의 역량을 갖추도록 효율적으로 훈련하는 방법을 연구합니다. 우리의 방법인 Branch-Train-MiX(BTX)는 시드 모델에서 시작하여, 병렬 처리 방식으로 전문가 모델을 분기하여 높은 처리량과 낮은 통신 비용으로 훈련합니다. 개별 전문가 모델이 비동기적으로 훈련된 후, BTX는 이들의 피드포워드 파라미터를 Mixture-of-Expert(MoE) 레이어의 전문가로 통합하고 나머지 파라미터를 평균화한 다음, 토큰 수준의 라우팅을 학습하기 위한 MoE 미세 조정 단계를 진행합니다. BTX는 두 가지 특수한 경우를 일반화합니다: 하나는 라우팅 학습을 위한 MoE 미세 조정 단계가 없는 Branch-Train-Merge 방법이고, 다른 하나는 전문가 모델을 비동기적으로 훈련하는 단계를 생략한 sparse upcycling입니다. 대안적인 접근법과 비교했을 때, BTX는 정확도와 효율성 간의 최적의 균형을 달성합니다.

English

We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as coding, math reasoning and world knowledge. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and reduced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage to learn token-level routing. BTX generalizes two special cases, the Branch-Train-Merge method, which does not have the MoE finetuning stage to learn routing, and sparse upcycling, which omits the stage of training experts asynchronously. Compared to alternative approaches, BTX achieves the best accuracy-efficiency tradeoff.

Branch-Train-MiX: 전문가 LLM을 Mixture-of-Experts LLM으로 통합하기

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

초록

Support