시간적으로 확장된 전문가 혼합 모델

기존 Mixture-of-Experts(MoE) 모델은 토큰마다 전문가를 교체하여 메모리 오프로딩 효율이 떨어지는 문제가 있었다. 이 논문은 강화학습의 Options Framework를 도입해 전문가 유지 시간을 늘림으로써 성능 저하 없이 메모리 사용량과 추론 지연 시간을 획기적으로 줄이는 방법을 제시한다.

Transformer 기반의 MoE 모델은 각 토큰이 입력될 때마다 Softmax 기반의 라우터를 통해 가장 적합한 전문가(Expert)를 선택한다. 하지만 이 방식은 거의 모든 토큰마다 활성화되는 전문가 세트가 바뀌는 극심한 변동(Churn)을 초래하며, 모델 크기가 GPU 메모리를 초과하여 외부 메모리에서 가중치를 불러와야 하는 상황에서는 심각한 병목 현상을 일으킨다.

이 논문은 전문가 선택을 단순한 분류 문제가 아닌, 일정 시간 동안 유지되는 '행동'으로 정의하는 강화학습의 Options Framework를 통해 해결한다. 마치 사람이 매 순간 결정을 내리는 대신 하나의 '옵션(예: 점심 먹기)'을 정하면 일정 시간 그 행동을 유지하는 것과 같다. 모델 내부에 탑재된 컨트롤러는 현재 전문가를 유지할 때의 이득과 새로운 전문가를 로드할 때 발생하는 지연 시간(Deliberation Cost)을 비교하여 최적의 교체 시점을 결정한다.

결과적으로 모델은 문맥적 흐름에 따라 필요한 전문가 세트를 한 번 로드하면 수십 개의 토큰 동안 이를 유지하게 된다. 이는 GPU 메모리에 상주시켜야 하는 전문가 수를 줄여주며, 전문가 교체 횟수를 획기적으로 낮추어 메모리 대역폭 한계를 극복하고 전체적인 추론 속도를 높이는 결과를 가져온다.

MoE 전문가 선택 과정을 semi-Markov Decision Process(s-MDP)로 정형화했다. 각 레이어의 전문가 마스크를 하나의 'Option'으로 간주하고, 전문가 로딩 지연 시간을 'Deliberation Cost'인 η로 설정하여 최적화한다.

컨트롤러는 DeepSets 인코더를 사용하여 현재 전문가 마스크를 임베딩하고, 이를 LLM의 Hidden State와 결합하여 종료 확률(Termination Probability) β를 계산한다. [Hidden State와 현재 마스크 임베딩 입력 → MLP 연산 → 0~1 사이의 확률 출력] 순으로 계산되며, 이 확률에 따라 현재 전문가를 유지할지(KEEP) 아니면 새로운 전문가 세트를 샘플링할지(SWITCH) 결정한다.

학습에는 Option-Critic 아키텍처를 사용하며, 기존 MoE 모델을 교사(Teacher)로 삼아 Student 모델이 교사의 확률 분포를 따르도록 하는 Self-distillation Reward를 부여한다. [교사 모델의 로그 확률 - 학생 모델의 로그 확률 → Reward 산출] 과정을 통해 성능 저하를 최소화하면서 교체율을 낮추도록 유도한다.

GPT-oss-20b 모델 실험 결과, 전문가 교체율(Switch Rate)이 기존 58.6%에서 η=0.02 설정 시 4.1%로 급감했다. 이는 전문가 세트가 한 번 선택되면 평균적으로 약 25토큰 동안 유지됨을 의미한다.

성능 측면에서 MATH 벤치마크는 베이스 모델 71.5% 대비 64.0%를 기록했고, MMLU는 79.5% 대비 72.5%를 유지했다. 전문가 수를 16개에서 8개로 더 줄였을 때도 기존 Pruning 기법들(Wanda, Reconstruction Loss 등)이 무너지는 것과 달리 유의미한 추론 능력을 유지했다.

VRAM 요구 사항 분석에서 GPT-oss-20b의 전문가 16개만 유지할 경우 약 4.7 GiB(37%), 8개 유지 시 약 7.1 GiB(55%)의 메모리 절감 효과가 있음을 확인했다.

각 레이어는 독립적인 컨트롤러를 가지며, 이는 Plackett-Luce 분포를 통해 전문가 세트를 샘플링한다. 학습 안정성을 위해 Gumbel-top-k 트릭을 사용하여 샘플링 과정을 벡터화하고 미분 가능하게 구현했다.

Intra-option Policy 업데이트 시에는 LoRA(rank=16)를 전문가 및 Attention 파라미터에 적용하여 효율적인 파인튜닝을 수행한다. Reward 설계 시 Student와 Teacher의 분포 혼합(Mixture) 샘플링을 통해 Reward Hacking과 모델 붕괴를 방지했다.

Deliberation Cost η는 하이퍼파라미터로 작용하며, 이 값을 조절함으로써 성능(Accuracy)과 효율성(Switch Rate) 사이의 Pareto Frontier를 사용자 요구에 맞게 조정할 수 있는 유연성을 제공한다.

현재 연구는 각 레이어가 독립적으로 전문가 교체 여부를 결정하므로, 모든 레이어가 동시에 교체되는 Cross-layer 최적화에 비해서는 시스템 구현 복잡도가 높을 수 있다. 또한 실제 하드웨어 로딩 시간을 η에 직접 매핑하는 실측 기반 연구가 향후 과제로 남아있다.