[论文解读] 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

一句话总结

本论文在测试时训练中使用 Masked Autoencoders (MAE) 通过自监督重建对每个测试输入进行自适应，从而提升在视觉基准上的对分布偏移的鲁棒性。它在 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.

研究动机与目标

• Motivate robust generalization under unseen distribution shifts and propose adapting models at test time using self-supervision.
• Leverage masked autoencoding as the self-supervised task to generate informative signals for per-sample adaptation.
• Evaluate TTT-MAE across diverse distribution-shift benchmarks (ImageNet-C, ImageNet-A, ImageNet-R, Portraits) and analyze its theoretical properties.
• Compare training-time design choices (fine-tuning, probing, joint training) and identify a practical, effective setup for test-time adaptation.

提出的方法

• Adopt MAE as the self-supervised component in a test-time training (TTT) framework with a Y-shaped architecture (encoder f, self-supervised head g, main task head h).
• Use MAE pre-trained encoder f0 and decoder g0; perform test-time optimization to minimize a self-supervised reconstruction loss on masked patches for each test input, yielding f_x and g_x before predicting with h∘f_x.
• Employ ViT probing (freeze f, train a head h) as the default training-time setup for a strong baseline, and compare with fine-tuning and joint training.
• Train with MAS (mask 75%) and non-mangled augmentations; use SGD for 20 TTT steps per test input, starting from f0, g0.
• Evaluate on ImageNet-C level-5 (and other levels in appendix) and report improvements over the baseline, without using corruption-specific augmentations.
• Provide theoretical insight in a linear setting showing that TTT with PCA-like autoencoding yields a better bias-variance trade-off than a fixed model.

实验结果

研究问题

• RQ1 Can test-time training with masked autoencoders improve robustness of vision models under various distribution shifts without relying on corruption-specific cues?
• RQ2 How does MAE-based TTT compare to rotation-prediction-based TTT and other training-time designs (fine-tuning, probing, joint training) across benchmarks?
• RQ3 What is the theoretical explanation for TTT-MAE’s effectiveness in terms of bias-variance trade-off in a linear model setting?

主要发现

• TTT-MAE significantly improves accuracy on ImageNet-C level-5 over the baseline ViT probing setup.
• TTT-MAE outperforms rotation-prediction based TTT and baseline models across most corruption types.
• Training-time design choice matters: ViT probing with MAE pretraining yields the strongest performance compared to fine-tuning or joint training.
• Using SGD for test-time optimization with a fixed 20-step budget yields ongoing improvements without validation-based early stopping.
• TTT-MAE also yields gains on ImageNet-A, ImageNet-R, and the Portraits dataset under distribution shifts.