Robust Generalization with Adaptive Optimal Transport Priors for Decision-Focused Learning

AI summary

Overview

• Research area: Few-shot learning, distributionally robust optimization (DRO), and optimal transport (OT), with a decision-focused framing.
• Technical level: Advanced. The paper leans on Sinkhorn distances, entropic optimal transport duality, Wasserstein ambiguity sets, and a convex one-dimensional dual problem in a scalar variable λ.
• Scope (one sentence): The paper proposes PG-DRO, a framework that replaces the single fixed reference distribution inside Sinkhorn DRO with class-adaptive Gaussian-mixture priors built by hierarchical optimal transport from base-class prototypes, and tests it on synthetic classification/regression and on CIFAR-10/CIFAR-100 and ImageNet-family few-shot benchmarks.

What This Paper Is About

Few-shot learners must generalize from only a handful of labeled examples, and in practice the test distribution is rarely the same as the training distribution, so small perturbations can push examples across decision boundaries. Existing Sinkhorn DRO methods give robustness guarantees but anchor the uncertainty set to one fixed reference distribution, which cannot adapt per class when support data are scarce. The goal of this work is to make the reference distribution itself class-specific and data-adaptive, so that robustness is aligned with transferable structure learned from abundant base classes.

Key Contributions

1. PG-DRO framework: A decision-making framework for low-data regimes that uses hierarchical optimal transport to construct class-adaptive priors inside an entropic (Sinkhorn) DRO formulation, enabling scalable robust learning under limited supervision.
2. Robust generalization in the few-shot regime: Demonstrates maintained accuracy under perturbations and distribution shifts, with the stated benefit falling particularly on few-shot minority classes.
3. Theoretical analysis: Provides conditions for robust optimization with adaptive priors, including a multi-class convexity result (Corollary 4.1), a contraction result for the robust value under the adaptive OT prior update (Theorem 4.2), and a consistency result (Theorem 4.3).
4. Empirical validation: Experiments on synthetic data and real benchmarks reporting consistent gains in both accuracy and robustness under distributional shifts.

Main Findings

• Fixed priors are the bottleneck: The paper argues that a single predefined reference distribution ν in Sinkhorn DRO "is agnostic to the downstream few-shot task," ignores class geometry, and is prone to misalignment when the support set is small.
• Class-adaptive priors are constructed by hierarchical OT: For each support example, a soft-min Sinkhorn cost is computed to every base class; the resulting cost matrix C ∈ ℝ^(B×N) is solved as an entropic OT problem to produce a transport plan T ∈ ℝ^(B×N); aggregating mass over each support class yields mixture weights w̃_bc and a Gaussian-mixture prior ν_c = Σ_b w̃_bc N(μ_b, Σ_b).
• One-dimensional per-class optimization: Plugging ν_c into the Sinkhorn kernel gives a class-specific robust logit V_c(x; θ) that is smooth and strictly convex in λ, solvable by a one-dimensional Newton step per class.
• Contraction under adaptive OT: Theorem 4.2 states that Δ_(t+1)(θ) ≤ (1 − η_t κ) Δ_t(θ) + O(N^(−1/2)), i.e., the robust value approaches the oracle prior value at a linear rate modulo a stochastic term from few-shot estimation; the gap vanishes as N → ∞ and Σ_t η_t = ∞.
• Consistency: Theorem 4.3 gives V_c^(N)(x; θ) → V_c^⋆(x; θ) and λ_c^(N)(x) → λ_c^⋆(x) under stated continuity, Lipschitz, compactness, and bounded-cost assumptions.
• Simulation results (Table 1, "different levels of disturbance"): PG-DRO reports the best average and worst-10% accuracy at every disturbance level shown. At λ_cov = 0.0: Pure ERM 0.07 ± 0.00 avg / 0.67 ± 0.02 worst-10%, Classical OT 74.06 ± 1.20 / 95.68 ± 0.15, PG-DRO 78.10 ± 0.06 / 98.93 ± 0.75. At λ_cov = 5.0: Pure ERM 21.68 ± 0.26 / 23.27 ± 0.95, Classical OT 17.43 ± 0.13 / 22.00 ± 0.21, PG-DRO 32.53 ± 0.04 / 26.60 ± 0.49.
• CIFAR-10 under Laplace noise (Table 2): With k = 1, 5, 10 supports, PG-DRO reports 20.88 ± 0.57 / 50.45 ± 0.54 / 61.54 ± 0.34 for Laplace1; 22.28 ± 0.71 / 50.21 ± 0.86 / 59.57 ± 0.45 for Laplace2; 17.71 ± 0.24 / 41.53 ± 0.39 / 50.58 ± 0.41 for Laplace5. The comparison methods were Few-shot, SAA, and W-DRO.
• CIFAR-10 under Gaussian noise (Table 3): PG-DRO reports 20.53 ± 1.06 / 53.62 ± 1.08 / 60.78 ± 0.44 for Gaussian1; 22.74 ± 1.44 / 51.00 ± 1.15 / 60.34 ± 0.48 for Gaussian2; 18.80 ± 0.34 / 41.86 ± 0.37 / 50.47 ± 0.48 for Gaussian5, with the same baseline set.
• Transport alignment sharpens with support size: The heatmap of normalized weights w̃_bc (Figure 2) shows alignments progressively concentrating along the diagonal as the number of support samples increases.
• Regression behavior: The paper states that PG-DRO is consistently competitive and that its relative advantage persists as noise grows, but the MSE, MAE, and worst-10% MSE values are not given as numbers in the provided text (they are shown in a figure).
• Setup details: Simulation classification used 3 to 8 randomly sampled labeled support instances per class; CIFAR-10 used S ∈ {1, 5, 10} supports per class. Source priors used Dir(1) and target priors used Dir(α_test) with α_test ≪ 1, producing a long-tailed head-tail few-shot setting. Experiments ran on a 13th Gen Intel Core i9-13900HX CPU (24 cores) and an NVIDIA A100 GPU (40 GB).

Methodology in Plain English

Standard Sinkhorn DRO asks: how bad could the loss get if the data distribution moved slightly away from a reference distribution, where movement is penalized by an entropy-regularized optimal transport cost? The answer is a one-dimensional convex problem in a scalar λ, which is fast to solve — but the reference distribution is fixed and, in few-shot settings, estimated from very few points.

PG-DRO keeps that machinery and swaps in a different reference for each class. It first looks at the abundant base data, computing summary statistics (mean μ_b and covariance Σ_b) for each base class. For each few-shot support example, it measures a smoothed (soft-min) transport cost to every base class. These costs form a matrix, and solving an entropic OT problem over that matrix produces a transport plan saying how much each base class contributes to each novel class. Those contributions become mixture weights, so each novel class gets its own Gaussian mixture prior built from base-class prototypes.

That class-specific prior is then plugged into the Sinkhorn DRO dual, producing a "robust logit" per class, which is solved with a single Newton step. Predictions take the argmax over these robust logits, and training minimizes ordinary softmax cross-entropy applied to them — so the method slots into normal training pipelines rather than requiring a special adversarial loop.

Why This Matters

• Research impact: It challenges the fixed-reference assumption shared across Wasserstein DRO, φ-divergence DRO, and MMD-based DRO, and shows one can adapt the ambiguity set per class while retaining Sinkhorn DRO's tractability. It also bridges hierarchical OT methods (previously used for distribution calibration or metric alignment) with robust optimization, using OT to produce robust logits rather than just corrected features.
• Real-world applications:
  • Rare or extreme operating conditions where abundant normal-condition data exists but only a few examples of the failure or anomaly regime are available.
  • Sensor-based deployment subject to noise, domain shift, or adversarial manipulation of inputs.
  • Long-tailed recognition where minority classes receive very few prototypes but must still be classified reliably.
  • Decision systems that must remain stable under unexpected environmental change, where the paper emphasizes decision-focused training.
• Industry relevance: Because the method reduces to a one-dimensional convex solve per class with a Newton step and reuses standard cross-entropy training, it is computationally practical for deployment; the experiments were run on a single A100 GPU (40 GB), and the code is stated to be available.

Future Directions

1. Close the empirical gaps in the provided text: The ImageNet-family results (mini-ImageNet, tiered-ImageNet) and the regression metrics are deferred to an appendix or figure; quantitifying how PG-DRO behaves there is a natural next step.
2. Relax the theoretical assumptions: The contraction guarantee assumes the OT map is locally contractive near w^⋆ and is stated with "wild and common assumptions"; testing behavior outside that neighborhood, and the O(N^(−1/2)) few-shot term, remain open.
3. Scaling and cost: Extending the per-class Newton solve and hierarchical OT to many more classes and larger base datasets is open, since the paper's experiments use CIFAR and ImageNet-family benchmarks.
4. Better robustness targets: The paper notes that informative outliers highlight which decision boundaries are most vulnerable to shift; systematically exploiting those to guide data collection or active support selection is a logical extension.

Target Audience

Researchers and practitioners working on few-shot learning, distributionally robust optimization, and optimal transport, particularly those who need robustness guarantees to survive real distribution shift rather than only clean-benchmark accuracy. It is most useful to readers with graduate-level background in convex duality and OT; applied machine learning engineers working on long-tailed or safety-critical deployment may benefit from the method's practical training recipe even if they skip the proofs.

Authors’ abstract

Few-shot learning requires models to generalize under limited supervision while remaining robust to distribution shifts. Existing Sinkhorn Distributionally Robust Optimization (DRO) methods provide theoretical guarantees but rely on a fixed reference distribution, which limits their adaptability. We propose a Prototype-Guided Distributionally Robust Optimization (PG-DRO) framework that learns class-adaptive priors from abundant base data via hierarchical optimal transport and embeds them into the Sinkhorn DRO formulation. This design enables few-shot information to be organically integrated into producing class-specific robust decisions that are both theoretically grounded and efficient, and further aligns the uncertainty set with transferable structural knowledge. Experiments show that PG-DRO achieves stronger robust generalization in few-shot scenarios, outperforming both standard learners and DRO baselines.