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: This paper proposes a Bayesian inference framework for test-time adaptation under covariate shift setting.  The setting discussed in the paper is, we have a model trained in a fully supervised way, and observe unlabelled test data (without labels) that might have come from a distribution different from the training data distribution. The goal is to (quickly) adapt the model for the test data without access to the test labels. The proposed method also tries to properly model the epistemic uncertainty of the model to be adapted for the test data, and for this purpose, the problem is formulated as a Bayesian hierarchical model with the data generating distribution and the parameter generating distributions are coupled. Here, an important difference from the previous approaches is that for test data we have a different data generating distribution (conditioning on a different parameter for data generating distribution). The dependence between the model parameter distribution and data generating distribution is expressed in terms of energy-based density whose energies are written as predictive entropies.  The proposed approach entitled Bayesian Adaptation for Covariate Shift (BACS) is demonstrated to be effective in various covariate shift benchmark datasets.