Distributionally Robust Classification for Multi-source Unsupervised Domain Adaptation

AI summary

Overview

Research area: Unsupervised domain adaptation (UDA) and distributionally robust optimization (DRO) for classification under distribution shift, with a focus on multi-source and single-source settings.

Technical level: Advanced. The paper builds a minimax ambiguity set over mixtures of source conditional distributions plus an infinite-order Wasserstein ball on the target inputs, and derives an alternating stochastic-gradient algorithm for it.

Scope: The paper proposes a distributionally robust classification framework that models uncertainty in both the target covariate distribution and the conditional label distribution, and evaluates it on digit classification and spurious-correlation benchmarks where target data are scarce.

What This Paper Is About

In unsupervised domain adaptation, a model is trained on labeled source data and unlabeled target data, and must generalize to the target domain even though the two distributions differ. Existing distribution-alignment and pseudo-labeling methods break down in two common situations: when very little unlabeled target data is available, and when the source domain contains spurious correlations that do not transfer. The paper's goal is a learning framework that stays robust in both situations by explicitly accounting for two kinds of uncertainty — which source conditionals to trust, and how far the target inputs may deviate from what was observed.

Key Contributions

1. A novel distributionally robust framework that jointly models uncertainty in the target covariate distribution and the conditional label distribution, defined in a learned feature space.
2. A tractable minimax optimization algorithm (Proposition 3.1 gives an upper-bounding surrogate objective) with alternating updates, designed so it can be combined with existing UDA methods such as CDAN and STAR.
3. A construction for applying a multi-source framework to single-source problems by treating random sub-samples with replacement as pseudo-sources, following the maximin-effect principle.
4. Extensive experiments across digit benchmarks (MNIST, SVHN, USPS) and spurious-correlation benchmarks (Waterbirds, CelebA, Colored MNIST) showing consistent improvements over strong baselines, particularly under extreme target-data scarcity.

Main Findings

• Digit adaptation with scarce targets: The method achieves the best reported accuracy in every task and target sample size tested. With 10² unlabeled target samples per class, the best variant (Ours (STAR)) reaches 94.4 ± 1.7 on SVHN → MNIST, 95.6 ± 1.0 on MNIST → USPS, and 97.3 ± 0.8 on USPS → MNIST. With only 10 samples per class, it reaches 91.3 ± 1.1, 91.2 ± 1.4, and 93.0 ± 2.8 respectively.
• Large gains over the base classifiers: Combining the method with CDAN improves classical CDAN by +29.1% on the SVHN → MNIST task.
• Baseline comparisons: On SVHN → MNIST with 10² samples, ERM (Src-only) scores 59.6 ± 1.8, DANN 66.0 ± 4.9, CDAN 63.4 ± 1.8, MK-MMD 50.0 ± 3.0, ATDOC 83.3 ± 9.1, STAR 76.4 ± 1.5, CORAL 75.4 ± 2.7, and MCD 79.1 ± 1.0, versus 92.0 ± 1.6 (Ours (ERM)), 92.5 ± 2.3 (Ours (CDAN)), and 94.4 ± 1.7 (Ours (STAR)).
• Robustness at extreme scarcity: The drop from 10² to 10 target samples per class is small for the proposed method, while several baselines degrade sharply — for example, CDAN falls from 80.8 ± 1.4 to 62.0 ± 1.8 on MNIST → USPS, and DANN falls from 82.0 ± 3.9 to 74.3 ± 5.7.
• Two mechanisms behind the gains: The ε₁ radius hedges against the empirical target input distribution being a poor estimate of the true one when target samples are few; using UDA methods (CDAN, STAR) to build the conditional estimators yields mixtures that align more closely with the true target distribution than mixtures built from ERM conditionals.
• Cross-validation variant: The tuning-free LODO-CV variant (marked with an asterisk) shows a moderate drop relative to the validation-tuned version but remains competitive and still outperforms baselines that rely on labeled target data for tuning.
• Spurious-correlation benchmarks: Target sample sizes are 56 (Waterbirds), 1,387 (CelebA), and 2,998 (CMNIST). Reported baseline values include ERM (Src-only) at 48.4 ± 0.9 on Waterbirds, 35.5 ± 0.6 on CelebA, and 0.9 ± 0.5 on CMNIST; DANN at 35.8 ± 4.5, 23.5 ± 2.1, and 0.9 ± 1.8; and CDAN at 46.2 ± 1.8 on Waterbirds and 24.6 on CelebA. The table in the provided content is truncated, so the proposed method's numbers on these three benchmarks are not reported here.

Methodology in Plain English

The method starts from a trained classifier and strips off its final layer to obtain a feature map. It then splits the labeled source data into K = 10 sub-samples, each of size N = N_sc/5 drawn with replacement, and treats each sub-sample as if it came from its own "source." On each sub-sample it trains a simple linear logistic regression in feature space (or, optionally, CDAN or STAR initialized with the learned features) to estimate that source's conditional distribution of the label given the input. If a natural grouping variable exists, it can be used instead of random sub-sampling.

These conditional estimates are combined into a mixture whose weights β lie on the probability simplex. The weights are not fixed: they are constrained to stay within a Euclidean distance ε₂ of a reference vector β̄ (set to uniform in the experiments), and the framework searches adversarially over them. Separately, the target input distribution is allowed to move inside an ε₁-radius infinite-order Wasserstein ball around the empirical target inputs, with the cost function defined as the Euclidean distance between feature representations.

Rather than optimizing this minimax problem directly, the authors prove an upper bound (Proposition 3.1) that replaces it with a tractable surrogate: an expectation over target inputs of a worst-case loss, where the worst case involves both a perturbed feature z' within ε₁ of the original feature and a soft pseudo-label y° formed as a convex combination of the source conditionals. The algorithm then alternates three steps — projected gradient ascent on the perturbed feature z', projected exponentiated gradient ascent on β (giving more weight to conditionals that cause higher loss under the current classifier), and ordinary stochastic gradient descent on the model parameters θ. The β update is the adversarial step that forces the classifier to be robust to whichever mixture of sources is hardest; the θ update then minimizes that adversarial objective.

Why This Matters

Impact on research: The paper extends maximin-effect estimation, previously studied mainly for regression where closed-form solutions are available, to classification, where the authors show the objective can be handled through a tractable surrogate and an alternating algorithm. It also offers a way to apply a multi-source robust framework in single-source settings by simulating pseudo-sources, and it addresses a setting that GroupDRO handles poorly: GroupDRO assumes access to group labels and does not explicitly use unlabeled target data, whereas this method uses neither group labels nor a target-labeled validation set (in its LODO-CV form).

Real-world applications:

• Medical imaging and healthcare, where models trained at one hospital or scanner must transfer to another site with little or no labeled data.
• Autonomous driving and robotics, where environmental changes, lighting, and weather create distribution shifts between training and deployment.
• Wildlife monitoring and other computer-vision tasks vulnerable to spurious background cues, as illustrated by the Waterbirds benchmark.
• Facial-attribute and demographic prediction, where demographic imbalance causes spurious correlations, as illustrated by CelebA.

Industry relevance: The method is designed as a reusable component rather than a replacement — it can be layered on top of existing UDA pipelines such as CDAN or STAR, which makes it relatively cheap to adopt in production systems that already use those methods. Its emphasis on the scarcity regime matters for industrial settings where collecting and annotating target data is expensive or slow.

Future Directions

• Providing theoretical guarantees for the classification setting beyond the surrogate upper bound in Proposition 3.1, which the paper presents as an approximation whose rationale is deferred to an appendix.
• Making the ambiguity-set radii ε₁ and ε₂ fully automatic: the current approach grid-searches them using a small labeled target validation set of 10 samples per class, with LODO-CV offered as a labeled-target-free alternative.
• Testing the framework on larger-scale and non-vision domains, given that the reported experiments cover digit benchmarks and three spurious-correlation datasets.
• A systematic study of how the number of simulated sources K (fixed at 10 here) and the sub-sample fraction (N = N_sc/5) affect robustness, especially when the source distribution's subpopulation structure is unknown.

Target Audience

Researchers and graduate students working on domain adaptation, distributionally robust optimization, or robustness to spurious correlations. Practitioners who already deploy UDA methods such as CDAN or STAR and want a robustness layer for settings with very few unlabeled target samples will also benefit, though the minimax formulation and exponentiated-gradient updates assume comfort with optimization-based machine learning.

Authors’ abstract

Unsupervised domain adaptation (UDA) is a statistical learning problem when the distribution of training (source) data is different from that of test (target) data. In this setting, one has access to labeled data only from the source domain and unlabeled data from the target domain. The central objective is to leverage the source data and the unlabeled target data to build models that generalize to the target domain. Despite its potential, existing UDA approaches often struggle in practice, particularly in scenarios where the target domain offers only limited unlabeled data or spurious correlations dominate the source domain. To address these challenges, we propose a novel distributionally robust learning framework that models uncertainty in both the covariate distribution and the conditional label distribution. Our approach is motivated by the multi-source domain adaptation setting but is also directly applicable to the single-source scenario, making it versatile in practice. We develop an efficient learning algorithm that can be seamlessly integrated with existing UDA methods. Extensive experiments under various distribution shift scenarios show that our method consistently outperforms strong baselines, especially when target data are extremely scarce.