[Paper Review] Test-Time Adaptation to Distribution Shift by Confidence Maximization and Input Transformation

Chaithanya Kumar Mummadi, Robin Hutmacher+4 more|arXiv (Cornell University)|Jun 28, 2021

Anomaly Detection Techniques and ApplicationsComputer Science39 references40 citations

TL;DR

The paper proposes fully test-time adaptation with non-saturating likelihood-ratio losses (HLR/SLR), a diversity regularizer with moving average, and a trainable input transformation module to improve robustness to distribution shifts without target labels.

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 the entropy by a non-saturating surrogate and adding a diversity regularizer based on batch-wise entropy maximization that prevents convergence to trivial collapsed solutions. Moreover, we propose to prepend an input transformation module to the network that can partially undo test-time distribution shifts. Surprisingly, this preprocessing can be learned solely using the fully test-time adaptation loss in an end-to-end fashion without any target domain labels or source domain data. We show that our approach outperforms previous work in improving the robustness of publicly available pretrained image classifiers to common corruptions on such challenging benchmarks as ImageNet-C.

Motivation & Objective

• Motivate robust performance of pretrained models under distribution shift without source/target labels.
• Introduce non-saturating regression-like losses to sustain learning from high-confidence predictions.
• Prevent collapsed solutions with a moving-average diversity regularizer over batches.
• Enhance adaptation by prepending a trainable input transformation to partially undo shifts.
• Demonstrate improved robustness on ImageNet-C and ImageNet-R across multiple pretrained backbones.

Proposed method

• Prepend a trainable input transformation module to the pretrained network to form g = f ∘ d, where d partially counteracts domain shift.
• Adapt only a subset of parameters (e.g., affine BN parameters) and update BN statistics on target data during test-time.
• Use two non-saturating, likelihood-ratio-based losses L_hlr and L_slr for self-supervised adaptation, avoiding vanishing gradients at high confidence.
• Incorporate a diversity regularizer L_div with a running-average update p_t(y) to prevent collapse to trivial predictions.
• Combine L_div with L_conf (non-saturating variants) to drive informative adaptation while maintaining diversity.

Experimental results

Research questions

• RQ1 Can fully test-time adaptation improve accuracy on corrupted or shifted distributions without source data?
• RQ2 Do non-saturating likelihood-ratio based losses (HLR/SLR) provide better gradient signals than entropy-based losses for high-confidence predictions during adaptation?
• RQ3 Does a moving-average diversity regularizer stabilize adaptation and prevent collapse?
• RQ4 Can an input transformation module learned at test-time further counteract distribution shifts?
• RQ5 How do the proposed methods perform across multiple pretrained architectures on ImageNet-C and ImageNet-R?

Key findings

• HLR and SLR outperform TENT and TENT+ across several pretrained models on ImageNet-C and ImageNet-R.
• Adaptation with the proposed losses improves mean corruption accuracy, notably when using robust backbones (e.g., DeepAugment+AugMix).
• A single epoch of test-time adaptation already yields substantial gains, with further gains up to epoch 5, while L_div stabilizes training.
• The input transformation (IT) module partially undoing shifts increases robustness for several distortions (e.g., Impulse, Contrast).
• Ground-truth supervised adaptation yields upper bounds, while the proposed method approaches these bounds in several settings.