MoE LLM 학습 성능을 개선하는 거의 공짜에 가까운 방법, 글로벌 배치 로드 밸런싱

Mixture-of-Experts (MoE) 아키텍처 학습 시 사용되는 기존의 **로드 밸런싱 손실(Load Balancing Loss)**은 주로 마이크로 배치(micro-batch) 단위로 계산된다. 하지만 특정 도메인의 데이터로만 구성된 마이크로 배치의 경우, 모델이 모든 전문가를 균등하게 사용하도록 강제하여 전문가의 **도메인 특화(specialization)**를 방해하는 문제가 발생한다.

이를 해결하기 위해 제안된 **글로벌 배치 로드 밸런싱(Global-batch load balance)**은 모든 병렬 그룹 간의 전문가 선택 빈도를 동기화하고, 전체 마이크로 배치에 걸쳐 손실을 집계한다. 전문가 선택 빈도는 단일 벡터 형태이므로 병렬 그룹 간 동기화 비용이 매우 적어 '거의 공짜(almost free lunch)'에 가까운 효율을 보인다.

실험 결과, 3.4B, 15B, 43B 규모의 다양한 MoE 모델에서 글로벌 배치 로드 밸런싱은 다음과 같은 성과를 거두었다:

• 모든 설정(모델, 데이터, 태스크)에서 기존 방식보다 우수한 성능 달성
• 특정 도메인에 따라 전문가가 활성화되는 뚜렷한 도메인 특화 현상 확인
• 더 큰 배치 사이즈에서도 안정적인 학습 가능

학습 효율 저하를 방지하기 위해 글로벌 배치 손실에 소량의 마이크로 배치 로드 밸런싱 손실을 추가하면, 모델 성능을 유지하면서도 계산 효율성을 높일 수 있다.