C3PO: Ottimizzazione Collaborativa del Percorso con Strato Critico e Nucleo Esperto per il Rimixaggio degli Esperti al Momento del Test

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

April 10, 2025

Autori: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) basati su Mixture-of-Experts (MoE) soffrono di percorsi di esperti fortemente sub-ottimali: il nostro studio rivela che la selezione ingenua degli esperti appresa durante il pre-training lascia un sorprendente divario di accuratezza del 10-20% da colmare. Motivati da questa osservazione, sviluppiamo una nuova classe di metodi di ottimizzazione al momento del test per ri-pesare o "ri-miscelare" gli esperti in diversi strati in modo congiunto per ogni campione di test. Poiché la verità di base del campione di test è sconosciuta, proponiamo di ottimizzare un obiettivo surrogato definito dai "vicini di successo" del campione, tratti da un insieme di riferimento di campioni. Introduciamo tre surrogati e algoritmi basati su ricerca della moda, regressione kernel e la perdita media di campioni/compiti di riferimento simili. Per ridurre il costo dell'ottimizzazione di interi percorsi, applichiamo i nostri algoritmi solo ai pesi di miscelazione degli esperti principali negli strati critici, ottenendo prestazioni simili ma risparmiando un calcolo significativo. Questo porta a "Ottimizzazione Collaborativa del Percorso con Esperti Principali negli Strati Critici (C3PO)". Applichiamo C3PO a due recenti LLM MoE e lo testiamo su sei benchmark ampiamente utilizzati. Migliora costantemente il modello base del 7-15% in accuratezza e supera di gran lunga i metodi di apprendimento al momento del test comunemente usati, come l'apprendimento in-context e il tuning di prompt/prefissi. Inoltre, C3PO consente a LLM MoE con 1-3 miliardi di parametri attivi di superare LLM con 7-9 miliardi di parametri, migliorando così i vantaggi di MoE in termini di efficienza. Il nostro studio di ablazione approfondito fornisce ulteriori nuove intuizioni su come ottenere miglioramenti al momento del test su 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.