研究者たちは、マルチモーダル混合専門家（MoE）モデルにおいて視覚トークンとテキストトークンのバランスを維持する、幾何学的なガイドに基づくトークンルーティング手法であるReBAルーティングを発表しました。初期の実験では、従来のルーターが通常つまずくような、画像解像度を大幅に引き上げたシナリオにおいても、この手法がトレーニングを安定させることが示されました。

なぜビジョン言語モデルに新しいルーターが必要なのか

ビジョン言語モデルは、数百に及ぶこともある画像パッチと、わずかな単語トークンという、大きく異なる2つのストリームを取り込みます。標準的なMoEルーターは、すべてのトークンを同じ種類のデータとして扱うため、画像パッチが少数のエキスパートに集中する一方で、テキストトークンはアイドル状態になってしまいます。既存の研究では、各エキスパートに割り当てられるトークンの総数のみを均等化する補助損失（auxiliary loss）を用いて負荷を調整しようとしています。しかし、膨大な画像負荷が微量なテキスト負荷を相殺してしまうため、パフォーマンスが低下するまで、この損失関数は真の不均衡を隠してしまいます。

ReBAがいかにルーティングのあり方を変えるか

ReBAは、ルーティングプロセスにモダリティ境界を追加します。パッチと単語を混ぜ合わせた単一のプールを使用する代わりに、画像トークンの幾何学的レイアウトを尊重した個別のパスを構築します。各画像インスタンスには同じ全体的な重みが割り当てられ、特定のエキスパートがボトルネックになるのを防ぎます。パッチの空間的配置に導かれることで、入力解像度が変化してもシステムは安定した状態を維持します。

著者が報告している主な利点は以下の通りです：

  • 視覚トークンと言語トークンの明確な分離を強制する。
  • バッチ内の各画像からの均等な寄与を保証する。
  • 特定のモダリティによってエキスパートが圧倒されるのを防ぐ。
  • 画像パッチの数が増加してもバランスを維持する。

いくつかのベンチマーク設定において、ReBAはバッチに追加されるパッチの数に関わらず、エキスパートプールを均等に活用し続け、従来の補助損失アプローチを上回る性能を示しました。

限界と未解決の課題

この研究では、速度やメモリ消費量に関する数値が示されていないため、追加のルーティングロジックが隠れたコストを生じさせているかどうかは不明です。また、著者は単一の画像のすべてのパッチが単一のユニットとして機能すると仮定していますが、この仮定は、異なる解像度の画像が混在するバッチや、一部がマスクされた領域を含むバッチでは崩れる可能性があります。

今後の注目点

  • パフォーマンスと効率性のトレードオフ: 今後の研究では、ReBAが追加するオーバーヘッドを定量化する必要があります。
  • 混合バッチの挙動: 高解像度画像と低解像度画像を組み合わせたバッチでのテストにより、モダリティ境界が維持されるかどうかが明らかになるでしょう。
  • 大規模パイプラインへの採用: ルーティングの安定性が、画像キャプション生成や視覚的質問応答（VQA）などのダウンストリームタスクの向上につながる場合、開発者は標準的な補助損失をReBAに置き換える可能性があります。

ビジョン言語MoEパイプラインを構築している場合、デフォルトのルーターをReBAに置き換えることで、画像解像度を高めても、より均一なエキスパート使用パターンを得られる可能性があります。変更後はトークン分布メトリクスに注目してください。分布が平坦になれば、モダリティ境界が適切に機能していることを示しています。