[論文レビュー] Discrete Key-Value Bottleneck

Frederik Träuble, Anirudh Goyal+5 more|arXiv (Cornell University)|Jul 22, 2022

Domain Adaptation and Few-Shot Learning被引用数 5

ひとこと要約

本論文では、事前学習済みエンコーダーを固定し、推論時にのみスパarsな学習可能な値コードのセットを更新することで、継続的学習を可能にする離散的キー・バリュー・ボトルネック機構を提案する。タスク固有の知識を離散的キー・バリュー対によって格納・取得することで、タスク境界なしのクラスインクリメンタル学習ベンチマークにおいて災難的忘却を低減し、さまざまな事前学習モデルで既存手法を上回る性能を達成する。

ABSTRACT

Deep neural networks perform well on classification tasks where data streams are i.i.d. and labeled data is abundant. Challenges emerge with non-stationary training data streams such as continual learning. One powerful approach that has addressed this challenge involves pre-training of large encoders on volumes of readily available data, followed by task-specific tuning. Given a new task, however, updating the weights of these encoders is challenging as a large number of weights needs to be fine-tuned, and as a result, they forget information about the previous tasks. In the present work, we propose a model architecture to address this issue, building upon a discrete bottleneck containing pairs of separate and learnable key-value codes. Our paradigm will be to encode; process the representation via a discrete bottleneck; and decode. Here, the input is fed to the pre-trained encoder, the output of the encoder is used to select the nearest keys, and the corresponding values are fed to the decoder to solve the current task. The model can only fetch and re-use a sparse number of these key-value pairs during inference, enabling localized and context-dependent model updates. We theoretically investigate the ability of the discrete key-value bottleneck to minimize the effect of learning under distribution shifts and show that it reduces the complexity of the hypothesis class. We empirically verify the proposed method under challenging class-incremental learning scenarios and show that the proposed model - without any task boundaries - reduces catastrophic forgetting across a wide variety of pre-trained models, outperforming relevant baselines on this task.

研究の動機と目的

• 学習データの分布が時間とともに変化する継続的学習の状況において、災難的忘却を解消すること。
• 事前学習済みエンコーダーの重みを微調整せずに、効率的で局所的なモデル更新を可能にすること。
• 離散的ボトルネックを用いることで仮説クラスの複雑さを低減し、共変量シフト下での一般化を向上させること。
• 事前学習済みエンコーダーからの知識を再利用しつつ、スパースで文脈依存の値コード更新によって新しいタスクに適応するモデルを設計すること。
• タスク境界のアノテーションなしで、複数の事前学習モデルにおけるクラスインクリメンタル学習において、手法の実証的妥当性を検証すること。

提案手法

• モデルは三段階のパイプライン（エンコード、離散的キー・バリュー・ボトルネックによる処理、デコード）を用いる。
• 事前学習済みエンコーダーが連続表現を出力し、それがC個のヘッドに投影されてキー・バリューの検索が行われる。
• 各ヘッドは学習可能なキー・コードブック内で最も近いキーを検索し、対応する連続的バリュー・コードを取得する。
• トレーニング中は、キーとエンコーダーは固定されたまま、バリュー・コードのみが更新されるため、局所的適応が可能になる。
• キー・バリュー・コードブックは、多様なデータ分布に対するEMAを用いて初期化され、特徴多様体全体を広くカバーするように設計される。
• キー（符号化）とバリュー（復号化）のコードを分離することで、仮説クラスの複雑さを低減し、モデル容量を制限することで一般化性能が向上する。

実験結果

リサーチクエスチョン

• RQ1 タスク境界なしのクラスインクリメンタル学習において、離散的キー・バリュー・ボトルネックが災難的忘却を低減できるか？
• RQ2 離散的ボトルネックアーキテクチャは、入力分布のシフト下での一般化にどのように影響するか？
• RQ3 パrameter更新をバリュー・コードに限定することで、過去のタスクからの知識がどの程度保持されるか？
• RQ4 非i.i.d.なデータストリーム下で、多様な事前学習エンコーダーに対して、モデルはどの程度の性能を示すか？
• RQ5 キー・バリュー対の利用パターンはどのようなものか？未使用容量は、利用不足を示すのか、それとも耐障害性を示すのか？

主な発見

• 提案手法は、クラスインクリメンタル学習において、全テスト済み事前学習モデル（ResNet50, ViT, DINO, SwAV, ConvMixer）で著しく災難的忘却を低減した。
• タスク境界のアノテーションやタスク固有のヘッド再初期化を必要とせず、関連するベースラインを上回る性能を示した。
• トレーニング中に70％以上のキー・バリュー対が少なくとも1回使用され、28％のキーが未使用のままであった。これは、依然として利用可能な容量があり、利用はスパースであることを示している。
• 初期化段階ではキー・コードがデータ多様体全体に広く分布していたが、ターゲットデータセットに適応するに従い、より凝集した分布を示し、ドメイン固有の特徴への適応を反映していた。
• 仮説クラスの複雑さが低減されたため、分布シフト下で標準モデルよりも良好な一般化境界を達成した。
• バリュー・コードの更新は非常に局所的であり、多くのキーが複数の入力から参照されていた。これにより、効率的で文脈依存の適応が可能になった。