UniPool: Mixture-of-Experts를 위한 전역 공유 전문가 풀

기존 MoE 모델이 레이어마다 별도의 전문가를 두어 파라미터가 낭비되던 문제를 전역 공유 풀(Global Shared Pool) 구조로 해결했다. 이를 통해 전문가 파라미터 수를 최대 60%까지 줄이면서도 기존 방식보다 더 낮은 손실값과 높은 정확도를 달성할 수 있음을 입증했다.

기존의 Mixture-of-Experts(MoE)는 각 층(Layer)이 자신만의 전용 전문가 팀을 가지고 있는 구조였다. 하지만 연구진이 분석한 결과, 깊은 층에 있는 전문가들은 서로 비슷한 기능을 학습하며 상당 부분 중복되어 있다는 사실이 밝혀졌다. 이는 마치 회사에서 부서마다 똑같은 기능을 하는 별도의 팀을 따로 두어 인력을 낭비하는 것과 같다.

UniPool은 이 부서별 전용 팀을 없애고 모든 부서가 필요할 때마다 호출할 수 있는 '전사 공용 전문가 풀'을 만들었다. 각 층에 있는 라우터(Router)는 입력된 데이터(Token)의 특징을 보고 공용 풀에서 가장 적합한 전문가를 골라낸다. 이때 층마다 데이터의 특성(Hidden-state scale)이 다르기 때문에, 이를 표준화하여 안정적으로 연결해주는 NormRouter 기법을 적용했다.

결과적으로 각 층은 자신에게 꼭 필요한 전문가만 골라 쓰면서도, 다른 층에서 이미 잘 학습된 전문가를 재사용할 수 있게 되었다. 이는 전체 파라미터 수를 획기적으로 줄이면서도 각 전문가가 더 다양한 데이터를 접하며 고도로 전문화되도록 유도하여 모델 전체의 지능을 높이는 결과를 가져왔다.

UniPool은 크게 세 가지 기술적 요소로 구성된다. 첫째, 전역 공유 전문가 풀(Global Shared Pool)이다. L개의 레이어가 각각 E개의 전문가를 갖는 대신, 총 M개의 전문가로 구성된 하나의 풀을 생성한다. 각 레이어 l은 독립적인 라우터 r_l을 유지하며 풀 내의 전문가 e_i를 선택한다.

둘째, 풀 레벨 보조 손실(Pool-level Auxiliary Loss)이다. 기존의 레이어별 손실 함수는 해당 레이어 내에서의 균형만 맞추려 하여 전역적인 전문가 활용을 방해한다. 이를 해결하기 위해 모든 레이어의 토큰 할당 비율 f_i를 평균 내어 [전체 레이어의 평균 할당량 → 전체 전문가 수 M과 곱함 → 전체 라우팅 확률 P_i와 결합]하는 연산을 거쳐 하나의 전역 손실값을 산출한다. 이 숫자가 낮아질수록 전체 풀의 전문가들이 골고루 사용된다.

셋째, NormRouter의 도입이다. 레이어마다 입력값의 크기(Scale)가 다르기 때문에 Softmax 대신 L2 정규화를 사용하는 NormRouter를 채택했다. [입력 로짓 z를 L2 노름으로 나눔 → 학습 가능한 스케일 인자 σ를 곱함 → ReLU를 통과시켜 희소한 점수 생성] 과정을 통해 레이어 간 간섭 없이 안정적인 라우팅이 가능하도록 설계했다.

LLaMA 아키텍처 기반의 182M부터 978M까지 5가지 규모 모델에서 실험한 결과, UniPool은 모든 규모에서 기존 MoE 대비 낮은 검증 손실(Validation Loss)과 퍼플렉서티(Perplexity)를 기록했다. 특히 830M 모델에서는 기존 MoE 대비 손실값을 0.0386 감소시키는 성과를 거두었다.

효율성 측면에서 UniPool은 기존 전문가 예산의 41.6%~66.7%만 사용하고도 레이어별 전문가를 사용하는 기존 MoE의 성능을 능가했다. 이는 전문가 파라미터가 모델 깊이에 따라 선형적으로 증가할 필요가 없음을 시사한다.

제로샷(Zero-shot) 벤치마크 평가에서도 ARC, PIQA, HellaSwag 등 주요 지표에서 기존 MoE보다 높은 정확도를 보였다. 특히 전문가 세분화(Granularity)를 높인 실험(32E/top-4)에서도 UniPool의 이점이 그대로 유지됨이 확인되었다.

UniPool의 핵심 차별점은 전문가 파라미터의 소유권(Ownership)을 레이어에서 아키텍처 전체로 전이시킨 것이다. 이는 레이어별로 독립적인 그래디언트 신호를 받던 전문가들이 이제 모든 레이어로부터 오는 신호를 통합하여 학습하게 됨을 의미한다. 수학적으로는 각 레이어 l의 라우팅 확률 P_i^(l)을 전체 레이어 L에 대해 평균낸 f_i를 기반으로 보조 손실을 정의함으로써, 레이어 로컬 최적화가 아닌 글로벌 최적화를 달성한다.

NormRouter의 채택은 공유 풀 구조에서 발생하는 수치적 불안정성을 해결하는 핵심 장치이다. Softmax는 로짓의 절대적 크기에 민감하여 레이어 깊이에 따른 히든 스테이트 스케일 변화에 취약하지만, NormRouter는 방향성(Direction)에 집중하고 스케일은 학습 가능한 파라미터 σ로 흡수한다. 또한 ReLU 활성화를 통해 자연스러운 희소성(Sparsity)을 유도하여 별도의 복잡한 희소화 연산 없이도 효율적인 라우팅을 구현했다.

본 연구는 182M에서 978M 사이의 모델 규모와 30B 토큰 학습 환경에서 검증되었다. 수십억 개 이상의 파라미터를 가진 초대형 모델(Billion-scale)에서의 성능 유지 여부와 더 긴 학습 시간에 따른 수렴 특성에 대해서는 추가적인 검증이 필요하다. 또한 전역 통계량을 수집하는 과정에서 발생하는 미세한 연산 오버헤드와 전문가 병렬화 효율성에 대한 상세 분석이 향후 과제로 남아있다.