MoE 환유기: 6, 최적 할당으로 균형 촉진

URL: https://kexue.fm/archives/11619

Title: MoE环游记：6、最优分配促均衡 - 科学空间|Scientific Spaces

우리는 부하 균형(Load Balance)이 MoE 아키텍처에서 기본적이면서도 핵심적인 요소이며, 모델의 효율과 성능에 직접적인 영향을 준다는 것을 알고 있다. 본 시리즈에서는 이미 두 편의 글에서 부하 균형을 구현하는 두 가지 주류 접근을 소개했다. 하나는 《MoE环游记：2、不患寡而患不均》에서 소개한 고전적 방법 Aux Loss이고, 다른 하나는 《MoE环游记：3、换个思路来分配》에서 DeepSeek가 제안한 Loss-Free 방법이다. 두 방법은 각각 장점이 있으나, 동시에 한계도 있다.

본 글에서는 세 번째 접근인 최적 할당을 탐구한다. 이는 부하 균형을 등식 제약 하의 선형계획 문제로 보고 푸는 방식이다. 최종 형태만 보면 여전히 Loss-Free 범주에 속하지만, 전혀 다른 원리에 기반하며 더 정확하고 하이퍼파라미터가 없는 업데이트 방식을 제공한다.

방법 회고#

기존 두 방법 중 Aux Loss의 발상은 비교적 소박하며, 핵심은 “불안정한 곳을 벌점으로 눌러준다”는 것이다. 정규화 항으로 부하 불균형에 패널티를 부여한다. 하지만 Aux Loss에는 두 가지 문제가 있다. 첫째, 패널티 계수를 맞추기 어렵다. 너무 크면 주 손실(main loss)의 최적화를 방해하고, 너무 작으면 균형 효과가 떨어진다. 둘째, Aux Loss의 배경에는 STE(Straight-Through Estimator)가 있어, 그라디언트가 차선(suboptimal)이라는 뜻이다. 이는 부하 균형 이외의 알 수 없는 영향을 가져올 수 있다.

이를 위해 DeepSeek는 두 번째 방안 Loss-Free를 제안했는데, 이는 추가적인 바이어스 항을 도입해 정렬을 보조한다. 식은 다음과 같다:

$y=∑i∈argtop k ρ ρ i e i→y=∑i∈argtop k ρ+b ρ i e i$

여기서 $b$는 Expert의 순위를 조정하는 데만 쓰이며, Expert에 곱해지는 것은 여전히 $ρ i$이므로, $b$는 모델의 계산에 직접 참여하지도 않고 방해 그라디언트도 만들지 않는다. 하지만 $b$에 그라디언트가 없다는 것은, $b$를 위한 별도의 업데이트 규칙을 설계해야 한다는 뜻이기도 하다. 발상은 역시 직관적이다: $b i$가 클수록 $i$번째 Expert가 선택될 확률이 커지므로, 현재 부하 분포 $F$를 통계로 구한 뒤, $F i$가 기대값 $1/n$을 넘으면 $b i$를 줄이고, 아니면 $b i$를 늘린다. 즉

$b←b−γ sign(F−1/n)$

전반적으로 Loss-Free는 모델에 대한 “침투(intrusion)”가 더 적어 확실히 더 간결하고 우아하다고 할 수 있지만, 완벽하진 않다. Aux Loss처럼 패널티 계수는 없지만, 대신 $γ$라는 파라미터를 조정해야 한다. 이는 사실상 $b$의 학습률이다. 논문에서는 $γ=10−3$을 추천하지만, 이는 $ρ$를 Sigmoid로 활성화하는 것과 강하게 결합되어 있어, 활성화 함수를 바꾸면 $γ$도 다시 조정해야 한다.

또한 Sigmoid 활성화를 쓰더라도 어떤 층에서는 $ρ$ 분포가 꽤 “기형적”일 수 있는데, 이때 모델은 $γ$에 민감해져서 상수 $γ$로 부하 균형을 이루기 어렵다. 이런 경우는 드물지 않다. 예컨대 모델 앞쪽 몇 개 층에서 MoE를 쓰면 균형이 잘 잡히지 않는 경우가 많아서 “first_k_dense” 같은 동작이 생겨났고, 또 모델이 크거나 Expert 총수 $n$이 클 때도 일부 층에서 균형이 어려운 현상이 나타나기 쉽다.

선형계획#

이제 풀고자 하는 문제를 더 정확히 기술해 보자. $m$개의 Token이 있고, $i$번째 Token의 Router가 $n$개의 Expert에 대해 매긴 점수가 $s i=(s i,1,s i,2,⋯,s i,n)$라고 하자. 그러면 총 $m n$개의 점수가 있다. 이 점수들은 양수일 수도 음수일 수도 있고, 미리 정해진 값 범위에 들어갈 필요도 없다. 우리는 이 점수들을 바탕으로 할당 방안을 정해 각 Token이 어떤 Expert를 활성화할지 결정하고자 한다.

각 Token은 $k$개의 Expert를 선택한다고 약속하자. 그러면 기본 방안은 점수가 가장 높은 $k$개의 Expert를 고르는 것이다. 하지만 이 방식은 부하 불균형을 초래할 수 있다. 즉 어떤 Expert는 활성화 횟수가 눈에 띄게 많거나 적을 수 있다. 그래서 추가로 각 Expert가 정확히 $m k/n$번 활성화되도록 약속하자. 이 두 제약 아래에서 총점이 최대가 되는 할당을 구하면, 다음과 같이 정식화된다.

$max x i,j∈{0,1}∑i,j x i,j s i,j s.t.∑j x i,j=k,∑i x i,j=m k n$

여기서 $x i,j=1$은 $i$번째 Token이 $j$번째 Expert를 선택했음을 의미하고, $x i,j=0$은 선택하지 않았음을 뜻한다. 또한 $m k/n$은 두 등식 제약이 엄밀히 성립하려면 정수여야 하므로, 우선 이것이 성립한다고 가정하자. 각 Expert가 엄밀히 $m k/n$번 활성화되는 것은 가장 이상적인 완전 균일 상태이며, 실제로는 느슨하게 해도 되지만 이론 전개에서는 가장 강한 제약을 사용한다.

위 식에서 $x i,j$는 0 또는 1만 될 수 있으므로 이는 정수계획(integer programming) 문제다. 이산 최적화는 보통 어렵기 때문에, 완화(relaxation) 버전을 고려하자.

$max x i,j∈[0,1]∑i,j x i,j s i,j s.t.∑j x i,j=k,∑i x i,j=m k n$

이제 $x i,j$는 $[0,1]$의 임의의 실수를 취할 수 있고, 나머지 조건은 동일하다. 목표함수와 제약 등식 모두 $x i,j$에 대해 선형이므로, 이는 유계 영역에서의 선형계획 문제가 된다.

극대-극소#

제약 최적화는 보통 쉽지 않으므로, 제약을 제거한 max-min 형태(즉 “라그랑주 승수법”)을 고려하자:

$max x i,j∈[0,1]min α i,β j∑i,j x i,j s i,j−∑i α i(∑j x i,j−k)−∑j β j(∑i x i,j−m k n)$

만약 $∑j x i,j=k$ 및 $∑i x i,j=m k/n$이면 위 식은 식(4)와 동치이며 최대값은 유한하다. 하지만 둘 중 하나라도 만족하지 않으면, $min$ 단계에서 음의 무한대까지 갈 수 있어 최댓값은 음의 무한대가 된다. 음의 무한대는 유한값보다 명백히 작으므로, 결국 전자의 경우만 가능해져 식(4)와 동치가 된다.

목표(5)는 $x i,j,α i,β j$에 대해 모두 선형이다. 선형함수는 볼록이면서 오목이고, $[0,1]$은 볼록집합이므로 Minimax theorem 조건을 만족한다. 따라서 $max$와 $min$의 순서를 바꿀 수 있고, 이는 다음과 같다.

$min α i,β j max x i,j∈[0,1]∑i,j x i,j(s i,j−α i−β j)+k∑i α i+m k n∑j β j$

위 식에서는 $x i,j,α i,β j$에 관한 항을 분리해 놓았다. 자세히 보면 $max$ 단계는 직접 구할 수 있다. $s i,j−α i−β j>0$이면 목표를 키우기 위해 $x i,j=1$을 취하고, $<0$이면 $x i,j=0$을 취한다. $=0$이면 $x i,j$를 무엇으로 해도 결과가 같다. 즉

{x∗i,j=1,s i,j−α i−β j>0 x∗i,j=0,s i,j−α i−β j<0 x∗i,j∈[0,1],s i,j−α i−β j=0

여기서 $s i,j−α i−β j=0$은 특수한 경우이며, 발생 확률이 무시 가능하다고 가정하면 $x∗i,j$는 0 아니면 1이다. 물론 $s i,j−α i−β j=0$이라도 $x∗i,j$의 임의성을 이용해 제약을 만족하도록 0 또는 1로 둘 수 있다. 이로부터, 형태(4)는 원문제(3)를 완화했음에도 최적해가 원문제의 최적해와 동일하여 둘은 완전히 동치임을 알 수 있다.

분할정복#

위의 $x∗i,j$를 식(6)에 대입하면 $x∗i,j(s i,j−α i−β j)=max(0,s i,j−α i−β j)$이므로, 목표(6)는 다음으로 단순화된다.

$min α i,β j∑i,j max(0,s i,j−α i−β j)+k∑i α i+m k n∑j β j$

이를 풀기 위해 교대 최소화(alternating minimization)를 사용하자. 즉 먼저 $β j$를 고정하고 $α i$를 구한 뒤, $α i$를 고정하고 $β j$를 구하는 과정을 번갈아 수행한다. $α i,β j$는 대칭성이 뚜렷하므로 두 단계는 사실상 같은 문제다.

먼저 $β j$를 고정하고 $α i$를 구하면, 문제는 다음과 동치다.

$min α i∑i,j max(0,s i,j−α i−β j)+k∑i α i$

또한 각 $α i$ 항은 서로 독립적으로 더해지므로, 이는 $m$개의 독립적인 부분문제로 분해된다. 즉 아래첨자 $i$를 잠시 생략하고 문제를

$min α k α+∑j max(0,s j−β j−α)$

로 단순화할 수 있다.

전체 $s j−β j$를 내림차순으로 정렬하여 $s σ 1−β σ 1≥s σ 2−β σ 2≥⋯≥s σ n−β σ n$이라 하자(즉 $j$번째로 큰 원소가 $s σ j−β σ j$). 그리고 $s σ l−β σ l≥α≥s σ l+1−β σ l+1$임을 안다고 가정하면, 목표함수는

k α+l∑j=1(s σ j−β σ j−α)={k∑j=1(s σ j−β σ j)+l∑j=k+1(s σ j−β σ j−α)⏟≥0,l≥k k∑j=1(s σ j−β σ j)−k∑j=l+1(s σ j−β σ j−α)⏟≤0,l≤k

가 된다.

이는 $l>k$이든 $l<k$이든 목표함수를 더 키우는 방향임을 뜻하므로, 최소값은 오직 $l=k$일 때만 달성된다. 이때 결과는 $α$의 구체적인 값에 의존하지 않는다. 즉 $α∗$는 $s j−β j$의 $k$번째 큰 값과 $k+1$번째 큰 값 사이의 임의의 수가 될 수 있으며, 관례적으로 $α∗$를 $k+1$번째 큰 값으로 잡는다.

교대 반복#

아래첨자 $i$를 되살리면 다음을 얻는다. 임의의 주어진 $i$에 대해 $α∗i$는 $s i,j−β j$를 내림차순 정렬했을 때의 $k+1$번째 원소다. 마찬가지로 $α i$를 고정하고 $β i$를 구하면: 임의의 주어진 $j$에 대해 $β∗j$는 $s i,j−β j$를 내림차순 정렬했을 때의 $m k/n+1$번째 원소가 된다. 이 두 단계를 교대로 수행하는 것이 최종 알고리즘이다.

충분히 정확한 $α∗$와 $β∗$를 얻었다고 가정하면, 앞선 분석에 의해 $x∗i,j$는 자동으로 0 또는 1이 되고 제약 조건도 만족한다. 또한 $x∗i,j=1$은 $s i,j−α∗i−β∗j>0$에 대응한다. 그리고 제약 $∑j x∗i,j=k$를 결합하면, 각 Token $i$가 선택하는 Expert는 반드시 $s i−β∗$의 Top-k가 된다.

이는 추론 단계에서 $β∗$만 필요하고 $α∗$는 풀이 과정의 중간 변수일 뿐이므로, 학습이 끝난 뒤에는 무시할 수 있음을 뜻한다. 이는 매우 중요하다. 왜냐하면 $β$의 크기는 고정된 $n$인 반면, $α$의 크기는 $m$이며 $m$은 전역 배치 크기(global batch size)로서 동적으로 변하기 때문에, 과학적인 추론 형식이 아니기 때문이다. 이 특성을 이용한 풀이 흐름은 아래와 같으며, 여기서 des_sort는 내림차순 정렬을 의미한다.

Quantile Balancing (QB): 문제(3)의 교대 풀이 알고리즘 입력: 점수 행렬 $s∈R m×n$ 출력: 할당 $x∈{0,1}m×n$ 1:Initialize $β=0 1×n$ 2:For $t=1,2,⋯,T$ do 3:$α←des_sort(s−β,axis=1)[:,k:k+1]$ 4:$β←des_sort(s−α,axis=0)[m k/n:m k/n+1]$ 5:Output $x i,j=1$ if $j∈argtop k s i−β$ else 0

또 하나의 개선점은 “분위수(Quantile)” 개념을 사용해 “$n$개 수 중 $k+1$번째 큰 원소, $m$개 수 중 $m k/n+1$번째 큰 원소”를 통일하는 것이다. 둘 다 각 축에서의 “$1−k/n$ 분위수”다. Numpy, Jax, Torch 등 수치 프레임워크에는 quantile 함수가 있어 이를 쓰면 전체 정렬을 피하고 복잡도를 줄일 수 있다.

이 때문에 우리는 이 알고리즘을 “Quantile Balancing (QB)”라고 부른다.

함정에 주의#

하지만 아직 축하하긴 이르다. 여기에는 눈에 잘 띄지 않는 함정이 하나 있는데, 빠지면 본질적으로 잘못된 결과를 얻을 수 있으니 특별히 주의해야 한다.

앞서 QB의 추론은 $β$만 쓰고, 학습 과정에서도 $β$만 저장하면 충분하다고 했다. 그러므로 Loss-Free 관점에서 QB는 새로운 Bias 업데이트 방법을 제공하며, 이를 “Quantile Bias”라고 불러도 무방하다. 기존의 SignSGD든 본문의 Quantile이든 전체 Token 점수에 의존하므로 순서를 절대 헷갈리면 안 된다: 반드시 옛 $β$로 현재 배치의 Expert를 먼저 선택하고, 그 다음에 $β$ 값을 업데이트해야 정보 누설 문제가 생기지 않는다.

독자는 “전방 계산에 직접 참여하지 않는 Bias 벡터가 뭘 누설하느냐?”고 의아해할 수 있다. 직관적으로 누설될 정보가 많지 않은 것은 맞지만, 위험은 결국 존재한다. 작은 모델이라면 누설 버전도 시험해 볼 수 있겠지만, 큰 모델 학습에서는 이런 위험을 감수하면 안 된다. 모델이 크고 능력이 강할수록 어떤 미세한 버그든 증폭시킬 수 있기 때문이다. 따라서 훈련-추론 일치(train-inference consistency)를 지키고, 어떤 형태의 정보 누설 가능성도 차단하는 것이 대규모 모델 학습의 기본 태도다.

실제 학습 시나리오를 고려하면 QB는 또 조정할 수 있다. 원래는 0으로 초기화한 뒤 $T$번 반복하지만, 이를 이전 단계의 Bias에서 시작해 매 스텝마다 1회만 반복하도록 바꿀 수 있다. 그러면 현재 배치에 과적합되는 것을 피하면서 계산량도 줄일 수 있다. $s−β$를 기반으로 Top-k를 뽑는 것은 원래 MoE가 하던 일이고, 이제는 Top-(k+1)을 뽑는 것으로 바뀌는 정도이므로 비용 증가가 거의 없다. 따라서 추가되는 단계는 $s−α$로 axis=0에서 “$1−k/n$ 분위수”를 찾는 것이다.

Quantile Balancing (QB) 실제 사용 형태 입력: 점수 행렬 $s∈R m×n$, 이전 단계 $β∈R n$ 출력: 할당 $x∈{0,1}m×n$, 새로운 $β∈R n$ 1:$x i,j=1$ if $j∈argtop k s i−β$ else 0 2:$α←des_sort(s−β,axis=1)[:,k:k+1]$ 3:$β←des_sort(s−α,axis=0)[m k/n:m k/n+1]$ 4:Output $x,β$

하지만 1회만 반복하더라도 이 추가 단계는 여전히 비싸다. $m$개의 원소에서 $m k/n+1$번째 큰 값을 찾아야 하기 때문이다. 여기서 $m$은 “전역 샘플 수 * 시퀀스 길이”이며 보통 백만 단위부터 시작한다. 각종 병렬화 전략과 gradient accumulation을 고려하면, 정확한 구현은 대개 받아들이기 어렵다. 타협안으로는, 샘플을 수용 가능한 최대 크기의 작은 미니배치로 나눠 각 미니배치에서 공식대로 $β$를 계산한 뒤, 이를 평균해 최종 결과로 쓰는 방법이 있다.

이 문제들을 해결하고 나면 QB는 거의 장점뿐이다. 첫째 학습률 같은 하이퍼파라미터를 조정할 필요가 없고, 둘째 균형이 매우 빠르게 잡힌다. 특히 극단적인 경우를 다루는 데 능하다. 예를 들어 완전 MoE 모델을 학습할 때 1층 MoE도 매우 균형적으로 바뀐다. 다만 원래 SignSGD 규칙만으로도 균형이 잡히는 층에서는 QB가 보통 더 큰 이점을 보이지는 않는다.

첫 번째 층 MoE의 MaxVio 비교 그림

데모 코드#

여기 데모 코드를 제공한다. 관심 있는 독자는 직접 체험해 볼 수 있다.

import numpy as np

def quantile_bias(s, k, T=5):
    """교대 quantile로 최적 bias를 구함
    원리：https://kexue.fm/archives/11619
    """
    m, n = s.shape
    beta = np.zeros((1, n))
    for _ in range(T):
        alpha = np.quantile(s - beta, 1 - k / n, axis=1, keepdims=True)
        # alpha = alpha.clip(0, np.inf)  # BIP会多出这一步
        beta = np.quantile(s - alpha, 1 - k / n, axis=0, keepdims=True)
        # beta = beta.clip(0, np.inf)  # BIP会多出这一步
    return beta

def max_min_avg_vio(s, k):
    """max_vio、min_vio、avg_vio를 계산
    여기서 max_vio ≥ 0, avg_vio ≥ 0, -1 ≤ min_vio ≤ 0이며,
    셋 모두 0에 가까울수록 더 균형적임
    """
    m, n = s.shape
    topk = np.argsort(-s, axis=1)[:, :k]
    f = np.bincount(topk.reshape(-1), minlength=n)
    f = f / f.sum() * n - 1
    return f.max(), f.min(), np.abs(f).mean()

m, n, k = 100000, 256, 8
s = np.random.rand(m, n) + np.random.rand(n)  # 불균일한 점수를 시뮬레이션

b = quantile_bias(s, k, 5)
max_min_avg_vio(s, k)  # top-k를 직접 취했을 때의 max_vio, min_vio, avg_vio
max_min_avg_vio(s - b, k)  # bias를 뺀 뒤 top-k의 max_vio, min_vio, avg_vio

관련 연구#

최적 할당 관점에서 MoE 부하 균형을 바라보는 접근은 논문 《BASE Layers: Simplifying Training of Large, Sparse Models》에서 가장 이르게 등장했지만, 일반해를 완전하게 제시한 것은 논문 《Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models》 (약칭 BIP)로 보인다. QB는 바로 BIP에서 개선한 것이다.

QB의 목표(3)와 달리 BIP는 등식 제약을 부등식 제약으로 바꾼다:

$max x i,j∈{0,1}∑i,j x i,j s i,j s.t.∑j x i,j≤k,∑i x i,j≤m k n$

그리고 앞서의 일련의 유도를 반복하면 max-min 형태 단계에서 $α i≥0,β j≥0$ 제약이 추가된다. 즉 다음을 고려해야 한다.

$max x i,j∈[0,1]min α i≥0,β j≥0∑i,j x i,j s i,j−∑i α i(∑j x i,j−k)−∑j β j(∑i x i,j−m k n)$

그래야 문제(12)의 완화 버전과 동치가 된다. 이후에도 $max$와 $min$은 여전히 교환할 수 있고 비슷한 유도 과정을 거치지만, 최종적으로 $α,β$의 반복 규칙에는 $max(0,⋅)$ 형태의 절단(truncation)이 추가되어 비음수임을 보장한다.

$α←max(0,des_sort(s−β,axis=1)[:,k:k+1])β←max(0,des_sort(s−α,axis=0)[m k/n:m k/n+1])$

하지만 필자의 테스트에 따르면 이 절단은 부하 균형 속도를 뚜렷하게 떨어뜨리고, “부자를 털어내리기(빈도가 매우 높은 Expert 억제)”는 되지만 “가난한 자를 구제하기(빈도가 매우 낮은 Expert 회복)”는 못하는 사례가 자주 나타났다. 이는 이 절단이 완전히 역방향 최적화임을 시사한다. 그래서 QB는 원 문제를 등식 제약으로 바꿔 $α,β$의 비음수 제약을 제거하여 풀이를 가장 단순화했다.

또한 BIP는 앞 절에서 언급한 실수를 범한 것으로 보인다. 즉 $β$를 먼저 업데이트한 뒤 Top-k를 뽑는데, 이는 인과 법칙(미래 정보 누설)과 훈련-추론 일치 원칙(샘플 간 불일치)을 위배한다. 그럼에도 BIP가 최적 할당 관점에서 부하 균형을 분석한 전체 프레임은 매우 시사적이다.

검색해 보면 BIP 이후에도 이 방향으로 탐색한 작업들이 몇 가지 있으며, 본문과 일부 겹치지만 완전히 같지는 않다. 여기서는 일일이 펼치지 않는다.

《Maximum Score Routing For Mixture-of-Experts》

《Selective Sinkhorn Routing for Improved Sparse Mixture of Experts》

《MicroMoE: Fine-Grained Load Balancing for Mixture-of-Experts with Token Scheduling》

《A Theoretical Framework for Auxiliary-Loss-Free Load Balancing of Sparse Mixture-of-Experts in Large-Scale AI Models》

경사하강#

마지막으로 Loss-Free와 QB의 중간쯤에 있는 또 하나의 풀이를 소개한다. QB 반복 형식에서 $α$ 계산은 비교적 저렴하지만 진짜 비싼 것은 $β$ 계산로, 전체 Token을 가로지르는 어떤 정렬이 필요하다. $α$가 주어졌다고 가정하면 $β$의 최적화 목표는

$min β j∑i,j max(0,s i,j−α i−β j)+m k n∑j β j⏟记为 ℓ$

이다.

Quantile로 최적해를 구하는 것 외에 더 저렴하게 근사해를 구할 방법이 있을까? 있다. 목표 함수 $ℓ$는 미분 가능하므로 경사하강을 고려할 수 있고, 그라디언트는

$∂ℓ∂β j=m k n−m∑i=1 χ(s i,j−α i−β j>0)$

가 된다. 여기서 $χ$는 지시함수로 $χ(True)=1,χ(False)=0$이다. 이 그라디언트 계산도 상대적으로 저렴하다. 그라디언트를 얻었으니 경사하강을 할 수 있고, Loss-Free와 맞추기 위해 SignSGD를 고려하면

$β j←β j−γ sign(∂ℓ∂β j)$

가 된다.

이를 QB에서 Quantile로 $β$를 업데이트하는 단계 대신 사용하면 Loss-Free와 비슷한 비용을 달성할 수 있으며, 실측 효과도 둘 사이(동일 조건에서 Loss-Free와 같은 Sigmoid 활성화 및 같은 $γ$를 쓸 때) 정도로 나온다.

글 요약#

본 글에서는 최적 할당 관점에서 MoE의 부하 균형 문제를 탐구하여, Aux Loss 없이 부하 균형을 달성하는 새로운 알고리즘 Quantile Balancing을 도출했다. 이는 기존 Loss-Free 방식보다 더 안정적이고 정확하며, 어떤 값 범위의 Router 점수에도 적용 가능하고, 추가 하이퍼파라미터를 조정할 필요가 없다.

재배포 시에는 본문 주소를 포함해 주세요: https://kexue.fm/archives/11619

더 자세한 재배포 관련 사항은 다음을 참고하세요: 《科学空间FAQ》