Expert-Choice 라우팅을 통한 확산 언어 모델의 적응형 연산 구현

확산 언어 모델(DLM)에서 기존의 토큰 선택 방식 대신 전문가 선택 방식을 도입하여 연산 효율을 극대화했다. 특히 디노이징 단계별로 전문가 처리 용량을 다르게 배정하는 기법을 통해 동일 연산량 대비 성능을 크게 개선하고 학습 속도를 2배 높였다.

기존의 Mixture-of-Experts(MoE) 모델은 각 토큰이 선호하는 전문가를 선택하는 Token-Choice(TC) 방식을 주로 사용한다. 하지만 이 방식은 특정 전문가에게 토큰이 몰리는 부하 불균형 문제를 해결하기 위해 복잡한 보조 손실 함수를 써야 하며, 이 과정에서 학습 효율이 저하되는 한계가 있다.

논문은 확산 언어 모델(DLM)이 한 번에 전체 시퀀스를 처리한다는 점에 주목하여, 전문가가 자신에게 적합한 토큰을 고정된 개수만큼 선택하는 Expert-Choice(EC) 방식을 제안한다. 이는 구조적으로 완벽한 부하 균형을 보장하므로 보조 손실 함수가 필요 없고 하드웨어 활용도를 극대화한다.

더 나아가 DLM의 디노이징 과정에서 마스크 비율이 낮아질수록(즉, 문맥 정보가 많아질수록) 토큰당 학습 효율이 10배 이상 높다는 사실을 발견했다. 이를 바탕으로 학습 효율이 높은 단계에 전문가 연산 용량을 집중적으로 배치함으로써 전체적인 모델 성능을 비약적으로 향상시켰다.

전체 시퀀스를 동시에 처리하는 DLM의 특성을 활용해 전문가가 상위 c개의 토큰을 선택하는 EC 라우팅을 적용한다. 각 전문가 j가 게이팅 점수 S_ij를 기준으로 상위 c개의 토큰 인덱스 집합 T_j를 선택하며, 출력 y_i는 선택된 전문가들의 FFN 결과값에 게이트 가중치 g_ij를 곱해 합산하는 방식으로 계산된다.

디노이징 단계 t에서의 마스크 비율 r에 따라 전문가 용량 k(r)을 결정하는 스케줄링 함수를 정의한다. k(r) = clamp(k_min + (k_max - k_min) * s(r), k_min, k_max) 수식을 사용하며, 여기서 s(r)은 마스크 비율에 따른 가중치 함수이다. [입력 r → s(r) 계산 → k_min/k_max 사이의 용량 결정 → 해당 단계의 전문가 연산량 확정] 순으로 연산이 진행되어 단계별로 차등화된 연산 자원을 배분한다.

학습 효율 분석을 위해 각 마스크 비율 구간별로 로그 손실의 시간에 따른 변화율인 수렴 속도 η_r = -d ln(L_r) / dt를 측정한다. [단계별 손실값 입력 → 시간 대비 로그 변화량 계산 → 수렴 속도 산출 → 학습 효율성 지표로 활용] 과정을 통해 마스크 비율이 낮은 구간이 모델의 주요 학습 전선임을 확인하고 연산량을 집중시켰다.

EC 라우팅은 TC 방식 대비 동일 손실값(3.75) 도달 시간을 20.7시간에서 10.6시간으로 단축하여 2.0배 빠른 수렴 속도를 기록했다. 하드웨어 처리량 측면에서도 EC는 52.1 TFLOP/s/GPU를 달성하여 TC 변체들(24.935.4 TFLOP/s)보다 1.52.1배 높은 효율을 보였다.

다양한 스케줄링 전략 중 마스크 비율이 낮을 때 연산량을 늘리는 Linear-Reverse 방식이 36.5 PPL을 기록해 정적 방식(37.1 PPL)이나 마스크 비율이 높을 때 연산량을 늘리는 방식(37.5 PPL)보다 우수한 성능을 보였다. 이는 8B 규모의 모델 실험에서도 MMLU 및 ARC-Challenge 정확도 향상으로 일관되게 나타났다.

기존 TC 기반 LLaDA-MoE 모델을 EC로 교체하여 미세 조정한 결과, HumanEval에서 Pass@1이 53.9%에서 58.6%로 상승했으며 추론 속도는 1.3~1.5배 빨라졌다. 이는 사후 교체만으로도 성능과 효율을 동시에 잡을 수 있음을 시사한다.

본 연구는 DLM의 비인과적(Non-causal) 어텐션 구조가 EC 라우팅의 전역적 토큰 선택 요구사항과 완벽히 일치함을 활용한다. TC 라우팅의 고질적 문제인 전문가 과부하 및 토큰 드롭 문제를 구조적으로 제거하여 보조 손실 함수 없이도 결정론적인 부하 분산을 실현했다.

핵심 기술적 차별점은 전문가 용량(Expert Capacity)을 고정된 상수가 아닌 디노이징 타임스텝의 함수로 취급한 점이다. 이는 DLM의 반복적 디노이징 루프가 각 단계마다 서로 다른 난이도와 학습 효율을 가진다는 특성에 기반한다. 실험을 통해 저마스크(Low-mask) 단계가 고마스크 단계보다 수렴 속도가 최대 20배 빠름을 정량적으로 분석했다.

구현 측면에서는 Megatron-LM 프레임워크를 기반으로 SwiGLU 활성화 함수와 WSD 학습률 스케줄을 사용했다. 8B 파라미터 규모(활성 파라미터 1B)에서 64개의 미세 전문가(Fine-grained experts)와 2개의 공유 전문가를 배치하여 실험적 타당성을 검증했다.

본 연구에서 탐구한 용량 스케줄링 함수는 선형, 코사인, 가우시안 등 수동으로 설계된 함수에 국한된다. 모델 규모나 데이터셋, 작업의 종류에 따라 최적의 스케줄이 달라질 수 있으며, 이를 자동으로 학습하는 경량 예측기나 강화학습 기반 정책 도입은 향후 과제로 남겨두었다.