Learning Predictive Ambiguity Sets for Decision-Focused Distributionally Robust Optimization

Abstract: Predict-then-optimize systems usually compress uncertainty into a point forecast and then solve a downstream optimization problem as if the forecast were reliable. Distributionally robust optimization (DRO) offers protection against misspecification, but the ambiguity set is often centered at historical samples and uses a fixed radius. We propose \emph{learned predictive ambiguity sets} (LPAS): a deep contextual model outputs a finite nominal scenario distribution, a state-dependent Wasserstein radius, and optionally an anisotropic ground metric. These outputs define a contextual ambiguity set that feeds a DRO decision layer. The radius is trained by a combination of conditional quantile calibration, size regularization, and downstream decision loss, so that robustness is adaptive rather than globally fixed. We derive the finite dual form used by the decision layer, present a staged training algorithm, and evaluate the method on distributionally robust portfolio optimization with 20 S&P 500 constituents from 2018--2026. The proposed method substantially improves over equal-weight, predict-then-optimize, and historical Wasserstein DRO baselines, achieving 26.28% annualized return, Sharpe ratio 1.30, final wealth 1.61, and lower tail loss than a deep fixed-radius DRO baseline while using a smaller average radius. The results show that learned ambiguity radii can recover most of the performance of strong fixed-radius DRO while reducing unnecessary conservatism and improving regime adaptivity.

• The paper presents a novel LPAS framework that learns context-dependent ambiguity set sizes to address limitations in static DRO methods.
• It integrates deep Transformer-based scenario generation with a calibrated radius network to optimize decision-focused performance under uncertainty.
• Empirical results in portfolio selection show LPAS achieves superior returns and risk management by adapting ambiguity radii to changing market volatility.

Learned Predictive Ambiguity Sets for Contextual Distributionally Robust Optimization

Context and Motivation

Predict-then-optimize paradigms in machine learning-driven decision systems typically rely on point forecasts when confronting real-world uncertainty. This approach insufficiently accounts for the amplification of model error through sensitive downstream optimization, particularly when errors align with decision-critical subspaces. Distributionally Robust Optimization (DRO)—with Wasserstein ambiguity sets—offers a principled alternative by safeguarding solutions against a ball of plausible distributions centered at empirical data with a prescribed radius. However, conventional DRO implementations use hand-tuned, globally fixed ambiguity radii, neglecting state-dependent uncertainty and sensitivity, and thus exhibit a misalignment between the deployed degree of robustness and the risk present in specific contexts.

Methodological Innovations

This work introduces Learned Predictive Ambiguity Sets (LPAS), a framework for end-to-end learning of both the center and size of the Wasserstein ambiguity set as direct contextual outputs of a deep model. Given observed contexts $z_t$, the model emits a finite-support scenario distribution and a state-dependent radius $\rho_\phi(z_t)$, with the ambiguity set

$$P_{\theta,\phi}(z_t) = \{ Q : W_c(Q, \widehat{P}_\theta(\cdot|z_t)) \leq \rho_\phi(z_t) \}.$$

This contextual ambiguity set directly informs a downstream DRO layer. Notably, the radius is not heuristically fixed or pre-calibrated; it is trained via a combination of objectives to maintain (i) predictive calibration (quantile-conformal style), (ii) regularization to avoid unnecessary conservatism, and (iii) direct optimization of realized decision-relevant loss. The framework supports the optional learning of state-dependent ground metrics but evaluates with the Euclidean metric for tractability.

Key architectural elements comprise:

• A Transformer-based scenario generator producing $N$ predicted scenarios and probabilities, enabling the construction of a nominal predictive distribution.
• A radius network $\rho_\phi(z_t)$, parameterized to ensure nonnegativity (optionally bounded), predicting the ambiguity set size as a function of context features.
• A decision layer solving the convex inner maximization implied by the Wasserstein-DRO dual, with closed forms available under affine loss and $\ell_2$ transport.

Numerical Performance and Empirical Analysis

LPAS is empirically validated in multi-asset portfolio selection with $d=20$ S&P 500 constituents over a rolling 8.5-year period. The predictive backbone is a deep Transformer. The experimental protocol compares LPAS against:

• Classical equal-weight portfolios (EW).
• Predict-then-optimize (P2O) methods (point forecast only).
• Historical Wasserstein DRO (Hist-WDRO) with empirical centers and validation-tuned fixed radii.
• Deep fixed-radius DRO with the same Transformer scenario generator but a constant radius.

Main findings include:

• LPAS achieves annualized return 26.28%, Sharpe ratio 1.30, and final wealth 1.61, consistently outperforming EW, P2O, and Hist-WDRO.
• Compared to fixed-radius DRO, LPAS uses a substantially smaller average radius (24.3 vs. 35.4) yet retains nearly all performance on standard risk-adjusted metrics, with notably improved downside tail behavior.

  Figure 1: Out-of-sample cumulative wealth. Robust methods stabilize wealth accumulation relative to predict-then-optimize.

Notably, the learned radius adapts sharply to volatility regimes, expanding in high-uncertainty windows and shrinking during tranquil market periods.

Figure 2: Learned LPAS-W radius across market-volatility regimes. The radius increases in uncertain periods and decreases in calmer periods.

Calibration diagnostics demonstrate that LPAS can achieve high empirical coverage with smaller mean radii versus constant-radius robust baselines, indicating a superior trade-off between protection and conservatism.

Figure 3: Coverage versus radius conservatism. LPAS-W attains high empirical coverage with a smaller mean ambiguity radius than fixed-radius robust optimization, illustrating the benefit of state-dependent radius learning.

Ablation studies reveal the critical role of the calibration objective for empirical coverage and downside protection. Elimination of the calibration term causes the coverage to collapse, leading to significant performance degradation and increased turnover.

Theoretical and Practical Implications

The introduction of contextually adaptive ambiguity radii directly addresses DRO's historical limitation of mismatched or globally-tuned robustness parameters. By fusing calibration-driven, decision-aware objectives, LPAS moves beyond risk-hedging via static regularization to a setting where the robustness level is learned in a regime-aware, data-driven manner. This offers strong theoretical implications for robust optimization under covariate shift and environment non-stationarity.

Practically, LPAS provides robust decision solutions with substantially reduced excess conservatism, which is crucial for high-stakes applications such as dynamic portfolio allocation, real-time logistics, and critical infrastructure operations. Empirical evidence demonstrates that adaptivity in ambiguity set size recovers nearly all the utility of strong robustification in turbulent regimes, without the attendant utility loss in stable periods.

Limitations and Future Directions

Current experiments are constrained to single rolling splits of a modest asset universe, and the learned anisotropic ground metric is not explored. The calibration achieved is generally robust but trades exact conditional coverage for improved decision performance; for applications requiring strict valid coverage, post-hoc conformal calibration can be incorporated. Richer universes, multi-seed evaluations, and integration with non-Euclidean or learnable ground metrics remain fruitful areas for extension.

Conclusion

LPAS demonstrates that state- and context-dependent learning of Wasserstein ambiguity sets yields robust yet non-conservative policies that adapt to regime-dependent risk, outperforming both classic empirical DRO and decision-agnostic robust baselines. The paradigm shift from static to learned, decision-sensitive ambiguity radii is substantiated both theoretically and empirically. This approach positions ambiguity set learning as a central task in robust machine learning and optimization, with implications for future research at the intersection of uncertainty quantification, adaptive risk management, and differentiable optimization.