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: The paper proposes a new loss to improve test-time BN adaptation for domain adaptation. The proposed loss consists of two components: the diversity maximization loss and the confidence maximization loss. Specifically, they use a running estimate for the diversity loss based on KL divergence. They propose the hard and soft likelihood ratio for the confidence loss which has large gradients for high confidence predictions.