Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

이 논문은 토큰이 각 전문가의 임계값을 독립적으로 초과하는지 확인하여 동적 계산 할당과 보조 손실 없이 부하 균형을 달성하는 '전문가 임계값 (Expert Threshold)' 라우팅을 제안하며, 이를 통해 TC-MoE 대비 더 낮은 손실과 적은 토큰으로 동등한 성능을 달성하는 것을 보여줍니다.

핵심

🏢 비유: 거대한 '지식 회사'와 '새로운 배정 시스템'

생각해 보세요. 거대한 AI 회사가 있다고 칩시다. 이 회사에는 **수백 명의 '전문가 (Expert)'**들이 있습니다. 어떤 전문가는 수학 문제를 잘 풀고, 어떤 전문가는 코딩을 잘하며, 또 어떤 전문가는 문법 교정을 잘합니다.

AI 가 문장을 하나 만들 때마다 (예: "안녕하세요"), 이 회사에서는 그 문장의 각 단어 (토큰) 에 맞는 전문가를 찾아 일을 시켜야 합니다.

1. 기존 방식의 문제점 (TC-MoE 와 EC-MoE)

기존에는 두 가지 방식이 있었는데, 둘 다 단점이 있었습니다.

• 방식 A: "정해진 인원 뽑기" (Token Choice)

  • 비유: 모든 단어가 "나는 3 명만 뽑아!"라고 외칩니다. 각 단어는 무조건 점수가 가장 높은 전문가 3 명을 뽑습니다.
  • 문제: 어떤 날은 모든 단어가 똑같은 '수학 전문가'만 뽑아서 그 분은 과로로 쓰러지고, '코딩 전문가'는 하루 종일 놀게 됩니다. (부하 불균형)
  • 해결책: 회사가 "너희는 고르게 분배해!"라고 강제로 지시하거나, 불균형하면 벌금을 매기는 복잡한 규칙을 만들어야 했습니다.

• 방식 B: "전문가가 직접 고르기" (Expert Choice)

  • 비유: 이번에는 전문가들이 "나에게 가장 잘 맞는 단어 100 개만 고르겠다!"라고 합니다.
  • 문제: 이 방식은 아주 효율적이지만, 시간 여행이 필요합니다. 전문가가 "내가 고를 단어 100 개를 정하려면, 미래에 나올 단어들도 다 봐야 해!"라고 말합니다.
  • 치명적 단점: AI 가 글을 쓸 때는 미래의 단어가 아직 존재하지 않습니다. ( autoregressive, 자기주도적 생성) 그래서 이 방식은 글을 실시간으로 쓰는 데는 쓸 수 없었습니다.

2. 이 논문이 제안한 해결책: "Expert Threshold (ET) Routing"

이 논문은 **"미래를 보지 않아도, 전문가가 스스로 적정선을 유지하게 하는 방법"**을 찾아냈습니다.

• 핵심 아이디어: "평균 점수 기준선 (Threshold)"
  • 비유: 각 전문가에게 **"오늘 내가 일할 수 있는 기준선"**을 정해줍니다. 이 기준선은 과거의 모든 데이터를 바탕으로 계산된 **'이동 평균 (EMA)'**입니다.
  • 작동 원리:
    1. 새로운 단어가 들어오면, 그 단어의 점수가 전문가의 기준선보다 높은지 확인합니다.
    2. 높다면: "일해!" (할당)
    3. 낮다면: "오늘은 쉬어." (할당 안 함)
  • 장점:
    • 미래 불필요: 현재 단어 하나만 보면 되므로, 실시간으로 글을 쓸 수 있습니다. (인과성 유지)
    • 자동 균형: 기준선은 과거 전체 데이터의 평균이므로, 장기적으로 보면 모든 전문가가 고르게 일을 하게 됩니다. 별도의 벌금이나 복잡한 규칙이 필요 없습니다.
    • 유연한 자원: 어려운 단어는 여러 전문가가, 쉬운 단어는 적은 전문가가 처리할 수 있어 계산 자원을 아낄 수 있습니다.

🚀 이 방법이 왜 대단한가요?

1. 더 똑똑해졌습니다: 실험 결과, 이 방법을 쓰면 기존 방식보다 오류가 0.067 만큼 줄어듭니다. 같은 실력을 내는 데 필요한 학습 데이터 양을 1.6 배나 줄일 수 있다는 뜻입니다. (더 빠르고 저렴하게 학습 가능)
2. 실시간 사용 가능: 미래의 단어를 알 수 없는 상황 (실시간 채팅, 글쓰기) 에서도 최고의 효율을 발휘합니다.
3. 균형 잡힌 업무: 어떤 전문가가 과부하 걸리거나, 어떤 전문가가 놀지 않도록 자연스럽게 조절됩니다.

📝 한 줄 요약

"AI 회사에서 각 전문가에게 '과거 평균'을 기준으로 스스로 일할지 말지 결정하게 했더니, 미래의 단어를 알 필요도 없이, 부하도 고르게 분배되고 더 똑똑해졌습니다."

이 기술은 앞으로 우리가 사용하는 AI 가 더 빠르고, 저렴하며, 똑똑하게 작동하는 데 큰 역할을 할 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

현재 MoE 아키텍처는 모델 용량을 늘리면서 계산 비용을 낮추기 위해 널리 사용되지만, 라우팅 메커니즘에는 근본적인 모순이 존재합니다.

• Token Choice (TC) 라우팅: 각 토큰이 고정된 수의 전문가 (Expert) 를 선택합니다.
  • 단점: 부하 불균형 (Load Imbalance) 을 해결하기 위해 보조 손실 함수 (Auxiliary Loss) 나 PID 제어기가 필요하며, 이는 라우팅을 복잡하게 만듭니다. 또한 계산 할당이 토큰당 고정되어 동적 할당이 어렵습니다.
• Expert Choice (EC) 라우팅: 각 전문가가 배치 (Batch) 내 토큰 중 상위 k 개를 선택합니다.
  • 장점: 완벽한 부하 균형과 동적 계산 할당이 가능합니다.
  • 치명적 단점 (인과성 위반): 특정 토큰을 선택하려면 **배치 내의 모든 토큰 (미래 토큰 포함)**의 점수를 비교해야 합니다. 이는 자회귀 생성 (Autoregressive Generation) 시 미래 토큰이 존재하지 않기 때문에 적용이 불가능하며, 훈련 시에는 미래 정보가 누출되는 (Information Leakage) 문제를 일으킵니다.

따라서, 인과성 (Causality) 을 유지하면서 동적 계산 할당과 부하 균형을 동시에 달성하는 새로운 라우팅 방식이 필요했습니다.

2. 제안 방법: Expert Threshold (ET) Routing

저자는 토큰당 희소성 (Sparsity) 제약과 배치 단위 부하 균형 제약을 완화하여, 기대값 (Expectation) 수준에서 부하 균형을 달성하는 방식을 제안합니다.

• 핵심 메커니즘:
  • 각 전문가 (Expert) 는 전역 토큰 분포를 기반으로 추정된 **지수 이동 평균 (EMA) 임계값 (Threshold, $c_i$)**을 유지합니다.
  • 각 토큰은 독립적으로 라우팅되며, 해당 토큰의 점수 ($r_{t,i}$) 가 전문가의 임계값 ($c_i$) 을 초과하면 ($r_{t,i} > c_i$) 해당 전문가로 라우팅됩니다.
  • 완전한 인과성 (Fully Causal): 라우팅 결정이 현재 토큰의 점수와 과거 통계 (EMA 임계값) 만에 의존하므로, 훈련과 추론 시 미래 토큰에 대한 접근이 불필요합니다.
• 부하 균형:
  • 배치 내에서의 완벽한 균형을 강제하지 않고, 전역 분포의 상위 $1/E$비율 (여기서$E$는 전문가 수) 에 해당하는 토큰이 선택되도록 임계값을 조정합니다.
  • 훈련 시 EMA 를 통해 임계값을 업데이트하고, 추론 시에도 동일한 임계값을 사용하여 훈련 - 추론 불일치 (Train-Inference Mismatch) 를 제거합니다.
• Warmup 전략:
  • 훈련 초기에는 토큰 분포가 불안정하여 임계값 추정이 어렵습니다. 이를 해결하기 위해 처음 4,000 스텝 동안은 기존 EC 라우팅을 사용하여 임계값이 안정화될 때까지 부하 균형을 유지한 후, ET 로 전환합니다.

3. 주요 기여 (Key Contributions)

1. 인과성 문제 해결: EC 라우팅의 인과성 위반 문제를 해결하여, MoE 를 자회귀 언어 모델에 안전하게 적용할 수 있게 했습니다.
2. 보조 손실 불필요: TC 라우팅에서 필요했던 보조 손실 함수 (Auxiliary Loss) 없이도 부하 균형을 달성합니다.
3. 동적 계산 할당: 토큰의 난이도나 중요도에 따라 활성화되는 전문가 수를 유연하게 조절할 수 있습니다.
4. 훈련 - 추론 일치: EMA 기반 임계값을 사용하여 훈련과 추론 시 라우팅 로직이 완전히 일치합니다.

4. 실험 결과 (Results)

FineWeb-Edu 데이터셋을 사용하여 24 억 (2.4B) 파라미터 (활성 파라미터 0.56B) 규모의 모델을 사전 학습 (Pretraining) 했습니다.

• 성능 향상:
  • TC 라우팅 대비: ET 는 교차 엔트로피 손실 (Cross-Entropy Loss) 에서 0.067만큼 더 낮은 값을 기록했습니다. 이는 동일한 성능을 달성하는 데 1.6 배 적은 토큰이 필요함을 의미합니다.
  • EC 라우팅 대비: 대용량 배치 (Large Batch) 를 사용한 EC 와 유사한 성능을 달성하면서도, 추론 시 배치 크기 조정이 필요 없습니다.
• 부하 균형: ET 는 TC 보다 훨씬 더 균일한 전문가 사용률을 보이며, EC 와 유사한 수준의 균형을 달성했습니다.
• 전문가 전문화 (Specialization): ET 는 EC 와 마찬가지로 특정 도메인 (코드, 수학 등) 에 특화된 전문가들을 형성하는 것을 확인했습니다.
• 안정성: Warmup 전략을 통해 초기 훈련 단계의 불안정성을 해결했으며, 임계값과 전문가 사용량이 안정적으로 수렴함을 보였습니다.

5. 의의 및 결론 (Significance)

• MoE 아키텍처의 확장: 이 연구는 MoE 모델이 대규모 자회귀 언어 모델 (LLM) 에 적용될 때 겪는 인과성 문제를 해결함으로써, MoE 의 확장성을 크게 높였습니다.
• 효율성과 성능의 균형: 보조 손실 없이도 부하 균형을 유지하며, 동적 계산 할당을 통해 모델의 효율성을 극대화합니다.
• 실용성: 추론 시 추가적인 배치 처리나 복잡한 조정 없이도 구현이 가능하여, 실제 서비스 환경에 적용하기 용이합니다.

결론적으로, Expert Threshold (ET) Routing은 MoE 기반 언어 모델의 라우팅 문제를 해결하는 강력한 대안으로, 인과성을 유지하면서도 동적 계산과 부하 균형을 동시에 달성하는 새로운 패러다임을 제시합니다.