[論文レビュー] Module-wise Adaptive Distillation for Multimodality Foundation Models

Liang Chen, Jiahui Yu+6 more|arXiv (Cornell University)|Oct 6, 2023

Multimodal Machine Learning Applications被引用数 4

ひとこと要約

本稿では、マルチモーダル基礎モデルにおけるモジュール単位の適応的蒸留手法OPTIMAを提案する。この手法は、各訓練ステップで最も寄与度の高いモジュール（例：画像、テキスト、マルチモーダルエンコーダー）を動的に選択するために、マルチアームバンディットフレームワークを用いる。各モジュールの損失低下量を追跡し、修正版トムソンサンプリングアルゴリズムを適用することで、均一または固定モジュール蒸留と比較して、より少ない蒸留ステップで優れた性能を達成する。画像キャプション生成および視覚的推論ベンチマークにおいて最先端の結果を示している。

ABSTRACT

Pre-trained multimodal foundation models have demonstrated remarkable generalizability but pose challenges for deployment due to their large sizes. One effective approach to reducing their sizes is layerwise distillation, wherein small student models are trained to match the hidden representations of large teacher models at each layer. Motivated by our observation that certain architecture components, referred to as modules, contribute more significantly to the student's performance than others, we propose to track the contributions of individual modules by recording the loss decrement after distillation each module and choose the module with a greater contribution to distill more frequently. Such an approach can be naturally formulated as a multi-armed bandit (MAB) problem, where modules and loss decrements are considered as arms and rewards, respectively. We then develop a modified-Thompson sampling algorithm named OPTIMA to address the nonstationarity of module contributions resulting from model updating. Specifically, we leverage the observed contributions in recent history to estimate the changing contribution of each module and select modules based on these estimations to maximize the cumulative contribution. We evaluate the effectiveness of OPTIMA through distillation experiments on various multimodal understanding and image captioning tasks, using the CoCa-Large model (Yu et al., 2022) as the teacher model.

研究の動機と目的

• 大規模なマルチモーダル基礎モデルにおける非効率な蒸留の課題に対処する。特に、均一なレイヤーワイズ蒸留では、最も情報量の多いモジュールを優先できない点を改善すること。
• ターゲットタスクの学生モデル性能に最も寄与するモジュールを同定し、優先的に蒸留すること。
• 固定された訓練予算のもとで、モジュール寄与度の探索と活用のバランスを取る動的で適応的な蒸留戦略を開発すること。
• 下流タスクにおける損失低下推定によるモジュールレベルの寄与度追跡が、固定または等重み蒸留戦略よりも効果的であることを実証すること。

提案手法

• モジュール蒸留をマルチアームバンディット（MAB）問題としてモデル化し、各モジュールをアームとし、固定ステップ数の蒸留後に得られる報酬を損失低下量とする。
• 訓練中に変化するモジュール寄与度に対応するため、修正版トムソンサンプリングアルゴリズムであるOPTIMAを導入する。
• OPTIMAは、最近の報酬履歴を用いて各モジュールの変化する寄与度を推定し、各訓練ラウンドで最も有益なモジュールを適応的に選択する。
• 各モジュールの寄与度は、固定ステップ数（1ラウンドあたりPステップ）の蒸留後に得られる損失低下の相対的減少量を測定することで評価する。
• 各モジュールの期待報酬（損失低下）に関する後方分布を維持し、これらの分布からのサンプリング値に基づいてアームを選択することで、探索と活用のバランスを取る。
• 本手法はCoCa-Largeに適用され、1ラウンドあたり一部のモジュールのみを蒸留することで、計算コストを削減しながら性能を維持または向上させる。

実験結果

リサーチクエスチョン

• RQ1 均一または固定モジュール蒸留と比較して、蒸留における動的で寄与度に基づくモジュール選択は、学生モデルの性能向上に寄与するか？
• RQ2 訓練中にモジュールの蒸留損失への寄与度はどのように変化するのか？ また、その変化を効果的にモデル化できるか？
• RQ3 トムソンサンプリングのような適応的サンプリング戦略を用いたマルチアームバンディットフレームワークは、モジュール単位の蒸留において探索と活用のバランスを効果的に取れるか？
• RQ4 各モジュールの損失低下量を追跡することは、下流タスクにおける最も有益なモジュールを特定する信頼できるシグナルを提供するか？

主な発見

• すべてのモジュールを均等に蒸留するのではなく、画像エンコーダーのみ、またはマルチモーダルデコーダーのみを蒸留する方が顕著に高い性能を示す。これは、特定のモジュールが下流タスクの精度により大きな寄与をしていることを示している。
• OPTIMAにおけるモジュールアームの報酬分布は時間とともに変化し、一部のモジュール（例：COCOにおけるテキストエンコーダー）は寄与度が増加する傾向を示す一方、他のモジュール（例：画像エンコーダー）は減少傾向を示す。
• モジュールアームの報酬分布の平均値と、そのモジュールを単独で蒸留した場合の最終テスト精度との間に強い正の相関が確認され、寄与度追跡機構の信頼性が裏付けられた。
• NLVR2およびMicrosoft COCO Captionタスクにおいて、OPTIMAはベースライン蒸留手法を上回る優れた性能を達成しており、「Img+Txt+Multi」が最も頻繁に選択されており、これが一貫して最高の累積的利益を提供していることが示された。
• 本手法は高い効率性を維持しており、標準的なレイヤーワイズ蒸留に加えて追加の計算コストやメモリオーバーヘッドを発生させない。1ラウンドあたり選択されたモジュールのみが更新されるため。