MoE Overview

1. MoE (Mixture-of-Expert) 란?

1.1. 개요

MoE (Mixture-of-Experts) 모델은 거대한 매개변수(파라미터)를 효율적으로 활용하기 위해 등장한 희소^sparse 모델 아키텍처입니다. 다수의 하위 신경망을 “전문가^expert”로 두고, 라우터^router 네트워크(또는 Gating network)가 입력 토큰별로 소수의 전문가만 선택하여 활성화합니다. 이를 통해 Forward pass에 필요한 계산량이 대폭 감소하는 스케일링 이점을 얻게 됩니다. 이를 간단히 요약하면 아래와 같습니다.

용량 증가: 모델이 훨씬 더 많은 수의 파라미터를 포함하게 되어, 토큰 당 계산 비용이 크게 증가하지 않으면서도 더 다양하고 전문화된 기능을 학습합니다.
상대적 일정한 계산량: 전체 모델 크기는 커지더라도, 활성화된 전문가 수 $k$ 를 고정함으로써 단일 토큰을 처리하는 데 필요한 계산 비용(FLOPs)은 대략 일정하게 유지되거나 훨씬 느리게 증가하게 됩니다.

아래 그림은 MoE 모델의 레이어 예시로, 라우터가 입력 토큰을 2개의 전문가 FFN에 라우팅하고 해당 전문가의 출력을 취합하는 Top-2 라우팅 방식을 보여줍니다. 이를 통해 한 레이어에서 다중 전문가 FFN들이 동시에 작동하면서도, 각 토큰은 소수 전문가의 출력만 활용하게 됩니다. Router는 간단한 피드포워드 네트워크(FFN; Feed Forward Network)로, 일반적으로 토큰의 입력 임베딩에 적용되는 단일 선형 레이어와 그 뒤에 오는 소프트맥스 함수로 구성됩니다. 이에 대해서는 1.2절에서 자세히 설명하겠습니다.

MoE 아키텍처는 전문가의 단위 크기와 라우팅 세분화 정도^granularity가 얼마나 세세한 단위까지 조정되는가에 따라 Coarse-grained MoE와 Fine-grained MoE로 구별합니다.

Coarse-grained MoE

Coarse-grained MoE는 비교적 큰 규모의 전문가 단위를 사용합니다. 예를 들어 하나의 MoE 레이어가 8개의 전문가로 구성되어 있고, 각 입력 토큰이 그중 상위 1~2개의 전문가만 선택되어 전달되는 구조입니다. 이 방식은 각 전문가가 담당하는 영역이 넓고 비교적 독립적으로 구조가 단순하고 안정성이 높으며 구현이 용이하다는 장점이 있으나, 전문가 간의 세밀한 협력이나 표현 조합이 어렵다는 한계가 있습니다.

Fine-grained MoE

Fine-grained MoE는 전문가를 훨씬 작은 단위로 세분화하기에, 각 토큰은 그 작은 전문가 단위들 사이에서 조합적으로 혹은 더욱 유연하게 라우팅됩니다. 이 방식은 라우팅 선택지가 폭넓어지고 전문가 간의 표현 조합 가능성이 커지는 반면, 라우팅 비용이 더 커지고 최적화 난이도도 상승합니다. DeepSeek, Qwen, GPT-OSS를 비롯한 최신 MoE에서 트랜드로 자리잡았습니다.

최근 MoE 연구는 단순히 전문가의 수를 늘리는 방향에서 벗어나, fine-grained MoE를 정교화하는 방향으로 발전하고 있습니다.

전문가 단위의 세분화 확대: Mixture of A Million Experts와 Scaling Laws for Fine-Grained Mixture of Experts 연구에서는 세분화 정도^granularity가 모델의 스케일링 효율과 성능 향상에 핵심적인 요인임을 보여주었습니다.
라우팅의 정교화 및 안정성 확보: Fine-grained 구조에서는 전문가 선택이 많아질수록 라우팅의 불안정성이 커질 수 있습니다. 이를 해결하기 위해 최근에는 전문가-토큰 친화도^affinity를 기반으로 한 라우팅, 토큰 간 관계를 고려한 라우팅 등이 Expert-Token Resonance MoE 논문을 통해 제안되고 있습니다.
동적 라우팅 및 사후 조정(Post-routing): 2025년 9월에 공개된 Ban&Pick 논문에서는 학습 이후에도 중요한 전문가를 동적으로 유지하고 불필요한 전문가를 제거하는 방식으로 추론 효율을 높였습니다. 즉, 학습 단계에서 결정된 라우팅을 그대로 두지 않고, 학습 이후에도 지속적으로 조정함으로써 효율성을 극대화하는 방향으로 발전하고 있습니다.
Chain-of-Experts (CoE): 기존 MoE는 전문가들이 병렬적으로 독립 작동하지만, 최근에는 전문가 간 상호작용을 허용하거나, 입력이 여러 전문가를 순차적으로 거치는 Chain-of-Experts (CoE) 구조가 등장했습니다. 이 방식은 각 전문가의 출력을 다음 전문가에 전달함으로써 전문가 조합의 다양성과 표현의 깊이를 동시에 확보할 수 있습니다.

1.2. 수학적 정의

$d$ 가 모델의 은닉^hidden 차원인 벡터 $x \in \mathbb{R}^d$ 로 표현되는 단일 입력 토큰을 고려해 보겠습니다. 라우터 자체는 간단한 선형 레이어로, $N$ 이 총 전문가 수인 훈련 가능한 가중치 행렬 $W_g \in \mathbb{R}^{d \times N}$ 으로 정의합니다. 이 레이어는 입력 토큰을 차원 $N$ 의 공간으로 투영^project하여 각 전문가에 대한 로짓^logit을 생성합니다.

\text{logits } H(x) = x \cdot W_g

$N$ : 레이어의 총 전문가 수로 각 전문가는 일반적으로 독립적인 FFN입니다. (가중치를 공유하지 않습니다.)
$W_g$ : 라우터의 학습 가능한 가중치 행렬

결과로 나온 벡터는 어떤 전문가들이 선택될지를 결정하는 점수들을 포함합니다. 희소성을 강제하기 위해, 이 벡터에 top-k 함수가 적용됩니다. 예를 들어, $k=2$ 인 경우 가장 높은 점수를 가진 두 명의 전문가만 선택됩니다. 나머지 전문가들의 점수는 모두 0으로 설정되어 현재 토큰에 대해 비활성화됩니다.

$H(x)$ 는 어떤 전문가들이 선택될지를 결정하는 $N$ 개 전문가의 raw score를 포함합니다. 이러한 점수를 확률 분포로 변환하기 위해 각 전문가에 대해 소프트맥스를 적용하여 확률 분포 $G(x)$ 로 변환합니다.

G(x) = \text{softmax}(H) = \frac{\exp(H)}{\sum_{j=1}^{N} \exp(H_j)}

출력 $G(x)$ 는 밀집^dense $N$ 차원 벡터로, 각 요소 $G_i(x)$ 는 토큰 $x$ 가 전문가 $i$ 에 할당될 라우터의 신뢰도를 나타냅니다. 소프트맥스 함수이기 때문에 $G(x)$ 의 모든 요소의 합은 1입니다.

이제 이러한 단계들을 결합하여 MoE 레이어의 최종 출력 $y(x)$ 를 정의할 수 있습니다. 라우터 $G(x)$ 는 입력 표현 $x$ (보통 트랜스포머의 셀프 어텐션 계층 출력)를 검사하여 어떤 전문가가 이 입력을 처리할지 나타내는 확률 또는 가중치를 생성합니다. 선택된 전문가들의 출력 $E_i(x)$ 는 라우터의 점수로 가중합되어 결합됩니다.

y(x) = \sum_{i=1}^{N} G_i(x) \cdot E_i(x)

각 전문가 $E_i$ 는 자체 매개변수 집합을 가집니다. 표준 2-layer FFN 전문가는 다음과 같이 표현 가능합니다.

E_i(x) = \text{ReLU}(x \cdot W_{1,i}) \cdot W_{2,i}

여기서 $W_{1,i}$ 와 $W_{2,i}$ 는 각각 전문가 $i$ 의 첫 번째와 두 번째 선형 레이어의 가중치 행렬입니다.

1.3. Sparse MoE (Top-k 라우팅)

위에서 산출한 밀집^dense 확률 벡터인 $G(x)$ 벡터는 모든 전문가가 출력에 기여한다는 것을 의미하며, 이는 MoE의 계산 효율성이 높지 않다는 것을 의미합니다. 희소성^sparse을 강제하기 위해 top-k 선택 메커니즘을 적용함으로써, 주어진 토큰에 대해 전문가의 일부만 활성화할 수 있습니다. 이는 모든 $N$ 개의 전문가를 사용하는 대신 로짓 벡터 $H$ 에서 가장 높은 점수를 가진 $k$ 개의 전문가를 선택하는 것입니다. $k$ 의 값은 대개 1 또는 2로 설정됩니다. (예컨대, $k=2$ 인 경우, 라우터는 가장 높은 로짓을 가진 두 전문가를 식별합니다.) 다른 모든 전문가는 이 특정 토큰에 대해 무시되어, 계산 비용이 총 전문가 수 $N$ 이 아닌 $k$ 에 비례합니다.

과정은 다음과 같이 진행됩니다.

점수 계산: $G(x)$ 를 계산합니다.
Top-k 선택: $G$ 에서 가장 큰 값 $k$ 개에 해당하는 인덱스 $\text{TopK}(H) = {i_1, i_2, ..., i_k}$ 를 식별합니다.
점수 재정규화 및 가중합 계산: 선택된 $k$ 명의 전문가에 대해서만 소프트맥스를 적용하여 점수(확률)를 재정규화합니다. 최종 출력은 (재정규화된) 게이팅 확률을 가중치로 사용하여 선택된 전문가들의 출력의 가중 합입니다.

y(x) = \sum_{i \in \text{TopK}(H)} \frac{\exp(H_i)}{\sum_{j \in \text{TopK}(H)} \exp(H_j)} \cdot E_i(x)

Sparse MoE 접근법은 모델 크기(매개변수 수)와 계산 비용(토큰당 FLOPs)을 분리하는 강력한 메커니즘을 제시합니다. 이를 통해 수조 개의 매개변수를 가진 초거대 모델을 구축하면서도 훈련과 추론 시 계산 용량을 관리 가능한 수준으로 유지할 수 있습니다.

1.4. Dense 모델과 Sparse MoE 모델의 계산 비용 (FLOPs) 비교

Dense 모델

밀집^dense 모델에서는 주어진 층의 모든 입력 토큰이 모든 파라미터에 의해 처리하기에 토큰당 FLOPs는 밀집 계층 파라미터의 크기에 비례합니다. 입력/출력 차원이 $d_{\text{model}}$ 이고 중간 차원이 $d_{\text{ffn}}$ 인 FFN의 경우, 단일 토큰에 필요한 파라미터 수와 계산량은 모델의 차원, 특히 $d_{\text{model}}$ 과 $d_{\text{ffn}}$ 에 정비례합니다.

이 식은 $d_{\text{ffn}}$ 을 증가시켜 모델을 더 똑똑하게 만들려면 방대한 계산 비용을 지불해야 한다는 것을 보여줍니다.

Llama3-70B 파라미터 예시

모델 구성
- 레이어 개수: 80
- Hidden dim: 8,192
- Vocab size: 100,000
- Query heads: 64개
- Key/Value heads: 8개 (8:1 그룹핑)
- Head dim: 128
레이어당 멀티헤드 어텐션 파라미터 개수
- Query projection: 8,192 × (64 × 128) = 8,192 × 8,192 ≈ 67M
- Key projection: 8,192 × (8 × 128) = 8,192 × 1,024 ≈ 8.4M
- Value projection: 8,192 × (8 × 128) = 8,192 × 1,024 ≈ 8.4M
- Output projection: 8,192 × 8,192 ≈ 67M
- 총 파라미터: 약 151M (A)
레이어당 FFN 파라미터 개수
- Intermediate size: 28,672: 표준 FFN은 4 × hidden_size = 32,768이나 SwiGLU는 gating 때문에 약간 작게 설정 (2/3 × 4 ≈ 2.67배)하고 메타에서 하드웨어 최적화를 위해 256의 배수로 조정하면서 조정한 최적 수치입니다.
- Gate projection: 8,192 × 28,672 ≈ 235M
- Up projection: 8,192 × 28,672 ≈ 235M
- Down projection: 28,672 × 8,192 ≈ 235M
- 총 파라미터 : 약 705M (B)
LayerNorm 파라미터 개수: 8,192 × 2 ≈ 16K (C)
레이어당 파라미터 개수: A + B + C ≈ 약 856M
80개 레이어이므로 856M × 80 ≈ 68.5B

FFN에서만 705M × 80 ≈ 58.4B 파라미터를 점유하며, 이는 총 파라메터의 85%입니다.

Sparse MoE 모델

Sparse MoE 모델의 경우 많은 전문가들에 걸쳐 방대한 수의 총 파라미터를 포함하고 있을지라도, 주어진 토큰에 대해 활성화되는 파라미터는 그 중 일부에 불과합니다. 예컨대 전문가를 16명에서 64명으로 늘려도 토큰당 FLOPs는 증가하지 않습니다.

MoE 레이어의 총 파라미터 수는 모든 전문가들의 합입니다.

\text{Parameters}_{\text{MoE}} = N_{\text{experts}} \times \text{Parameters}_{\text{expert}}

하지만 FLOPs는 활성화된 전문가 수 $k$ 에만 비례합니다.

\text{FLOPs}_{\text{MoE}} \approx k \times (2 \times \text{batch\_size} \times \text{seq\_len} \times d_{\text{model}} \times d_{\text{expert}})

이것이 MoE 아키텍처의 핵심 장점입니다. 전문가 수 $N_{\text{experts}}$ 를 늘려 모델의 용량^capacity(총 파라미터 수)을 극적으로 증가시키면서도, 토큰당 계산 비용은 $k$ 를 고정함으로써 일정하게 유지할 수 있습니다. 이는 밀집 모델보다 훨씬 작은 계산 예산으로도 수조 개의 파라미터를 가진 모델을 훈련할 수 있게 합니다.

2. Router 네트워크 훈련

직관적으로 특정 유형의 토큰을 특정 전문가에게 라우팅하는 것이 지속적으로 모델의 loss를 줄일 수 있게, 오류역전파를 통해 $W_g$ 를 업데이트하여 미래에 특정 전문가 할당의 확률을 증가시키는 피드백 루프를 수행합니다. 그러나 이 과정에서 큰 난제가 발생하는데 게이팅 네트워크가 소수의 전문가들만 선호하여 다른 전문가들을 훈련되지 않은 상태로 남겨두는 불균형 라우팅 문제(또는 전문가 붕괴^{expert collapse})가 발생한다는 것입니다. 이를 완화하기 위해 보조 손실^{auxiliary loss}(또는 로드 밸런싱 소실), 전문가 용량^{capacity factor}, 라우터 노이즈, 전문가 선택 라우팅 등의 기법을 적용합니다.

보조 손실 (Auxiliary Loss): 로드 밸런싱 손실에 보조 손실을 추가함으로써 전문가의 분포를 균등하게 합니다.
라우터 노이즈: top-k 전문가를 선택하기 전에 게이팅 네트워크의 로짓에 소량의 무작위 노이즈를 추가하여 선택될 확률이 낮은 전문가도 가끔 탐색하도록 강제합니다.
전문가 용량 (Expert Capacity): 전문가가 수용할 수 있는 토큰 수에 대한 버퍼 크기를 조정합니다.
전문가 선택 라우팅 (Expert Choice Routing): 전문가가 자신에게 도달한 토큰 점수를 기준으로 상위 $k$ 개의 토큰을 선택하고, 나머지는 버립니다.

2.1. 전문가 붕괴 (Expert Collapse)

전문가 붕괴는 라우터가 대부분의 토큰을 작고 선호되는 전문가 부분집합으로 라우팅하는 것을 학습하는 반면, 나머지 전문가들은 적거나 아예 토큰을 받지 못할 때 발생합니다. 이러한 활용도가 낮은 전문가들은 의미 있는 전문화를 학습하지 못하고, 효과적으로 "죽은" 매개변수가 됩니다.

MoE 모델 훈련 중에는 전문가 붕괴^{expert collapse} 현상이 자연스럽게 발생하는 경우가 많습니다:

초기화: 게이팅 네트워크의 초기 무작위 가중치가 특정 전문가들을 본질적으로 선호할 수 있습니다.
초기 훈련 역학^dynamics: 초기 그래디언트가 초기에 설정된 라우팅 패턴을 강화하여, 약간 더 잘 수행하는 일부 전문가들이 점점 더 많은 토큰을 받는 피드백 루프를 만들 수 있습니다.
데이터 분포: 특정 특징이나 패턴이 훈련 데이터에 더 많이 존재할 수 있어, 최적화가 잘 된 게이팅 네트워크가 이를 특정 전문가에게 라우팅하게 되고, 해당 패턴이 흔할 경우 전문가들이 과부하될 수 있습니다.
명시적 메커니즘 부족: 부하 균형을 장려하는 메커니즘이 없으면 부하를 고르게 분배할 동기가 없습니다.

전문가 붕괴 문제를 해결하는 것은 단순한 최적화 세부사항이 아니라, 크고 성능 좋은 MoE 모델을 성공적으로 훈련하기 위한 필수 절차입니다.

전문가 붕괴를 식별하려면 훈련 과정 전반에 걸쳐 라우터 행동을 모니터링해야 합니다. 단순히 총 손실을 관찰하는 것만으로는 불충분하며, 근본적인 불균형을 드러내지 못할 수 있습니다. 진단을 위해 아래의 주요 지표를 활용합니다.

전문가 활용도 (Expert Utilization): 설정된 훈련 단계 수에 걸쳐 각 전문가에게 전달된 토큰 수를 추적합니다. 전문가당 토큰 수의 히스토그램이 심하게 치우쳐 있다면 전문과 붕괴임을 의미합니다.
로드 밸런싱 손실 (Load Balancing Loss): 높은 값은 라우터가 로드 균형을 맞추는 데 어려움을 겪고 있음을 시사합니다. 그러나 낮은 값이 항상 좋은 결과를 보장하는 것은 아닙니다. 라우터가 토큰을 드롭하여 낮은 손실을 달성할 수 있기 때문입니다.
변동 계수 (CV; Coefficient Variation): 전문가당 토큰 분포의 변동 계수로 표준편차를 평균으로 나눈 값으로 정의됩니다. 0에 가까운 CV는 완벽한 균형을 나타내고, 높은 CV는 상당한 불균형을 나타냅니다. 쉽게 말해, 각 전문가의 중요도 점수가 많이 다르면 CV가 높고 모든 전문가의 중요도 점수가 비슷하면 CV가 낮아집니다.

\text{CV} = \frac{\sqrt{\frac{1}{N} \sum_{i=1}^{N} (L_i - \bar{L})^2}}{\bar{L}}

$N$ : 전문가 수, $L_i$ : 훈련 단계에서 전문가 $i$ 의 로드(토큰 수), $\bar{L}$ : 평균 로드

2.2. 보조 손실 (Auxiliary Loss)

보조 손실의 목표는 라우터의 불균형에 대해 페널티를 주는 손실 항을 만드는 것입니다. Sparsely-Gated MoE 논문에서 도입된 가장 일반적인 접근법은 배치에서 토큰과 라우터 확률 분포를 기반으로 값을 계산하는 것입니다.

보조 손실 함수

보조 손실 $L_{aux}$ 는 이 두 벡터의 내적으로 계산되며, 전문가 수 $N$ 으로 스케일됩니다.

L_{aux} = N \cdot \sum_{i=1}^{N} f_i \cdot P_i

전문가당 토큰 비율 ( $f_i$ ): 전문가 $i$ 로 보내지는 토큰의 수를 측정합니다. 배치에 $B$ 개의 토큰이 있고 top-1 게이팅을 사용한다면, 이는 단순히 전문가 $i$ 로 라우팅된 토큰 수를 $B$ 로 나눈 것입니다.
전문가당 평균 라우터 확률 ( $P_i$ ): 배치의 모든 토큰에 걸쳐 라우터가 전문가 $i$ 에 할당하는 평균 확률(또는 게이팅 점수)입니다. 라우터가 전문가 $i$ 에게 부여하는 "중요도"을 나타냅니다.

이 손실을 최소화하기 위해, 모델은 단일 전문가 $i$ 가 높은 토큰 비율 $f_i$ 와 높은 평균 확률 $P_i$ 를 동시에 갖는 것을 방지해야 합니다. 손실은 곱 $f_i \cdot P_i$ 가 작고 고르게 분포될 때 가장 낮으며, 이는 모든 전문가에 걸쳐 로드가 균형을 이룰 때 발생합니다.

보조 손실 함수는 보다 다양한 형태로 정의할 수 있습니다. 예컨대 DeepSpeed-MoE는 아래와 같은 보조 손실 함수를 정의합니다.

${L}_{aux} = \sum_{i=1}^N |f_i - \frac{1}{N}|$

또한, 불균형한 토큰 분포의 영향을 더 완화하기 위해 DeepSpeed-MoE는 선택된 전문가의 출력물을 밀집 FFN의 출력과 결합하는Residual-MoE 아키텍처를 도입했습니다: $y=\text{FFN}(x) + g \cdot E(x)$

Router z-손실

MoE 구조에서는 라우터가 softmax 또는 sparse gating 연산을 해서 전문가 선택 확률을 계산하므로, 라우터의 입력 로짓^logits(정규화 이전 점수)이 매우 커지면 지수^exp 연산에서 오버플로우, 정밀도 손실^{round-off error}, 불안정한 기울기 등 문제가 발생할 수 있습니다. router-z 손실은 전문가 선택 확률과는 직접적인 관계 없는, 라우터의 로짓 값이 너무 커지는 걸 억제하기 위한 보조 손실 항을 추가하여 수치 안정성과 훈련 안정성을 확보하는 기법으로, ST-MoE 논문에서 처음 정의되었고 여러 MoE 아키텍처에서 활용하고 있습니다.

L_z = \frac{1}{B} \sum_{i=1}^B \left( \log \sum_{j=1}^N \exp\bigl(x_j^{(i)}\bigr) \right)^2

$B$ : 배치 내 토큰 수 또는 샘플 수
$x_j^{(i)}$ : $i$ 번째 토큰/샘플의 라우터가 전문가 $j$ 에 대해 계산한 로짓 (softmax 연산 이전 점수)
$\log \sum_{j} \exp(x_j^{(i)})$ : softmax 계산의 정규화^{normalization} 정규화 항과 관련이 있는데, 이 값이 지나치게 크면 라우터 로짓 전반이 “스케일이 커진다”는 의미입니다. 그걸 제곱된 항으로 패널티를 주면, 라우터가 무분별하게 큰 값을 쓰지 못하도록 억제하는 역할을 합니다. 그렇게 하면 숫자적으로 안정해지고, softmax 지수부 연산이 과도한 오버플로우나 수치 불안정성을 덜 일으키게 됩니다.

최종 손실 함수

보조 손실은 작업 손실( $L_{\text{task}} = \text{CrossEntropy}(y_{\text{pred}}, y_{\text{true}})$ )과 결합되어 역전파에 사용되는 최종 손실 함수를 형성합니다.

L_{total} = L_{task} + \alpha \cdot L_{aux} + c_z \cdot L_z

하이퍼파라미터 $\alpha$ 는 종종 load_balance_loss_coef라고 불리며, 보조 손실의 강도를 제어하는 작은 스칼라 값입니다.

$\alpha$ 가 너무 작으면, 균형 힘이 너무 약해서 전문가 붕괴를 방지할 수 없습니다.
$\alpha$ 가 너무 크면, 모델이 주요 작업의 성능을 희생하면서 완벽한 로드 밸런싱을 우선시하여 전체적인 정확도가 떨어질 수 있습니다.

하이퍼파라미터 $c_z$ 는 종종 router_z_loss_coef라고 불리며, router z-손실의 강도를 제어하는 작은 스칼라 값입니다.

$c_z$ 가 너무 크면, 훈련 초기 안정성이 떨어지고 loss 진동이 커지거나 라우팅이 지나치게 한 전문가에 몰리는 현상이 발생할 수 있습니다.
$c_z$ 가 너무 크면, 라우터가 거의 균등한 확률을 뿌리는 수준에서 굳어 버려, MoE의 장점인 표현 다양성과 전문가 전문화 효과를 떨어뜨릴 위험이 있습니다.

적합한 값을 찾는 것은 MoE 모델의 하이퍼파라미터 튜닝 과정이며, 권장 베이스라인은 $\alpha=0.01, c_z=0.001$ 입니다. (참조: NVIDIA NeMo Framework User Guide)

스케줄링은 모델과 데이터에 따라 가변적으로 두시는 편이 좋습니다. 현업에서 흔한 전략은 초반에는 규제를 조금 강하게, 이후 완만히 낮추는 방식입니다. 예를 들어 초기에는 보조 손실 가중치 계수를 조금 큰 값으로 시작해서 안정화 구간에 진입할 때 점차 감소시키는 스케줄링을 적용하면, 초기 붕괴를 막으면서 후반에는 전문가 특화 여지를 돌려줄 수 있습니다. 최근 대규모 MoE 훈련 노하우에서도 “부하가 이미 잘 균형을 이루면 보조 손실 계수를 낮추고, 불균형이 크면 일시적으로 올려서 강제 균형을 준다”는 상태 기반 조정을 권합니다.

파인튜닝 단계에서는 보조 손실 가중치 계수를 크게 낮추거나 0으로 두어도 성능 손실이 크지 않다는 연구 결과도 있습니다.

2.3. 전문가 용량 (Expert Capacity)

전문가 용량^{expert capacity}은 일반적으로 전문가당 이상적인 균일 부하를 기준으로 전문가 용량 계수^{capacity factor}를 조정하여 설정합니다. $T$ 가 처리되는 토큰 수일 때, 이상적인 완벽하게 균형 잡힌 라우터는 각 전문가에게 정확히 $\frac{T}{N}$ 개의 토큰을 보냅니다. 용량 계수 $C$ 는 통계적 변동성에 대한 완충 역할을 제공하며 다음과 같이 설정합니다.

\text{Expert Capacity} = \left( \frac{\text{Tokens per Batch}}{N_{\text{experts}}} \right) \times C

$C$ 의 일반적인 값 범위는 1.0에서 2.0 사이입니다. (예컨대 $C$ 가 1.25라는 것은 각 전문가가 평균보다 25% 더 많은 토큰을 처리할 수 있음을 의미합니다.)

1.0에 가까운 값(예: 1.25)은 계산 오버헤드를 최소화하지만, 부하가 약간만 불균형해도 토큰이 손실될 위험이 증가합니다.
더 큰 값(예: 1.5, 2.0)은 토큰 손실 가능성을 줄이지만, 전문가가 종종 $C$ 개의 토큰보다 적게 처리할 수 있기 때문에 필요 이상으로 계산 자원이 할당될 수 있습니다.

전문가의 용량이 초과되면 어떻게 될까요? 전문가가 용량에 도달했을 때, 생성된 토큰들은 Top-2일 경우 다음 전문가에게 전달되며, 다음 전문가의 용량이 초과된다면 현재 계층에서의 토큰을 버립니다. (이를 토큰 오버플로^{token overflow}라고 합니다.) Top-1은 전달한 다른 전문가가 없으므로 전문가 용량 초과 시 곧바로 토큰을 버립니다. (즉, 어떤 전문가에 의해서도 처리되지 않습니다.) 대신, 잔차 연결^{residual connection}에서의 표현이 다음 층으로 직접 전달됩니다. 이는 MoE 계층 내에서 토큰이 항등 함수^{identity function}를 통과하는 것과 같습니다. 토큰 드롭이 해로워 보일 수 있지만, 실험 결과 잘 조정된 용량 계수와 보조 손실^{auxiliary loss}을 사용하면 버려지는 토큰의 비율이 종종 낮으며(<1%), 전체 모델 성능에 큰 영향을 미치지 않는 것으로 나타났습니다.

하지만 토큰 드롭이 여전히 모델 정확도에 영향을 줄 수 있다는 우려가 있기에, 토큰 리라우팅^re-routing 기법으로 초과 토큰을 다른 한가한 전문가에게 재할당하여 드롭 없이 처리할 수도 있습니다. DeepSpeed-MoE의 경우 훈련 중 전문가 용량을 넘는 토큰들을 드롭 대신 덜 바쁜 다른 전문가로 동적으로 재분배하여 처리하고 추가로 Residual MoE 구조를 도입해 밀집 경로와 전문가 경로 출력을 합치는 방식을 사용합니다. 이를 통해 일부 토큰이 전문가 처리를 못 받아도 밀집 경로로 보완하고, 전체적으로 전문가 활용의 안정성을 높입니다.

2.4. 라우터 노이즈 (Noisy Top-k Gating)

전문가 로드 밸런싱의 또다른 기법은 라우터 노이즈입니다. 원시 게이팅 로짓^{raw gating logits}만을 기반으로 상위 $k$ 명의 전문가를 선택하는 대신, 선택 과정 전에 노이즈를 추가합니다. 일반적인 방법으로는 가우시안 노이즈를 추가합니다.

\text{noisy\_logits} = \text{logits} + \mathcal{N}(0, \sigma^2)

$\mathcal{N}(0, \sigma^2)$ 는 분산이 $\sigma^2$ , 평균이 0인 가우시안 분포에서 추출된 샘플을 의미합니다. 그런 다음 이 noisy_logits에 대해 TopK 선택이 수행됩니다. 노이즈 분산 $\sigma^2$ 는 종종 훈련 가능한 파라미터로 구현되거나 훈련 스케줄링에 의해 조절됩니다.

조금 더 수학적으로 표현해 보겠습니다. 만약 $W_g$ 가 게이팅 가중치 행렬이고 $x$ 가 입력 토큰 표현이라면, 표준 로짓은 $\text{logits} = x \cdot W_g$ 입니다. Noisy top-k의 경우, 다음과 같이 계산합니다.

$\epsilon \sim \mathcal{N}(0, 1)$ 는 토큰별로 샘플링된 표준 가우시안 노이즈이며, $W_{noise}$ 는 입력에 기반해 노이즈를 곱셈적으로 스케일링하는 또 다른 훈련 가능한 가중치 행렬입니다. softplus 함수는 노이즈 스케일링 인자가 양수가 되도록 보장합니다.

이렇게 추가된 잡음은 라우팅 결정에 확률성을 도입합니다. 이는 점수가 낮은 전문가들도 선택될 기회를 제공하여 탐색을 촉진하고, 라우터가 너무 일찍 최적이 아닌 할당 패턴에 고착되는 것을 방지할 수 있습니다.

Noisy Top-K 라우터 구현 예시

class NoisyTopKRouter(nn.Module):
    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        
        # Layer to generate logits
        self.gate = nn.Linear(d_model, num_experts)
        # Layer to generate noise scaling factor
        self.noise_net = nn.Linear(d_model, num_experts)

    def forward(self, x):
        # x shape: [TOKENS_PER_BATCH, d_model]
        logits = self.gate(x)
        
        # Add noise during training
        if self.training:
            noise_logits = self.noise_net(x)
            # Use softplus to ensure the noise scaling factor is positive
            noise = torch.randn_like(logits) * F.softplus(noise_logits)
            logits = logits + noise
            
        # Get the top-k logits and their indices
        # top_k_logits shape: [TOKENS_PER_BATCH, top_k]
        # top_k_indices shape: [TOKENS_PER_BATCH, top_k]
        top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)
        
        # Apply softmax to the top-k logits to get the weights
        gating_scores = F.softmax(top_k_logits, dim=-1)
        
        # We also need a mask for the load balancing loss
        # This mask has a 1 for each token-expert pair that was selected
        # B = TOKENS_PER_BATCH, E = NUM_EXPERTS
        # zero_mask shape: [B, E]
        zero_mask = torch.zeros(x.size(0), self.num_experts, device=x.device)
        # router_mask shape: [B, E]
        router_mask = zero_mask.scatter(1, top_k_indices, 1)

        return top_k_indices, gating_scores, router_mask

# --- Example Usage ---
# Configuration
NUM_EXPERTS = 8
D_MODEL = 512
BATCH_SIZE = 4
SEQ_LEN = 1024
TOKENS_PER_BATCH = BATCH_SIZE * SEQ_LEN

noisy_router = NoisyTopKRouter(D_MODEL, NUM_EXPERTS, top_k=2)
noisy_router.train() 

input_tokens = torch.randn(TOKENS_PER_BATCH, D_MODEL)
indices, scores, mask = noisy_router(input_tokens)

print("Noisy Top-k Router Output:")
print("Indices Shape:", indices.shape) # [TOKENS_PER_BATCH, 2]
print("Scores Shape:", scores.shape)   # [TOKENS_PER_BATCH, 2]
print("Mask Shape:", mask.shape)       # [TOKENS_PER_BATCH, 8]

2.5. 전문가 선택 (Expert Choice)

2022년에 공개된 전문가 선택^{Expert Choice} 라우팅 기법은 전문가가 주체가 되어 자신에게 도달한 토큰의 점수를 기준으로 상위 $k$ 개의 토큰을 선택하고, 나머지는 버리는 방식입니다. 이 방식은 전문가별 처리량을 균등하게 유지하여 자연스럽게 부하를 분산시키며, 결과적으로 분산 학습 및 추론 과정에서 연산 효율성을 높입니다.

전문가 선택 라우팅의 동작 과정은 다음과 같습니다. 우선 게이트 네트워크가 모든 <토큰, 전문가> 쌍에 대해 점수를 산출합니다. 이후 각 전문가는 자신의 열에 해당하는 점수를 정렬한 뒤 상위 $k$ 개 capacity에 해당하는 토큰만 선택합니다. 선택된 토큰은 해당 전문가의 FFN을 통과하여 처리되고, 최종적으로 토큰 순서에 맞추어 결과가 다시 집계됩니다. 이 과정에서 capacity를 초과한 토큰은 처리되지 않고 드롭될 수 있으므로, 게이트 네트워크가 전문가 간 점수를 균등하게 분포하도록 학습되는 것이 모델 성능에 중요합니다.

조금 더 수학적으로 표현해 보겠습니다. 입력 벡터 $x \in \mathbb{R}^{d}$ 가 주어졌을 때, 게이트 네트워크는 $N$ 개의 전문가에 대한 점수 벡터를 산출합니다.

H = x \cdot W_g\in \mathbb{R}^{N}, \qquad G = \mathrm{softmax}(H) \in \Delta^{N-1}

전문가 선택 라우팅 방식에서는 각 전문가가 자신의 점수 집합에서 상위 $k$ 개 토큰을 선택하므로, 단일 토큰 $x$ 는 전문가 $i$ 에 의해 선택될 확률적 이벤트로 표현됩니다. ( $G_i(x)$ 는 토큰 $x$ 가 전문가 $i$ 에 할당될 적합도입니다.)

R_i(x) = \begin{cases} 1, & G_i(x) \text{가 전문가 \(i\)의 상위 \(k\) 토큰에 속할 때} \\ 0, & \text{그 외의 경우} \end{cases}

이를 indicator 함수로 표현하면 아래와 같습니다.

\underbrace{R_i(x;\tau_i)}_{\text{route indicator}} = \mathbf{1}\!\left\{\,p_i(x)\;\ge\;\tau_i\,\right\}\in\{0,1\}

$\tau_i$ 는 전문가 $i$ 가 상위 $k$ 개의 토큰만 받도록 하는 토큰 점수의 최솟값입니다. 토큰당 최대 전문가 수를 $b$ 로 제한하려면 $\sum_{i=1}^{N} R_i(x;\tau_i)\le b$ 의 제약을 추가합니다.

그 후 선택된 전문가들에 대해만 FFN $E_i(x)$ 를 적용하고, 가중합으로 출력을 계산합니다.

y(x) = \sum_{i=1}^{N} R_i(x;\tau_i) \cdot G_i(x) \cdot {E}_i(x)

이 식에서 중요한 점은 각 전문가별로 최대 $k$ 개의 토큰만 선택되므로, $x$ 는 최대 $b$ 개의 전문가에 라우팅될 수 있으며( $b \le N$ ), 어떤 전문가에도 선택되지 않을 가능성(드롭)이 존재합니다.

이 방식의 주요 장점은 균형 잡힌 부하 분산과 추론 지연의 안정성입니다. 전문가별 처리량이 일정하므로 통신 및 동기화 병목이 줄어들고, 대규모 GPU/TPU 환경에서 지연 시간^latency 예측이 용이합니다. 또한 전문가 단위의 capacity 제어가 가능해 메모리 관리 및 연산 자원 배분이 효율적입니다. 다만 전문가가 선택할 수 있는 토큰의 수가 제한적이므로 일부 토큰이 버려질 위험이 있고, <토큰, 전문가> 쌍의 전체 스코어를 계산하고 정렬해야 하므로 계산 비용이 증가합니다. 또한 이 방식으로 MoE 모델을 훈련할 시에는 미분 불가이므로(indicator 함수이기 때문에) soft-approximation(예: 시그모이드^sigmoid 스무딩)으로 대체합니다.

3. MoE 최적화 베이스라인

아래 베이스라인을 기반으로 특정 애플리케이션과 하드웨어 환경에 맞게 성능, 전문화, 계산 가능성의 균형을 맞춘 MoE 아키텍처를 도출하는 것을 권장합니다.

3.1. 전문가 수 결정

전체 전문가 수

전체 전문가 수 $N$ 은 주요 아키텍처 선택 사항입니다. 이 값을 늘리면 다음과 같은 장점이 있습니다:

더 많은 파라미터 수: MoE 계층은 모델 크기를 효율적으로 확장합니다. 전문가를 추가하면 전체 파라미터 수가 크게 증가하지만, 각 토큰이 라우팅되는 전문가 수 $k$ 가 작게 유지된다면 계산 비용(FLOPs per token)은 적당히만 증가합니다.
더 세밀한 전문화: 더 많은 전문가는 모델이 다양한 입력 유형이나 문맥에 대해 더 구별되고 특화된 기능을 학습할 수 있게 합니다.
하드웨어 매핑: 전문가 병렬 처리^{Expert Parallelism}를 사용하는 분산 환경에서는 $N$ 이 보통 사용 가능한 처리 장치(GPU 등)의 수에 따라 결정되며, 각 장치에 하나 이상의 전문가가 할당됩니다. 대형 모델에서 흔히 사용되는 $N$ 값은 8, 16, 64, 128 또는 그 이상입니다.

하지만 $N$ 을 늘리면 다음과 같은 문제도 발생합니다:

통신 오버헤드: 전문가 병렬 처리에서는 All-to-All 통신이 필요하여 장치 간 토큰을 섞어야 합니다. 전문가 수가 많아질수록 통신 오버헤드가 커져 병목 현상이 발생합니다.
부하 균형의 어려움: 더 많은 전문가가 모두 효과적으로 활용되도록 라우터와 보조 손실 함수를 설계하는 것이 더 어려워집니다.
이득 감소: 데이터가 자연스럽게 그렇게 많은 구별된 전문 분야로 분해되지 않거나, 라우터가 토큰을 효과적으로 할당하지 못하면 전문가를 더 추가해도 모델 품질이 크게 향상되지 않을 수 있습니다.

레이어별 전문가 수

더 깊은 레이어일수록 더 많은 전문가의 이점이 크다는 관찰을 활용하여, DeepSpeed-MoE는 후반 레이어에 더 많은 전문가를 배치합니다. 하지만 레이어별로 전문가 수가 달라지기에 전문가 로드 밸런싱과 병렬화 측면에서의 난제가 대두됩니다.

DeepSpeed-MoE는 전문가 병렬도를 각 레이어의 전문가 수에 맞추면 병목 현상을 피하고 자원 활용도를 최대화할 수 있다는 가설에 기인하여 레이어별 병렬도^parallelism를 동적으로 조정함으로써 이 난제를 개선합니다. 예를 들어, 128장의 GPU로 MoE 모델 훈련 시, 32개의 전문가를 가진 레이어는 32-way 전문가 병렬화와 4-way 데이터 병렬화를 적용하고, 128개의 전문가를 가진 레이어는 128-way 전문가 병렬화와 1-way 데이터 병렬화를 적용합니다. 이를 통해 각 GPU는 레이어별로 정확히 한 개의 전문가만 처리하게 됩니다.

3.2. Top-k 라우팅과의 상호작용

대부분의 최신 MoE 구현에서는 top-k 라우팅을 사용하며, 일반적으로 $k=1$ 또는 $k=2$ 를 선택합니다.

k=1: 각 토큰은 단일 전문가에게 라우팅됩니다. 배치/그룹 전체의 토큰 할당 수는 $T$ 입니다. 총 가용 용량은 $N \times C$ 입니다. 부하 분산은 $T$ 개의 할당을 $N$ 명의 전문가들 사이에 대략 균등하게 분배하는 것을 목표로 합니다.
k=2: 각 토큰은 두 명의 전문가에게 라우팅됩니다. 토큰 할당 수는 총 $2T$ 가 됩니다. 총 가용 용량은 여전히 $N \times C$ 입니다. 이는 시스템에 훨씬 더 큰 부담을 줍니다. 과도한 토큰 손실을 피하기 위해서는 다음이 필요할 수 있습니다:
- 더 높은 전문가 용량^{expert capacity}
- 보조 손실^{auxiliary loss} 을 통한 더 효과적인 부하 분산
- 토큰 배치 크기( $T$ )에 비해 더 많은 전문가 수( $N$ )

$k=2$ 를 사용하는 것은 토큰이 여러 전문화된 기능의 혜택을 받을 수 있게 하여 모델 품질을 향상시키는 경우도 있지만, 토큰당 두 명의 전문가를 활성화함으로써 계산량 증가와 잠재적으로 더 높은 통신 및 용량 요구 사항이라는 비용이 따르기 때문에 신중히 검토해야 합니다.

3.3. Upcycling MoE

MoE 모델 파라미터를 처음부터 훈련하는 대신(처음부터 MoE로 학습하는 것은 계산 자원, 통신 인프라, 데이터 양 면에서 매우 부담이 큽니다!), 이미 훈련된 밀집^dense 모델의 체크포인트를 살려서 전문가 구조를 덧붙이거나 변형하는 방식입니다 Sparse Upcycling 논문에서 제안되었고 Qwen을 비롯한 대다수 MoE에서 적용 중인 기법입니다. Megatron-Core에서도 --moe-use-upcycling와 --moe-upcycling-granularity 파라미터로 업사이클링을 지원하고 있습니다.

기본(upcycling 기본 전략; --moe-upcycling-granularity = 1)에서는 dense 모델의 MLP (혹은 FFN) 가중치를 그대로 복제하여 여러 전문가로 나누는 방식입니다. 즉, 전문가 hidden size는 dense 모델의 hidden size와 동일하게 유지되고, 여러 복제본을 만드는 것입니다
granular upcycling 전략에서는, 각 전문가의 hidden size를 dense의 hidden size를 --moe-upcycling-granularity만큼 나눈 값으로 작게 설정합니다. (예를 들어 granularity 값을 2로 하면 각 전문가가 dense의 절반 크기로 세분화됩니다.) 전문가의 크기를 작게 함으로써, 전문가 수를 늘리거나 표현의 다양성을 확보하는 fine-grained 특성을 강화할 수 있게 됩니다.

Upcycling MoE 적용 시 고려해야 할 핵심 요소들은 다음과 같습니다.

전문가 수와 세분화 정도 (granularity) 결정: Dense 모델의 FFN 층을 MoE로 바꿀 때 몇 개의 전문가로 나눌지, 전문가 하나당 파라미터 크기를 어떻게 조정할지 결정해야 합니다. (Megatron-Core 사용 시, --moe-upcycling-granularity 조정)
전문가 초기화 전략: dense 모델의 가중치를 그대로 복제하거나 일부 변형하여 여러 전문가로 초기화하는 방식이 기본입니다. 그러나 전문가 간 차별성을 주지 않으면 모든 전문가가 동일하게 동작할 수 있어 specialization이 일어나지 않을 수 있습니다. Upcycling 논문에서는 “virtual group initialization” 기법을 제안해 dense 체크포인트에서 MoE 초기화를 할 때 전문가들이 라우터 상에서 top-k 후보군에 골고루 분포하도록 유도합니다. 또 다른 기법으로는 parameter merging, partial reinitialization 등을 병행하여 전문가 다양성을 유도하는 방법도 있습니다. 예를 들어 Drop-Upcycling 논문에서는 일부 파라미터를 통계적으로 재초기화하여 전문화 경로를 강화하는 전략을 제시합니다.
라우터 도입 및 사전 보정 (warm-up / router pre-training): Dense 모델에는 라우터가 없기 때문에 새로 라우터를 설계해야 합니다. 이때 라우터 파라미터이므로 초기 무작위 초기화가 가능하나, 너무 무작위이면 학습 초기에 routing 붕괴^collapse가 잘 일어납니다. 따라서 소량의 seed 데이터를 사용해 각 새로운 전문가가 잘 작동하도록 라우터를 미리 보정(pre-optimize)하는 단계를 도입하기도 합니다.
가중치 스케일링 (weight scaling): Dense → MoE 전환 시, 여러 전문가로 쪼갤 때 가중치 분산이 달라지고 표현 강도가 약해질 수 있어서, 적절한 스케일 보정 (예: 각 전문가 가중치에 상수 계수 곱하기) 전략이 필요합니다. Upcycling 논문에서는 weight scaling을 통해 손실 개선 효과가 있다는 보고가 있습니다.
점진적 확장 또는 단계적 업사이클링: 처음부터 Dense → 아주 많은 전문가로 일괄 전환하기보다, 적은 전문가 수로 시작하고 점차 전문가 수를 늘리면서 upcycling을 점진적으로 수행하는 방식도 효과적입니다. Upcycling Instruction Tuning (UpIT) 논문에서는 전문가 확장 스테이지^{expert expansion stage}를 두고 점진적으로 전문가를 늘려 가는 전략을 사용합니다.
Drop-Upcycling 전략: Drop-Upcycling 논문은 dense에서 upcycling한 모델이 초기에는 좋은 성능을 보이나 장기 학습 시 속도가 느려지거나 specialization이 약해지는 문제가 있다고 보고하며, 일부 파라미터를 무작위 재초기화하는 전략을 병행해 학습 성능과 효율을 모두 개선하는 방식을 제안했습니다.

3.4. 도메인별 토큰 분포를 MoE 훈련에 반영

도메인별 토큰 분포을 MoE 모델 학습 및 라우팅 설계에 반영하려면 아래 전략들을 고려하시는 것이 좋습니다.

도메인 태깅 및 라우팅 보조 입력: 입력 텍스트에 도메인 태그를 포함시키고, 라우터가 이 태그까지 고려하도록 만드는 방식입니다. 예: “[domain_news] … text …” 형태로 토큰 앞에 삽입하거나 도메인 임베딩을 라우터 입력에 합침
도메인별 로드 밸런싱 항 및 가중치: 일반 로드 밸런싱 손실뿐 아니라 도메인별 균등화를 유도하는 보조 항을 둡니다. 예컨대 도메인 A, B, C가 있을 때 각 도메인별 전문가 사용률 차이를 줄이기 위한 항 추가
토큰 분포 스무딩 / 업샘플링: 특정 도메인 토큰의 표현이 너무 희소해 전문가가 학습하기 어렵다면, 해당 도메인의 데이터를 약간 업샘플링하거나, 토큰 빈도가 낮은 부분을 보정해서 라우팅이 너무 치우치지 않도록 조정할 수 있습니다.

전문가 구성은 “공유 전문가 + 도메인 경향 전문가”의 혼합이 좋습니다. 모든 입력을 무난히 처리하는 1–2개의 공유 전문가를 두고, 나머지는 특정 도메인(예: 자연어, 코드, 수학)에 각각 경향성만 주는 수준으로 설계합니다. 처음부터 도메인 고정 전문가로 만들기보다는, 데이터와 라우팅이 자연스럽게 특화를 유도하게 두는 편이 수렴이 매끄럽기 때문입니다. 이미 보유한 dense 모델을 upcycling해 시작한다면, FFN 가중치를 복제·분할하여 전문가 초기화를 하되, 일부 전문가에는 미세한 재초기화(스케일 조정, layer-norm 재설정 등)를 가해 초기에 동일 동작으로 수렴하는 것을 피하세요. Upcycling 직후에는 라우터만 소량 데이터로 짧게 워밍업해 라우팅 편향을 예방하고, 그 다음 전체 업데이트를 수행하는 2-stage 절차가 실패 확률을 줄여 줍니다. 자세한 내용은 DeepSeekMoE 논문을 참조 바랍니다.

도메인 혼합의 경계에서 불안정이 잦다면 토큰-전문가 친화도^affinity를 추가로 보정하는 방식도 효과적입니다. 예컨대 토큰 임베딩과 전문가 키(전문가당 요약 벡터)의 코사인 유사도를 얕게 섞어 게이트 점수를 살짝 재정렬하면, 도메인 급전환 구간에서의 흔들림이 줄어듭니다.

3.5. 보조 손실 가중치 계수 조정

적절한 값은 보통 경험적으로^heuristic 찾아야 하지만, 일반적인 방법론은 다음과 같습니다.

추천 범위: 값은 보통 $10^{-2}$ 또는 $10^{-3}$ 처럼 작지만, 최적 값은 모델 아키텍처, 작업, 전문가 수, 선택한 보조 손실 공식에 크게 의존합니다. 지나치게 큰 가중치 계수는 로드 밸런싱을 우선시하여 전체적인 정확도가 떨어질 수 있고 지나치게 작은 가중치 계수는 라우터 불균형을 개선하지 않다는 점을 기억해 주세요.
모니터링: 훈련 중에는 아래와 같은 중요한 지표들을 관찰하세요. 모니터링은 Weights & Biases나 텐서보드 같은 툴킷이 매우 유용합니다.
1. 전문가 활용도: 각 배치 또는 여러 단계에 걸쳐 각 전문가가 처리한 토큰의 수나 비율을 추적합니다. 이상적으로는 전문가들이 대략 비슷한 활용도를 가져야 하지만, 완벽한 균일성이 항상 필요하거나 최적은 아닙니다.
2. 보조 손실 크기: 보조 손실 값 자체를 관찰합니다. 일반적으로 훈련이 진행되고 전문가 균형이 개선됨에 따라 감소해야 합니다. 보조 손실이 $L_{task}$ 에 비해 어느 정도인지도 모니터링합니다.
3. 작업 손실 및 검증 성능: 보조 손실 계수를 증가시킬 때 모델이 주 작업을 학습하는 능력이 지나치게 손상되지 않도록 합니다. $L_{task}$ 와 관련된 검증 지표(예: 정확도, BLEU 등)를 추적하세요.
스케줄링: 모델과 데이터에 따라 가변적으로 두시는 편이 좋습니다. 현업에서 흔한 전략은 초반에는 규제를 조금 강하게, 이후 완만히 낮추는 방식입니다. 예를 들어 초기에는 균형 강제를 위해 보조 손실 가중치 계수를 조금 큰 값으로 시작해서 안정화 구간에 진입할 때 점차 감소시키는 스케줄링을 적용하면, 초기 붕괴를 막으면서 후반에는 전문가 특화 여지를 돌려줄 수 있습니다. 최근 대규모 MoE 훈련 노하우에서도 “부하가 이미 잘 균형을 이루면 보조 손실 계수를 낮추고, 불균형이 크면 일시적으로 올려서 강제 균형을 준다”는 상태 기반 조정을 권합니다.

3.6. 실용적 지침 및 모니터링

MoE 모델 개발 및 훈련 중에 추적해야 할 중요한 지표 몇 가지입니다.

드롭된 토큰 비율: 용량이 충분한지 여부를 나타내는 주요 지표입니다. 높은 값(1-2% 이상)은 문제가 발생했음을 의미합니다.
전문가 활용도: 전문가 용량이나 전문가 할당 시각화와 같은 지표를 모니터링합니다. 이는 라우터 문제를 진단하고 부하 균형 손실을 조정하는 데 도움이 됩니다.
계산 비용(FLOPs) 및 훈련 처리량: 다양한 구성의 실제 성능 영향을 측정합니다.
전체 모델 성능: LLM의 BLEU/ROUGE나 다운스트림 작업의 정확도와 같은 표준 지표를 추적합니다. 아키텍처 선택이 최종 성능 향상으로 이어지는지 확인합니다.
메모리 사용량: 특히 분산 환경에서 중요하며, 선택한 구성이 장치 메모리 제약 내에 맞는지 확인합니다.

모니터링은 도메인 인식형으로 설계해야 합니다. 훈련 루프에서 전문가별 토큰 점유율을 도메인별로 분해해 기록하고, 전문가 라우팅 확률의 엔트로피, 미사용(혹은 저사용) 전문가 비율, capacity 초과율, 드롭률을 함께 추적하세요. 코드와 수학 도메인의 길이 분포가 자연어와 다르므로, 시퀀스 길이 구간별로도 나눠서 보시면 병목 지점을 빨리 찾을 수 있습니다. 라우팅 변동성이 큰 경우에는 인접 토큰의 게이트를 약하게 참조하는 토큰 유사도 기반 스무딩 기법도 도움이 됩니다.

3.7. 옵티마이저 및 하이퍼파라미터

MoE 훈련을 위한 최적의 옵티마이저 설정과 하이퍼파라미터 조합을 찾는 것은 일반적으로 경험적인 과정이만, 아래의 베이스라인부터 시작해보는 것을 권장합니다.

기본값으로 시작하기: 표준 AdamW 설정 $\beta_1=0.9, \beta_2=0.999$ , 잘 알려진 학습률 스케줄(warmup + decay), 그리고 적당한 가중치 decay(예: 0.01 또는 0.1)부터 시작하세요.
학습률^{learning rate}과 보조 손실 계수를 함께 조정하기: 학습률과 보조 손실 계수 $\alpha$ 는 보통 가장 민감한 파라미터입니다. 두 값을 함께 변화시키면서 작업 손실과 부하 균형 지표를 모니터링하는 실험을 수행하세요.
Gradient Clipping: 그래디언트 폭주^{exploding gradients}를 방지하기 위해 gradient norm clipping을 안전장치로 적용합니다.
Longer Warmup: 초기 안정성에 문제가 있다면 warmup 단계를 연장해 보세요.

References

PreviousMoE (Mixture-of-Experts)NextMoE 모델 비교 및 주요 기법 정리

Last updated 18 days ago