Zhenyu (Zach) Wang
Home Blog Teaching
← All posts

Classifying reliably in a domain you've never labelled: CG-DRO

Jul 2025, revised Jan 2026 · with Zijian Guo, Yifan Hu, Francis Bach · arXiv:2507.09905

The short version

Setting
Labelled data from \(L\) source domains and only unlabelled covariates from the target. Both the features and the label mechanism may shift.
Idea
Hold the observable target covariate law \(\mathbb Q_X\) fixed and treat the target's conditional as some mixture of the sources' conditionals. Minimize the worst-case cross-entropy over that set.
The catch
The worst-case weight \(\widehat\gamma\) can be pinned to the boundary of the simplex, or swing wildly between vertices. In either case the classifier is not asymptotically normal, and both normal-approximation intervals and the bootstrap break.
The fix
A perturbation-resampling scheme that gives uniformly valid confidence intervals of parametric length — and runs in about 20 seconds, because it never refits a machine-learning model.

The problem

A hospital wants to deploy a diagnostic classifier on its own patients. It has richly labelled records from several other hospitals, and plenty of its own patient records — but none of those local records carry the outcome label yet. This is multi-source unsupervised domain adaptation: labelled data from source domains \(1,\dots,L\), and only unlabelled covariates from the target, typically many more of them than labelled source records.

The domains genuinely differ. Both the feature distribution and the way the label depends on features can shift. So a classifier trained on any single source, or naively pooled across all of them, bakes in assumptions about the target that you cannot check without target labels. And the target's conditional \(\mathbb Q_{Y\mid X}\) is not identifiable from unlabelled data, no matter how much of it you collect.

The idea: fix what you can see, hedge over what you cannot

Since \(\mathbb Q_{Y\mid X}\) is unobservable, we refuse to guess it. Instead we hold the (observable) target covariate distribution \(\mathbb Q_X\) fixed and treat the conditional as some mixture of the sources' conditionals, giving the uncertainty set

\[ \mathcal C \;=\; \Big\{\, (\mathbb Q_X,\, \mathbb T_{Y\mid X}) \ :\ \mathbb T_{Y\mid X} = \textstyle\sum_{l=1}^{L} \gamma_l\, \mathbb P^{(l)}_{Y\mid X},\ \ \gamma \in \Delta^L \,\Big\}. \]

Conditional Group DRO then learns the classifier minimizing the worst-case risk over that set. With the cross-entropy loss \(\ell\), the model \(\theta^\star\) solves

\[ \theta^\star \;=\; \arg\min_{\theta}\ \max_{\gamma \in \Delta^L}\ \sum_{l=1}^{L} \gamma_l\; \mathbb E_{X\sim \mathbb Q_X}\, \mathbb E_{Y\sim \mathbb P^{(l)}_{Y\mid X}}\, \ell(X, Y, \theta). \]

Whatever the true target mixture turns out to be — provided it lies in this class — the classifier's loss is controlled.

Cross-entropy makes this tractable in a way few losses do. Deploying a multinomial-logit model with \(\theta = (\theta_1^\top,\dots,\theta_K^\top)^\top\), the objective splits into a part linear in \(\gamma\) and a log-sum-exp part free of \(\gamma\) entirely:

\[ \phi(\theta,\gamma) \;=\; \sum_{l=1}^{L}\gamma_l\, \theta^\top\mu^{(l)} \;+\; S(\theta), \qquad S(\theta) = \mathbb E_{X\sim\mathbb Q_X}\log\Big[1+\textstyle\sum_{c=1}^{K}\exp(\theta_c^\top X)\Big], \]

where each source contributes only through a single vector of first moments,

\[ \mu_c^{(l)} \;=\; -\,\mathbb E_{X\sim\mathbb Q_X}\ \mathbb E_{Y\sim\mathbb P^{(l)}_{Y\mid X}}\big[\mathbf 1(Y=c)\,X\big]. \]

Everything the \(L\) sources have to say is compressed into the \(dK \times L\) matrix \(\mathcal U = (\mu^{(1)},\dots,\mu^{(L)})\). The problem is strictly convex in \(\theta\) and linear in \(\gamma\), and we solve it with Mirror Prox: gradient descent on \(\theta\), entropic mirror ascent on \(\gamma\) (a closed-form simplex update), with an extra-gradient correction that buys an \(O(1/T)\) duality gap instead of the \(O(1/\sqrt T)\) of one-step gradient descent–ascent.

Under covariate shift the \(\mu^{(l)}\) are not sample means, because the expectation over \(X\) is under the target law while the labels come from the source. Writing \(f_c^{(l)}(x) = \mathbb P^{(l)}(Y=c\mid x)\) and \(\omega^{(l)} = d\mathbb Q_X/d\mathbb P^{(l)}_X\), the estimator is a plug-in on target covariates plus a source-side correction:

\[ \widehat\mu_c^{(l)} = -\frac1N\sum_{j=1}^{N}\widehat f_c^{(l)}\big(X_j^{\mathbb Q}\big)X_j^{\mathbb Q} \;-\; \frac{1}{n_l}\sum_{i=1}^{n_l}\widehat\omega^{(l)}\big(X_i^{(l)}\big)\Big(\mathbf 1\big(Y_i^{(l)}=c\big) - \widehat f_c^{(l)}\big(X_i^{(l)}\big)\Big)X_i^{(l)} . \]

This is double machine learning, cross-fitted: both nuisances may be any flexible learner (multinomial logistic, random forest, XGBoost, a net), and their errors enter only through the product \(\|\widehat\omega/\omega - 1\|\cdot\|\widehat f - f\|\). If that product is \(o(n^{-1/2})\) — which slow-converging ML estimators can satisfy jointly even when neither is \(\sqrt n\)-consistent alone — then \(\max_l\|\widehat\mu^{(l)}-\mu^{(l)}\|_2 \lesssim \sqrt{d/n}\). The usual DML idea, applied not to a scalar parameter but to a whole risk function.

The same uncertainty set, applied to regression with a worst-case reward instead of a worst-case cross-entropy, is the subject of an earlier paper. There the minimax problem collapses to a quadratic program and the analysis stops at convergence rates. Cross-entropy is harder, and inference is the whole point here.

How this differs from Group DRO

Classical Group DRO mixes the full joint distributions \(\mathbb P^{(l)}\), and so never looks at the target at all. CG-DRO keeps \(\mathbb Q_X\) fixed and mixes only the conditionals, tailoring the guarantee to the covariates you will actually deploy on. When there is no covariate shift the two uncertainty sets coincide and the two methods agree — in the paper's simulation both land at a worst-case loss of ≈ 0.687. Under shift they separate: Group DRO reaches ≈ 0.697–0.700 while CG-DRO reaches ≈ 0.685–0.691, and Group DRO's distance to \(\theta^\star\) is roughly 0.10 against CG-DRO's 0.06.

0.60.81.01.21.40.10.30.50.70.91.351.120.72ERM (pooled)Group DRO ≈ 0.89CG-DRO ≈ 0.71fraction of the labelled sample drawn from source 1worst-case target risk
Figure 1. Worst-case risk on the target as the labelled data is mixed differently across two sources (\(L=2\), \(n_1+n_2=4000\)). ERM's risk depends entirely on an accident of sampling: it is 1.35 when source 1 is scarce, dips to 0.72 near a balanced mix, and climbs back to 1.12 at the other extreme. It happens to do well in the middle — and you have no way to know whether your sample sits there. Both DRO methods are flat, which is the point of a worst-case objective, but Group DRO sits at ≈ 0.89 while CG-DRO sits at ≈ 0.71, because CG-DRO alone uses the target's covariates. Values read from the paper's Figure 2.

Getting the rate: a detour through two surrogate problems

Before inference, one needs a rate — and the obvious route gives a bad one. Standard M-estimation arguments fail here because after taking expectations, the inner maximum over \(\gamma\) is concave but not strongly concave: the maximizer need not be unique. Pushed through naively, that yields only \(\|\widehat\theta - \theta^\star\|_2 \lesssim (d/n)^{1/4}\), a rate no one wants.

The repair is to replace the log-sum-exp term \(S(\theta)\) with its second-order Taylor expansion about \(\widehat\theta\), giving a pair of surrogate minimax problems — one population, one empirical — whose duals are strongly concave quadratics in \(\gamma\). Two facts make this more than a trick. First, the empirical surrogate's solution is not an approximation at all: \(\widehat\theta_{\mathrm{ap}} = \widehat\theta\) exactly. Second, the population surrogate is within \(\|\widehat\theta-\theta^\star\|_2^2\) of \(\theta^\star\), a second-order distance. Chaining the two recovers the parametric rate,

\[ \|\widehat\theta - \theta^\star\|_2 \;\lesssim\; \frac{1}{\sigma_L^2(\mathcal U)}\sqrt{\frac{d}{n}}, \qquad \|\gamma^\star_{\mathrm{ap}} - \gamma^\star\|_2 \;\lesssim\; \Big(\frac{d}{n}\Big)^{3/4}, \]

where \(\sigma_L(\mathcal U)\) is the smallest singular value of the source-summary matrix. That quantity is the paper's measure of cross-source heterogeneity, and it is not a technicality: if two sources are identical their \(\mu^{(l)}\) coincide, \(\sigma_L(\mathcal U) = 0\), and the worst-case weight is simply not identified. Robustness across sources requires that the sources actually differ. Empirically the rate holds even with gradient-boosted nuisances — the paper fits \(1.87\,n^{-1/2}\) in a linear-score setting and \(1.35\,n^{-1/2}\) in a strongly nonlinear one.

The statistical twist: two ways normality dies

Worst-case optimizers are statistically awkward, and it pays to be precise about why. Write \(\widehat\theta = \arg\min_\theta \widehat\phi(\theta, \widehat\gamma)\). If the worst-case weight \(\gamma^\star\) were known, \(\widehat\theta\) would be an ordinary M-estimator and asymptotically normal. The trouble is entirely in \(\widehat\gamma\), and it arrives in two distinct ways:

One source dominates. When the shift between sources is sparse but pronounced, the maximizing weight sits on the boundary of the simplex — one domain takes essentially all of it. The estimator \(\widehat\gamma_1\) then piles up at zero, a point mass rather than a bell curve.

Sources are nearly tied. When two source conditionals are almost identical, the inner maximizer is nearly non-unique, and \(\widehat\gamma_1\) swings between the vertices 0 and 1 depending on noise. The histogram is bimodal.

In both regimes \(\widehat\theta_1\) inherits the pathology: it is biased, non-Gaussian, and normal-approximation confidence intervals are simply invalid. So is the bootstrap, which resamples the same broken object.

regular00.51interior optimum: a bell curveone source dominates00.51mass piles onto the boundarysources nearly tied00.51bimodal across the vertices
Figure 2. A schematic of the sampling distribution of the estimated worst-case weight \(\widehat\gamma_1\) across simulation runs. Only the first regime is well behaved, with the estimate scattered around an interior optimum (dashed). In the second the mass collapses onto the simplex boundary; in the third it splits between the two vertices, so no single Gaussian describes it. These are the three shapes the paper's Figure 4 exhibits for the classifier coordinate \(\widehat\theta_1\) itself; the histograms here are drawn to illustrate them, not simulated from the paper's design (whose three panels use different \(L\) and \(d\)).

The fix: perturb the troublesome part, keep the rest

The randomness decouples cleanly. Given the weights, the classifier is a standard, asymptotically normal M-estimator; all the nonstandard behaviour lives in \(\widehat\gamma\). So we resample only that piece:

  1. Estimate the source summaries \(\widehat\mu^{(l)}\) once, by double ML, together with their covariance estimates \(\widehat{\mathbf V}^{(l)}_\mu\).
  2. Draw \(M\) Gaussian perturbations \(\widehat\mu^{(l,m)} \sim \mathcal N\big(\widehat\mu^{(l)},\, \widehat{\mathbf V}^{(l)}_\mu + \tfrac{1}{n_l}\mathbf I\big)\), exploiting the asymptotic normality of \(\widehat\mu^{(l)}\). The inflation term keeps the draw non-degenerate. This is cheap, and needs no ML refitting.
  3. For each draw, maximize the perturbed quadratic \(\widehat F^{[m]}(\gamma)\) over the simplex to get a weight \(\widehat\gamma^{[m]}\). With enough draws, some \(m^\star\) lands within \(o(n^{-1/2})\) of the true \(\gamma^\star\) — we just don't know which.
  4. Discard extreme draws (those whose standardized perturbation exceeds \((1+\eta_0)\,z_{\alpha_0/(dKL)}\), with defaults \(\eta_0 = 0.1\), \(\alpha_0 \le 0.01\)), then for each survivor solve a plain convex problem for \(\widehat\theta^{[m]}\) and form its ordinary sandwich-based normal interval at level \(\alpha' = \alpha - \alpha_0\).
  5. Report the union of those intervals.

Because some perturbation recovers the true weight, the union covers \(\theta^\star_1\) whether or not the ties or the boundary case actually occur — coverage is uniformly valid across regular, boundary, and near-tied regimes. And despite being a union of many intervals, its length is \(O(n^{-1/2})\): no rate is lost. The same construction tests any linear contrast, and rejects \(H_0 : \theta^\star_1 = 0\) exactly when 0 falls outside the union.

The one price is that the uniform-validity theorem needs the heterogeneity \(\sigma_L^2(\mathcal U)\) to be not merely positive but bounded below — roughly \(\sigma_L \gg n^{-1/16}\). Perfect robustness across near-identical sources is not merely hard to certify; it is ill-posed.

Does it work?

The paper's evidence is simulation — there is no real-data application in it — but the simulations target exactly the failure modes above, and the comparison is stark. The benchmark "OBA" is an oracle bias-aware interval that knows the true bias and cannot be run in practice; the "Normal approximation" row is itself generous, using oracle standard errors computed across the 500 replications.

Coverage of nominal 95% intervalsBoundary regimeNear-tied regime
Perturbation (ours)≈ 95%≈ 97%
Normal approximation≈ 65%≈ 79%
Bootstrap≈ 55%≈ 82%
Oracle bias-aware (not implementable)≈ 94%≈ 94%

Hardest configuration of each regime, 500 replications, values read from the paper's Figure 8. In the regular regime all methods sit near 95%, as they should — the failures appear only as the regime degenerates. Note that the two invalid methods fail in opposite orders across regimes: the bootstrap is the worse of the two when one source dominates (55%), the normal approximation when sources are nearly tied (79%).

Validity is not bought with uselessly wide intervals. In the boundary regime the perturbation interval grows from 0.67 to 0.90 as the shift \(\delta\) increases, converging toward the oracle's length (≈ 0.55 → 0.86) and nearly coinciding with it once the shift is pronounced; the invalid intervals stay near 0.51–0.63, which is precisely why they miss. A more pointed comparison appears in the paper's fully nonparametric setting, where nuisances are fit by cross-validated XGBoost: there the bootstrap is omitted entirely, because refitting 500 ML models per replication is computationally prohibitive. The perturbation scheme, which refits nothing, runs in about twenty seconds.

≈ 20 sFull uncertainty quantification with \(M = 500\) perturbations — no ML refitting, unlike the bootstrap
\(\sqrt{d/n}\)Estimation rate for \(\widehat\theta\), recovered via two surrogate problems and confirmed empirically with gradient-boosted nuisances
55%Bootstrap's actual coverage of its nominal 95% intervals when one source dominates

Why it matters

CG-DRO delivers both halves of what unsupervised domain adaptation needs: a classifier with a guarantee that holds under distribution shift, and honest uncertainty quantification for it — even in the non-regular regimes where standard asymptotics break down. Testing whether a coordinate \(\theta^\star_j\) is zero becomes a test of whether a feature is robustly predictive: predictive across every plausible target, not just on the sources you happened to collect. That is what lets a model trained on other hospitals be deployed on your own patients with a confidence interval you can defend.

The guarantee is relative to the mixture class: if the target's conditional is not representable as a mixture of the sources', the worst case need not cover it, and the result can be conservative — the supplement shows how prior knowledge tightens \(\mathcal C\). The method needs the sources to be genuinely heterogeneous: with identical sources the summaries \(\mu^{(l)}\) become collinear, \(\sigma_L(\mathcal U) = 0\), and the weight is unidentified. The deployed classifier is a multinomial logistic model in \(X\), with all the flexibility living in the ML nuisance estimates rather than the final model. The theory also asks for a lot of data in the right places — \(n \gg d^3\), and unlabelled target data dominating the sources, \(N \gg n\log N\) — and the confidence-interval results are stated for fixed dimension. Finally, all evidence is simulated; the paper contains no real-data application.

© 2026 Zhenyu (Zach) Wang 王振宇