Expert Upcycling: Mixture-of-Experts의 연산 효율적 경계를 확장하는 방법

대규모 언어 모델의 주류가 된 MoE 아키텍처를 처음부터 거대하게 학습시키는 대신, 작은 모델에서 전문가 수를 늘려가며 점진적으로 확장하는 효율적인 방법론을 제시한다. 이를 통해 동일한 성능의 모델을 훨씬 적은 GPU 시간으로 확보할 수 있어 모델 학습의 경제성을 크게 개선한다.

MoE 모델은 입력 토큰마다 전체 전문가 중 일부(Top-K)만 활성화하여 연산 효율을 챙기지만, 학습 시에는 모든 전문가의 가중치와 그래디언트를 관리해야 하므로 전문가 수가 늘어날수록 메모리와 통신 비용이 급증한다. 기존에는 처음부터 목표하는 전문가 수를 정해두고 학습했으나, 이는 초기 학습 단계에서 불필요하게 높은 비용을 지불하게 만든다.

이 논문은 Transformer의 전문가 계층이 학습 초기에는 서로 유사한 표현을 배우다가 점차 전문화된다는 점에 착안했다. 먼저 적은 수의 전문가(E)로 모델을 학습시켜 기초적인 표현 능력을 확보한 뒤, 이를 복제하여 더 많은 전문가(mE)를 가진 모델로 확장한다. 이때 복제된 전문가들은 동일한 가중치에서 시작하는 Warm Initialization 효과를 누리며, 이후 추가 학습(CPT) 과정에서 발생하는 stochastic gradient의 다양성에 의해 서로 다른 전문 분야로 분화된다.

결과적으로 7B 모델에서 시작해 13B 모델로 확장했을 때, 처음부터 13B로 학습한 모델과 동일한 성능을 내면서도 전체 GPU 학습 시간은 약 32% 절감하는 성과를 거두었다. 이는 모델의 크기를 키우는 과정이 반드시 처음부터 시작할 필요가 없음을 시사한다.

전체 프로세스는 세 단계로 구성된다. Step 1에서는 E개의 전문가를 가진 MoE 모델을 일정 단계(τ)까지 사전 학습한다. Step 2에서는 Upcycling Operator(Um)를 적용하여 전문가를 복제하고 라우터를 확장한다. 이때 각 전문가 e는 re번 복제되며, 라우터 가중치는 기존 값을 복사하되 복제된 슬롯에 미세한 노이즈를 추가하여 대칭성을 깬다. Step 3에서는 확장된 mE 모델로 남은 기간 동안 Continued Pre-training(CPT)을 수행하여 전문가들의 전문화를 유도한다.

핵심인 Utility-based Selection은 전문가의 중요도를 측정하기 위해 Squared gradient norm uG(e) = ||ge||²₂를 계산한다. [전문가 파라미터에 대한 손실 함수의 그래디언트 ge를 입력으로] → [L2 노름의 제곱 연산을 수행하여] → [스칼라 값인 중요도 점수를 얻고] → [이 점수가 높을수록 해당 전문가가 손실 감소에 민감하게 기여함을 의미]하므로 해당 전문가를 우선적으로 복제한다.

이론적으로는 Upcycling 모델의 손실(L_up)과 처음부터 학습한 모델의 손실(L_fs) 차이를 분석한다. [학습률 ηt와 각 모델의 최적 손실 차이를 입력으로] → [시간에 따른 가중 평균 연산을 수행하여] → [Capacity gap과 Initialization gain이라는 두 항을 얻고] → [Initialization gain이 음수일 때, 즉 복제된 초기화가 랜덤 초기화보다 최적점에 가까울 때 품질 격차가 해소됨을 증명]한다.

7B에서 13B 파라미터 규모로 확장하는 실험에서, Upcycling 모델은 100% CPT 예산을 투입했을 때 처음부터 학습한 Fixed-64 모델과 거의 동일한 검증 손실(1.263 vs 1.267)을 기록했다. 11개 주요 벤치마크 평균 정확도에서도 56.4%를 기록하여 Fixed-64의 56.7%에 근접했으며, 이 과정에서 GPU 시간은 약 32% 절약되었다.

전문가 선택 전략 비교에서는 Utility-based 방식이 Uniform 방식보다 우수함이 증명되었다. 특히 CPT 예산이 25%로 제한된 상황에서 Utility 방식은 성능 격차 해소 효율을 Uniform 대비 3배 이상(26.5% vs 8.2%) 높였다. 또한, Dense 모델을 MoE로 바꾸는 Sparse Upcycling보다 MoE에서 더 큰 MoE로 가는 Expert Upcycling이 활성화 비율이 낮아질수록 훨씬 더 효율적임이 확인되었다.

본 연구는 MoE의 확장 법칙(Scaling Law)이 활성화 파라미터가 아닌 전체 파라미터 수에 지배받는다는 점을 활용한다. 아키텍처적으로는 20개 레이어 중 일부만 MoE를 사용하는 Interleaved MoE 구조를 채택하여 통신 오버헤드를 줄였다. 학습 시에는 Loss-free load balancing 기법을 적용하여 복제된 전문가들이 고르게 그래디언트 신호를 받아 분화될 수 있도록 보장했다.

수학적으로는 Online Convex Optimization(OCO) 프레임워크를 사용하여 Upcycling의 수렴 경계를 도출했다. 특히 전문가 복제 시 가중치 공간에서의 Taylor 전개를 통해 Squared gradient norm이 손실 민감도를 측정하는 최적의 1차 근사치임을 이론적으로 뒷받침했다. 이는 단순한 휴리스틱이 아닌 최적화 이론에 기반한 전문가 선택 기준임을 보여준다.

전문가 수 확장 계수(m)가 2인 경우를 중심으로 검증되었으며, 더 큰 확장 계수나 사전 학습과 CPT 사이의 데이터 분포 변화가 극심한 환경에서의 안정성은 추가 연구가 필요하다. 또한 라우터 붕괴(Router collapse)나 밸런싱의 취약성 문제가 더 큰 규모에서 나타날 가능성을 언급했다.