슬림가능 전문가 혼합으로 효율적인 언어 모델 구현

초록

MoSE는 기존 MoE 구조에 슬림가능(가변 폭) 전문가를 도입해, 라우터가 선택한 전문가마다 실행 폭을 조절한다. 이를 통해 하나의 사전학습 모델로 연산량과 정확도 사이의 연속적인 트레이드오프를 제공하며, 테스트‑타임 훈련을 이용해 라우터 확신도에 기반한 폭 할당을 자동화한다. 실험 결과, MoSE는 동일 FLOPs에서 표준 MoE보다 낮은 퍼플렉시티와 향상된 제로샷 성능을 달성한다.

상세 분석

MoSE는 “Mixture‑of‑Experts”(MoE) 아키텍처의 핵심 한계를 보완한다. 기존 MoE는 토큰당 몇 개의 전문가만을 선택하고, 선택된 전문가를 전폭(full‑width)으로 실행한다. 이 설계는 라우터가 토큰별로 전문가를 선택하는 이산적 조건부 연산은 제공하지만, 선택된 전문가 내부의 연산량을 조절할 수 없다는 제약을 남긴다. MoSE는 각 전문가를 슬림가능(slimmable) 구조로 설계함으로써, 전문가 내부의 중간 차원(FFN의 확장 비율 4)을 폭 multiplier w∈A 로 슬라이스한다. 이렇게 하면 하나의 전문가 파라미터 집합이 여러 폭을 공유하는 중첩 네트워크가 되고, 라우터가 선택한 전문가마다 실행 폭을 다르게 지정할 수 있다.

학습 단계에서는 두 개의 폭을 동시에 샘플링한다. 하나는 항상 최대 폭 w_max, 다른 하나는 w_min과 w_max 사이의 균등 분포에서 무작위로 선택된 폭이다. 두 폭에 대해 각각 순전파·역전파를 수행하고 손실을 평균해 최적화한다. 이 방식은 기존 슬림가능 네트워크에서 사용되는 “random width training”을 MoE의 sparse routing과 결합한 것으로, 추가적인 연산 오버헤드가 거의 없으며, 전문가가 모든 폭에서 안정적으로 동작하도록 만든다. 라우터 자체는 기존 MoE와 동일하게 top‑k 방식을 사용하고, 로드 밸런싱 손실과 router‑z‑loss를 그대로 적용해 전문가 붕괴를 방지한다.

추론 단계에서는 세 가지 폭 할당 전략을 제시한다. (1) 전역 단일 폭을 모든 활성 전문가에 적용하는 uniform‑width 모드; (2) 라우터 확률 p_i를 직접 폭에 매핑하는 normalized‑probability 모드(γ 파라미터로 집중도를 조절); (3) 테스트‑타임 트레이닝(TTT)으로 γ(또는 레이어별 γ_l)를 짧은 캘리브레이션 데이터에 대해 최적화해, 사전 정의된 연산 예산 Γ 하에서 가장 효율적인 폭 분배를 학습한다. TTT는 모델 파라미터를 고정한 채 스칼라 γ만 조정하므로 매우 가볍고, 배포 환경에서 실시간 적응이 가능하다.

실험에서는 GPT‑2‑Small(55 M), GPT‑2‑Standard(322 M), GPT‑2‑Medium(1 B) 규모의 모델을 OpenWebText에 3 B~15 B 토큰으로 사전학습했다. 라우팅 구성(E8A2, E8A4, E16A4)과 다양한 FLOPs 예산에서 MoSE는 표준 MoE와 동등하거나 더 나은 퍼플렉시티를 기록했으며, 특히 TTT를 적용한 경우 동일 FLOPs 대비 1‑2 % 정도 퍼플렉시티가 감소했다. 또한 LAMBADA와 WSC 같은 제로샷 벤치마크에서도 정확도가 상승했다. 중요한 점은 MoSE가 전폭 실행(w=1.0)만 사용해도 MoE와 성능 차이가 없으며, 슬림가능 구조 자체가 표현력을 해치지 않는다는 점이다.

이 논문은 두 가지 혁신을 제공한다. 첫째, 전문가 내부의 연산량을 가변적으로 조절함으로써 “전문가 선택”과 “전문가 용량”이라는 두 축의 조건부 연산을 동시에 활용한다. 둘째, 라우터 확신도와 연산 예산을 연결하는 경량 테스트‑타임 학습 메커니즘을 제시해, 실제 서비스 환경에서 동적 비용‑품질 최적화를 가능하게 한다. 이러한 접근은 대규모 언어 모델을 클라우드·엣지 양쪽에서 효율적으로 운영하려는 현업에 큰 시사점을 제공한다.