C3PO: Оптимизация критического слоя, ключевых экспертов и совместных путей для повторного смешивания экспертов во время тестирования

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

April 10, 2025

Авторы: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

Аннотация

Модели больших языковых моделей (LLM) с архитектурой Mixture-of-Experts (MoE) страдают от существенно неоптимальных путей использования экспертов — наше исследование показывает, что наивный выбор экспертов, усвоенный в процессе предварительного обучения, оставляет удивительный разрыв в точности на 10–20%, который можно улучшить. Вдохновленные этим наблюдением, мы разработали новый класс методов оптимизации на этапе тестирования, которые пересчитывают или "пересмешивают" веса экспертов в различных слоях совместно для каждого тестового примера. Поскольку истинное значение для тестового примера неизвестно, мы предлагаем оптимизировать суррогатную цель, определенную на основе "успешных соседей" примера из эталонного набора данных. Мы вводим три суррогата и алгоритма, основанные на поиске моды, ядерной регрессии и среднем значении потерь для похожих эталонных примеров/задач. Чтобы снизить затраты на оптимизацию всего пути, мы применяем наши алгоритмы только к весам смешения ключевых экспертов в критических слоях, что обеспечивает схожую производительность, но значительно экономит вычисления. Это приводит к "Оптимизации Совместного Пути Ключевых Экспертов в Критических Слоях (C3PO)". Мы применяем C3PO к двум недавним MoE LLM и тестируем на шести широко используемых бенчмарках. Метод стабильно улучшает базовую модель на 7–15% по точности и значительно превосходит популярные методы обучения на этапе тестирования, такие как обучение в контексте и настройка промтов/префиксов. Более того, C3PO позволяет MoE LLM с 1–3 млрд активных параметров превосходить LLM с 7–9 млрд параметров, тем самым усиливая преимущества MoE в эффективности. Наше детальное исследование также дает новые инсайты о достижении улучшений на этапе тестирования для 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.