Bayesian Adaptation for Covariate Shift

1 Introduction . Modern deep learning methods can provide high accuracy in settings where the model is evaluated on data from the same distribution as the training set , but accuracy often degrades severely when there is a mismatch between the training and test distributions [ Hendrycks and Dietterich , 2019 , Taori et al. , 2020 ] . In safety-critical settings , effectively deploying machine learning models requires not only high accuracy , but also requires the model to reliably quantify uncertainty in its predictions in order to assess risk and potentially abstain from making dangerous , unreliable predictions . Reliably estimating uncertainty is especially important in settings with distribution shift where inaccurate predictions are more prevalent , but the reliability of uncertainty estimates often also degrades along with the accuracy as the shifts become more severe [ Ovadia et al. , 2019 ] . In real-world applications , distribution mismatch at test time is often inevitable , thus necessitating methods that can robustly handle distribution shifts , both in terms of retaining high accuracy but also in providing meaningful uncertainty estimates . As it can be difficult to train a single model to be robust to all potential distribution shifts we might encounter at test time , we can instead robustify models by allowing for adaptation at test time , where we finetune the network on unlabeled inputs from the shifted target distribution , thus allowing the model to specialize in the particular shift it encounters . Since test-time distribution shifts often can not be anticipated during training time , we restrict the adaptation procedure to operate without any further access to the original training data . Prior work on test-time adaptation [ Wang et al. , 2020a , Sun et al. , 2019b ] focused on improving accuracy , and found that simple objectives like entropy minimization capable of providing substantial improvements under distribution shift [ Wang et al. , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2020a ] . However , these prior works do not consider uncertainty estimation , and an objective like entropy minimization can quickly make predictions overly confident and less calibrated , leaving us without reliable uncertainty estimates and thus unable to quantify risks when making predictions . Our goal in this paper is to design a test-time adaptation procedure that can not only improve predictive accuracy under distribution shift , but also provide reliable uncertainty estimates . While a number of prior papers have proposed various heuristic methods for test-time adaptation [ Wang et al. , 2020a , Sun et al. , 2019a ] , it remains unclear what precisely unlabeled test data under covariate shift can actually tell us about the optimal classifier . In this work , we take a Bayesian approach to this question , and explicitly formulate a Bayesian model that describes how unlabeled test data from a different domain can be related to the classifier parameters . Such a model requires introducing an additional explicit assumption , as the classifier parameters are conditionally independent of unlabeled data in the standard model for discriminative classification [ Seeger , 2000 ] . The additional assumption we introduce intuitively states that the data generation process at test-time , though distinct from the one at training time ( hence , under covariate shift ) is still more likely to produce inputs that have a single unambiguous labeling , even if that labeling is not known . We argue that this assumption is reasonable in practice , and leads to an appealing graphical model where approximate inference corresponds to a Bayesian extension of entropy minimization . We propose a practical test-time adaptation strategy , Bayesian Adapatation for Covariate Shift ( BACS ) , which approximates Bayesian inference in this model and outperforms prior adaptive methods both in terms of increasing accuracy and improving calibration under distribution shift . Our adaptation strategy is simple to implement , requires minimal changes to standard training procedures , and outperforms prior test-time adaptation techniques on a variety of benchmarks for robustness to distribution shifts . 2 Bayesian Adaptation for Covariate Shift . In this section , we will devise a probabilistic graphical model that describes how unlabeled data in a new test domain can inform our posterior about the model , and then describe a practical deep learning algorithm that can instantiate this model in a system that enables test-time adaptation . We will begin by reviewing standard probabilistic models for supervised and discuss why such models are unable to utilize unlabeled data . We then discuss a probabilistic model proposed by Seeger [ 2000 ] that does incorporate unlabeled data in a semisupervised learning ( SSL ) setting , and propose an extension to the model to account for distribution shift . Finally , we discuss the challenges in performing exact inference in our proposed model and describe the approximations we introduce in order to derive a tractable inference procedure suitable for test-time adaptation . 2.1 Probabilistic Model for Covariate Shift In the standard Bayesian model for supervised learning ( Figure 1 ) , we assume that the inputs X are sampled i.i.d . from a generative model with parameters φ , while the corresponding labels are sampled from a conditional distribution p ( Y |X , θ ) parameterized by θ . The parameters θ and φ are themselves random variables , sampled independently from prior distributions p ( φ ) and p ( θ ) . We observe only a dataset D = { Xi , Yi } ni=1 , and then perform inference over the parameters θ using Bayes rule : p ( θ|D ) ∝ p ( θ ) n∏ i=1 p ( Yi|Xi , θ ) . ( 1 ) To make predictions on a new input xn+1 , we marginalize over the posterior distribution of the classifier parameters θ to obtain the predictive distribution p ( Yn+1|Xn+1 , D ) = ∫ p ( Yn+1|Xn+1 , θ ) p ( θ|D ) dθ . Note that observing the ( unlabeled ) test input xn+1 ( or more generally , any number of test inputs ) does not affect the posterior distribution of parameters θ , and so within this probabilistic model , there is no benefit to observing multiple unlabeled datapoints before making predictions . Inference for semi-supervised learning . By treating the parameters θ , φ as a-priori independent , the standard model assumes there is no relationship between the input distribution and the process that generates labels , leading to the inability to utilize unlabeled data in inferring θ . To utilize unlabeled data , we can introduce assumptions about the relationship between the labeling process and input distribution using a model of the form in Figure 2 [ Seeger , 2000 ] . The assumption we use has a simple intuitive interpretation : we assume that the inputs that will be shown to our classifier have an unambiguous and clear labeling . This assumption is reasonable in many cases : if the classifier discriminates between automobiles and bicycles , it is reasonable that it will be presented with images that contain either an automobile or bicycle . Of course , this assumption is not necessarily true in all settings , but this intuition often agrees with how discriminatively trained models are actually used . This simple intuition can be formalized in terms of a prior belief about the conditional entropy of labels conditioned on the inputs . Similar to Grandvalet and Bengio [ 2004 ] , we can encode this belief using the additional factor p ( θ|φ ) ∝ µ ( θ ) exp ( −αHθ , φ ( Y |X ) ) ( 2 ) = µ ( θ ) exp ( αEX∼p ( X|φ ) , Y∼p ( Y |X , θ ) [ log p ( Y |X , θ ) ] ) , ( 3 ) where µ ( θ ) is a prior over the parameters θ that is agnostic of the input distribution parameter φ . Now , observing unlabeled test inputs U = { x̃1 , . . . , x̃m } provides information about the parameter φ governing the input distribution , which then allows us to update our belief over the learned parameters θ through Equation 2 and thus allows inference to utilize unlabeled data within a Bayesian framework . Extension to covariate shift . The previous probabilistic model for SSL assumes all inputs were drawn from the same distribution ( given by φ ) . However , our goal is use unlabeled data to adapt our model to a different test distribution , so we extend the model to incorporate covariate shift . We now assume we have two input-generating distributions ; φ specifies the input distribution for our labeled training set , and φ̃ specifies the shifted distribution of unlabeled test inputs which we aim to adapt to . Under the assumption of covariate shift , the same classifier θ is used to generate labels in both the train and test domains , leading to the model in Figure 3 . We argue that our prior belief of low aleatoric uncertainty is reasonable for both the distribution induced by φ and the one induced by φ̃ ; even if there is some distribution shift from the training set , we can still often expect the test data we are shown to have unambiguous labels . We can then incorporate both φ and φ̃ into our belief over θ with the factor p ( θ|φ , φ̃ ) ∝ µ ( θ ) exp ( −αHθ , φ ( Y |X ) ) exp ( −α̃Hθ , φ̃ ( Y |X ) ) , ( 4 ) with the two scalar hyperparameters α , α̃ controlling how to weight the entropies from each distribution with the likelihoods of the labeled training set . While we can extend this model to include more labeled or unlabeled input distributions , we focus on simply having one training distribution and one test distribution as it is the most relevant for our problem setting of adapting to a particular distribution shift at test time . 2.2 Approximations for Tractable Inference . Performing inference over θ in the above model can be challenging in our test-time adaptation setting for several reasons . First of all , inference over θ would also require performing inference over the parameters of the generative models φ and φ̃ to evaluate likelihoods in Equation 4 . This is difficult in practice for two reasons : First , the inputs x might be high-dimensional and difficult to model , as in the case of images . Second , the amount of unlabeled data might be fairly small , and insufficient to accurately estimate the parameters φ̃ if we employ a highly expressive generative model . As such , we would much prefer to avoid explict generative modeling and to only perform inference over the discriminative model parameters θ instead . Another issue is that performing inference would require access to both the labeled training data and the unlabeled test data at the same time , while our test-time adaptation setting assumes that we can no longer look at the training set when it is time to adapt to the test distribution . We will discuss how to address these issues and describe our method , Bayesian Adaptation for Covariate Shift ( BACS ) , which provides a practical instantiation of inference in this Bayesian model in a computationally tractable test-time adaptation setting . Plug-in approximation with empirical Bayes . We first propose to avoid explicit generative modeling by using with a plug-in empirical Bayes approach . Empirical Bayes is a common approximation for hierarchical Bayesian models that simply uses a point estimate of φ rather than marginalizing over φ ’ s ( similarly for φ̃ ) , which would reduce the computation to estimating only a single generative model for each input distribution given by parameters φ∗ and φ̃∗ . To eliminate the need to train the parameters φ∗ of a generative model of the inputs altogether , we note that p ( θ|φ , φ̃ ) only depends on φ and φ̃ through the input distributions p ( x|φ ) , p ( x̃|φ̃ ) . We can then approximate Equation 4 by plugging in the empirical distributions of x and x̃ in place of p ( x|φ∗ ) , p ( x̃|φ̃∗ ) , resulting in p ( θ|φ̂ , φ̃ ) ∝ µ ( θ ) exp ( −α n n∑ i=1 H ( Yi|xi , θ ) ) exp ( − α̃ m m∑ i=1 H ( Yi|x̃i , θ ) ) . Now , given a labeled training set D and unlabeled test points U = { x̃i } mi=1 , the new posterior distribution over parameters now has log probabilities ( up to an additive normalizing constant ) log p ( θ|D , U ) = logµ ( θ ) + n∑ i=1 log p ( yi|xi , θ ) − α n n∑ i=1 H ( Y |xi , θ ) − α̃ m m∑ j=1 H ( Y |x̃j , θ ) . ( 5 ) For convenience , we simply set α = 0 , as additionally minimizing entropy on the labeled training set is unnecessary when we are already maximizing the likelihood of the given labels with highly expressive models . Now , to infer θ given the observed datasets D and U , the simplified log-density is log p ( θ|D , U ) = logµ ( θ ) + n∑ i=1 log p ( yi|xi , θ ) − α̃ m m∑ j=1 H ( Y |x̃j , θ ) . ( 6 ) Computing the maximum-a-posteriori ( MAP ) solution of this model corresponds to simply optimizing model parameters θ that on the supervised objective on the labeled training set in addition a minimum entropy regularizer on the unlabeled input . However , we should not necessarily expect the MAP solution to provide reasonable uncertainty estimates , as the learned model is being encouraged to make confident predictions on the test inputs , and so any single model will likely provide overconfident predictions . Marginalizing the predictions over the posterior distribution over parameters is thus essential to recovering meaningful uncertainty estimates , as the different models , though each being individually confident , can still express uncertainty through their combined predictions when the models predict different labels . Approximate inference for test-time adaptation . We now discuss how to perform inference in the above model in a way suitable for test-time adaptation , where we adapt to the test data without any further access to the original training set . To enable this , we propose to learn an approximate posterior q ( θ ) ≈ p ( θ|D ) = log µ ( θ ) + n∑ i=1 log p ( yi|xi , θ ) ( 7 ) during training time , and then use this approximate training set posterior q ( θ ) in place of the training set when performing inference on the unlabeled test data . This gives us an approximate posterior with log density log p ( θ|D , U ) = log q ( θ ) − α̃ m m∑ j=1 H ( Y |x̃j , θ ) ( 8 ) Algorithm 1 Bayesian Adaptation for Covariate Shift ( BACS ) Input : Ensemble size k , Entropy weight α̃ , Training Data : x1 : n , y1 : n , Test Data : x̃1 : m Output : Predictive distributions p ( y|x̃j ) for each test input x̃j Training : For each i ∈ ( 1 , . . . , k ) compute an approximate density qi ( θ ) ≈ p ( θ|x1 : n , y1 : n ) Adaptation : for all ensemble members i ∈ ( 1 , . . . , k ) do Compute adapted parameters θ̂i = argmaxθ α̃ m ∑m j=1−H ( Y |x̃j , θ ) + log qi ( θ ) end for For each test input x̃j , marginalize over ensemble p ( y|x̃j ) = 1k ∑k i=1 p ( y|x̃j , θ̂i ) . Unlike the full training set , the learned approximate posterior density q ( θ ) , which can be as simple as a diagonal Gaussian distribution , can be much easier to store and optimize over for test-time adaptation , and the training time procedure would be identical to any approximate Bayesian method that learns a posterior density . In principle , we can now instantiate test-time adaptation by running any approximate Bayesian inference algorithm , such as variational inference or MCMC , to sample θ ’ s from the density in Equation 8 , and average predictions from these samples to compute the marginal probabilities for each desired test label . Practical instantiation with ensembles . As previously mentioned , marginalizing over different models that provide diverse labelings of the test set is crucial to providing uncertainty estimates after adaptation via entropy minimization . We thus propose to use an ensembling approach [ Lakshminarayanan et al. , 2016 ] as a practical method to adapt to the test distribution while maintaining diverse labelings . Deep ensembles simply train multiple models from different random initializations , each independently optimizing the target likelihoods , and averages together the models ’ predicted probabilities at test time . They are able to provide effective approximations to Bayesian marginalization due to their ability to aggregate models across highly distinct modes of the loss landscape [ Fort et al. , 2019 , Wilson and Izmailov , 2020 ] . Our method , BACS , summarized in Algorithm 1 , trains an ensemble of k different models on the training set , each with their approximate posterior qi ( θ ) that captures the local loss landscape around each mode in the ensemble . Then at test time , we independently optimize each of the k models by minimizing Equation 8 ( using the corresponding qi ( θ ) for each ensemble member ) , and then average the predictions across all adapted ensemble members .

Summary: The paper extends test-time entropy minimization through a novel Bayesian criterion for adaptation. It allows to improve performance during adaptation without degrading the model calibration, a shortcoming of standard entropy minimization. The model is evaluated on CIFAR-10C/-100C/ImageNet-C and compared to batch norm adaptation and test-time entropy minimization.