[論文レビュー] Test-Time Training with Masked Autoencoders

Yossi Gandelsman, Yu Sun+2 more|arXiv (Cornell University)|Sep 15, 2022

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

ひとこと要約

本論文は test-time training に MAE (Masked Autoencoders) を用い、自己监督再構成によって各テスト入力へモデルを適応させ、視覚ベンチマークにおける分布シフトへの頑健性を向上させる。ImageNet-C など他データセットで経験的な利得を示し、アプローチの線形モデルにおけるバイアス・バリアンス分析を提供する。

ABSTRACT

Test-time training adapts to a new test distribution on the fly by optimizing a model for each test input using self-supervision. In this paper, we use masked autoencoders for this one-sample learning problem. Empirically, our simple method improves generalization on many visual benchmarks for distribution shifts. Theoretically, we characterize this improvement in terms of the bias-variance trade-off.

研究の動機と目的

• unseen distribution shifts に対する頑健な汎化を動機づけ、テスト時に自己監視でモデルを適応させることを提案。
• 自己監視タスクとして masked autoencoding を用い、サンプルごとの適応に有用な信号を生成。
• ImageNet-C, ImageNet-A, ImageNet-R, Portraits など多様な分布シフトベンチマークで TTT-MAE を評価し、その理論的特性を分析。
• 訓練時の設計選択（微調整、プロービング、結合訓練）を比較し、テスト時適応の実用的で効果的な設定を特定。

提案手法

• TTT フレームワークに MAE を自己監視コンポーネントとして採用し、Y 形状のアーキテクチャ（エンコーダ f、自己監視ヘッド g、主タスクヘッド h）を用いる。
• MAE 事前学習済みのエンコーダ f0 とデコーダ g0 を用い、マスクされたパッチに対する自己監視再構成損失を各テスト入力で最小化するテスト時最適化を実施し、得られた f_x と g_x を h∘f_x で予測する前に得る。
• ViT プロービング（f を凍結、ヘッド h を訓練）をデフォルトの訓練時設定として強力なベースラインとし、微調整や結合訓練と比較する。
• MAS（マスク率 75%）と非 mangled 増強を用いて訓練；テスト入力ごとに 20 ステップの SGD で開始点 f0, g0 から訓練する。
• ImageNet-C レベル5（および付録の他のレベル）で評価し、汚染特異的増強を用いずにベースラインより改善を報告する。
• 線形設定で、PCA に似たオートエンコーディングを用いたTTT が、固定モデルよりバイアス-分散のトレードオフを改善する理論的洞察を提供する。

実験結果

リサーチクエスチョン

• RQ1 MAE ベースの TTT は、腐敗特異的手掛かりに依存せず、分布シフト下で視覚モデルの頑健性を改善できるか？
• RQ2 MAE ベースの TTT は、回転予測に基づく TTT や他の訓練時設計（微調整、プロービング、結合訓練）と比較して、ベンチマーク全体でどのように比較されるか？
• RQ3 線形モデル設定におけるバイアス-分散の観点から、TTT-MAE の有効性を説明する理論的根拠は何か？

主な発見

• TTT-MAE はベースラインの ViT プロービング設定より ImageNet-C レベル-5 で精度を大きく向上させる。
• TTT-MAE は回転予測ベースの TTT やベースラインモデルより、ほとんどの汚染タイプで上回る。
• 訓練時の設計選択は重要で、MAE 事前学習済みの ViT プロービングが、微調整や結合訓練より最も強力な性能をもたらす。
• テスト時最適化を 20 ステップの予算固定で SGD を用いて行っても、検証ベースの早期停止なしに継続的な改善が得られる。
• TTT-MAE は分布シフト下で ImageNet-A、ImageNet-R、Portraits データセットにも利得をもたらす。