MoE 로드 밸런싱과 그래디언트 간섭

부하를 맞추는 신호를 손실 함수에 넣느냐 라우팅 점수에 넣느냐가 MoE 학습 품질을 가릅니다.

top-K 라우팅과 routing collapse

희소 전문가 혼합(Sparse Mixture-of-Experts, MoE) 레이어는 피드포워드 블록 하나를 $N$개의 전문가로 쪼갭니다. 게이트 함수가 토큰마다 전문가별 점수를 매기고, 상위 $K$개만 활성화해 나머지 전문가의 계산을 건너뜁니다. 전체 파라미터 수는 유지한 채 토큰당 연산량만 줄이는 구조입니다.

문제는 이 선택이 자기 강화된다는 점입니다. 특정 전문가가 반복해서 뽑히면 그 전문가만 그래디언트를 받아 좋아지고, 좋아진 만큼 다시 더 자주 뽑힙니다. 나머지 전문가는 학습에서 소외되고, 결국 소수만 일하는 routing collapse에 빠집니다.

두 번째 비용은 학습 인프라 쪽에서 나옵니다. 분산 학습에서는 전문가마다 처리할 토큰 수(capacity)에 상한을 두는데, 과부하 전문가로 몰린 토큰은 상한을 넘는 순간 드롭됩니다. 부하 편중은 모델 품질 문제이면서 동시에 연산 낭비 문제입니다.

용어는 하나 구분해 둡니다. 여기서 로드 밸런싱은 전문가 사이의 토큰 분배를 뜻하며, 서버 앞단의 L4/L7 네트워크 로드 밸런싱과는 이름만 같습니다. 균형을 맞추는 대상도, 신호를 주는 방법도 다릅니다.

auxiliary loss의 간섭 그래디언트

Switch Transformer 이후의 표준 해법은 보조 손실(auxiliary loss)입니다. 전문가 $i$가 선택된 빈도 $f_i$와 그 전문가에 매겨진 평균 게이팅 점수 $p_i$를 곱해 모두 더하고, 계수 $\alpha$를 붙여 주 손실에 얹습니다.

$\mathcal{L}_{aux} = \alpha \, N_E \sum_i f_i p_i$

빈도와 점수가 한 전문가에 동시에 몰릴수록 이 값이 커집니다. 따라서 이 항을 줄이는 방향이 곧 균등 분배 방향이고, 계수 $\alpha$ 하나가 균형을 얼마나 세게 밀지 정합니다.

그 계수가 두 목표를 동시에 쥐고 있다는 점이 문제입니다.

근원은 보조 손실이 라우팅 점수의 그래디언트 계산에 직접 개입한다는 데 있습니다. 역전파는 다음 토큰 예측에서 온 신호와 부하 균형에서 온 신호를 구분하지 못하고 합쳐서 흘려보냅니다. Wang, Gao, Zhao, Sun, Dai(arXiv:2408.15664)는 이렇게 섞여 드는 성분을 간섭 그래디언트(interference gradient)라고 부릅니다.

bias 업데이트로 선택만 바꾸기

같은 논문이 제안한 Loss-Free Balancing은 손실 항을 없애고 개입 지점을 옮깁니다. 전문가마다 스칼라 bias $b_i$를 하나씩 두고, top-K를 고를 때만 라우팅 점수에 더합니다. 선택이 끝난 뒤 출력을 가중합할 때는 bias를 뺀 원래 점수를 씁니다.

# 매 스텝(또는 매 배치)마다
for expert i in 1..N:
    e_i = c_avg - c_i          # 평균 부하 - 전문가 i의 실제 부하
    b_i = b_i + u * sign(e_i)  # u: 업데이트 속도(논문 권장값 0.001)
 
# top-K 선택에만 쓰는 점수
g_i,t = s_i,t + b_i            # s_i,t: 원래 라우팅 점수
 
# 출력 가중합에는 bias를 더하지 않은 s_i,t를 사용

bias가 선택에만 관여하므로 모델 파라미터로 흐르는 그래디언트는 그대로 남습니다. 부하가 평균보다 적은 전문가는 bias가 올라가 다음 스텝에 뽑힐 확률이 커지고, 과부하 전문가는 반대로 내려갑니다. 균형 신호가 손실을 거치지 않고 선택 경계만 밀어내는 셈입니다.

갱신 비용도 가볍습니다. 스텝마다 전문가 수만큼만 더하면 되므로 배치당 $O(N)$이고, 토큰 수 $T$에 비례하는 보조 손실의 $O(N \times T)$보다 작습니다.

논문의 대조 실험은 두 지표를 함께 봅니다. MaxVio는 가장 부하가 큰 전문가와 평균 부하의 편차를 정규화한 값으로, 낮을수록 균형이 좋습니다.

두 규모 모두에서 perplexity가 내려가는 동시에 MaxVio가 90% 이상 줄었습니다. 균형을 얻으려면 성능을 내줘야 하던 관계 자체가 사라졌습니다.

sign 업데이트가 수렴하는 이유

부호 하나로 정의된 규칙이 왜 발산하지 않는지는 원 논문에서 실험으로만 확인됐습니다. Han과 Zhong(arXiv:2512.03915)은 토큰-전문가 할당을 정수계획법(Integer Programming, IP)으로 정식화해 이 규칙의 정체를 밝힙니다. 목적함수는 $\max \sum_{i,k} \gamma_{ik} x_{ik}$입니다. 제약은 각 토큰이 정확히 $K$개 전문가에, 각 전문가가 최대 $L$개 토큰에 배정되는 것입니다.

전문가 용량 제약을 쌍대 변수 $p_k$로 완화하면 라그랑지안이 다음 형태가 됩니다.

$\mathcal{L}(x,p) = \sum_{i,k}(\gamma_{ik}+p_k)x_{ik} - L\sum_k p_k$

여기서 $\gamma_{ik}+p_k$는 bias를 더한 라우팅 점수와 정확히 같은 자리에 놓입니다. 쌍대 변수는 부하가 용량에 못 미치면 올리고 넘치면 내리는 규칙으로 갱신합니다.

$p_k^{(n+1)} \leftarrow p_k^{(n)} + \varepsilon_k^{(n)}(L - A_k^{(n)}), \quad \varepsilon_k^{(n)} = \frac{u}{|L - A_k^{(n)}|}$

$A_k^{(n)}$은 $n$번째 스텝에서 전문가 $k$가 받은 토큰 수입니다. 스텝 크기가 편차의 절대값을 그대로 상쇄하므로 갱신량에는 부호와 $u$만 남습니다. DeepSeek의 sign 업데이트와 정확히 일치하는 형태입니다.

같은 논문은 이 대응에서 세 가지 보장을 끌어냅니다.

확률적 온라인 설정에서는 쌍대 목적함수의 강볼록성(strong convexity)을 이용합니다. 스텝 크기를 $u/n$으로 줄이면 $O(\log T)$ 기대 후회(regret) 상한이 증명됩니다. 경험적으로 찾아낸 트릭이 할당 문제의 표준 쌍대 상승법(dual ascent)이었다는 사후 해석입니다.

micro-batch 균형과 전문가 특화

균형 신호를 어디에 넣느냐만큼 부하를 어느 단위에서 재느냐도 결과를 바꿉니다. Qiu 등(arXiv:2501.11873)은 대부분의 MoE 학습 프레임워크가 병렬 처리 편의상 로드 밸런싱 손실을 micro-batch 단위로 계산한다고 지적합니다. 대규모 모델의 micro-batch에는 시퀀스가 몇 개 들어가지 않습니다.

시퀀스가 적으면 그 손실은 사실상 시퀀스 레벨 균형을 강제합니다. 코드 한 편으로 채워진 시퀀스의 토큰조차 모든 전문가에 고르게 흩어야 손실이 줄어듭니다. 도메인별로 전문가가 갈리는 특화(specialization)가 이 지점에서 막힙니다.

해법은 측정 단위를 넓히는 것입니다. 여러 micro-batch의 전문가 선택 빈도를 동기화해 global-batch 단위로 손실을 계산하면, 균형 요구가 배치 전체 수준으로 올라갑니다. 여러 도메인이 섞인 배치 전체만 고르면 되므로, 개별 시퀀스 안에서는 전문가가 특정 도메인에 몰릴 여지가 생깁니다. 선택 빈도는 전문가 수만큼의 원소를 가진 벡터 하나뿐이라 동기화 통신 비용도 거의 없습니다.

3.4B(활성 0.6B)부터 43B(활성 6.6B)까지의 실험에서 global-batch 방식이 사전학습 perplexity와 다운스트림 성능을 함께 개선했습니다. balance batch size를 2에서 128로 키울수록 개선폭이 커졌습니다.

세 방식의 트레이드오프

세 접근을 나란히 두면 선택 기준이 드러납니다.

두 축은 배타적이지 않습니다. 신호를 bias로 줄지 손실 항으로 줄지, 부하를 micro-batch에서 잴지 global-batch에서 잴지는 따로 고를 수 있습니다. 최근 오픈소스 MoE 모델들은 loss-free bias 업데이트와 global-batch 통계 집계를 함께 쓰는 쪽으로 모이고 있습니다.

남는 경계도 분명합니다. 여기까지는 모두 훈련 단계에서 그래디언트를 지키며 균형을 맞추는 이야기입니다. 배포된 모델은 서비스 트래픽의 프롬프트 분포에 따라 특정 전문가에 요청이 몰리는 hot expert 문제를 다시 겪습니다. 추론 단계의 편중은 vLLM이나 TensorRT-LLM 같은 서빙 프레임워크의 토큰 단위 전문가 그룹핑과 all-to-all 통신 최적화가 담당하는 영역입니다.

정리

loss-free 방식의 요지는 부하 균형의 초점을 손실 함수에서 라우팅 선택으로 옮긴다는 데 있습니다. 보조 손실은 균형 신호를 역전파에 섞어 계수 $\alpha$ 하나로 균형과 성능을 맞바꾸게 만들지만, bias 업데이트는 top-K 선택에만 개입해 그래디언트를 건드리지 않습니다. Wang 등의 3B/200B 토큰 실험에서 perplexity가 7.97에서 7.92로, MaxVio가 0.52에서 0.04로 함께 내려간 것이 그 결과입니다. Han과 Zhong의 정식화는 이 sign 업데이트가 토큰-전문가 할당 문제의 쌍대 상승법과 동치임을 보여, 휴리스틱에 수렴 근거를 붙였습니다. 다만 부하를 어느 단위에서 잴지는 여전히 별개 선택이고, 서빙 시점의 전문가 편중은 서빙 프레임워크가 다뤄야 할 문제로 남습니다.