동적 연산 할당 및 부하 분산을 지원하는 자기회귀 언어 모델링을 위한 전문가 임계값 라우팅

기존 MoE 모델은 토큰별 연산량을 고정하거나 미래 정보를 참조해야 하는 한계가 있었다. 이 논문은 전문가별 동적 임계값을 도입해 인과성을 유지하면서도 토큰의 중요도에 따라 연산 자원을 지능적으로 배분하는 새로운 표준을 제시했다.

MoE(Mixture of Experts)는 입력 토큰을 수많은 전문가 중 일부에게만 보내 연산량을 줄이는 구조다. 기존의 Token-choice(TC) 방식은 토큰마다 무조건 N개의 전문가를 고르게 강제하는데, 이는 쉬운 토큰에는 연산 낭비를, 어려운 토큰에는 연산 부족을 초래하는 한계가 있다. 이를 해결하기 위해 제안된 Expert-choice(EC)는 전문가가 토큰을 고르게 하여 토큰별 연산량을 조절하지만, 배치 전체의 토큰 점수를 비교해야 하므로 미래 토큰을 알 수 없는 실시간 생성 환경에서는 사용할 수 없었다. ET(Expert Threshold) 라우팅은 각 전문가가 과거에 들어온 토큰들의 점수 분포를 바탕으로 상위 1/E에 해당하는 점수 커트라인(임계값)을 스스로 관리하게 한다. 새로운 토큰이 들어오면 이 고정된 커트라인과 비교만 하면 되므로, 다른 토큰을 기다릴 필요 없이 즉시 라우팅이 가능하다. 결과적으로 어려운 토큰은 여러 전문가의 임계값을 통과해 집중 케어를 받고, 쉬운 토큰은 적은 전문가만 거치게 되어 전체적인 연산 자원을 지능적으로 배분하면서도 실시간 생성이 가능한 구조를 완성했다.

ET 라우팅의 핵심은 각 전문가 i가 유지하는 지수 이동 평균(EMA) 임계값 c_i이다. 라우터가 계산한 토큰 t와 전문가 i 사이의 점수 r_{t,i}가 c_i보다 크면 해당 전문가를 활성화(z_{t,i}=1)한다. [토큰 표현 벡터와 전문가 가중치 행렬을 내적하여 점수 r_{t,i}를 구하고] → [이를 현재 전문가의 임계값 c_i와 비교하여] → [이진 활성화 여부 z_{t,i}를 결정하며] → [이는 토큰이 독립적으로 전문가를 선택함을 의미한다]. 학습 과정에서 임계값 c_i는 매 스텝 업데이트된다. 배치 내에서 해당 전문가에 대한 점수들 중 상위 k번째(목표 활성화율에 해당) 점수를 찾아 현재 EMA 값과 결합한다. [c_i = beta * c_i + (1-beta) * kth-largest(r) 연산을 수행하여] → [과거의 점수 분포를 반영한 새로운 임계값을 도출하고] → [이를 다음 스텝의 기준으로 삼아] → [전체 학습 과정에서 전문가 부하가 목표치에 수렴하도록 유도한다]. 초기 학습의 불안정성을 해결하기 위해 Warmup 전략을 사용한다. 초기 4,000 스텝 동안은 표준 EC 라우팅(배치 내 상위 k개 선택)을 사용하여 임계값이 유의미한 통계치에 도달할 때까지 전문가 굶주림(Starvation) 현상을 방지하며 안정적인 통계치를 확보한다.

FineWeb-Edu 데이터셋을 이용한 2.4B 파라미터(활성 파라미터 561M) 모델 실험 결과, ET 라우팅은 TC-MoE 대비 Cross-Entropy Loss에서 0.067의 성능 향상을 보였다. 이는 동일한 손실값에 도달하기 위해 필요한 토큰 수를 1.6배 절감한 것과 같은 효과다. 벤치마크 평가인 CORE Eval에서도 ET는 25.14점을 기록하여 Dense 모델(20.43) 및 TC aux(22.31)를 크게 상회했으며, 비인과적 방식인 EC(24.98)와 대등하거나 소폭 앞서는 성능을 증명했다. 전문가 전문화(Expert Specialization) 분석 결과, ET는 HumanEval(코드) 및 GSM8K(수학) 도메인에서 특정 전문가가 특정 도메인 토큰을 집중적으로 처리하는 뚜렷한 패턴을 보였다. 이는 대규모 배치 EC와 유사한 수준의 전문화가 독립적 라우팅만으로도 가능함을 시사한다.

ET 아키텍처는 Nanochat 기반의 Transformer 구조를 따르며, 16개의 라우팅 전문가와 1개의 공유 전문가(Shared Expert)를 혼합하여 사용한다. 공유 전문가는 모든 토큰을 처리하여 글로벌 지식을 캡처하고, 라우팅 전문가는 ET 메커니즘을 통해 선택적으로 활성화된다. 수학적으로 ET는 E[z]=1/E라는 확률적 기댓값 제약 조건을 만족하도록 설계되었다. 이는 배치 단위의 엄격한 부하 분산 대신 모집단(Population) 수준의 부하 분산을 추구함으로써 인과성을 확보하는 접근법이다. 구현 측면에서 Sigmoid Gate를 사용하여 Softmax의 상호 의존성을 제거했다. 또한 GPU 메모리 부족을 방지하기 위해 전문가 용량 계수(Capacity Factor) C=0.5를 설정하여 학습 시 토큰 드롭이나 패딩을 관리한다. 연구팀은 유한 정밀도 부동 소수점을 사용할 경우 임계값 전달을 통한 미래 정보 누설이 상수로 제한됨을 이론적으로 증명하여, ET가 실질적으로 완전한 인과성을 유지함을 뒷받침했다.

초기 학습 단계에서 임계값이 안정화되기 전까지 전문가 활용도가 낮아지는 전문가 굶주림(Starvation) 현상이 발생할 수 있으며, 이를 방지하기 위해 초기 수천 스텝 동안의 Warmup 과정이 필수적으로 요구된다.