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

Чому моделям зі сприйняттям зору та мови потрібен новий маршрутизатор

Моделі зі сприйняттям зору та мови обробляють два дуже різні потоки: сотні патчів зображень та лише кілька текстових токенів. Стандартні маршрутизатори MoE ставляться до кожного токена так, ніби це дані одного типу, тому патчі зображень перевантажують кількох експертів, тоді як текстові токени залишаються невикористаними. Існуючі підходи намагаються виправити навантаження за допомогою допоміжної функції втрат (auxiliary loss), яка лише вирівнює загальну кількість токенів, призначених кожному експерту. Оскільки величезне навантаження від зображень може нівелювати мізерне навантаження від тексту, функція втрат маскує реальний дисбаланс до моменту, поки продуктивність не почне падати.

Як ReBA змінює правила гри в маршрутизації

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

Основні переваги, про які повідомляють автори:

  • Забезпечує чіткий розподіл між візуальними та лінгвістичними токенами.
  • Гарантує рівний внесок кожного зображення в пакеті (batch).
  • Запобігає перевантаженню експертів однією модальністю.
  • Підтримує баланс при збільшенні кількості патчів зображення.

У кількох тестових сценаріях ReBA забезпечувала рівномірне використання пулу експертів незалежно від кількості патчів, доданих до пакета, перевершивши традиційний підхід із використанням допоміжної функції втрат.

Обмеження та відкриті питання

Дослідження не надає даних щодо швидкості або споживання пам'яті, тому ми не знаємо, чи додає додаткова логіка маршрутизації приховані витрати. Автори також припускають, що всі патчі одного зображення діють як єдине ціле; це припущення може бути хибним у пакетах, що змішують зображення різної роздільної здатності або містять частково замасковані області.

На що звернути увагу далі

  • Компроміс між продуктивністю та ефективністю: майбутні роботи мають кількісно оцінити будь-які накладні витрати, які додає ReBA.
  • Поведінка у змішаних пакетах: тестування на пакетах, що поєднують зображення високої та низької роздільної здатності, покаже, чи зберігається межа модальності.
  • Впровадження у масштабні конвеєри: якщо стабільність маршрутизації призведе до покращення результатів у подальших завданнях — таких як опису зображень або візуальних відповідей на запитання — розробники можуть замінити стандартну допоміжну функцію втрат на ReBA.

Якщо ви розробляєте конвеєр MoE для сприйняття зору та мови, заміна стандартного маршрутизатора на ReBA може забезпечити більш рівномірне використання експертів при підвищенні роздільної здатності зображень. Після внесення змін стежте за метриками розподілу токенів: більш плаский розподіл свідчитиме про те, що межа модальності виконує свою роботу.