MEMO: Test Time Robustness via Adaptation and Augmentation
Anonymous Author(s)
AfﬁliationAddress email
Abstract
While deep neural networks can attain good accuracy on in-distribution test points, many applications require robustness even in the face of unexpected perturbationsin the input, changes in the domain, or other sources of distribution shift.
We studythe problem of test time robustiﬁcation , i.e., using the test input to improve modelrobustness.
Recent prior works have proposed methods for test time adaptation, however, they each introduce additional assumptions, such as access to multiple test points, that prevent widespread adoption.
In this work, we aim to study anddevise methods that make no assumptions about the model training process and are broadly applicable at test time.
We propose a simple approach that can be used in any test setting where the model is probabilistic and adaptable: when presented with a test example, perform different data augmentations on the data point, and then adapt (all of) the model parameters by minimizing the entropy of the model’saverage, or marginal , output distribution across the augmentations.
Intuitively, this objective encourages the model to make the same prediction across differentaugmentations, thus enforcing the invariances encoded in these augmentations,while also maintaining conﬁdence in its predictions.
In our experiments, weevaluate two baseline ResNet models, two robust ResNet-50 models, and a robustvision transformer model, and we demonstrate that this approach achieves accuracygains of 1-8% over standard model evaluation and also generally outperforms prior augmentation and adaptation strategies.
For the setting in which only one test pointis available, we achieve state-of-the-art results on the ImageNet-C, ImageNet-R,and, among ResNet-50 models, ImageNet-A distribution shift benchmarks.
Introduction
Deep neural network models have achieved excellent performance on many machine learningproblems, such as image classiﬁcation but are often brittle and susceptible to issues stemmingfrom distribution shift .
For example, deep image classiﬁers may degrade precipitously in accuracywhen encountering input perturbations, such as noise or changes in lighting [12] or domain shiftswhich occur naturally in real world applications [21].
Therefore, robustiﬁcation of deep modelsagainst these test shifts is an important and active area of study.
Most prior works in this area have focused on techniques for training time robustiﬁcation, includingutilizing larger models and datasets [34], various forms of adversarial training [39, 48], and aggressivedata augmentation [51, 13, 24, 14].
Employing these techniques requires modifying the trainingprocess, which may not be feasible if, e.g., it involves heavy computation or non public data.
Furthermore, these techniques do not rely on any information about the test points that the modelmust predict on, even though these test points may provide signiﬁcant information for improvingmodel robustness.
Recently, several works have proposed methods for improving accuracy via adaptation after seeing the test data, typically by updating a subset of the model’s weights [44, 45,18], normalization statistics [40], or both [46, 52].
Though effective at handling test shifts, thesemethods sometimes still require specialized training procedures, and they typically rely on extractinginformation via batches or even entire sets of test inputs, thus introducing additional assumptions.
In this work, we focus on methods for test time robustness , in which the speciﬁc test input may beleveraged in order to improve the model’s prediction on that point. We are interested in studyingand devising methods for improving model robustness that are “plug and play”, i.e., they can bereadily used with a wide variety of pretrained models and test settings.
We also want methods thatsynergize with other robustiﬁcation techniques, in order to achieve greater performance than usingeither in isolation.
With these goals in mind, we devise a novel test time robustness method basedon adaptation and augmentation.
As illustrated in Figure 1, when presented with a test point, weadapt the model by augmenting the test point in different ways while encouraging the model to makeconsistent predictions, thus respecting the invariances encoded in the data augmentations.
We furtherencourage the model to make conﬁdent predictions, thus arriving at the proposed method: minimizethe marginal entropy of the model’s predictions across the augmented versions of the test point.
We refer to the proposed method as m arginal e ntropy m inimization with o ne test point (MEMO), andthis is the primary contribution of our work.
MEMO makes direct use of pretrained models withoutany assumptions about their particular training procedure or architecture, while requiring only asingle test input for adaptation.
In Section 4, we demonstrate empirically that MEMO consistentlyimproves the performance of ResNet [11] and vision transformer [7] models on several challenging
ImageNet distribution shift benchmarks, achieving several new state-of-the-art results for thesemodels in the setting in which only one test point is available.
MEMO consistently outperforms nonadaptive marginal distribution predictions (between 1-10% improvement) on the ImageNet-C [12]and ImageNet-R [14] test sets, indicating that adaptation plays a crucial role in improving predictiveaccuracy.
MEMO encourages both invariance across augmentations and conﬁdent predictions, and anablation study in Section 4 shows that both components are important for maximal performance gains.
Also, MEMO is, to the best of our knowledge, the ﬁrst adaptation method to improve performance(by 1-4% over standard model evaluation) on the ImageNet-A test set [15].
Related work
Distribution shift has been studied under a number of frameworks [36], including domain adaptation[41, 6, 47], domain generalization [5, 32, 9], and distributionally robust optimization [4, 16, 39].
Theseframeworks typically leverage additional training or test assumptions in order to make the distributionshift problem more tractable.
Largely separate from these frameworks, various empirical methodshave also been proposed for dealing with shift, such as increasing the model and training dataset size or using heavy training augmentations [34, 51, 14].
The focus of this work is complementary tothese efforts: MEMO is applicable to a wide range of pretrained models, including those trained viarobustness methods, and can achieve further performance gains via test time adaptation.
Prior test time adaptation methods generally either make signiﬁcant training or test time assumptions.
Some methods update the model using batches or even entire datasets of test inputs, such as bycomputing batch normalization (BN) statistics on the test set [25, 19, 33, 40], computing classprototypes [18], or minimizing the (conditional) entropy of model predictions across a batch oftest data [46].
The latter approach is closely related to MEMO.
The differences are that MEMOminimizes marginal entropy using single test points and data augmentation and adapts all of themodel parameters rather than just those associated with normalization layers, thus not requiringmultiple test points or speciﬁc model architectures.
Other test time adaptation methods can be appliedto single test points but require speciﬁc training procedures or models [44, 17, 40, 1, 3].
Test timetraining (TTT) [44] requires a specialized model with a rotation prediction head and a differentprocedure for training this model.
Schneider et al.
[40] show that BN adaptation can be effective evenwith only one test point.
As we discuss in Section 3, MEMO synergizes well with this techniqueof “single point” BN adaptation. Mao et al.
[29] propose a test time adaptation method based oninput perturbations for robustness to adversarial attacks.
Concurrently with our work, Sivaprasadand Fleuret [43] propose a similar method for test time adaptation by encouraging invariance to dataaugmentations, and they test their method on the corrupted CIFAR and VisDA [35] datasets.
A number of works have noted that varying forms of strong data augmentation on the training set canimprove the resulting model’s robustness [51, 13, 24, 14].
Data augmentations are also sometimesused on the test data directly by averaging the model’s outputs across augmented copies of the testpoint [23, 42], i.e., predicting according to the model’s marginal output distribution.
When usingcropping as the augmentation, this technique is often referred to as multicrop evaluation [22].
Weinstead use the term test time augmentation (TTA), as we use additional augmentations beyondcropping [13].
TTA has been shown to be useful both for improving model accuracy and calibration[2] as well as handling distribution shift [31].
We take this idea one step further by explicitly adaptingthe model such that its marginal output distribution has low entropy.
This extracts an additionallearning signal for improving the model, and furthermore, the adapted model can then make its ﬁnalprediction on the clean test point rather than the augmented copies.
We empirically show in Section 4that these differences lead to improved performance over this non adaptive TTA baseline.
Augmenting and Adapting at Test Time
Data augmentations are typically used to train the model to respect certain invariances – e.g., changesin lighting or viewpoint do not change the underlying class label – but, especially when faced withdistribution shift, the model is not guaranteed to obey the same invariances at test time.
In this section,we introduce MEMO, a method for test time robustness that adapts the model such that it respectsthese invariances on the test input.
We use “test time robustness” speciﬁcally to refer to techniquesthat operate directly on pretrained models and single test inputs – single point BN adaptation and
TTA, as described in Section 2, are examples of prior test time robustness methods.
In the test time robustness setting, we are given a trained model f ✓ with parameters ✓ 2 ⇥ .
We donot require any special training procedure and do not make any assumptions about the model, exceptthat ✓ is adaptable and that f ✓ produces a conditional output distribution p ✓ ( y | x ) that is differentiablewith respect to ✓ .
1 All standard deep neural network models satisfy these assumptions.
A singlepoint x 2 X is presented to f ✓ , for which it must predict a label ˆ y 2 Y immediately.
Note that this is 114 precisely identical to the standard test time inference procedure for regular supervised learning models– in effect, we are simply modifying how inference is done, without any additional assumptions onthe training process or on test time data availability.
This makes test time robustness methods asimple “slot-in” replacement for the ubiquitous and standard test time inference process.
We assumesampling access to a set of augmentation functions A , { a 1 , . . .
, a M } that can be applied to the testpoint x .
We use these augmentations and the self-supervised objective detailed below to adapt themodel before it predicts on x .
When given a set of test inputs, the model adapts and predicts on eachtest point independently.
We do not assume access to any ground truth labels.
3.1 Marginal Entropy Minimization with One test point
Given a test point x and set of augmentation functions A , we sample B augmentations from A andapply them to x in order to produce a batch of augmented data ˜ x 1 , . . . , ˜ x B .
The model’s average, ormarginal , output distribution with respect to the augmented points is given by
[Equation]
where the expectation is with respect to uniformly sampled augmentations a ⇠ U ( A ) .
What properties do we desire from this marginal distribution?
To answer this question, consider therole that data augmentation typically serves during training.
For each training point ( x train , y train ) , themodel f ✓ is trained using multiple augmented forms of the input ˜ x train 1 , . . .
, ˜ x train E .
f is trained to obeythe invariances between the augmentations and the label – no matter the augmentation on x train , fshould predict, with conﬁdence, the same label y train .
We seek to devise a similar learning signalduring test time, without any ground truth labels.
That is, after adapting:
(1) the model f ✓ predictions should be invariant across augmented versions of the test
point, and(2) the model f ✓ should be conﬁdent in its predictions, even for heavily augmented versions of thetest point, since all versions have the same underlying label.
Optimizing the model for more conﬁdent predictions can be justiﬁed from the assumption that thetrue underlying decision boundaries between classes lie in low density regions of the data space [8].
With these two goals in mind, we propose to adapt the model using the entropy of its marginal outputdistribution over augmentations (Eq. 1), i.e.,
[Equation]
Note that this objective is not the same as optimizing the average conditional entropy of the model’spredictive distributions across augmentations, i.e.,
[Equation]
A model which predicts conﬁdently but differently across augmentations would minimize Eq.
3 butnot Eq. 2.
Optimizing Eq.
2 encourages both conﬁdence and invariance, since the entropy of ¯ p ✓ ( ·| x )is minimized when the model outputs the same (conﬁdent) prediction regardless of the augmentation.
Algorithm 1 presents the overall method MEMO for test time adaptation.
Though prior test timeadaptation methods must carefully choose which parameters to adapt in order to avoid degeneratesolutions [46], our adaptation procedure simply adapts all of the model’s parameters ✓ (line 3).
Giventhat p
✓ ( y | x ) is differentiable with respect to ✓ , we can directly use gradient based optimization toadapt ✓ according to Eq.
We use only one gradient step per test point, because empirically wefound this to be sufﬁcient for improved performance while being more computationally efﬁcient.
After this step, we use the adapted model f ✓ 0 to predict on the original test input x (line 4).
3.2 Composing MEMO with Prior Methods
An additional beneﬁt of MEMO is that it synergizes with other approaches for handling distributionshift.
In particular, MEMO can be composed with prior methods for training robust models andadapting model statistics, thus leveraging the performance improvements of each technique.
Pretrained robust models.
Since MEMO makes no assumptions about, or modiﬁcations to, themodel training procedure, performing adaptation on top of pretrained robust models, such as thosetrained with heavy data augmentations, is as simple as using any other pretrained model.
Crucially,we ﬁnd that, in practice, the set of test augmentations A does not have to match the augmentationsthat were used to train the model.
For simplicity and efﬁciency, we use augmentations that can beeasily sampled and are applied directly to the model input x .
These properties do not hold for, e.g.,data augmentation techniques based on image translation models, such as DeepAugment [14], orfeature mixing, such as moment exchange [24].
However, we can still use models trained with thesedata augmentation techniques as our starting point for adaptation, thus allowing us to improve upontheir state-of-the-art results.
As noted above, using pretrained models is not as easily accomplishedfor adaptation methods which require complicated or specialized training procedures and modelarchitectures, such as TTT [44] or ARM [52].
In our experiments, we use AugMix as our set ofaugmentations [13], as it satisﬁes the above properties and still yields signiﬁcant diversity whenapplied, as depicted in Figure 2.
Note that AugMix explicitly does not use augmentations that aresimilar to the corruptions in the CIFAR-10-C and ImageNet-C test sets [12, 13].
Adapting BN statistics.
Schneider et al.
[40] showed that, even when presented with just asingle test point, partially adapting the estimated mean and variance of the activations in eachbatch normalization (BN) layer of the model can still be effective in some cases for handlingdistribution shift.
In this setting, to prevent overﬁtting to the test point, the channelwise mean andvariance [ µ test ,  2 test ] estimated from this point are mixed with the the mean and variance [ µ train ,  2 train ]computed during training according to a prior strength N .
That is, for ⌫ 2 { µ ,  2 } ,
[Equation]
This technique is also straightforward to combine with MEMO: we simply use the adapted BNstatistics whenever computing the model’s output distribution.
We ﬁnd in our experiments that thistechnique never degrades, and generally improves, the performance of test time adaptation, thus wecombine MEMO with this technique by default whenever applicable.
Following the suggestion in Schneider et al.
[40], we set N = 16 for all of our experiments in the next section.
Experiments
Our experiments aim to answer the following questions:
(1) How does MEMO compare to prior methods for test time adaptation and test time robustness?(2)
Can MEMO be combined with a wide range of model architectures and pretraining methods?(3)
Which aspect of MEMO, the adaptation or augmentation, is the most important?
We evaluate MEMO on a total of ﬁve distribution shift benchmarks.
We conduct CIFAR-10 [22]experiments on the CIFAR-10-C [12] and CIFAR-10.1 [37] test sets, and we conduct ImageNet [38]experiments on the ImageNet-C [12], ImageNet-R [14], and ImageNet-A [15] test sets.
To answer question (1), we compare to test time training (TTT)
[44] in the CIFAR-10 experiments,for which we train ResNet-26 models following their protocol and specialized architecture.
We donot compare to TTT for the ImageNet experiments due to the computational demands of trainingstate-of-the-art models and because Sun et al.
[44] do not report competitive ImageNet results.
Forthe ImageNet experiments, we compare to Tent [46] and BN adaptation, which can be used withpretrained models but require multiple test inputs (or even the entire test set) for adaptation.
Weprovide BN adaptation with 256 test inputs at a time and set the prior strength N = 256 [40].
For Tent, we use test batch sizes of 64 and, for ResNet-50 models, test both “online” adaptation– where the model adapts continually through the entire evaluation – and “episodic” adaptation –where the model is reset after each test batch [46].
Note that the evaluation protocols are different forthese two methods: whereas MEMO is tasked with predicting on each test point immediately afteradaptation, BN adaptation predicts on a batch of 256 test points after computing BN statistics onthe batch, and Tent predicts on a batch of 64 inputs after adaptation but also, in the online setting,continually adapts throughout evaluation.
In all experiments, we further compare to single point BNadaptation [40] and the TTA baseline that simply predicts according to ¯ p ✓ ( y | x ) (Eq. 1) [23, 2].
Fulldetails on our experimental protocol are provided in Appendix A.
To answer question (2), we apply MEMO on top of multiple pretrained models with differentarchitectures, trained via several different procedures.
For CIFAR-10, we train our own ResNet26 [11] models.
For ImageNet, we use the best performing ResNet-50 robust models from priorwork, which includes those trained with DeepAugment and AugMix augmentations [14] as well asthose trained with moment exchange and CutMix [24].
To evaluate the generality of prior test timerobustness methods and MEMO, we also evaluate the small robust vision transformer (RVT ⇤ -small),which provides superior performance on all three ImageNet distribution shift benchmarks compared tothe robust ResNet-50 models [30].
Finally, we evaluate ResNext-101 models [50, 28] on ImageNet-A,as these models previously achieved the strongest results for this test set [14].
Finally, to answer (3), we conduct ablative studies in subsection 4.2: ﬁrst to determine the relative importance of maximizing conﬁdence (via entropy minimization) versus enforcing invariant predictionsacross augmented copies of each test point, and second to determine the importance of the particularaugmentation functions used.
The comparison to the non adaptive TTA baseline also helps determinewhether simply augmenting the test point is sufﬁcient or if adaptation is additionally helpful.
In
Appendix B, we provide further experiments ablating the augmentation component speciﬁcally, alongwith experiments that vary the number of augmented samples per inference.
4.1 Main Results
We summarize results for CIFAR-10, CIFAR-10.1, and CIFAR-10-C in Table 1, with full CIFAR-10C results in Appendix C. We use indentations to indicate composition, e.g., TTT is performed at testtime on top of their specialized joint training procedure.
Across all corruption types in CIFAR-10-C,
MEMO consistently improves test error compared to the baselines, non adaptive TTA, and TTT.
MEMO also provides a larger performance gain on CIFAR-10.1 compared to TTT.
We ﬁnd that thenon adaptive TTA baseline is competitive for these relatively simple test sets, though it is worsethan MEMO for CIFAR-10-C. Of these three test sets, CIFAR-10-C is the only benchmark thatexplicitly introduces distribution shift, which suggests that adaptation is useful when the test shifts aremore prominent.
Both TTA and MEMO are also effective at improving performance for the original
CIFAR-10 test set where there is no distribution shift, providing further support for the widespreaduse of augmentations in standard evaluation protocols [23, 2].
[Table] [Table]
[Table]
4 (  22 . 3) 57 .
7 (  6 . 2) 99 .
(  0 . 2) + Tent (episodic) 64 .
7 (  12 . 0) 61 .
(  2 . 9) 99 .
7 (  0 . 3) + DeepAugment+AugMix {+[14] 53 .
+}
.+
TTA 55 .
2 (+1 . 6) 51 .
(  2 . 2) 93 .
5 (  2 . 6) + Single point BN 51 .
3 (  2 . 3) 51 .
2 (  2 . 0) 95 .
4 (  0 .
+ MEMO (ours) 49 .
8 (  3 . 8 ) 49 .
[Table]
4 . 0 ) 94 .
8 (  1 . 3) + BN ( N = 256 , n = 256 ) 45 .
[Table]
[Table] [Table]
[Table]
[Table]
+
[Table]
+ Tent (online) 43 .
5 (
[Table]
[Table]
[Table]
[Table]
+ Tent (episodic) 47 .
1 (  6 .
[Table] [Table]
[Table]
[Table]
+ MoEx+CutMix [24] 74 {+.
+}
.91 .+ TTA 75 .
[Table]
+0 . 9) 62 .
7 (
[Table]
5 (  2 . 4) + Single point BN 71 . {+
+}
(  3 .
6 (  1 . 9) 91 .
(  0 .
8) + MEMO
(ours) 69 .
1 (  5 . 7) 59 .
4 (  {+3 . 3) 89 . 
+}
(  2 . 9 ) + BN ( N = 256 , n = 256 ) 60 .
(  13 . 9) 61 .
6 (  2 . 9) 93 .
9 (+2 . 0)
+ Tent (online)
[Table] [Table]
[Table]
7 (
[Table]
We summarize results for ImageNet-C, ImageNet-R, and ImageNet-A in Table 2, with complete
ImageNet-C results in Appendix C. We again use indentations to indicate composition, e.g., thebest results on ImageNet-C for our setting are attained through a combination of starting from amodel trained with DeepAugment and AugMix [14] and using MEMO on top.
For both ImageNet-Cand ImageNet-R, and for both the ResNet-50 and RVT ⇤ -small models, combining MEMO withrobust training techniques leads to new state-of-the-art performance among methods that observeonly one test point at a time.
We highlight in gray the methods that require multiple test points foradaptation, and we list in bold the best results from these methods which outperform the test timerobustness methods.
As Table 2 and prior work both show [40, 46], accessing multiple test pointscan be powerful for benchmarks such as ImageNet-C and ImageNet-R, in which inferred statistics from the test input distribution may aid in prediction.
However, these methods do not help, andoftentimes even hurt, for ImageNet-A.
Furthermore, we ﬁnd that these methods are less effective withthe RVT ⇤ -small model, which may indicate their sensitivity to model architecture choices.
Therefore,for this model, we also test a modiﬁcation of Tent which adapts all parameters, and we ﬁnd that thisversion of Tent works better for ImageNet-C but is signiﬁcantly worse for ImageNet-R.
MEMO also results in substantial improvement for ImageNet-A.
No prior test time adaptationmethods have reported improvements on ImageNet-A, and some have reported explicit negativeresults [40].
As discussed, it is reasonable for adaptation methods that rely on multiple test points toachieve greater success on other benchmarks such as ImageNet-C, in which a batch of inputs providessigniﬁcant information about the speciﬁc corruption that must be dealt with.
In contrast, ImageNet-Adoes not have such obvious characteristics associated with the input distribution, as it is simply acollection of images that are difﬁcult to classify.
As MEMO instead extracts a learning signal fromsingle test points, it is, to the best of our knowledge, the ﬁrst test time adaptation method to reportsuccessful results on this testbed.
We view the consistency with which MEMO outperforms the bestprior methods, which change across different test sets, as a major advantage of the proposed method.
[Table]
TTA is the most competitive priormethod for ImageNet-A, e.g., it resultsin larger improvements than MEMO forthe
RVT ⇤ -small model.
MEMO, however, achieves state-of-the-art performanceamong ResNet-50 models.
To furthercompare MEMO to TTA, in Table 3, weevaluate whether MEMO can successfullyadapt ResNext-101 models [50] and furtherimprove performance on this challengingtest set.
We evaluate both a ResNext-101(32x8d) baseline model pretrained on ImageNet, as well as the same model pretrainedwith weakly supervised learning (WSL) onbillions of Instagram images [28].
For the
WSL model, we did not use single point BN adaptation for MEMO as we found this technique to beactually harmful to performance, and this corroborates previous ﬁndings [40].
From the results, wecan see that, although both TTA and MEMO signiﬁcantly improve upon the baseline model evaluation,
MEMO ultimately achieves the best accuracy by a signiﬁcant margin as it is more successful atadapting the WSL model.
This suggests that MEMO may synergize well with large scale pretraining,and further exploring this combination is an interesting direction for future work.
4.2 Ablative Study
MEMO uses both adaptation and augmentations.
In this section, we ablate the adaptation procedure,and in Appendix B we ablate the choice of augmentations as well as the number of augmentations.
From the results above, we conclude that adaptation generally provides additional beneﬁts beyondsimply using TTA to predict via the marginal output distribution ¯ p ✓ ( y | x ) .
However, we can disentangle two distinct self-supervised learning signals that may be effective for adaptation: encouraginginvariant predictions across different augmentations of the test point, and encouraging conﬁdence viaentropy minimization.
The marginal entropy objective in Eq.
2 encapsulates both of these learningsignals, but it cannot easily be decomposed into these pieces.
Thus, we instead use two ablativeadaptation methods that each only make use of one of these learning signals.
First, we consider optimizing the pairwise cross entropy between each pair of augmented points, i.e.,` PCE ( ✓ ; x ) ,
[Equation]
Where ˜ x i again refers to the i -th sampled augmentation applied to x .
Intuitively, this loss functionencourages the model to adapt such that it produces the same predictive distribution for all augmentations of the test point, but it does not encourage the model to produce conﬁdent predictions.
Conversely, as an objective that encourages conﬁdence but not invariance, we also consider optimizing the conditional entropy objective detailed in Eq.
This ablation is effectively a version of
[Table]
the episodic variant of Tent [46] that produces augmented copies of a single test point rather thanassuming access to a test batch.
We ﬁrst evaluate these ablations on the CIFAR-10 test sets.
We usethe same adaptation procedure and hyperparameters, with ` replaced with the above objectives.
The results are presented in Table 4.
We see that MEMO, i.e., marginal entropy minimization,generally performs better than adaptation with either of the alternative objectives.
This supports thehypothesis that both invariance across, and conﬁdence on, the augmentations are important learningsignals for self-supervised adaptation.
When faced with CIFAR-10.1, we see poor performance fromthe pairwise cross entropy based adaptation method.
On the original CIFAR-10 test set and CIFAR-10C, the ablations perform nearly identically and uniformly worse than MEMO.
To further test the ` CEablation, which is the stronger of the two ablations, we also evaluate it on the ImageNet test sets forthe RVT ⇤ -small model.
We ﬁnd that, similarly, minimizing conditional entropy generally improvesperformance compared to the baseline evaluation.
MEMO is more performant for ImageNet-C and
ImageNet-R. Adaptation via ` CE performs slightly better for ImageNet-A, though for this problemand model, TTA is still the best method.
Thus, MEMO results in relatively small, but consistent,performance gains compared to only maximizing conﬁdence on the augmentations.
Discussion
We presented MEMO, a method for test time robustiﬁcation again distribution shift via adaptationand augmentation.
MEMO does not require access or changes to the model training procedure and isthus broadly applicable for a wide range of model architectures pretrained in a number of differentways.
Furthermore, MEMO adapts at test time using single test inputs, thus it does not assume accessto multiple test points as in several recent methods for test time adaptation [40, 46].
On a range of
CIFAR-10 and ImageNet distribution shift benchmarks, and for ResNet, vision transformer, and, toan extent, ResNext models, MEMO consistently improves performance at test time and achievesseveral new state-of-the-art results for these models in the single test point setting.
Inference via MEMO is more computationally expensive than standard model inference due to itsaugmentation and adaptation procedure – though, as Appendix B shows, more favorable tradeoffsbetween efﬁciency and accuracy are possible with smaller values of B , the number of augmentationsper test point.
One interesting direction for future work is to develop techniques for selectivelydetermining when to adapt the model in order to achieve more efﬁcient inference.
For example,with well calibrated models [10], we may run simple “feedforward” inference when the predictionconﬁdence is over a certain threshold, thus achieving better efﬁciency.
Additionally, it would beinteresting to explore MEMO in the test setting where the model is allowed to continually adapt asmore test data is observed.
In our preliminary experiments in this setting, MEMO tended to lead todegenerate solutions, e.g., the model predicting a constant label with maximal conﬁdence regardlessof the input.
This failure mode may potentially be rectiﬁed by carefully choosing which parametersto adapt, such as only adapting the parameters in BN layers [46], or regularizing the model such thatit does not change too drastically from the pretrained model [26].