C3PO: 테스트 타임 전문가 재조합을 위한 핵심 계층, 코어 전문가, 협력적 경로 최적화

C3PO: Critical-Layer, Core-Expert, Collaborative Pathway Optimization for Test-Time Expert Re-Mixing

April 10, 2025

저자: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

초록

전문가 혼합(Mixture-of-Experts, MoE) 대형 언어 모델(Large Language Models, LLMs)은 심각하게 최적화되지 않은 전문가 경로 문제를 겪고 있습니다. 우리의 연구는 사전 학습을 통해 학습된 단순한 전문가 선택이 놀랍게도 10-20%의 정확도 향상 여지를 남겨둔다는 것을 밝혀냈습니다. 이러한 관찰에 동기를 받아, 우리는 각 테스트 샘플에 대해 서로 다른 계층의 전문가들을 재가중하거나 "재혼합"하는 새로운 테스트 시점 최적화 방법론을 개발했습니다. 테스트 샘플의 실제 정답을 알 수 없기 때문에, 우리는 참조 샘플 집합에서 해당 샘플의 "성공적인 이웃"을 기반으로 한 대리 목적 함수를 최적화하는 방식을 제안합니다. 우리는 모드 탐색, 커널 회귀, 그리고 유사한 참조 샘플/태스크의 평균 손실을 기반으로 한 세 가지 대리 목적 함수와 알고리즘을 소개합니다. 전체 경로를 최적화하는 비용을 줄이기 위해, 우리는 알고리즘을 핵심 전문가들의 혼합 가중치에만 적용하며, 이는 비슷한 성능을 유지하면서도 상당한 계산 비용을 절약합니다. 이를 통해 "중요 계층, 핵심 전문가, 협력적 경로 최적화(Critical-Layer, Core-Expert, Collaborative Pathway Optimization, C3PO)"를 도출했습니다. 우리는 C3PO를 최근의 두 MoE LLM에 적용하고, 널리 사용되는 여섯 가지 벤치마크에서 이를 검증했습니다. C3PO는 기본 모델의 정확도를 7-15% 향상시키며, 테스트 시점 학습의 대표적인 기법들(예: 컨텍스트 내 학습, 프롬프트/프리픽스 튜닝)을 큰 차이로 능가했습니다. 더 나아가, C3PO는 1-3B 활성 파라미터를 가진 MoE LLM이 7-9B 파라미터의 LLM을 능가하도록 하여, MoE의 효율성 장점을 더욱 강화했습니다. 우리의 철저한 제거 연구는 MoE에서 테스트 시점 개선을 달성하는 데 대한 새로운 통찰을 제공합니다.

English

Mixture-of-Experts (MoE) Large Language Models (LLMs) suffer from severely sub-optimal expert pathways-our study reveals that naive expert selection learned from pretraining leaves a surprising 10-20% accuracy gap for improvement. Motivated by this observation, we develop a novel class of test-time optimization methods to re-weight or "re-mixing" the experts in different layers jointly for each test sample. Since the test sample's ground truth is unknown, we propose to optimize a surrogate objective defined by the sample's "successful neighbors" from a reference set of samples. We introduce three surrogates and algorithms based on mode-finding, kernel regression, and the average loss of similar reference samples/tasks. To reduce the cost of optimizing whole pathways, we apply our algorithms merely to the core experts' mixing weights in critical layers, which enjoy similar performance but save significant computation. This leads to "Critical-Layer, Core-Expert, Collaborative Pathway Optimization (C3PO)". We apply C3PO to two recent MoE LLMs and examine it on six widely-used benchmarks. It consistently improves the base model by 7-15% in accuracy and outperforms widely used test-time learning baselines, e.g., in-context learning and prompt/prefix tuning, by a large margin. Moreover, C3PO enables MoE LLMs with 1-3B active parameters to outperform LLMs of 7-9B parameters, hence improving MoE's advantages on efficiency. Our thorough ablation study further sheds novel insights on achieving test-time improvement on MoE.