Test-Time Training with Masked Autoencoders
Anonymous Author(s)
AffiliationAddress email
Abstract
Prior work has shown masked autoencoding as an effective self-supervised taskacross many visual distributions with sufficient training data.
We use maskedautoencoding to train on each unlabeled test sample as it arrives at test time, beforemaking a prediction.
This simple method improves generalization of predictivemodels on many visual benchmarks of unknown distributions, where input imagesare not covered by any training data.
Our theoretical analysis explains how test-timetraining with autoencoding helps a linear model under distribution shifts.
Introduction
The success story of deep learning has been about general methods that leverage data and computationinstead of domain knowledge. When the test distribution is not restricted to be the same as thatof training, test-time training [50, 22, 1, 49] is a paradigm of methods that takes this philosophybeyond the training data.
It performs self-supervised learning for each unlabeled test input, beforemaking a prediction on this input itself.
Instead of anticipating a specific test distribution throughprior knowledge, this paradigm generalizes to diverse kinds of distribution shifts by leveraging theextra input and computation at test time. However, the generality of test-time training is still limited by its choice of the self-supervised task.
Take the rotation prediction task used in [50], which hopes to form semantic features by learningthe orientation of rotated images.
For natural scenes, image rotation can be trivially predicted,without scene understanding, from the horizon’s orientation alone.
Even for centered objects, manyclasses, e.g. frogs and lizards, have mostly top-down views, for which every orientation looks equallyplausible.
In both cases, rotation prediction is unlikely to form semantic features as hoped for. Latest advances in self-supervised learning suggest masked autoencoding as a potentially generalpurpose task.
This task is simple: remove parts of the data, and predict the removed content.
Innatural language processing, masked autoencoding [12] and the closely related task of autoregressivelanguage modeling [39, 42, 4] have long been a staple of the community. Recent research in computervision also shows that predicting missing patches on images can lead to state-of-the-art performanceon a variety of downstream tasks and datasets [23, 2, 55]. Motivated by these successes, we experiment with masked autoencoding (MAE) as self-supervisionfor test-time training on four datasets.
Our method leads to improvements in all of them.
We alsotheoretically characterize the relationship between our method and the well known bias-variancetrade-off in statistics.
Related
Work
2.1 Masked Autoencoders and Self-Supervised Learning
Key to the success of deep learning is the formation of semantic features that transfer across manytasks; the goal of self-supervised pre-training is to learn these transferable features without humanannotations, by making up inputs and labels to an auxiliary pretext task.
Given an original inputimage x , the pretext task of masked autoencoding [23] splits the image into many non-overlappingpatches (e.g. 196), randomly masks out majority of the patches (e.g. 75%), and trains a model toreconstruct the missing patches, by optimizing the mean squared error between the reconstructed andthe original pixels.
This model, called a masked autoencoder (MAE), is composed of two parts: a large encoder anda small decoder, both based on Vision Transformers [15].
For a downstream task, e.g. objectrecognition, the encoder takes in all the image patches and outputs the features, which then becomeinput to a linear projection head.
There are two common ways to combine the pre-trained encoder anduntrained head.
1) Fine-tuning: both the encoder and head are trained, end-to-end, for the downstreamtask.
Linear probing: the encoder is frozen as a fixed feature extractor, and only the head is trained.
We will refer back to these training options in Section 3.
2.2 Generalization under Distribution Shifts
Traditionally, training and test data are drawn i.i.d. from the same distribution; deep learninggeneralizes well by interpolating the training set with rather limited prior knowledge [56]. Whentraining and test distributions are different, generalization is intrinsically hard since interpolationcannot cover many of the test samples [43, 44, 46, 14]. One way to turn extrapolation into interpolation again is to make the training distribution wider, andthe training set larger, to cover the test distribution.
This brute force approach, however, naturally fallsshort of scaling up to all possible distribution shifts in the world, even if they could be anticipated attraining time.
Training on many distribution shifts at once causes underfitting [57, 13].
The robustnessobtained by training or fine-tuning on one type of distribution shift (e.g. Gaussian noise) also doesnot transfer to another (e.g. salt-and-pepper noise) [53, 18], even for visually similar corruptions.
To avoid the hardness of extrapolation, researchers work on alternative problem settings that anticipatethe test distribution.
For domain generalization [33, 32, 45, 40, 5, 17, 38, 21], a meta distributiongenerates many distribution shifts, some of which are used for training and others for testing.
Fordomain adaptation [35, 36, 9, 6, 48, 28, 7, 20], a potentially unlabeled dataset from the test distributionis available for training.
Solutions motivated by these settings share the same flaw as augmenting thetraining set: they are only effective on the particular meta / test distribution that was trained for.
2.3 Test-Time Training
Test-time training (TTT) [50] is a paradigm of methods that generalizes to distribution shifts withoutanticipating them, through self-supervised learning on each test input x .
In [50], the main task isobject recognition, and the self-supervised task is rotation prediction [18]: x is simply rotated in theimage plane by multiples of 90 degrees to create a four-way classification problem.
The network for TTT is Y-shaped: a feature extractor f , followed simultaneous by a self-supervisedhead g and a main task head h .
During training, labels are available for both the self-supervisedtask, using g ◦ f , and the main task, using h ◦ f ; so both branches are trained jointly by summingtheir losses together for gradient descent.
During testing, for each x , only the self-supervised task islabeled; so g ◦ f is trained, end-to-end, on this single input.
This produces a different model, call it
f x(and g x ), initialized from f (and g ), locally optimized for each test input x .
After making a predictionon x , now with h ◦ f x , the algorithm never uses f x (and g x
[Figure]
TTT as explained above is not a new problem setting; it is a proposed solution to the canonicalproblem of generalization.
Our paper stay in the same setting.
Other followups have worked on newsettings such as a batch [54] or dataset [34] of samples from the test distribution.
Before [50], single-image training has been performed for super-resolution [47].
The philosophy of
TTT is also preceded by that of local learning [3] and transductive learning [52], although usuallyperformed for support vector machines with multiple test samples [8, 29].
Since [50], TTT hasbeen applied to many fields with domain-specific self-supervised tasks: vision-based reinforcementlearning [22], model-based legged locomotion
16], natural languagequestion answering [1], and even medical imaging [31].
Method
At a high level, our method simply substitutes MAE [23] for the self-supervised part of TTT [50].
In practice, making this work involves many design choices.
Architecture.
Like in [50], our architecture is Y-shaped.
The feature extractor f is exactly theencoder of MAE, and g the decoder; each is a Vision Transformer (ViT)
We intentionally avoidmodifying them to make clean comparisons with results in [23].
Choice of the main task (objectrecognition) head h presents more opportunities.
[23] uses a linear projection, from the encoderfeature dimension to the number of classes; the authors discuss this as mostly a historic artifact ofhow people have been evaluating the quality of features in unsupervised learning.
Therefore we alsoexperiment with a more expressive main task head – an entire ViT-Base – to strengthen our baseline.
Training-time training.
Following standard practice, we start from the MAE model provided bythe authors, with a ViT-Large encoder, pre-trained for reconstruction on ImageNet-1k [11].
There arethree ways to combine it with the untrained main task head: 1) Fine-tuning: train h ◦ f end-to-end.2)
Probing: train h only, with f frozen.
3) Joint training: train both h ◦ f and g ◦ f , by summingthe respective losses together.
The first two are practiced by [23], only for linear h (see Subsection2.1).
The third is practiced by [50], for linear h , as well as a neural network (see Subsection 2.3).
We experiment with all three, and choose probing with a ViT-Base head, a.k.a.
ViT probing , as ourdefault setup, based on the following considerations:
1) Fine-tuning does not work with TTT.
As discussed in [23], reconstruction and recognition require
very different features.
Fine-tuning makes the encoder specialize to recognition, only to find it
training for reconstruction again at test time and losing the recognition-only features that the head
learned to rely on.
This agrees with the intuition in [50] for joint training, instead of training only forrecognition, as preparation for TTT.
[Figure]
2) ViT probing is much more lightweight than both fine-tuning and joint-training.
It trains 3.5 timesfewer parameters than even linear fine-tuning (86M vs. 306M).
It also obtains higher accuracy thanlinear fine-tuning without aggressive augmentations on the ImageNet validation set, and accuracysimilar to linear fine-tuning with augmentations, which [23] optimized heavily with many hyperparameters that prevent overfitting.
Altogether, we think that ViT probing strikes the right balancebetween the highly constrained linear probing and highly expressive linear / ViT fine-tuning.3)
Joint training is the recommended design in [50], but its accuracy on the ImageNet validation set,to the best of our ability, has been worse than the linear / ViT probing counterparts.
This is most likelybecause we have not explored enough of the hyper-parameter space.
As discussed in 2), fine-tuning isalready heavyweight and hard to train
; joint training is only heavier, and has more hyper-parameters.
If future work finds a reliable recipe for joint training, our method will likely benefit from its setup.
Meanwhile, we consider it our contribution to find ViT probing as a a lightweight substitute.
Augmentations.
Our default setup, during training-time training, is to only use image croppingand horizontal flips for augmentation, follow the protocol in [23] for pre-training and linear probing.
Fine-tuning in [23] (and [2, 55]), however, adds aggressive augmentations on top of the two above,including random changes in brightness, contrast, color and sharpness 2 .
Many of them are in factcorruptions / distribution shifts contained in our evaluation benchmarks.
Training with them isanalogous to training on the test set, for the purpose of studying generalization to truly unknowndistributions.
Therefore, we choose not to use these augmentations, even though they would improveour results at face value.
Test-time training.
The choice of optimizer for TTT is straightforward in [50]: it simply takesexactly the same optimizer setting as during the last epoch of training-time training of the selfsupervised task.
This choice, however, is not available for us, because the learning rate scheduleof MAE reaches zero by the end of pre-training.
We experiment with various learning rates for AdamW [37] and stochastic gradient descent (SGD) with momentum.
Performance of both, usingrespectively the best learning rate, is shown in Figure 2.
AdamW is used in [23], but for TTT it hurts performance with too many iterations.
On the otherhand, more iterations with SGD consistently improve performance on all distribution shifts.
Testaccuracy keeps improving even after 20 iterations.
This is very desirable for TTT: the single test b r i g h t c o n t r a s t d e f o c u s e l a s t i c f o g f r o s t g a u ss g l a ss i m p u l s e j p e g m o t i o n p i x
e l a t e s h o t s n o w z oo m
[Figure]
image is all we know about the test distribution, so there is no validation set to tune hyper-parameters,or monitor performance for early stopping.
With SGD, we can simply keep training.
Empirical Results
4.1 Implementation Details
In all experiments we use the MAE based on ViT-Large, configured as in [23], pre-trained on
ImageNet-1k [11] for 800 epochs.
Our ViT-Base head takes as input the image features from thepre-trained MAE encoder.
There is a linear layer in between that resizes those features to fit as inputsto the head, just like between the encoder and the decoder in [23].
A linear layer is also appended tothe final class token of the head to produce the classification logits.
TTT is performed with SGD, as discussed, for 20 steps, using a momentum of 0.9, weight decay of0.2, batch size of 128, and fixed learning rate of 5e-3.
The choice of 20 steps is purely computational;more steps will likely improve performance marginally, judging from the trend observed in Figure 2.
Most experiments are performed on one NVIDIA A100 GPU; hyper-parameter sweeps are ran on anindustrial cluster with mostly V100 GPUs.
Like in pre-training, only tokens for non-masked patches are given to the encoder during TTT,whereas the decoder takes all tokens.
Each image in a TTT batch has a different random mask.
Weuse the same masking ratio as MAE [23] – 75%.
We do not use any augmentation on top of randommasking for TTT.
Specifically, every iteration of TTT is performed on the same 224 × 224 centercrop of the image as we later use to make a prediction; the only difference being that predictions aremade on images without masking.
We optimize both the encoder and decoder weights during TTT, together with the class token and themask token.
We have experimented with freezing the decoder and found that the difference, even formultiple iterations of SGD, is negligible; see the appendix for detailed comparison.
This is consistentwith the observations in [50].
We also tried reconstruction loss both with and without normalizedpixels, as done in [23].
Our method works well for both, slightly better for normalized pixels becausethe baseline is slightly better.
We include these results in the appendix.
4.2 ImageNet-C
ImageNet-C [26] is a benchmark for object recognition under distribution shifts.
It contains copies ofthe ImageNet [11] validation set with 15 types of corruptions, independent across images, each with5 levels of severity.
In the spirit of treating these distribution shifts as truly unknown, we do not use training data, priorknowledge, or data augmentations derived from these corruptions.
This complies with the statedrule of the benchmark, that the corruptions should be used only for evaluation, and the algorithmsshould not be corruption-specific [26].
This rule, in our opinion, is helpful for community progress: [Table] as explained in Subsection 2.2, the numerous distribution shifts outside of research evaluation cannotbe anticipated, and algorithms that rely on information specific to the test distribution cannot scale.
Our main results on ImageNet-C appear in Figure 3.
Following the convention in [50], we plot theaccuracy only on the level-5 (most severe) corruptions.
Results on the other levels are in the appendix.
Our baseline is the fixed model without TTT, produced by ViT probing on a pre-trained MAE encoderas discussed, i.e. the default training method throughout this paper.
Performing TTT-MAE on ourbaseline significantly improves performance.
Reconstruction vs. rotation prediction.
The baseline for TTT-Rot, taken from [50], is called Joint
Train in Figure 3.
This is a ResNet [24] with 16-layers, trained jointly for rotation prediction andobject recognition.
Because our baseline is much more advanced than the ResNet baseline, theformer already outperforms the reported results of TTT-Rot, for all corruptions except contrast.
3 Socomparison to [50] is more meaningful in relative terms: TTT-MAE has higher performance gains inall corruptions than TTT-Rot, on top of their respective baselines.
Rotation invariant classes.
As discussed in the introduction, many self-supervised tasks likerotation prediction [18] are designed for specific domains.
This can be especially limiting when thegoal is, in fact, to generalize to unknown domains.
Here we show examples in ImageNet where theassumptions behind rotation prediction do not hold.
We find rotation invariant classes through the following: 1) Take the ImageNet pre-trained ResNet-18from the official PyTorch website [41].
2) Run it on the original validation set, and one copy of itwith 90-degree rotation.
3) Filter out classes for which the original accuracy is lower than 60%. 4)
For each class, compare accuracy for the original and the rotate images, and choose the 5 classes forwhich the normalized difference is the smallest.
Sample images from these 5 classes are shown in Table 1.
Not surprisingly, these images are usuallytaken from top-down views; rotation prediction on them can only memorize the generated labels, [Table] without forming semantic features as intended.
TTT-Rot of [50] to harm performance, seen in Table 1.
Also not surprisingly, TTT-MAE is agnostic to rotation invariance and still helps
Training-time training.
In the synonymous paragraph in Section 3, we discussed our choice of ViT probing instead of fine-tuning or joint training.
Table 2 compares accuracy of these three choices under our default setting.
As discussed earlier, joint training does not achieve satisfactoryperformance in most categories; while fine-tuning is initially better than probing, it is not amenable toour method.
In the end, TTT-MAE on top of probing outperforms fine-tuning for all the corruptions.
Numbers in the last two rows are the same as in Figure 3.
4.3 Other ImageNet Variants
We evaluated TTT-MAE on two other popular benchmarks for distribution shifts.
ImageNet-A [27]contains real-world, unmodified, and naturally occurring examples where popular ImageNet modelsperform poorly.
ImageNet-R [25] contains renditions of ImageNet classes, e.g. art, cartoons andorigami.
Both TTT-MAE and the baseline are the same models and algorithms, with exactly thesame hyper-parameters as for ImageNet-C (Subsection 4.2).
Our method continues to outperform thebaseline by large margins, see Table 4.3.
4.4 Portraits Dataset
The Portraits dataset [19] contains American high school yearbook photos labeled by gender, takenover more than a century.
It contains real-world distribution shifts in appearances over the years.
We sort the entire dataset by year and split it into four equal parts, with 5062 images each.
Sampleimages from the four splits are shown in Figure 4.
Between the four splits, there are visible low-leveldifferences such as blurriness and contrast, as well as high-level change like hair-style and smile.
We create three experiments: for each we train on one of the first three splits and test on the fourth.
Performance is in terms of accuracy in binary gender classification.
As expected, performanceincreases slightly as the training split gets closer to test in time.
Our model is the same ImageNet pre-trained MAE + ViT head (with sigmoid for binary classification),and probing is performed on the training split.
We use exactly the same hyper-parameters for TTT.
Our results are shown in 4: our method improves on the baseline for all three experiments. [Table] [Figure] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Table]
Theoretical Results
Why does TTT help?
The theory of [50] gives an intuitive but shallow explanation: because theself-supervised task happens to propagate gradients that correlate with those of the main task.
Butthis evasive theory only delegates one unknown – the test distribution, to another – the magicalself-supervised task, without really answering the question. In this paper, we do not want to assume an oracle, i.e. domain expert, on the supposedly unknowntest distribution, handing over a self-supervised task with the desirable properties.
Algorithmically,we have experimented with masked autoencoding to demonstrate its generality. Theoretically, in thissection, we explain why TTT with autoencoding helps in the linear world, in simple mathematicalterms.
This analysis is possible thanks to the well known equivalence between linear autoencodingand principle component analysis (PCA).
The most illustrative insight, in our opinion, is that under distribution shifts, test-time training findsa better trade-off between bias and variance than applying a fixed model.
And in our context, biasmeans dataset bias in its colloquial sense [51].
We begin by setting up the problem.
Problem setup.
Throughout this section, we denote a i as the i th column of matrix A .
Let x, y ∼ P ,the training distribution, and assume that the population covariance matrix Σ
= Cov( x ) is known.
PCA performs spectral decomposition on Σ = UDU ⊤ , and takes the top k eigenvectors u 1 , . .
, u kto project x ∈ R d to R k .
Assume that for each x , the ground truth y is a linear function of the PCA projection with k = 1 :
[Equation]
for some known weight w ∈ R .
For mathematical convenience, we also assume the following aboutthe eigenvalues, i.e. the diagonal entries of D :
[Equation]
At test time, nature draws a new x, y ∼ P , but we can only see ˜ x , a corrupted version of x . We modela corruption as an unknown orthogonal transformation R , and ˜ x = Rx .
Algorithms.
The fixed model baseline is to blithely apply our old procedure and predict
[Equation]
as an estimate of y .
Intuitively, this will be inaccurate if corruption is severe i.e. R is far from I .
Andhere is the PCA version of TTT.
We form a new (and rather trivial) covariance matrix with ˜ x , theonly piece of information we have about the corruption.
Let α ∈ [0 , 1] be a hyper-parameter, and thelinear combination of Σ with our new covariance matrix be
[Equation]
Denote its spectral decomposition as M ( α ) =
V ( α ) · S ( α ) · V ⊤ ( α ) .
We then predict
[Equation]
Bias-variance trade-off.
α = 0 is equivalent to the baseline, and α = 1 means that we exclusiveuse the one test input and completely forget about the training features.
In statistical terms, α presentsa bias-variance trade-off: α ↓ 0 means more bias, α ↑ 1 means more variance.
Theorem.
Define the prediction risk as E [ | ˆ y − y | ] , where the expectation is taken over the corruptedtest distribution.
This risk is strictly dominated when α = 0 .
That is, TTT with some hyper-parameterchoice α > 0 is, on average, strictly better than the baseline.
The proof is given in the appendix.
Remarks on assumptions.
The assumption in Equation 1 is mild; it basically just says that PCA isat least helpful on the training distribution.
It is also mild to assume that Σ and w are known in thiscontext, since the estimated covariance matrix and weight asymptotically approach the populationground truth as the training set grows larger.
The assumption on the eigenvalues in Equation 4 greatlysimplifies the math, but a more complicated version of our analysis should exist without it.
Lastly,we believe that our general statement should still hold for invertible linear transformation instead oforthogonal transformations, since they share the same intuition.
Discussion
Limitations.
Due to the additional computation at test time, our method is slower than the baseline.
Inference speed has not been the focus of this paper, and it can potentially be improved through betteroptimizers and hyper-parameters.
In general, test-time training has been limited by the deep learningtoolbox.
Most tools, for architectural design, regularization and optimization, have been refined fortraining-time training on large datasets, especially for large models like ours; this is another directionfor future work.
Broader impact.
In his book Thinking, Fast and Slow , Nobel Prize Winner Daniel Kahneman [30]describes a dichotomy between two kinds of thinking: System 1, the reactive and intuitive, and
System 2, the deliberate and logical.
Consensus holds that deep learning systems today are almostexclusively System 1.
Pushing them closer to System 2 has the potential to make artificial intelligencemore reliable, interpretable and transparent, and better vehicles to understand human intelligence.
Webelieve that our work is a small step towards this goal, by allowing machine learning model to slowlyrefine their predictions beyond a reactive feed-forward process.
Choosing a general self-supervisedlearning task makes this kind of “thinking” less affected by biases of human experts, and moreattentive to each individual at test time.