C3PO: Kritische-Laag, Kern-Expert, Collaboratief Padoptimalisatie voor Testtijdige Expert Her-Mixen

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

April 10, 2025

Auteurs: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

Samenvatting

Mixture-of-Experts (MoE) Large Language Models (LLMs) lijden onder sterk suboptimale expertpaden—onze studie toont aan dat de naïeve expertselectie die tijdens het vooraf trainen wordt geleerd, een verrassende nauwkeurigheidskloof van 10-20% achterlaat voor verbetering. Gemotiveerd door deze observatie ontwikkelen we een nieuwe klasse van optimalisatiemethoden tijdens de testfase om de experts in verschillende lagen gezamenlijk te herwegen of "her-mixen" voor elk testsample. Omdat de grondwaarheid van het testsample onbekend is, stellen we voor om een surrogaatdoel te optimaliseren dat wordt gedefinieerd door de "succesvolle buren" van het sample uit een referentieset van samples. We introduceren drie surrogaten en algoritmen gebaseerd op modusvinden, kernelregressie en het gemiddelde verlies van vergelijkbare referentiesamples/taken. Om de kosten van het optimaliseren van hele paden te verminderen, passen we onze algoritmen alleen toe op de menggewichten van de kernexperts in kritieke lagen, wat vergelijkbare prestaties oplevert maar aanzienlijke rekenkracht bespaart. Dit leidt tot "Critical-Layer, Core-Expert, Collaborative Pathway Optimization (C3PO)". We passen C3PO toe op twee recente MoE LLMs en onderzoeken het op zes veelgebruikte benchmarks. Het verbetert het basismodel consistent met 7-15% in nauwkeurigheid en overtreft veelgebruikte leerbaselines tijdens de testfase, zoals in-context leren en prompt/prefix tuning, met een grote marge. Bovendien stelt C3PO MoE LLMs met 1-3B actieve parameters in staat om LLMs van 7-9B parameters te overtreffen, waardoor de efficiëntievoordelen van MoE worden verbeterd. Onze grondige ablatiestudie werpt verder nieuw licht op het bereiken van verbetering tijdens de testfase op 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.