Hierarchical Copula-Gumbel-Top-K: управление зависимостями маршрутизации в MoE
Исследователи предложили новый метод маршрутизации для моделей Mixture-of-Experts (MoE), позволяющий контролировать совместное распределение выбора экспертов разными токенами при сохранении фиксированных per-token законов. Метод основан на иерархической copula-модели и расширяет Gumbel-Top-K router.
Модели Mixture-of-Experts (MoE) стали ключевым инструментом для масштабирования нейронных сетей: вместо одного большого слоя используется множество «экспертов», а специальный маршрутизатор решает, к каким экспертам отправить каждый токен. Один из популярных подходов — стохастический Gumbel-Top-K router, который задаёт для каждого токена вероятностное распределение над списками экспертов и весами смеси — так называемый per-token routing law. Однако до сих пор оставался открытым вопрос: какие совместные распределения (joint distributions) между маршрутизационными решениями разных токенов достижимы, если каждый токен подчиняется своему фиксированному закону? Новая работа на arXiv предлагает элегантное решение на основе иерархической copula-модели.
Что произошло
В статье «Hierarchical Copula-Gumbel-Top-K Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws» (arXiv:2607.28670) авторы вводят метод, который позволяет явно контролировать двустороннюю зависимость (two-sided dependence) между выборами экспертов разными токенами, не нарушая при этом per-token законы маршрутизации. Ключевая идея — использовать иерархическую копулу (hierarchical copula) для моделирования корреляций между Gumbel-шумами, которые лежат в основе Gumbel-Top-K. Это даёт возможность «заморозить» сам маршрутизатор (frozen MoE), то есть оставить неизменными индивидуальные вероятности выбора эксперта для каждого токена, но при этом гибко настраивать, насколько схожими или разными будут эксперты, выбираемые разными токенами в одном батче.
Почему это важно
Способность контролировать совместное распределение маршрутизации открывает новые возможности для балансировки нагрузки между экспертами. В стандартных MoE часто возникает дисбаланс: многие токены выбирают одних и тех же экспертов, что приводит к неэффективному использованию ресурсов и перегрузке отдельных модулей. Предложенный метод позволяет, например, ввести отрицательную корреляцию между токенами, чтобы они равномернее распределялись по экспертам, или наоборот — положительную корреляцию для задач, где соседние токены должны обрабатываться одинаково. При этом per-token законы остаются неизменными, то есть не требуется переобучать или донастраивать сам маршрутизатор — достаточно изменить параметры копулы.
«Мы спрашиваем, какие совместные распределения над маршрутизационными выборами разных токенов достижимы, в то время как каждый токен подчиняется фиксированному per-token закону. Иерархическая copula-Gumbel-Top-K даёт полный контроль над этой зависимостью, сохраняя вычислительную простоту исходного Gumbel-Top-K.» — из аннотации работы.
Что дальше
Предложенный подход может быть интегрирован в существующие MoE-архитектуры (например, в Mixtral или Switch Transformer) без изменения процесса обучения — достаточно модифицировать только этап инференса. В перспективе это позволит динамически адаптировать степень зависимости между токенами под конкретную задачу или домен. Кроме того, авторы отмечают, что иерархическая копула допускает эффективную реализацию на GPU, что делает метод практичным для крупномасштабных моделей. Дальнейшие исследования могут быть направлены на изучение влияния различных копульных структур на качество предсказаний и на автоматический подбор параметров зависимости в процессе обучения.
Материал подготовлен с помощью ИИ на основе открытых источников и может содержать неточности. Нашли ошибку — напишите на hello@tolk.media.