Technical Summary: Statistical Analysis of Conditional Group Distributionally Robust Optimization with Cross-Entropy Loss

1. Problem Formulation

The paper addresses multi-source unsupervised domain adaptation for multi-category classification. In this setting, labeled data are available from $L$ source domains, while only unlabeled covariate data are observed from a target domain. The core challenge is distributional heterogeneity, where both the covariate distributions ($P_X$ vs. $Q_X$) and the conditional outcome distributions ($P_{Y|X}$ vs. $Q_{Y|X}$) may shift between sources and the target.

Standard Empirical Risk Minimization (ERM) often fails to generalize under such shifts. While Group Distributionally Robust Optimization (Group DRO) offers a solution by minimizing worst-case risk over mixtures of joint source distributions, it does not leverage the specific covariate information available in the unlabeled target domain.

The authors propose the Conditional Group Distributionally Robust Optimization (CG-DRO) framework. Unlike Group DRO, which considers uncertainty over joint distributions, CG-DRO constructs an uncertainty set over conditional outcome distributions (mixtures of source conditionals) while holding the target covariate distribution fixed. The goal is to learn a classifier $\theta^*$ that minimizes the worst-case cross-entropy loss over this uncertainty set:
$$\theta^* = \arg \min_{\theta} \left[ \max_{\gamma \in \Delta_L} \sum_{l=1}^L \gamma_l \cdot \mathbb{E}_{X \sim Q_X, Y \sim P^{(l)}_{Y|X}} [\ell(X, Y, \theta)] \right]$$
where $\ell$ is the cross-entropy loss and $\gamma$ represents the mixture weights.

2. Methodology

2.1 Optimization Algorithm (Mirror Prox with DML)

To solve the empirical CG-DRO problem, the authors employ a Mirror Prox algorithm, which performs gradient descent on the model parameters $\theta$ and mirror ascent on the mixture weights $\gamma$.

A critical component is the estimation of the risk function components under covariate shift. Since the target covariate distribution $Q_X$ differs from source distributions $P^{(l)}_X$, direct empirical averages from source data are biased. The authors integrate Double Machine Learning (DML) to construct an unbiased estimator $\hat{\mu}^{(l)}$ of the conditional expectation terms.

• Step 1: Fit a machine learning model (e.g., neural networks, random forests) to estimate the conditional probability $f^{(l)}(x) = P^{(l)}(Y|X)$ using source data.
• Step 2: Estimate the density ratio $\omega^{(l)}(x) = dQ_X/dP^{(l)}_X$.
• Bias Correction: Construct a DML estimator that corrects the bias of the plug-in estimator using the density ratio and the residual of the conditional probability model. This ensures that errors in nuisance estimation (the ML models) enter the final risk estimator only at a higher-order rate.

2.2 Statistical Convergence Analysis

The paper establishes fast statistical convergence rates for the empirical estimator $\hat{\theta}$.

• Challenge: Standard M-estimation theory yields slow rates ($n^{-1/4}$) here because the objective involves a maximization over $\gamma$ after the expectation, introducing non-smoothness and potential non-strong concavity.
• Solution: The authors construct two surrogate minimax optimization problems (approximate population and empirical problems) using a second-order Taylor expansion of the log-sum-exp term.
  • These surrogates allow the derivation of a closed-form expression for the approximate solution.
  • By analyzing the discrepancy between the original and surrogate problems, they prove that the empirical estimator $\hat{\theta}$ converges to the population optimum $\theta^*$ at a fast parametric rate of $O(\sqrt{d/n})$, provided the source conditional distributions are sufficiently distinct (linear independence of $\mu^{(l)}$ vectors).

2.3 Uncertainty Quantification (Perturbation-Based Inference)

A major contribution is addressing the nonstandard asymptotics of the estimator. The empirical CG-DRO estimator $\hat{\theta}$ often fails to converge to a standard normal distribution due to:

1. Boundary effects: When the optimal mixture weight $\gamma^*$ lies on the boundary of the simplex (e.g., one source dominates).
2. System instability: When source conditional distributions are highly similar, causing the optimal weight to fluctuate wildly.

To enable valid inference (confidence intervals and hypothesis testing) in these non-regular regimes, the authors propose a perturbation-based inference procedure:

• Idea: Instead of relying on the asymptotic normality of $\hat{\theta}$, the method generates a collection of perturbed weight vectors $\{\hat{\gamma}^{[m]}\}$ by resampling the estimated nuisance parameters ($\hat{\mu}^{(l)}$).
• Filtering: A filtering step excludes extreme perturbations that deviate significantly from the observed data structure.
• Union of Intervals: For each filtered perturbation, a "semi-oracle" estimator $\hat{\theta}^{[m]}$ is computed (treating the perturbed weight as fixed). Confidence intervals are constructed for each $\hat{\theta}^{[m]}$ using standard Gaussian quantiles (since the weight is fixed, the estimator becomes asymptotically normal).
• Final CI: The final confidence interval is the union of these individual intervals. Theoretically, this union is guaranteed to contain the true parameter with the desired coverage probability, even if the original estimator is non-normal.

3. Key Contributions

1. Framework: Proposal of the CG-DRO framework for cross-entropy loss, which leverages target covariates to define a conditional uncertainty set, distinct from standard Group DRO.
2. Algorithm & Estimation: Development of a Mirror Prox algorithm enhanced with a Double Machine Learning estimator to handle covariate shifts and nuisance parameter estimation errors efficiently.
3. Convergence Theory: Establishment of a fast $O(\sqrt{d/n})$ convergence rate for the empirical CG-DRO estimator by constructing surrogate minimax problems, overcoming the limitations of standard M-estimation in minimax settings.
4. Inference: Introduction of a perturbation-based inference method that yields uniformly valid confidence intervals with parametric length, even when the estimator exhibits nonstandard (non-normal) limiting distributions due to boundary effects or system instability.

4. Results

• Convergence: Simulations confirm that the estimation error $\|\hat{\theta} - \theta^*\|_2$ decays at the rate of $1/\sqrt{n}$, aligning with the theoretical bounds.
• Robustness: In numerical comparisons, CG-DRO consistently achieves lower worst-case risk than both Group DRO and ERM under covariate shifts.
• Inference Validity:
  • In regular settings (distinct sources), standard methods work, but CG-DRO maintains performance.
  • In non-regular settings (sparse shifts causing boundary weights) and unstable settings (similar sources causing weight fluctuation), standard Normality and Bootstrap methods fail to maintain nominal coverage (e.g., dropping below 95%).
  • The proposed perturbation-based method maintains uniform coverage across all regimes, with interval lengths comparable to an oracle benchmark.

5. Significance

The paper claims to fill a critical gap in the theoretical understanding of Distributionally Robust Optimization (DRO) under cross-entropy loss, a setting where closed-form solutions are unavailable and standard asymptotic theory often breaks down.

The significance lies in providing a principled framework for inference in minimax optimization problems where the estimator's limiting distribution is non-standard. By decoupling the uncertainty of the nuisance parameters (weights) from the structural parameters (classifier), the proposed perturbation method offers a computationally efficient and theoretically sound alternative to bootstrap methods that fail in non-regular regimes. This enables reliable uncertainty quantification for robust models deployed in environments with distributional shifts, such as multi-source medical or geographic data analysis.