Distributionally Robust Learning for Uncertainty Calibration under Domain Shift

We propose a framework for learning calibrated uncertainties under domain shifts . We consider the case where the source ( training ) distribution differs from the target ( test ) distribution . We detect such domain shifts through the use of binary domain classifier and integrate it with the task network and train them jointly endto-end . The binary domain classifier yields a density ratio that reflects the closeness of a target ( test ) sample to the source ( training ) distribution . We employ it to adjust the uncertainty of prediction in the task network . This idea of using the density ratio is based on the distributionally robust learning ( DRL ) framework , which accounts for the domain shift through adversarial risk minimization . We demonstrate that our method generates calibrated uncertainties that benefit many downstream tasks , such as unsupervised domain adaptation ( UDA ) and semi-supervised learning ( SSL ) . In these tasks , methods like self-training and FixMatch use uncertainties to select confident pseudo-labels for re-training . Our experiments show that the introduction of DRL leads to significant improvements in cross-domain performance . We also demonstrate that the estimated density ratios show agreement with the human selection frequencies , suggesting a positive correlation with a proxy of human perceived uncertainties . 1 INTRODUCTION . Uncertainty estimation is an important machine learning problem that is central to trustworthy AI ( Tomsett et al. , 2020 ; Antifakos et al. , 2005 ) . In addition , many important downstream applications rely on the correct estimation of uncertainties . This includes unsupervised domain adaptation ( UDA ) ( Zou et al. , 2019 ) and semi-supervised learning ( SSL ) ( Sohn et al. , 2020 ) where they are used to solicit confident pseudo-labels for re-training . In these applications , reliable pseudo-labels help avoid error propagation and catastrophic failures in early iterations ( Kumar et al. , 2020 ) . Obtaining reliable uncertainty estimation is challenging . In contrast to human annotation of labels , obtaining the ground-truth uncertainties from real-world data can be costly or even infeasible . It is also known that commonly used uncertainty proxies in deep neural networks , such as the softmax output , tend to give overconfident estimates ( Guo et al. , 2017a ) . This overconfidence is further amplified under domain shifts , where the target ( test ) domain and the source training domain differ significantly . Such distributional shifts tend to aggravate the existing issues in uncertainty estimation , leading to wrong but overconfident predictions on unfamiliar samples ( Li & Hoiem , 2020 ) . Many methods have been proposed to calibrate the confidence of deep learning models so that the uncertainty level of a model prediction reflects the likelihood of the true event ( Guo et al. , 2017a ) . Label smoothing is a popular approach to reduce overconfidence and to promote more uniform outputs ( Szegedy et al. , 2016 ) . Temperature scaling is another method where the logit scores are rescaled by a calibrated temperature ( Platt et al. , 1999 ) . Approaches such as Monte-Carlo sampling ( Gal & Ghahramani , 2016 ) and Bayesian inference ( Blundell et al. , 2015 ; Riquelme et al. , 2018 ) model uncertainties from a Bayesian perspective but are computationally more expensive . Even though these methods lead to more calibrated uncertainties , recent studies show that their results can not be fully trusted under domain shift ( Snoek et al. , 2019 ) . Our approach . To handle domain shifts , we characterize the “ overlap ” between the source training data and the test data . Intuitively , if the test sample is highly unlikely in the training distribution , then the resulting confidence levels should be lowered . We incorporate this idea through the density ratio between the two domains and employ it for confidence calibration . To make our idea more concrete , recall that the probability output of a standard neural network for classification can be expressed as P ( y|x ) ∝ exp ( θy · φ ( x ) ) , where x is the predictor input , φ ( x ) is the data feature , and θy is the model parameter of the y-th class . Instead , we propose the following predictive form for our neural network : P ( y|x ) ∝ exp ( Ps ( x ) Pt ( x ) θy · φ ( x ) ) , ( 1 ) where Ps ( x ) and Pt ( x ) are the densities of a data sample under the source and target distributions , respectively . When a target sample is close to the source domain ( large Ps ( x ) /Pt ( x ) ) , the prediction is confident . However , when a target sample x is far away from the source distribution ( small Ps ( x ) /Pt ( x ) ) , the confidence is lowered and the prediction is closer to a uniform distribution . This intuition is analogous to incorporating a sample-wise temperature to adjust the confidence according to the closeness of a test sample to the training distribution . Eq . ( 1 ) is based on the distributionally robust learning ( DRL ) framework . DRL is an adversarial risk minimization framework that involves a two-player minimax game between a predictor and an adversary ( Grünwald et al. , 2004 ) . While many previous DRL methods ( Liu et al. , 2020 ; Nakka et al. , 2020 ) operate in low dimensional spaces using kernel density estimators for the density ratio estimation , we develop a DRL method to scale to real-world computer vision tasks , which is able to produce calibrated uncertainties under domain shift . Our contributions are : 1 ) We propose a DRL method for uncertainty estimation under domain shift . We introduce a binary domain classifier network which learns to predict the density ratios between source and target domains ( See Fig . 1 ( a ) ) . The domain classifier and the target classifier are trained simultaneously in an end-to-end fashion . We also introduce a regularized DRL framework to further promote smoothed model prediction and improve the calibration . 2 ) We show that the estimated density ratio reflects the distance of a test sample from both training and test distributions ( Fig . 1 ( b ) ) . Our experiments further empirically show that these estimates are also correlated with human selection frequency , based on the available ground-truth labels in ImageNetV2 ( Recht et al. , 2019a ) , which can be regarded as a proxy of human uncertainty perception . 3 ) We empirically show that the top-1 class predictions of DRL are more calibrated than the empirical risk minimization ( ERM ) and the temperature scaling on Office31 ( Saenko et al. , 2010 ) , Office-Home ( Venkateswara et al. , 2017 ) , and VisDA ( Peng et al. , 2017 ) . We measure the level of calibration using expected calibration error ( ECE ) , Brier Score and reliability plots . 4 ) We integrate our method as a plug-in module in downstream applications such as unsupervised domain adaption and semi-supervised learning , leading to significant improvements . For example , incorporating self-training ( Zou et al. , 2019 ) with DRL leads to state-of-the-art performance on VisDA-2017 ( Peng et al. , 2017 ) and a 6 % improvement on hard examples . Incorporating Fix- match ( Sohn et al. , 2020 ) with DRL improves the original Fixmatch by a relative 17 % increase in accuracy under the cross-domain setting . 2 DISTRIBUTIONALLY ROBUST LEARNING . In this section , we first review the preliminaries of DRL ( Sec . 2.1 ) , followed by a proposed variant of DRL with class regularization ( Sec . 2.2 ) . We then propose an instantiation of DRL with a differentiable density ratio estimation network ( Sec . 2.3 ) . Finally , we show the application of DRL in UDA and SSL tasks ( Sec . 2.4 ) . 2.1 PRELIMINARIES . Notations and definitions : We denote the input and labels by random variables X and Y . We also use x ∈ Rd and X to represent the realization and sample space ofX . Our goal is to find a predictor F : Rd 7→ RC , { f ( x ) |x ∈ X , f ( x ) ∈ RC ∩∆ } ( 2 ) close to the true underlying Pt ( Y |X ) . Here , d , C and ∆ denote the input dimension , class number , and probabilistic simplex , respectively . We consider the problem with labeled data sampled from a source distribution Ps ( X , Y ) and unlabeled data sampled from a target distribution Pt ( X ) , and use ePs ( X , Y ) to represent the empirical source distribution . In this work , we consider an important form of domain shift with the covariate shift assumption Ps ( X ) 6= Pt ( X ) , Ps ( Y |X ) = Pt ( Y |X ) . Motivation : Traditional empirical risk minimization ( ERM ) frameworks tend to fail under covariate shift since ERM empirically learns a predictor P̂s ( Y |X ) from the finite source data that usually fails to generalize to the target distribution . DRL was previously proposed to overcome this issue . DRL can be formulated as a two-player adversarial risk minimization game ( Grünwald et al. , 2004 ) with the predictor player minimizing a loss , while the adversary player maximizing the loss . The adversary is allowed to perturb the labels , subject to certain feature-matching constraints to ensure data-compatibility . Formulation : Under covariate shift , DRL ( Liu & Ziebart , 2014 ) deals with the mismatch between the expected loss and the training data , and is defined on the target data distribution : P̂t ( Y |X ) = argmin F max G∈Σ Ex∼Pt ( X ) L ( f ( x ) , g ( x ) ) , ( 3 ) where f ( x ) , g ( x ) ∈ RC are the conditional label distributions given an input x and F , G are the entire distribution over all the input . Ex∼PtL ( · ) is an expected log loss on the target input defined : Ex∼Pt ( X ) L ( f ( x ) , g ( x ) ) , Ex∼Pt ( X ) [ −g ( x ) · log f ( x ) ] . ( 4 ) F is the predictor player minimizing the loss function while G is the adversary maximizing the loss function . After solving this game , F is our estimate of P̂t ( Y |X ) , which we will use for the classification on the target domain . Eq . 3 is defined on the target domain only . How could a predictor be properly trained while there are no target labels available ? The answer is that the adversary G is implicitly constrained by the source features . We use the following constraints to make sure thatG is close to Ps ( Y |X ) : Σ = { G|∑igyφ ( xi ) = ∑ iI [ yi = y ] φ ( xi ) , ∀y } , ( 5 ) where xi ∼ Ps ( X ) and gy is the y-th dimension of g. Eq . 5 is a necessary but not sufficient condition for G = Ps ( Y |X ) , thus serving as an implicit constraint for G to be close to the true Pt ( Y |X ) under the covariate shift assumption . Given a predefined feature function φ , when the adversary perturbs the conditional label distribution , certain aggregate function of φ on g should equal to the counterpart on the empirical source data . From DRL to density ratio : When using the expected target logloss in Eq . 3 , Eq . 1 is derived by solving the predictor F . Here we refer the derivation details to ( Liu & Ziebart , 2014 ) but emphasize the important properties of the prediction : representation-level conservativeness . Representation-level conservativeness : The predictions have higher certainty for inputs closer to the source domain when Ps ( x ) /Pt ( x ) is large . On the contrary , when Ps ( x ) /Pt ( x ) is small , the prediction is more uncertain . This property reflects the model ’ s ability to convey information about what it does not know through the model uncertainty .

Summary: This paper focuses on the problem of uncertainty calibration under distribution shift. By using a domain classifier in the distributional robust learning (DRL) framework, the authors estimate the density ratio between the source and target domain to achieve well-calibrated predictions under domain shift. A regularized DRL framework is further proposed to promote smoothed model prediction and improve the calibration. Experiments on Office31, Office-Home, and VisDA-2017 demonstrate the superiority of DRL over empirical risk minimization (ERM) and the temperature scaling method measured by expected calibration error (ECE), Brier Score, and reliability plots.