[Paper Review] Training on Test Data with Bayesian Adaptation for Covariate Shift

Aurick Zhou, Sergey Levine|arXiv (Cornell University)|Sep 27, 2021

Domain Adaptation and Few-Shot Learning45 references4 citations

TL;DR

This paper proposes Bayesian Adaptation for Covariate Shift (BACS), a test-time adaptation method that improves both accuracy and uncertainty calibration under distribution shift by formulating a Bayesian model linking unlabeled test data to model parameters. It uses regularized entropy minimization via approximate Bayesian inference, outperforming prior methods on image classification benchmarks with significant gains in accuracy and uncertainty reliability.

ABSTRACT

When faced with distribution shift at test time, deep neural networks often make inaccurate predictions with unreliable uncertainty estimates. While improving the robustness of neural networks is one promising approach to mitigate this issue, an appealing alternate to robustifying networks against all possible test-time shifts is to instead directly adapt them to unlabeled inputs from the particular distribution shift we encounter at test time. However, this poses a challenging question: in the standard Bayesian model for supervised learning, unlabeled inputs are conditionally independent of model parameters when the labels are unobserved, so what can unlabeled data tell us about the model parameters at test-time? In this paper, we derive a Bayesian model that provides for a well-defined relationship between unlabeled inputs under distributional shift and model parameters, and show how approximate inference in this model can be instantiated with a simple regularized entropy minimization procedure at test-time. We evaluate our method on a variety of distribution shifts for image classification, including image corruptions, natural distribution shifts, and domain adaptation settings, and show that our method improves both accuracy and uncertainty estimation.

Motivation & Objective

• To address the challenge of distribution shift in deep learning, where model accuracy and uncertainty calibration degrade under test-time distribution shifts.
• To develop a principled Bayesian framework that enables meaningful use of unlabeled test data for model adaptation, despite conditional independence in standard models.
• To improve both predictive accuracy and uncertainty estimation under covariate shift, avoiding overconfidence from heuristic entropy minimization.
• To provide a practical adaptation method that requires minimal changes to standard training and works without access to training data at test time.
• To demonstrate that Bayesian marginalization during adaptation leads to better generalization and calibration than heuristic approaches.

Proposed method

• Proposes a Bayesian probabilistic model that explicitly links unlabeled test data under distribution shift to model parameters via an assumption of unambiguous labeling.
• Introduces a conditional independence assumption that the test data generation process favors single, unambiguous labels, enabling inference on model parameters.
• Derives an approximate inference procedure that corresponds to regularized entropy minimization, with a prior term from the training posterior to prevent overfitting.
• Uses the model's approximate posterior as a regularizer during test-time adaptation, constraining the network to stay near the initial training solution.
• Applies the method via a simple optimization objective combining entropy minimization and a log-posterior term from the training data, enabling end-to-end adaptation.
• Implements the method with minimal architectural changes, adapting the full network while maintaining calibration through Bayesian regularization.

Experimental results

Research questions

• RQ1 Can unlabeled test data under covariate shift provide meaningful information about model parameters in a principled Bayesian framework?
• RQ2 How can we design a test-time adaptation method that improves both accuracy and uncertainty calibration under distribution shift?
• RQ3 What is the impact of incorporating a training-data-derived posterior as a regularizer during test-time adaptation?
• RQ4 How does Bayesian regularization compare to heuristic adaptation methods like TENT in terms of robustness and calibration?
• RQ5 In what scenarios does the training posterior constraint limit adaptation, and when is it beneficial?

Key findings

• BACS achieves 47.31% top-1 accuracy on ImageNet-R, outperforming all baselines including TENT (45.75%) and vanilla models (40.38%).
• On ImageNet-C, BACS reduces NLL from 3.288 (vanilla) to 2.565 and Brier score from 0.7602 to 0.6625, indicating improved calibration and uncertainty estimation.
• The MAP variant of BACS (without ensembling) outperforms TENT, demonstrating that full-network adaptation with Bayesian regularization is more effective than partial adaptation without regularization.
• Ablation studies show that removing the posterior term leads to degenerate solutions with near-zero accuracy on some corruptions, highlighting the necessity of the regularization term.
• On SVHN-to-MNIST transfer, removing the posterior term improves performance, indicating that the training posterior can be overly restrictive in large domain shifts.
• BACS improves ECE (Expected Calibration Error) from 0.05667 (vanilla) to 0.04270 on ImageNet-R, showing better uncertainty calibration.