Generative models for decision-making under distributional shift

Generative Models for Decision-Making under Distributional Shift

Introduction and Motivation

The paper "Generative models for decision-making under distributional shift" (2604.04342) advances a mathematically principled framework for leveraging deep generative models—specifically, flow-based and score-based architectures—to construct decision-relevant distributions beyond the nominal distributions identified by historical data. The central premise is that in many operations research (OR) and data-driven decision problems, the nominal data distribution $P$ is not representative of the uncertainty encountered during deployment; rather, the true deployment distribution $Q$ must be constructed or adapted, accounting for distributional shift, context dependency, stress scenarios, or partial observability.

Contrary to typical machine learning paradigms that focus on modeling $P$ for prediction or imitation, this paper emphasizes algorithmic and theoretical mechanisms—such as transport maps, velocity fields, and stochastic dynamics—for constructing $Q$ as required by decision objectives like robust optimization, scenario generation, or Bayesian inference under uncertainty.

Mathematical Formulation and Theoretical Underpinnings

The authors ground their framework in the formalism of optimal transport and probability flows. The deployment distribution $Q$ is constructed as a pushforward of $P$ via an appropriately chosen transport map $T$, i.e., $Q = T_\#P$. The Wasserstein-2 distance and its associated dynamic and Monge formulations are invoked to capture the geometry of distributional shift, while ODEs, SDEs, and Fokker–Planck equations describe the continuous evolution and perturbation of probability densities.

The generative construction is made concrete through several algorithmic frameworks:

• Score-based diffusion models, which learn the time-dependent score field $\nabla \log \rho_t$ via score matching, enabling reverse-time generative sampling equivalent to probability flow ODE integration.
• Normalizing and continuous-time flows, realized via invertible maps or neural ODEs parameterizing the dynamics of probability distributions.
• Flow Matching, which avoids expensive likelihood computations by directly regressing neural velocity fields onto target particle flows.
• Consistency models, designed for efficient one- or few-step generation via distilled transport maps.

All of these frameworks support rigorous sample-based optimization, with flow-matching providing a unifying computational approach for lifting discrete particle evolutions to continuous-time, Eulerian generative models.

Structured Distributional Operations for OR

The paper details how these machinery underpin various OR-relevant decision problems, including:

• Distribution-to-distribution transport and domain adaptation, enabling generative adaptation from nominal to deployment regimes given only empirical samples of source and target distributions.
• Worst-case/stress scenario generation, formalized via Wasserstein distributionally robust optimization (DRO) and cast as a min-max optimization in map-space. The worst-case distribution is the output of an adversarial transport map maximizing expected loss under feasible distributional perturbations.
• Conditional and posterior construction, where generative mechanisms are conditioned on observed variables or side information, or adapted in latent space to realize Bayesian posterior sampling with learned generative priors.
• Mean-field games and equilibrium distributions, where generative dynamics are coupled across interacting agents and solved as fixed point problems in probability space.

A universal particle-based implementation template is given: starting with empirical samples, particles are evolved according to the relevant decision-driven optimization, with final distributions (and the associated flows/maps) learned via flow matching. This provides reusable, generalizable generative representations beyond the original finite set of samples.

Applied Illustrations

Power Outage Scenario Generation

The capacity of generative models to support scenario generation under distributional shift is demonstrated with power outage data from Atlanta area counties. The flow-based model produces synthetic distributions that match empirical marginals and cross-county dependence structures, as shown in the following figures:

Figure 1: Empirical and generated marginal outage distributions for representative counties, demonstrating fit across regimes with varying outage frequencies.

Figure 2: Generated and empirical cross-county correlation matrices, highlighting accurate dependence structure reproduction in the transformed space.

The reported Maximum Mean Discrepancy (MMD) of $0.0150$ underscores the model's fidelity in multivariate scenario construction—a core requirement in realistic OR planning.

Stress Testing and Robust Optimization

A robust portfolio optimization example highlights generative models' potential for explicit stress distribution generation. Here, the adversarially perturbed return distributions, optimized via the Wasserstein-penalized min-max framework in map-space, result in portfolios with improved out-of-sample performance under stress:

Figure 3: Compared marginal distributions for nominal and worst-case generated returns under increasing regularization, illustrating controlled movement into adverse regions.

Figure 4: Out-of-sample cumulative wealth curves for nominal and robust portfolios; robust optimization under $Q$0 achieves the most favorable stressed performance.

Strong numerical results show that model-driven stress distributions yield robust decisions outperforming nominal optimization, particularly under adversity.

Theoretical Guarantees

The paper provides template theoretical guarantees tailored to each construction:

• Iterative flows in Wasserstein space benefit from contraction rates, with logarithmic convergence in KL divergence and total variation distance between the generate-and-reverse process.
• Min-max optimization in transport-map space (adversarial DRO) admits first-order stationarity guarantees, directly realizable via gradient descent-ascent schemes in the function space of transport maps.
• Posterior sampling with generative priors possesses error-transfer bounds, where the total variation distance between the true and generative-based posterior is controlled by prior approximation and numerical sampling errors.

These results clarify the correspondence between algorithmic construction, operational performance, and mathematical structure.

Implications, Limitations, and Future Directions

The implications for AI and decision science are significant. By connecting generative modeling to task- and decision-driven distribution construction, the field advances beyond black-box data imitation toward rigorous, interpretable, and actionable probabilistic representations:

• Scenario generation is now flexible and tailored, informed directly by decision requirements beyond matching historical frequencies.
• Robust optimization and stress testing leverage adversarial generative maps to probe model vulnerabilities in a controlled, operationally meaningful manner.
• Posterior and conditional generative modeling enables integrated Bayesian inference, filtering, and adaptive decision-making in high-dimensional, nonparametric settings.

The framework remains subject to open issues. The gap between idealized distributional guarantees and neural implementations persists, especially in high dimensions or under sequential, constrained feedback. There is demand for sharper generalization analysis, stability results, and robustification in the presence of finite-sample or architectural limitations. As generative architectures advance, the synthesis between optimization in probability space and deep learning will continue to reshape the algorithmic and theoretical landscape.

Conclusion

"Generative models for decision-making under distributional shift" (2604.04342) systematizes the operational use of flow- and score-based generative models for constructing, perturbing, and updating probability distributions as demanded by OR and AI objectives. By foregrounding the centrality of transportation, adversarial and conditional operations, and tying them to optimization in probability space, this work provides a technical and conceptual foundation for advanced methods in robust and context-aware decision-making under uncertainty.