Выбор каналов в MoE-моделях обходится без калибровочных данных

Выбор каналов в MoE-моделях обходится без калибровочных данных

В Mixture-of-Experts (MoE) моделях при интерактивной генерации с малыми батчами узкое место, загрузка весов маршрутизируемых экспертов из памяти. Большинство промежуточных каналов эксперта почти не влияет на отдельный токен, поэтому селекторы каналов предсказывают, какие каналы важны, и читают только их веса. Обычно для этого оставляют каналы выше порога, который калибруют на офлайн-данных.

Авторы статьи (в виде аннотации на Hugging Face Papers) утверждают, что калибровочные данные для выбора каналов не нужны. Первый довод: откалиброванный порог не фиксирует, сколько именно читает селектор, потому что число каналов, прошедших порог, меняется от входа к входу. Пример: селектор Prox с порогами, откалиброванными под бюджет чтения 18%, в среднем по бенчмаркам авторов реально читает только 12,9% весов маршрутизируемых экспертов.

Второй довод: top-k оставляет для каждого токена k каналов с наивысшей оценкой и калибровочных данных не требует. При откалиброванном бюджете top-k-версия селектора не хуже пороговой (совпадает или лучше по той же метрике) в 40 из 43 сравнений при бюджетах 13% и 18%, включая все 23 сравнения на задачах генерации. Кроме того, top-k-версии избегают зацикливания генерации (коллапса с повторами), которое вызывают откалиброванные пороги на Gemma 4. Вывод авторов: калибровка порога не даёт ни того, ради чего её делают, заранее известного бюджета чтения, ни стабильного выигрыша в точности при этом бюджете.

Кроме того, авторы предлагают собственный считыватель, целиком построенный на весах модели: он отсеивает каналы с помощью тернарного эскиза (ternary sketch), а пропущенные каналы восстанавливает низкоранговым отображением (low-rank map). На трёх дообученных (post-trained) MoE-моделях он сохраняет 98,1% точности полной модели при 17,8% чтений весов маршрутизируемых экспертов. Реализованные на Triton ядра выполняют блоки MoE в 1,60-1,90 раза быстрее оптимизированного vLLM при батче 1 и до 2,69 раза быстрее при больших батчах, на тех же трёх моделях.

Ключевые факты

  • Откалиброванный порог не фиксирует бюджет чтения: Prox, настроенный на 18%, в среднем по бенчмаркам авторов читает лишь 12,9% весов экспертов.
  • Top-k-версия селектора не хуже пороговой в 40 из 43 сравнений при бюджетах 13% и 18%, включая все 23 на задачах генерации; top-k также избегает зацикливания генерации, которое пороги вызывают на Gemma 4.
  • Предложенный считыватель, построенный только на весах модели (тернарный эскиз и низкоранговое отображение), сохраняет 98,1% точности полной модели при 17,8% чтений весов экспертов на трёх дообученных MoE-моделях.
  • Ядра на Triton ускоряют блоки MoE в 1,60-1,90 раза относительно оптимизированного vLLM при батче 1 и до 2,69 раза при больших батчах.

Почему это важно

При интерактивной генерации с малыми батчами MoE-модели упираются в чтение весов экспертов из памяти, и ускорять приходится именно его. Авторы утверждают, что привычный шаг, калибровка порога на офлайн-данных, не даёт ни предсказуемого бюджета чтения, ни стабильного выигрыша в точности. Если это верно, из конвейера можно убрать целый этап подготовки и заодно получить заранее известный объём чтения.

Кому это важно

Инженерам по инференсу MoE-моделей, разработчикам ядер и движков вроде vLLM, исследователям разреженного и приближённого вычисления в экспертах. Тем, кто запускает MoE локально или с малыми батчами, где главная задача, снизить объём читаемых из памяти весов.

Как это применить

Идея проста для проверки: вместо порога, калиброванного на данных, оставлять для каждого токена k лучших по оценке каналов и сравнить точность и скорость на своих задачах. Предложенный считыватель использует только веса модели, то есть не требует калибровочного набора. Выложен ли код или веса, в аннотации не указано, поэтому воспроизведение придётся строить по описанию подхода.

Можно ли доверять

Это аннотация статьи на Hugging Face Papers, полный текст в пересказе не использован. Все цифры, заявления самих авторов. В аннотации не названы авторы и организации, не перечислены три использованные MoE-модели и 43 сравнения, не указано, на каком GPU измерены ускорения Triton-ядер, и не раскрыто, что такое Prox и кто его создал. Три сравнения из 43, где top-k не оказался не хуже порога, не описаны. Независимого подтверждения в источнике нет.

Риски и подводные камни

Результаты относятся к конкретным бенчмаркам и трём моделям, которые в аннотации не названы, поэтому перенос на другие MoE-архитектуры не гарантирован. Скорость измерена относительно оптимизированного vLLM при батче 1 и «больших батчах» без указания их размера. Точность 98,1%, это доля от точности полной модели, а не абсолютная метрика, а 17,8%, доля чтений весов экспертов, а не сокращение времени. В 3 из 43 сравнений top-k не оказался не хуже порога.