Экспертное пороговое маршрутизирование для авторегрессионного языкового моделирования с динамическим распределением вычислений и балансировкой нагрузки
Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing
March 12, 2026
Авторы: Hanchi Sun, Yixin Liu, Yonghui Wu, Lichao Sun
cs.AI
Аннотация
Токен-выборочная смесь экспертов (TC-MoE) направляет каждый токен к фиксированному числу экспертов, что ограничивает динамическое распределение вычислений и требует вспомогательных функций потерь для поддержания баланса нагрузки. Мы предлагаем маршрутизацию по порогу эксперта (ET), при которой каждый эксперт поддерживает экспоненциально взвешенное скользящее среднее (EMA) пороговое значение, оцениваемое на основе глобального распределения токенов. Как на этапе обучения, так и на этапе вывода каждый токен независимо направляется к эксперту, если его оценка превышает порог этого эксперта, что позволяет динамически распределять вычисления и достигать баланса нагрузки без вспомогательных потерь. Данный полностью причинный механизм устраняет зависимость от других токенов в пакете, что делает его хорошо подходящим для авторегрессионного языкового моделирования. В экспериментах по предварительному обучению с масштабированием до 2.4 млрд параметров на FineWeb-Edu подход ET демонстрирует перекрестную энтропию на 0.067 ниже, чем TC-MoE, что эквивалентно достижению той же производительности при использовании в 1.6 раза меньшего количества токенов.
English
Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6times fewer tokens.