Learning Calibrated Uncertainties for Domain Shift: A Distributionally Robust Learning Approach

TL;DR

This work tackles uncertainty estimation under domain shift by formulating a distributionally robust learning (DRL) framework that learns a differentiable density-ratio to calibrate predictions. The end-to-end approach jointly optimizes a density-ratio estimator with a target classifier, incorporating class-regularization to produce conservative, well-calibrated probabilities, and yields a predictive form $P(y|x) \propto \exp\left(\frac{P_s(x)}{P_t(x)} \mathbf{\theta}_y \cdot \mathbf{\phi}(x) \right)$. The method advances unsupervised domain adaptation and cross-domain semi-supervised learning by providing calibrated uncertainties for pseudo-label selection, leading to improved cross-domain accuracy and calibration metrics (ECE, Brier score) across Office31, Office-Home, VisDA2017, and ImageNet-based analyses. Empirical results show that density ratios align with human uncertainty proxies and that DRL-based self-training (DRST) and SSL (DRSSL) yield substantial gains, especially on harder target examples. Overall, the approach offers a practical, end-to-end plug-in for robust learning under covariate shift with improved reliability in downstream tasks.

Abstract

We propose a framework for learning calibrated uncertainties under domain shifts, where the source (training) distribution differs from the target (test) distribution. We detect such domain shifts via a differentiable density ratio estimator and train it together with the task network, composing an adjusted softmax predictive form concerning domain shift. In particular, the density ratio estimation 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 by adversarial risk minimization. We show that our proposed method generates calibrated uncertainties that benefit downstream tasks, such as unsupervised domain adaptation (UDA) and semi-supervised learning (SSL). On 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 show that the estimated density ratios align with human selection frequencies, suggesting a positive correlation with a proxy of human perceived uncertainties.

Figures (6)

• Figure 1: (a) Architecture for end-to-end training of our DRL framework (see Sec. \ref{['sec:endtoend']}). (b) Examples for category 'Train' in VisDA. The estimated density ratios for the easy and hard target images are shown. The DRL framework gives higher uncertain predictions for the harder example ($\bm{x}_2$) that is more cluttered and hence not well-represented in the source domain.
• Figure 2: Formulation of the pseudo label based UDA or SSL methods with DRL. The unsupervised losses represent the loss imposed on the unlabeled target data. DRST conducts this procedure multiple iterations, while DRSSL minimizes the unsupervised losses on the augmented target data.
• Figure 3: Brier score (top) and reliability diagrams (bottom) on Office31, Office-Home and VisDA. DRL generates more calibrated uncertainties than source-only and temperature scaling and VADA. Brier score measures the mean squared difference between the predicted probability and the actual outcome. For a fully calibrated classifier, the confidence should match the accuracy across the full range of confidence. Thus the closer the lines are to the dashed line, the more calibrated the method is. Our method gets more and more calibrated as the confidence increases. Note that in the first row, TS's Brier scores are much larger and excluded to not effect the scale.
• Figure 4: Density ratios vs HSF on ImageNetV2.
• Figure 5: (a)-(b) Results on VisDA-17 (performed with 5 random seeds) with test accuracy and Brier score. DRST outperforms the baselines significantly. (c) We adopt distribution gap $P_s(\bm{\phi}(\bm{x}))/P_t(\bm{\phi}(\bm{x})) - P_s(\bm{\phi}(\bm{x}), y)/P_t(\bm{\phi}(\bm{x}), y)$ as a proxy of covariate shift. DRST helps further reduce this gap with self-training.

Theorems & Definitions (2)

• Lemma 1
• Theorem 1