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 training data.
We also theoretically characterize the relationshipbetween our method and the well known bias-variance trade-off in statistics.
Introduction
Generalization is the central theme of supervised learning, and a hallmark of intelligence. Whilemost of the research focuses on generalization when the training and test data are drawn from thesame distribution, this is rarely the case for real world deployment [47], and is certainly not true forenvironments where intelligence has emerged.
Most models today are fixed during deployment, even when the test distribution changes. As aconsequence, the trained models need to be robust to all potential distribution shifts that could happenin the future [24, 14, 49, 38]. This turns out to be quite difficult because being ready for all possiblefutures limits the model’s capacity to be good at any particular one. But only one of these futures isactually going to happen.
This motivates an alternative perspective on generalization: instead of being ready for everything,one simply adapts to the future once it arrives. Test-time training (TTT) is one line of work thattakes this perspective [46, 36, 20, 1]. The key insight is that each test input gives a hint about thetest distribution. We modify the model at test time to take advantage of this hint by setting up aone-sample learning problem .
The only issue is that the test input comes without a ground truth label. But we can generate labels from the input itself thanks to self-supervised learning.
At training time,
TTT trains on both the main task (e.g. classification) and the self-supervised task. Then at test time,it adapts the model on the self-supervised task alone for each test input, before making a predictionon the main task.
The choice of the self-supervised task is critical.
On one hand, this task must be general and broadenough to work on a wide range of potential test distributions.
But on the other hand, it should not beinvariant to the changes across these test distributions, otherwise test-time training on a new samplewill not provide the learning signal needed to improve the main task. What self-supervised task isboth general and difficult enough to fit these requirements? We turn to a fundamental property sharedby natural visual data – spatial smoothness, i.e. the local redundancy of information in xy space.
Thetask of spatial prediction – remove parts of the data, and predict the removed content – forms the basis of some of the most successful self-supervised approaches, e.g. denoising autoencoder [50],context encoder [41], and masked autoencoder (MAE) [21]. In this paper, we should that MAE [21] turns out to be extremely well-suited for test-time training.
Our method leads to substantial improvements on four different object recognition datasets.
We alsooffer a theoretical justification for the use of test-time training based on a simplified linear formulation. The code and the models will be made public upon acceptance.
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 [21] 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 [12].
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
When training and test distributions are different, generalization is intrinsically hard.
Training onmany distribution shifts at once causes underfitting [53, 11].
The robustness obtained by training orfine-tuning on one type of distribution shift (e.g. Gaussian noise) also does not transfer to another(e.g. salt-and-pepper noise) [14, 49], even for visually similar corruptions.
To avoid the hardness of extrapolation, researchers work on alternative problem settings that anticipatethe test distribution.
For domain generalization [30, 29, 42, 39, 4, 15, 37, 19], a meta distributiongenerates many distribution shifts, some of which are used for training and others for testing.
Fordomain adaptation [33, 34, 8, 5, 44, 26, 6, 18], 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 they were trained for.
2.3 Test-Time Training
Test-time training (TTT) [46] is a paradigm of methods that generalizes to distribution shifts withoutanticipating them, through self-supervised learning on each test input x .
In [46], the main task isobject recognition, and the self-supervised task is rotation prediction [16]: 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 ) again
[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 [51] or dataset [32] of samples from the test distribution.
Before [46], single-image training has been performed for super-resolution [43].
The philosophy of
TTT is also preceded by that of local learning [3] and transductive learning [48], although usuallyperformed for support vector machines with multiple test samples [7, 27].
Since [46], TTT hasbeen applied to many fields with domain-specific self-supervised tasks: vision-based reinforcementlearning [20], model-based legged locomotion 45], tracking
13], natural languagequestion answering [1], and even medical imaging [28].
Method
At a high level, our method simply substitutes MAE [21] for the self-supervised part of TTT [46].
In practice, making this work involves many design choices.
Architecture.
Like in [46], our architecture is Y-shaped.
The feature extractor f is exactly theencoder of MAE, and g the decoder; each is a Vision Transformer (ViT) [12].
We intentionally avoidmodifying them to make clean comparisons with results in [21].
Choice of the main task (objectrecognition) head h presents more opportunities.
[21] 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 [10].
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 [21], only for linear h (see Subsection2.1).
The third is practiced by [46], 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 [21], 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 [46] 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 [21] optimized heavily with many hyperparameters that prevent overfitting.
Altogether, we believe 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 [46], 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.
Formally, denote the encoder produced by MAE pre-training as f 0 (and the decoder, used later for
TTT, as g 0 ).
ViT probing produces a trained main task head h 0 :
[Equation] The summation is over the training set with n samples, each consisting of input x i and label y i . The main task loss l m in our case is the cross entropy loss for classification.
Note that we are onlyoptimizing the main task head, while the pre-trained encoder f 0 is frozen.
Augmentations. Our default setup, during training-time training, only uses image cropping andhorizontal flips for augmentations, follow the protocol in [21] for pre-training and linear probing.
Fine-tuning in [21] (and [2, 52]), however, adds aggressive augmentations on top of the two above,including random changes in brightness, contrast, color and sharpness 1 .
Many of them are in factdistribution shifts in our evaluation benchmarks.
Training with them is analogous to training on thetest set, for the purpose of studying generalization to new test distributions.
Therefore, we choose notto use these augmentations, even though they would improve our results at face value.
Test-time training.
At test time, we start from the main task head h 0 produced by ViT probing, aswell as the MAE pre-trained encoder f 0 and decoder g 0 . Once each test input x arrives, we optimizethe following loss for TTT: [Equation] [Figure]
i g
[Figure]
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]
The self-supervised reconstruction loss l s computes the pixel-wise squared error of the decodedpatches with the ground truth.
After TTT, we make a prediction on x as h ◦ f x ( x ) .
Note thatgradient-based optimization for Equation 2 always starts from f 0 and g 0 .
When evaluating on a testset, we always discard f x and g x after making a prediction on each test input x , and reset the weightsto f 0 and g 0 for the next test input.
By test-time training on the test inputs independently, we do notassume that they come from the same distribution.
The choice of optimizer for TTT is straightforward in [46]: it simply takes exactly the same optimizersetting as during the last epoch of training-time training of the self-supervised task.
This choice,however, is not available for us, because the learning rate schedule of MAE reaches zero by the endof pre-training.
We experiment with various learning rates for AdamW [35] and stochastic gradientdescent (SGD) with momentum.
Performance of both, using the best learning rate respectively, isshown in Figure 2.
AdamW is used in [21], 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 testimage 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 [21], pre-trained on
ImageNet-1k [10] 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 [21].
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 four NVIDIA A100 GPUs; hyper-parameter sweeps are ran onan industrial cluster with mostly V100 GPUs.
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 [46].
We also tried reconstruction loss both with and without normalized pixels, as done in [21].
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 [24] is a benchmark for object recognition under distribution shifts.
It contains copiesof the ImageNet [10] validation set with 15 types of corruptions, independent across images, eachwith 5 levels of severity.
Due to space constraints, all results in the main text are on level 5, the mostsevere, unless stated otherwise.
Results on the other four levels are in the appendix.
Sample imagesfrom this benchmark are also in the appendix in Figure 4.
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 [24].
This rule, in our opinion, is helpful for community progress: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 [46], we plotthe accuracy only on the level-5 (most severe) corruptions.
Results on the other levels are in theappendix.
Our baseline (green) in Figure 3 is the default training method throughout this paper: takea pre-trained MAE encoder, perform ViT probing for object recognition on the original ImageNettraining set, then apply the fixed model to the corrupted test sets.
Performing TTT-MAE on ourbaseline significantly improves performance.
Reconstruction vs. rotation prediction.
The baseline for TTT-Rot, taken from [46], is called Joint
Train in Figure 3.
This is a ResNet [22] 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.
2 Socomparison to [46] 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 [16] 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 [40].
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,without forming semantic features as intended.
This causes TTT-Rot of [46] to harm performance, asseen 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.
The first three rows of Table 2 compare accuracyof these three designs.
As discussed earlier, joint training does not achieve satisfactory performance [Table] [Table] third row, ViT probing, is our default baseline throughout the paper; it has the same numbers the baseline in
Figure
The last row
TTT MAE ViT probing. It achieves the best performance types it numbers
MAE in Figure 3.
in most categories.
While fine-tuning is initially better than ViT probing, it is not amenable to TTT.
Therefore, ViT probing is chosen as the default baseline throughout the paper.
All three of the aboveare only for training-time training, after which a fixed model is applied during testing.
The last row is
TTT-MAE on a model trained with ViT probing, and performs the best across all corruption types.
Numbers in the last two rows are the same as for the baseline and TTT-MAE in Figure 3.
4.3 Other ImageNet Variants
We evaluated TTT-MAE on two other popular benchmarks for distribution shifts.
ImageNet-A [25]contains real-world, unmodified, and naturally occurring examples where popular ImageNet modelsperform poorly.
ImageNet-R [23] 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 [17] 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.
Wesort the entire dataset by year and split it into four equal parts, with 5062 images each.
Between thefour splits, there are visible low-level differences such as blurriness and contrast, as well as high-levelchange 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.
Theoretical Results
Why does TTT help?
The theory of [46] gives an intuitive but shallow explanation: when 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 with correlated gradients, without really answering the question, or showing anyconcrete example of such a self-supervised task. In this paper, we show that autoencoding, i.e. reconstruction, is a self-supervised task that makes
TTT help. To do so with minimal mathematical overhead, we restrict ourselves to the linear world,where our models are linear and the distribution shifts are produced by linear transformations.
Weanalyze autoencoding with dimensionality reduction instead of masking, as we believe the two areclosely related in essence.
The most illustrative insight, in our opinion, is that under distribution shifts, TTT finds a better biasvariance trade-off than applying a fixed model.
The fixed model is biased because it is completelybased on biased training data, which do not represent the test sample under distribution shifts. Theother extreme is to completely forget the training data, and train a new model from scratch on eachtest input. This is also undesirable because the single test input is high variance, although unbiasedby definition. A sweet spot of the trade-off can be found by performing TTT while remembering thetraining data in some capacity: in practice by initializing with a model trained on them, and in thissection by using part of their data covariance matrix.
It is well known that linear autoencoding is equivalent to principle component analysis (PCA). This equivalence simplifies the mathematics by giving us closed form solutions to the optimizationproblems both during training and test-time training. For the rest of the section, we use the term PCAinstead of linear autoencoding, following convention of the theory community.
Problem setup.
Let x, y ∼ P , the training distribution, and assume that the population covariancematrix Σ
= Cov( x ) is known.
PCA performs spectral decomposition on Σ = UDU ⊤ , and takes thetop k eigenvectors u 1 , . .
, u k to project x ∈ R d to R k .
Throughout this section, we denote u i as thei th column of the matrix U , and likewise for other matrices.
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]
[Table] (train) [Table]
Algorithms.
The fixed model baseline is to blithely apply our old procedure and predict
[Equation]
[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 data.
In statistical terms, α presents abias-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 3 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 6 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. Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] Code and data
will be released upon acceptance.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running ex-
periments multiple times)?
[N/A] Following the convention of [21] and many other large-scale experiments. (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets? [Yes] In appendix (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]