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