C3PO: Otimização Colaborativa de Caminho com Camada Crítica e Especialista Central para Recombinação de Especialistas em Tempo de Teste

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

Resumo

Modelos de Linguagem de Grande Escala (LLMs) baseados em Mistura de Especialistas (MoE) sofrem com caminhos de especialistas severamente subótimos — nosso estudo revela que a seleção ingênua de especialistas aprendida durante o pré-treinamento deixa uma lacuna surpreendente de 10-20% de precisão para melhoria. Motivados por essa observação, desenvolvemos uma nova classe de métodos de otimização em tempo de teste para re-ponderar ou "re-misturar" os especialistas em diferentes camadas de forma conjunta para cada amostra de teste. Como a verdade fundamental da amostra de teste é desconhecida, propomos otimizar um objetivo substituto definido pelos "vizinhos bem-sucedidos" da amostra a partir de um conjunto de referência de amostras. Introduzimos três substitutos e algoritmos baseados em busca de moda, regressão de kernel e a perda média de amostras/tarefas de referência semelhantes. Para reduzir o custo de otimizar caminhos inteiros, aplicamos nossos algoritmos apenas aos pesos de mistura dos especialistas principais em camadas críticas, que apresentam desempenho semelhante, mas economizam computação significativa. Isso resulta na "Otimização Colaborativa de Caminhos de Especialistas Principais em Camadas Críticas (C3PO)". Aplicamos o C3PO a dois LLMs MoE recentes e o avaliamos em seis benchmarks amplamente utilizados. Ele melhora consistentemente o modelo base em 7-15% em precisão e supera baselines amplamente utilizados de aprendizado em tempo de teste, como aprendizado em contexto e ajuste de prompt/prefixo, por uma grande margem. Além disso, o C3PO permite que LLMs MoE com 1-3B parâmetros ativos superem LLMs de 7-9B parâmetros, melhorando assim as vantagens de eficiência do MoE. Nosso estudo de ablação detalhado ainda traz novos insights sobre como alcançar melhorias em tempo de teste em 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.