C3PO: Optimización Colaborativa de Rutas Críticas con Expertos Centrales para la Remezcla de Expertos en Tiempo de Prueba

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

April 10, 2025

Autores: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs) basados en Mezcla de Expertos (MoE) presentan rutas de expertos severamente subóptimas: nuestro estudio revela que la selección ingenua de expertos aprendida durante el preentrenamiento deja un sorprendente margen de mejora del 10-20% en precisión. Motivados por esta observación, desarrollamos una nueva clase de métodos de optimización en tiempo de prueba para re-ponderar o "re-mezclar" a los expertos en diferentes capas de manera conjunta para cada muestra de prueba. Dado que la verdad fundamental de la muestra de prueba es desconocida, proponemos optimizar un objetivo sustituto definido por los "vecinos exitosos" de la muestra dentro de un conjunto de referencia. Introducimos tres sustitutos y algoritmos basados en la búsqueda de modas, regresión de kernel y la pérdida promedio de muestras/tareas de referencia similares. Para reducir el costo de optimizar rutas completas, aplicamos nuestros algoritmos únicamente a los pesos de mezcla de los expertos principales en capas críticas, lo que ofrece un rendimiento similar pero ahorra un cálculo significativo. Esto da lugar a la "Optimización Colaborativa de Rutas de Expertos Principales en Capas Críticas (C3PO)". Aplicamos C3PO a dos LLMs MoE recientes y lo evaluamos en seis benchmarks ampliamente utilizados. Consistemente mejora el modelo base en un 7-15% en precisión y supera a los métodos de aprendizaje en tiempo de prueba comúnmente utilizados, como el aprendizaje en contexto y la sintonización de prefijos/prompt, por un amplio margen. Además, C3PO permite que los LLMs MoE con 1-3B parámetros activos superen a LLMs de 7-9B parámetros, mejorando así las ventajas de eficiencia de MoE. Nuestro exhaustivo estudio de ablación ofrece además nuevas perspectivas sobre cómo lograr mejoras en tiempo de prueba en 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.