[논문 리뷰] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

Lean Wang, Huazuo Gao외 3명|arXiv (Cornell University)|2024. 08. 28.

Distributed Sensor Networks and Detection Algorithms인용 수 6

한 줄 요약

이 논문은 Loss-Free Balancing을 도입하여 동적으로 전문가 라우팅 점수에 편향을 주어 MoE 부하 균형을 제어하는 보조 손실 없이 perplexity와 부하 균형을 개선하고, 그래디언트에 간섭하지 않습니다.

ABSTRACT

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.

연구 동기 및 목표

• MoE에서 라우팅 붕괴와 과도한 계산을 피하기 위한 전문가 라우팅의 균형 필요성 동기 부여.
• 간섭 그래디언트를 도입하지 않고 부하 균형을 제어하기 위해 Loss-Free Balancing 제안.
• 이 방법이 1B 및 3B MoE 모델에서 보조 손실 기반 방법보다 더 나은 검증 perplexity와 부하 균형을 산출함을 입증.
• 전문가 병렬화와의 호환성을 보여주고 대안 방법과 비교한 이론적 누설을 분석합니다."],
• method':['상위-K 라우팅 전에 게이팅 점수에 추가되는 전문가별 편향을 도입합니다.','관찰된 전문 로드에 따라 각 바이어스 b_i를 반복적으로 업데이트하여 무거운 로드를 가진 전문가를 억제하고 가벼운 로드를 가진 전문가를 높입니다.','주요 목적에 그라디언트를 추가하지 않고 학습 배치마다 바이어스를 업데이트하기 위해 알고리즘 1을 사용합니다.','벤치마크 MoE 설정에서 보조 손실로 제어된 부하 균형 및 Expert Choice와의 Loss-Free Balancing 비교.','검증 perplexity와 전역 부하 균형 지표 MaxVio(전역 및 배치 버전)를 사용하여 평가합니다.','전문가 병렬화와의 호환성을 논의하고 미래 토큰 누설 우려를 분석합니다.'],
• research_questions':['게이팅 점수의 보조 손실 없는 바이어스가 보조 손실보다 더 나은 부하 균형을 산출합니까?','Loss-Free Balancing이 1B 및 3B MoE 모델에서 모델 성능(퍼플렉시티)과 부하 균형을 모두 개선할 수 있습니까?','이 방법은 전문가 병렬화와 호환되며 간섭 그래디언트로부터 자유로운가?','Loss-Free Balancing은 미래 토큰 누설 및 학습 안전 측면에서 Expert Choice와 어떻게 비교됩니까?'],
• key_findings':['Loss-Free Balancing은 1B 및 3B 모델 모두에서 보조 손실 제어 방법보다 더 낮은 검증 perplexity를 달성합니다.','전역 MaxVio 부하 균형 지표는 Loss-Free Balancing에서(0.04)가 보조 손실보다 현저히 낮습니다(1B에서 0.72; 3B에서 0.52).','Loss-Free Balancing은 MaxVio_batch 곡선에서 지속적인 이점을 보여 학습 전반에 걸쳐 더 나은 부하 균형을 유지합니다.','이 방법은 전문가 병렬화와의 호환성을 유지하며 계산-배치 크기가 커질수록 균형이 개선됩니다.','가산적 전문가 바이어스가 곱셈 바이어스보다 균형 및 성능에서 우수합니다.','소프트맥스 게이트 실험에서 해당 설정에서 Loss-Free Balancing은 더 나은 균형과 보조 손실보다 약간 더 나은 perplexity를 보여줍니다.'],
• table_headers':['모델 크기','부하 균형 방법','검증 perplexity','MaxVio_global'],
• table_rows':[[
• 1B
• Loss-Controlled
• 9.56
• 0.72

제안 방법

• 1B
• Loss-Free
• 9.50
• 0.04

실험 결과

연구 질문

• RQ1 3B
• RQ2 Loss-Controlled
• RQ3 7.97
• RQ4 0.52

주요 결과

모델 크기	부하 균형 방법	검증 perplexity	MaxVio_global
1B	Loss-Controlled	9.56	0.72
1B	Loss-Free	9.50	0.04
3B	Loss-Controlled	7.97	0.52
3B	Loss-Free	7.92	0.04

• 3B
• Loss-Free
• 7.92
• 0.04