C3PO : Optimisation collaborative des voies critiques et des experts centraux pour le réassemblage des experts au moment du test

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

Résumé

Les modèles de langage de grande taille (LLM) à mélange d'experts (MoE) souffrent de chemins d'experts fortement sous-optimaux - notre étude révèle que la sélection naïve des experts apprise lors du pré-entraînement laisse un écart de précision surprenant de 10 à 20 % à combler. Motivés par cette observation, nous développons une nouvelle classe de méthodes d'optimisation au moment du test pour ré-équilibrer ou "re-mélanger" les experts dans différentes couches de manière conjointe pour chaque échantillon de test. Comme la vérité terrain de l'échantillon de test est inconnue, nous proposons d'optimiser un objectif de substitution défini par les "voisins réussis" de l'échantillon provenant d'un ensemble de référence. Nous introduisons trois substituts et algorithmes basés sur la recherche de mode, la régression par noyau et la perte moyenne des échantillons/tâches de référence similaires. Pour réduire le coût de l'optimisation des chemins entiers, nous appliquons nos algorithmes uniquement aux poids de mélange des experts centraux dans les couches critiques, ce qui offre des performances similaires mais permet d'économiser des calculs significatifs. Cela conduit à "l'Optimisation Collaborative des Chemins par Experts Centraux dans les Couches Critiques (C3PO)". Nous appliquons C3PO à deux récents LLM MoE et l'évaluons sur six benchmarks largement utilisés. Il améliore systématiquement le modèle de base de 7 à 15 % en précision et surpasse les méthodes d'apprentissage au moment du test couramment utilisées, comme l'apprentissage en contexte et le réglage par prompt/préfixe, avec une large marge. De plus, C3PO permet aux LLM MoE avec 1 à 3 milliards de paramètres actifs de surpasser les LLM de 7 à 9 milliards de paramètres, améliorant ainsi les avantages de MoE en termes d'efficacité. Notre étude d'ablation approfondie apporte également de nouvelles perspectives sur l'amélioration au moment du test pour 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.