MEMO: Test Time Robustness via Adaptation and Augmentation
Authors: Marvin Zhang, Sergey Levine, Chelsea Finn
Abstract
While deep neural networks can attain good accuracy on in-distribution test points, many applications require robustness even in the face of unexpected perturbations in the input, changes in the domain, or other sources of distribution shift. We study the problem of test time robustification, i.e., using the test input to improve model robustness. 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 and devise 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's average, or marginal, output distribution across the augmentations. Intuitively, this objective encourages the model to make the same prediction across different augmentations, thus enforcing the invariances encoded in these augmentations, while also maintaining confidence in its predictions. In our experiments, we demonstrate that this approach consistently improves robust ResNet and vision transformer models, achieving accuracy gains of 1-8% over standard model evaluation and also generally outperforming prior augmentation and adaptation strategies. We achieve state-of-the-art results for test shifts caused by image corruptions (ImageNet-C), renditions of common objects (ImageNet-R), and, among ResNet-50 models, adversarially chosen natural examples (ImageNet-A).
The papers around it, by date and closeness
Higher is closer to this paper; left is older. Lines join papers close to each other: groups of lines are sub-topics. Select a point to see why it is there. Back to the summary of this paper
Papers around this one
Ranked by how close their title and abstract are to this paper’s. Select some to read them together.
Closest papers
-
Test time Adaptation through Perturbation RobustnessPrabhu Teja Sivaprasad, François Fleuret · 2021 · arXiv.org · preprint, not peer reviewed · cited 47 times · Not summarized yet
Test time Adaptation through Perturbation Robustness enforces prediction consistency around each test image to handle distribution shift, as MEMO does through augmented views and adaptation.
Abstract
Data samples generated by several real world processes are dynamic in nature \textit{i.e.}, their characteristics vary with time. Thus it is not possible to train and tackle all possible distributional shifts between training and inference, using the host of transfer learning methods in literature. In this paper, we tackle this problem of adapting to domain shift at inference time \textit{i.e.}, we do not change the training process, but quickly adapt the model at test-time to handle any domain shift. For this, we propose to enforce consistency of predictions of data sampled in the vicinity of test sample on the image manifold. On a host of test scenarios like dealing with corruptions (CIFAR…
-
Learning Robustness at Test-Time from a Non-Robust TeacherStefano Bianchettin, Giulio Rossolini, Giorgio C. Buttazzo · 2026 · Image and Vision Computing · Not summarized yet
Learning Robustness at Test-Time from a Non-Robust Teacher studies adversarially robust test-time adaptation and cites MEMO as prior work.
Abstract
Nowadays, pretrained models are increasingly used as general-purpose backbones and adapted at test-time to downstream environments where target data are scarce and unlabeled. While this paradigm has proven effective for improving clean accuracy on the target domain, adversarial robustness has received far less attention, especially when the original pretrained model is not explicitly designed to be robust. This raises a practical question: \emph{can a pretrained, non-robust model be adapted at test-time to improve adversarial robustness on a target distribution?} To face this question, this work studies how adversarial training strategies behave when integrated into adaptation schemes for th…
-
Test-Time Adaptation to Distribution Shift by Confidence Maximization and Input TransformationChaithanya Kumar Mummadi, Robin Hutmacher, K. Rambach et al. · 2021 · arXiv.org · preprint, not peer reviewed · cited 87 times · Not summarized yet
Test-Time Adaptation to Distribution Shift by Confidence Maximization and Input Transformation improves test-time adaptation with a confidence-based loss and input transformations.
Abstract
Deep neural networks often exhibit poor performance on data that is unlikely under the train-time data distribution, for instance data affected by corruptions. Previous works demonstrate that test-time adaptation to data shift, for instance using entropy minimization, effectively improves performance on such shifted distributions. This paper focuses on the fully test-time adaptation setting, where only unlabeled data from the target distribution is required. This allows adapting arbitrary pretrained networks. Specifically, we propose a novel loss that improves test-time adaptation by addressing both premature convergence and instability of entropy minimization. This is achieved by replacing …
-
Approaching Test Time Augmentation in the Context of Uncertainty Calibration for Deep Neural NetworksPedro Conde, Tiago Barros, Rui L. Lopes et al. · 2023 · Computer Vision and Pattern Recognition · preprint, not peer reviewed · cited 9 times · AI summary available
Approaching Test Time Augmentation in the Context of Uncertainty Calibration uses augmented test views to improve uncertainty calibration, rather than MEMO’s robustness through model adaptation.
Abstract
With the rise of Deep Neural Networks, machine learning systems are nowadays ubiquitous in a number of real-world applications, which bears the need for highly reliable models. This requires a thorough look not only at the accuracy of such systems, but also to their predictive uncertainty. Hence, we propose a novel technique (with two different variations, named M-ATTA and V-ATTA) based on test time augmentation, to improve the uncertainty calibration of deep models for image classification. Unlike other test time augmentation approaches, M/V-ATTA improves uncertainty calibration without affecting the model's accuracy, by leveraging an adaptive weighting system. We evaluate the performance o…
-
Fully Test-time Adaptation by Entropy MinimizationDequan Wang, Evan Shelhamer, Shaoteng Liu et al. · 2020 · Machine Learning · preprint, not peer reviewed · cited 71 times · AI summary available
Fully Test-time Adaptation by Entropy Minimization adapts model parameters by minimizing prediction entropy on unlabeled shifted test data.
Abstract
Faced with new and different data during testing, a model must adapt itself. We consider the setting of fully test-time adaptation, in which a supervised model confronts unlabeled test data from a different distribution, without the help of its labeled training data. We propose an entropy minimization approach for adaptation: we take the model's confidence as our objective as measured by the entropy of its predictions. During testing, we adapt the model by modulating its representation with affine transformations to minimize entropy. Our experiments show improved robustness to corruptions for image classification on CIFAR-10/100 and ILSVRC and demonstrate the feasibility of target-only domai…
-
Parameter-free Online Test-time AdaptationMalik Boudiaf, Romain Mueller, Ismail Ben Ayed et al. · 2022 · Computer Vision and Pattern Recognition · cited 283 times · AI summary available
Parameter-free Online Test-time Adaptation evaluates adaptation methods across pretrained models and scenarios, addressing the practical test-time setting shared with MEMO.
Abstract
Training state-of-the-art vision models has become prohibitively expensive for researchers and practitioners. For the sake of accessibility and resource reuse, it is important to focus on adapting these models to a variety of downstream scenarios. An interesting and practical paradigm is online test-time adaptation, according to which training data is inaccessible, no labelled data from the test distribution is available, and adaptation can only happen at test time and on a handful of samples. In this paper, we investigate how test-time adaptation methods fare for a number of pre-trained models on a variety of real-world scenarios, significantly extending the way they have been originally ev…
-
PatAug: Augmentation of Augmentation for Test-Time AdaptationXin-Yao Li, Dan Zhang, Zhe-Kai Du et al. · 2025 · ACM Multimedia · cited 6 times · Not summarized yet
PatAug develops augmentation strategies for vision-language-model test-time adaptation and cites MEMO as related augmentation-based adaptation work.
Abstract
The rich pretrained knowledge in vision-language models (VLMs) endows them with the ability to discriminate common objects given only category names, but may be challenged by out-of-distribution unlabeled samples. To address this limitation, test-time adaptation (TTA) dynamically adjusts VLMs to target distributions during inference. Current TTA frameworks rely heavily on unsupervised data augmentations to enhance sample informativeness, but remain vulnerable to naive augmented views. This work introduces Patch Augmentation (PatAug), a pixel-level perturbation framework that optimizes the benefits of informative augmentations and mitigates negative transformation impacts. Implemented as trai…
-
Continual Test-Time Adaptation in Computer Vision: Methods, Benchmarks, and Future DirectionsSarthak Kumar Maharana, Shambhavi Mishra, Yun-Bei Zhang et al. · 2026 · Trans. Mach. Learn. Res. · cited 2 times · Not summarized yet
Continual Test-Time Adaptation in Computer Vision surveys methods and benchmarks for adapting to ongoing distribution shifts, including work citing MEMO.
Abstract
Deep neural nets achieve remarkable performance when training and test data share the same distribution, but this assumption frequently breaks in real-world deployment, where data undergoes continual distributional shifts. Continual Test-Time Adaptation (CTTA) addresses this challenge by adapting pretrained models to non-stationary target distributions on-the-fly, without access to source data or labeled targets, while mitigating two critical failure modes: catastrophic forgetting of source knowledge and error accumulation from noisy pseudo-labels over extended time horizons. In this comprehensive survey, we formally define the CTTA problem, analyze the diverse continual domain shift pattern…
Prior work it builds on
-
The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution GeneralizationDan Hendrycks, Steven Basart, Norman Mu et al. · 2020 · IEEE International Conference on Computer Vision · cited 2,632 times · Not summarized yet
The Many Faces of Robustness introduces real-world distribution-shift datasets and evaluates augmentation-based robustness methods used as context for MEMO.
Abstract
We introduce four new real-world distribution shift datasets consisting of changes in image style, image blurriness, geographic location, camera operation, and more. With our new datasets, we take stock of previously proposed methods for improving out-of-distribution robustness and put them to the test. We find that using larger models and artificial data augmentations can improve robustness on real-world distribution shifts, contrary to claims in prior work. We find improvements in artificial robustness benchmarks can transfer to real-world distribution shifts, contrary to claims in prior work. Motivated by our observation that data augmentations can help with real-world distribution shifts…
-
Benchmarking Neural Network Robustness to Common Corruptions and PerturbationsDan Hendrycks, Thomas Dietterich · 2019 · Machine Learning · cited 5,008 times · AI summary available
Benchmarking Neural Network Robustness to Common Corruptions and Perturbations establishes ImageNet-C, a corruption benchmark on which MEMO reports results.
Abstract
In this paper we establish rigorous benchmarks for image classifier robustness. Our first benchmark, ImageNet-C, standardizes and expands the corruption robustness topic, while showing which classifiers are preferable in safety-critical applications. Then we propose a new dataset called ImageNet-P which enables researchers to benchmark a classifier's robustness to common perturbations. Unlike recent robustness research, this benchmark evaluates performance on common corruptions and perturbations not worst-case adversarial perturbations. We find that there are negligible changes in relative corruption robustness from AlexNet classifiers to ResNet classifiers. Afterward we discover ways to enh…
-
Towards Robust Vision TransformerXiaofeng Mao, Ge-Ge Qi, Yue-Feng Chen et al. · 2021 · Computer Vision and Pattern Recognition · cited 263 times · Not summarized yet
Towards Robust Vision Transformer studies transformer robustness to adversarial examples, corruptions, and shifts, providing robust-model context for MEMO.
Abstract
Recent advances on Vision Transformer (ViT) and its improved variants have shown that self-attention-based networks surpass traditional Convolutional Neural Networks (CNNs) in most vision tasks. However, existing ViTs focus on the standard accuracy and computation cost, lacking the investigation of the intrinsic influence on model robustness and generalization. In this work, we conduct systematic evaluation on components of ViTs in terms of their impact on robustness to adversarial examples, common corruptions and distribution shifts. We find some components can be harmful to robustness. By leveraging robust components as building blocks of ViTs, we propose Robust Vision Transformer (RVT), w…
-
A Fourier Perspective on Model Robustness in Computer VisionDong Yin, Raphael Gontijo Lopes, Jonathon Shlens et al. · 2019 · Neural Information Processing Systems · cited 694 times · Not summarized yet
A Fourier Perspective on Model Robustness analyzes corruption-specific trade-offs from data augmentation and adversarial training, a neighboring robustness question.
Abstract
Achieving robustness to distributional shift is a longstanding and challenging goal of computer vision. Data augmentation is a commonly used approach for improving robustness, however robustness gains are typically not uniform across corruption types. Indeed increasing performance in the presence of random noise is often met with reduced performance on other corruptions such as contrast change. Understanding when and why these sorts of trade-offs occur is a crucial step towards mitigating them. Towards this end, we investigate recently observed trade-offs caused by Gaussian data augmentation and adversarial training. We find that both methods improve robustness to corruptions that are concen…
-
MT3: Meta Test-Time Training for Self-Supervised Test-Time AdaptionAlexander Bartler, Andre Bühler, Felix Wiewel et al. · 2021 · International Conference on Artificial Intelligence and Statistics · cited 107 times · Not summarized yet
MT3 learns to adapt to unseen test distributions using meta-learning and self-supervision, addressing test-time adaptation through a different training strategy.
Abstract
An unresolved problem in Deep Learning is the ability of neural networks to cope with domain shifts during test-time, imposed by commonly fixing network parameters after training. Our proposed method Meta Test-Time Training (MT3), however, breaks this paradigm and enables adaption at test-time. We combine meta-learning, self-supervision and test-time training to learn to adapt to unseen test distributions. By minimizing the self-supervised loss, we learn task-specific model parameters for different tasks. A meta-model is optimized such that its adaption to the different task-specific models leads to higher performance on those tasks. During test-time a single unlabeled image is sufficient to…
Papers that cite it
-
R-TPT: Improving Adversarial Robustness of Vision-Language Models through Test-Time Prompt TuningLi-Jun Sheng, Jian Liang, Zi-Lei Wang et al. · 2025 · Computer Vision and Pattern Recognition · cited 51 times · Not summarized yet
R-TPT adapts vision-language models at test time for adversarial robustness and cites MEMO among prior test-time adaptation methods.
Abstract
Vision-language models (VLMs), such as CLIP, have gained significant popularity as foundation models, with numerous fine-tuning methods developed to enhance performance on downstream tasks. However, due to their inherent vulnerability and the common practice of selecting from a limited set of open-source models, VLMs suffer from a higher risk of adversarial attacks than traditional vision models. Existing defense techniques typically rely on adversarial fine-tuning during training, which requires labeled data and lacks of flexibility for downstream tasks. To address these limitations, we propose robust test-time prompt tuning (R-TPT), which mitigates the impact of adversarial attacks during …
-
A Comprehensive Survey on Test-Time Adaptation Under Distribution ShiftsJian Liang, R. He, Tien-Ping Tan · 2023 · International Journal of Computer Vision · cited 693 times · Not summarized yet
A Comprehensive Survey on Test-Time Adaptation Under Distribution Shifts reviews adaptation methods for unlabeled test data and cites MEMO.
Abstract
Machine learning methods strive to acquire a robust model during the training process that can effectively generalize to test samples, even in the presence of distribution shifts. However, these methods often suffer from performance degradation due to unknown test distributions. Test-time adaptation (TTA), an emerging paradigm, has the potential to adapt a pre-trained model to unlabeled data during testing, before making predictions. Recent progress in this paradigm has highlighted the significant benefits of using unlabeled data to train self-adapted models prior to inference. In this survey, we categorize TTA into several distinct groups based on the form of test data, namely, test-time do…
-
Robust Test-Time Adaptation in Dynamic ScenariosLonghui Yuan, Binhui Xie, Shuang Li · 2023 · IEEE Computer Society Conference on Computer Vision and Pattern Recognition/Proceedings - IEEE Computer Society Conferen · cited 135 times · Not summarized yet
Robust Test-Time Adaptation in Dynamic Scenarios develops adaptation for temporally changing test streams and cites MEMO as prior work.
Abstract
Test-time adaptation (TTA) intends to adapt the pretrained model to test distributions with only unlabeled test data streams. Most of the previous TTA methods have achieved great success on simple test data streams such as independently sampled data from single or multiple distributions. However, these attempts may fail in dynamic scenarios of real-world applications like autonomous driving, where the environments gradually change and the test data is sampled correlatively over time. In this work, we explore such practical test data streams to deploy the model on the fly, namely practical test-time adaptation (PTTA). To do so, we elaborate a Robust Test-Time Adaptation (RoTTA) method against…
-
Self-Bootstrapping for Versatile Test-Time AdaptationShuai-Cheng Niu, Guo-Hao Chen, Pei-Lin Zhao et al. · 2025 · International Conference on Machine Learning · cited 15 times · Not summarized yet
Self-Bootstrapping for Versatile Test-Time Adaptation uses prediction consistency between original and deteriorated views and cites MEMO.
Abstract
In this paper, we seek to develop a versatile test-time adaptation (TTA) objective for a variety of tasks - classification and regression across image-, object-, and pixel-level predictions. We achieve this through a self-bootstrapping scheme that optimizes prediction consistency between the test image (as target) and its deteriorated view. The key challenge lies in devising effective augmentations/deteriorations that: i) preserve the image's geometric information, e.g., object sizes and locations, which is crucial for TTA on object/pixel-level tasks, and ii) provide sufficient learning signals for TTA. To this end, we analyze how common distribution shifts affect the image's information pow…
-
Back to the Source: Diffusion-Driven Test-Time AdaptationJin Gao, Jialing Zhang, Xihui Liu et al. · 2022 · Machine Learning · cited 112 times · AI summary available
Back to the Source: Diffusion-Driven Test-Time Adaptation shifts test inputs toward the source domain with diffusion models instead of adapting model parameters.
Abstract
Test-time adaptation harnesses test inputs to improve the accuracy of a model trained on source data when tested on shifted target data. Existing methods update the source model by (re-)training on each target domain. While effective, re-training is sensitive to the amount and order of the data and the hyperparameters for optimization. We instead update the target data, by projecting all test inputs toward the source domain with a generative diffusion model. Our diffusion-driven adaptation method, DDA, shares its models for classification and generation across all domains. Both models are trained on the source domain, then fixed during testing. We augment diffusion with image guidance and se…
Go further
How are these papers found?
We compare the title and abstract of this paper with those of every paper summarized on this site, and with the papers Semantic Scholar recommends, those it cites and those citing it (for a much-cited paper, its most recent citations and, from OpenAlex, the most cited ones). Among the prior work and the follow-ups, the papers that use each other in depth, a method it uses and the much-cited follow-ups come a little higher. The levels (very close, close, related) come from the similarity of the texts, not from a judgment of quality; * marks a paper without an abstract, compared by its title. The sentence under a paper is written by an AI from the two abstracts: check it in the papers. A paper marked not relevant by three readers leaves the list.
Updated on 6 November 2026.