Test-Time Adaptation to Distribution Shifts by Confidence Maximization and Input Transformation

1 INTRODUCTION . Deep neural networks achieve impressive performance on test data , which has the same distribution as the training data . Nevertheless , they often exhibit a large performance drop on test ( target ) data which differs from training ( source ) data ; this effect is known as data shift ( Quionero-Candela et al. , 2009 ) and can be caused for instance by image corruptions . There exist different methods to improve the robustness of the model during training ( Geirhos et al. , 2019 ; Hendrycks et al. , 2019 ; Tzeng et al. , 2017 ) . However , generalization to different data shifts is limited since it is infeasible to include sufficiently many augmentations during training to cover the excessively wide range of potential data shifts ( Mintun et al. , 2021a ) . Alternatively , in order to generalize to the data shift at hand , the model can be adapted during test-time . Unsupervised domain adaptation methods such as Vu et al . ( 2019 ) use both source and target data to improve the model performance during test-time . In general source data might not be available during inference time , e.g. , due to legal constraints ( privacy or profit ) . Therefore we focus on the fully test-time adaptation setting ( Wang et al. , 2020 ) : model is adapted to the target data during test time given only the arbitrarily pretrained model parameters and unlabeled target data that share the same label space as source data . We extend the work of Wang et al . ( 2020 ) by introducing a novel loss function , using a diversity regularizer , and prepending a parametrized input transformation module to the network . We show that our approach outperform previous works and make pretrained models robust against common corruptions on image classification benchmarks as ImageNet-C ( Hendrycks & Dietterich , 2019 ) and ImageNet-R ( Hendrycks et al. , 2020 ) . Sun et al . ( 2020 ) investigate test-time adaptation using a self-supervision task . Wang et al . ( 2020 ) and Liang et al . ( 2020 ) use the entropy minimization loss that uses maximization of prediction confidence as self-supervision signal during test-time adaptation . Wang et al . ( 2020 ) has shown that such loss performs better adaptation than a proxy task ( Sun et al. , 2020 ) . When using entropy minimization , however , high confidence predictions do not contribute to the loss significantly anymore and thus provide little self-supervision . This is a drawback since high-confidence samples provide the most trustworthy self-supervision . We mitigate this by introducing two novel loss functions that ensure that gradients of samples with high confidence predictions do not vanish and learning based on self-supervision from these samples continues . Our losses do not focus on minimizing entropy but on minimizing the negative log likelihood ratio between classes ; the two variants differ in using either soft or hard pseudo-labels . In contrast to entropy minimization , the proposed loss functions provide non-saturating gradients , even when there are high confident predictions . Figure 1 provides illustration of the losses and the resulting gradients . Using these new loss functions , we are able to improve the network performance under data shifts in both online and offline adaptation settings . In general , self-supervision by confidence maximization can lead to collapsed trivial solutions , which make the network to predict only a single or a set of classes independent of the input . To overcome this issue a diversity regularizer ( Liang et al. , 2020 ; Wu et al. , 2020 ) can be used , that acts on a batch of samples . It encourages the network to make diverse class predictions on different samples . We extend the regularizer by including a moving average , in order to include the history of the previous batches and show that this stabilizes the adaptation of the network to unlabeled test samples . Furthermore we also introduce a parametrized input transformation module , which we prepend to the network . The module is trained in a fully test-time adaptation manner using the proposed loss function , and without using source data or target labels . It aims to partially undo the data shift at hand and helps to further improve the performance on image classification benchmark with corruptions . Since our method does not change the training process , it allows to use any pretrained models . This is beneficial because any good performing pretrained network can be readily reused , e.g. , a network trained on some proprietary data not available to the public . We show , that our method significantly improves performance of different pretrained models that are trained on clean ImageNet data . In summary our main contributions are as follows : we propose non-saturating losses based on the negative log likelihood ratio , such that gradients from high confidence predictions still contribute to test-time adaptation . We extend diversity regularizer to its moving average to include the history of previous batch samples to prevent the model collapsing to trivial solutions . We also introduce an input transformation module , which partially undoes the data shift at hand . We show that the performance of different pretrained models can be significantly improved on ImageNet-C and ImageNet-R . 2 RELATED WORK . Common image corruptions are potentially stochastic image transformations motivated by realworld effects that can be used for evaluating a model ’ s robustness . One such benchmark , ImageNetC ( Hendrycks & Dietterich , 2019 ) , contains simulated corruptions such as noise , blur , weather effects , and digital image transformations . Additionally , Hendrycks et al . ( 2020 ) proposed three data sets containing real-world distribution shifts , including Imagenet-R . Most proposals for improving robustness involve special training protocols , requiring time and additional resources . This includes data augmentation like Gaussian noise ( Ford et al. , 2019 ; Lopes et al. , 2019 ; Hendrycks et al. , 2020 ) , CutMix ( Yun et al. , 2019 ) , AugMix ( Hendrycks et al. , 2019 ) , training on stylized images ( Geirhos et al. , 2019 ; Kamann et al. , 2020 ) or against adversarial noise distributions ( Rusak et al. , 2020a ) . Mintun et al . ( 2021b ) pointed out that many improvements on ImageNet-C are due to data augmentations which are too similar to the test corruptions , that is : overfitting to ImageNet-C occurs . Thus , the model might be less robust to corruptions not included in the test set of ImageNet-C. Unsupervised domain adaptation methods train a joint model of source and target domain by crossdomain losses to find more general and robust features , e. g. optimize feature alignment ( QuiñoneroCandela et al. , 2008 ; Sun et al. , 2017 ) between domains , adversarial invariance ( Ganin & Lempitsky , 2015 ; Tzeng et al. , 2017 ; Ganin et al. , 2016 ; Hoffman et al. , 2018 ) , shared proxy tasks ( Sun et al. , 2019 ) or adapt entropy minimization via an adversarial loss ( Vu et al. , 2019 ) . While these approaches are effective , they require explicit access to source and target data at the same time , which may not always be feasible . Our approach works with any pretrained model and only needs target data . Test-time adaptation is a setting , when training ( source ) data is unavailable at test-time . It is related to source free adaptation , where several works use generative models , alter training ( Kundu et al. , 2020 ; Li et al. , 2020b ; Kurmi et al. , 2021 ; Yeh et al. , 2021 ) and require several thousand epochs to adapt to the target data ( Li et al. , 2020b ; Yeh et al. , 2021 ) . Besides , there is another line of work ( Sun et al. , 2020 ; Schneider et al. , 2020 ; Nado et al. , 2021 ; Benz et al. , 2021 ; Wang et al. , 2020 ) that interpret the common corruptions as data shift and aim to improve the model robustness against these corruptions with efficient test-time adaptation strategy to facilitate online adaptation . such settings spare the cost of additional computational overhead . Our work also falls in this line of research and aims to adapt the model to common corruptions efficiently with both online and offline adaptation . Sun et al . ( 2020 ) update feature extractor parameters at test-time via a self-supervised proxy task ( predicting image rotations ) . However , Sun et al . ( 2020 ) alter the training procedure by including the proxy loss into the optimization objective as well , hence arbitrary pretrained models can not be used directly for test-time adaptation . Inspired by the domain adaptation strategies ( Maria Carlucci et al. , 2017 ; Li et al. , 2016 ) , several works ( Schneider et al. , 2020 ; Nado et al. , 2021 ; Benz et al. , 2021 ) replace the estimates of Batch Normalization ( BN ) activation statistics with the statistics of the corrupted test images . Fully test time adaptation , studied by Wang et al . ( 2020 ) ( TENT ) uses entropy minimization to update the channel-wise affine parameters of BN layers on corrupted data along with the batch statistics estimates . SHOT ( Liang et al. , 2020 ) also uses entropy minimization and a diversity regularizer to avoid collapsed solutions . SHOT modifies the model from the standard setting by adopting weight normalization at the fully connected classifier layer during training to facilitate their pseudo labeling technique . Hence , SHOT is not readily applicable to arbitrary pretrained models . We show that pure entropy minimization ( Wang et al. , 2020 ; Liang et al. , 2020 ) as well as alternatives such as max square loss ( Chen et al. , 2019 ) and Charbonnier penalty ( Yang & Soatto , 2020 ) results in vanishing gradients for high confidence predictions , thus inhibiting learning . Our work addresses this issue by proposing a novel non-saturating loss , that provides non-vanishing gradients for high confidence predictions . We show that our proposed loss function improves the network performance through test-time adaptation . In particular , performance on corruptions of higher severity improves significantly . Furthermore , we add and extend the diversity regularizer ( Liang et al. , 2020 ; Wu et al. , 2020 ) to avoid collapse to trivial , high confidence solutions . Existing diversity regularizers ( Liang et al. , 2020 ; Wu et al. , 2020 ) act on a batch of samples , hence the number of classes has to be smaller than the batch size . We mitigate this problem by extending the regularizer to a moving average version . Li et al . ( 2020a ) also use a moving average to estimate the entropy of the unconditional class distribution but source data is used to estimate the gradient of the entropy . In contrast , our work does not need access to the source data since the gradient is estimated using only target data . Prior work Tzeng et al . ( 2017 ) ; Rusak et al . ( 2020b ) ; Talebi & Milanfar ( 2021 ) transformed inputs by an additional module to overcome domain shift , obtain robust models , and also to learn to resize . In our work , we prepend an input transformation module to the model , but in contrast to former works , this module is trained purely at test-time to partially undo the data shift at hand to aid the adaptation .

Summary: Test-time adaptation by entropy minimization can help models adapt to dataset shifts like corruptions without altering training. This work extends tent, an entropy minimization method, by proposing alternative non-saturating losses, adding a diversity regularizer, and adapting the input data along with the model parameters. The input is adapted by applying a convolutional image transformation model between the input and classification model. These extensions do not need more optimization iterations or supervision than the baselines: the method adapts online and efficiently without auxiliary supervision. Experiments on the corruption benchmark ImageNet-C and the newer benchmark ImageNet-R report reduced generalization error. The improvements are there but marginal, and they are consistent across multiple baseline architectures (ResNet, DenseNet, MobileNet, etc.). However, the clean accuracy reduced, so the proposed method does not strictly dominate prior work.