[論文レビュー] Fine-tuning Image Transformers using Learnable Memory

Mark Sandler, Andrey Zhmoginov+2 more|arXiv (Cornell University)|Mar 29, 2022

Advanced Neural Network Applications被引用数 4

ひとこと要約

この論文では、各層に追加の学習可能な埋め込み（メモリトークン）を挿入することで、ビジョントランスフォーマーモデルに学習可能なメモリを拡張し、下流タスクにおける効率的でパラメータ効率の良い微調整を可能にする手法を提案する。アテンションマスクを用いることで、新しいタスクに拡張しても以前のタスクの性能を維持でき、精度の低下を最小限に抑えられる。この手法により、各層あたり数個の追加パラメータで、完全微調整に近い性能を達成する。

ABSTRACT

In this paper we propose augmenting Vision Transformer models with learnable memory tokens. Our approach allows the model to adapt to new tasks, using few parameters, while optionally preserving its capabilities on previously learned tasks. At each layer we introduce a set of learnable embedding vectors that provide contextual information useful for specific datasets. We call these "memory tokens". We show that augmenting a model with just a handful of such tokens per layer significantly improves accuracy when compared to conventional head-only fine-tuning, and performs only slightly below the significantly more expensive full fine-tuning. We then propose an attention-masking approach that enables extension to new downstream tasks, with a computation reuse. In this setup in addition to being parameters efficient, models can execute both old and new tasks as a part of single inference at a small incremental cost.

研究の動機と目的

• ビジョントランスフォーマーにおける完全微調整の高コストおよび不安定性、特に継続的学習の文脈での課題を解決すること。
• 最小限のパラメータ更新で新しい下流タスクへの効率的適応を可能にすること。
• 新しいタスクにモデルを拡張する際、以前に学習したタスクの高い性能を維持すること。
• 1つのモデル内で複数のタスクの計算再利用と共同推論を可能にすること。
• 視覚トランスフォーマーにおける生涯学習のための、パラメータ効率的でスケーラブルかつ安定した手法を検討すること。

提案手法

• 各ViTエンコーダー層に、アテンション計算の前に入力シーケンスに追加する小さな学習可能なメモリトークンを追加する。
• メモリトークンが他のトークンに注目しないように変更されたアテンションメカニズムを採用し、クラストークンとパッチトークンのみが層間で更新されるようにする。
• 古いタスクのトークンが新しい追加メモリトークンやクラストークンに注目しないように制限するアテンションマスク戦略を導入し、推論時に干渉を防ぐ。
• 新しいタスクを段階的に訓練する際、入力シーケンスに新しいタスク固有のメモリトークンとクラストークンを追加しつつ、アテンション制約を維持することで、古いタスクの性能を保つ。
• アテンションマスクを用いて、独立に訓練されたモデルを連結することでモデル拡張を可能にし、計算コストを段階的に増加させるだけで、すべてのタスクの共同推論を実現する。
• 全タスクに共通のバックボーンを用い、タスク固有のメモリと分類ヘッドを設けることで、パラメータ共有と効率的な推論を実現する。

実験結果

リサーチクエスチョン

• RQ1 学習可能なメモリトークンは、最小限のパラメータ更新でビジョントランスフォーマーの微調整精度を顕著に向上させることができるか？
• RQ2 アテンションマスクにより、新しいタスクにモデルを拡張した際、以前に学習したタスクの性能劣化を防げるか？
• RQ3 メモリ拡張微調整の性能は、ヘッドのみ微調整および完全微調整と比べてどうか？
• RQ4 異なるタスクで訓練された複数のモデルを、最小限の精度損失と段階的な計算コストで1つのモデルに統合できるか？
• RQ5 1層あたりのメモリトークン数はモデル性能に影響を及ぼすか？また、その効果に飽和点は存在するか？

主な発見

• 1層あたり数個のメモリトークンを追加するだけで、ヘッドのみ微調整より顕著に精度が向上し、完全微調整に近い性能にまで到達する。
• アテンションマスクを用いることで、複数のタスクを段階的に訓練したモデルは、すべての過去のタスクで高い精度を維持でき、単純な連結手法で見られる20%～40%の精度低下を回避できる。
• SUN-397からPlaces-365へ、またはその逆に段階的にタスクを追加したモデルは、劣化を伴わず安定した性能を示し、独立に訓練・連結したモデルを上回る性能を達成した。
• メモリトークンを備えたモデルは、完全微調整に必要なパラメータの一部にとどまる状況でも、ベースラインモデルを上回る性能を示した。
• 1層あたり10個を超えると、メモリサイズの増加に伴う性能向上の効果が鈍り、限界に達する傾向が見られた。
• アテンションマスクアプローチにより、計算再利用が効率的に行えるようになり、新しいタスクごとにわずかな追加コストで、すべてのタスクの共同推論が可能になった。