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

TL;DR

This paper introduces PG-DRO, a novel framework that leverages hierarchical optimal transport to learn class-specific priors, enhancing robust generalization in few-shot learning scenarios with limited data.

Contribution

It proposes a new PG-DRO method that learns adaptive priors from base data and integrates them into Sinkhorn DRO, improving robustness and transferability.

Findings

PG-DRO outperforms standard learners in few-shot tasks.

It achieves stronger robustness against distribution shifts.

The method effectively incorporates structural knowledge into decision-making.

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.