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

Почему мультимодальным моделям нужен новый маршрутизатор

Мультимодальные модели (vision-language models) обрабатывают два очень разных потока данных: сотни патчей изображений и лишь несколько текстовых токенов. Стандартные маршрутизаторы MoE относятся к каждому токену как к однотипным данным, из-за чего патчи изображений перегружают несколько экспертов, в то время как текстовые токены простаивают. Существующие подходы пытаются сбалансировать нагрузку с помощью вспомогательной функции потерь (auxiliary loss), которая лишь уравнивает общее количество токенов, назначенных каждому эксперту. Поскольку огромный объем данных от изображений может нивелировать малый объем текста, функция потерь маскирует реальный дисбаланс до тех пор, пока не начнет падать производительность.

Как ReBA меняет правила игры в маршрутизации

ReBA добавляет границу модальностей (modality boundary) в процесс маршрутизации. Вместо единого пула, смешивающего патчи и слова, она выстраивает отдельные пути, учитывающие геометрическое расположение токенов изображения. Каждому экземпляру изображения присваивается одинаковый общий вес, что предотвращает превращение отдельного эксперта в «узкое место». Благодаря учету пространственного расположения патчей система остается стабильной даже при изменении разрешения входных данных.

Основные преимущества, отмеченные авторами:

  • Обеспечивает четкое разделение между визуальными и лингвистическими токенами.
  • Гарантирует равный вклад каждого изображения в батче.
  • Предотвращает перегрузку экспертов одной модальностью.
  • Поддерживает баланс при росте количества патчей изображения.

В различных бенчмарках ReBA поддерживала равномерную загрузку пула экспертов независимо от количества патчей, добавленных в батч, превзойдя традиционный подход с использованием вспомогательной функции потерь.

Ограничения и открытые вопросы

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

За чем стоит следить в будущем

  • Компромисс между производительностью и эффективностью: будущим работам потребуется количественно оценить любые накладные расходы, которые вносит ReBA.
  • Поведение в смешанных батчах: тестирование на батчах, сочетающих изображения высокого и низкого разрешения, покажет, сохраняется ли граница модальностей.
  • Внедрение в крупномасштабные конвейеры: если стабильность маршрутизации приведет к улучшению результатов в прикладных задачах — таких как описание изображений (image captioning) или визуальный вопрос-ответ (VQA) — разработчики могут заменить стандартную вспомогательную функцию потерь на ReBA.

Если вы строите конвейер vision-language MoE, замена стандартного маршрутизатора на ReBA может обеспечить более равномерное использование экспертов при повышении разрешения изображений. Следите за метриками распределения токенов после изменений: более «плоское» распределение будет сигнализировать о том, что граница модальностей работает эффективно.