Diffusion-DFL: Decision-focused Diffusion Models for Stochastic Optimization

## Abstract

Decision-focused learning (DFL) integrates predictive modeling and optimization by training predictors to optimize the downstream decision target rather than merely minimizing prediction error. To date, existing DFL methods typically rely on deterministic point predictions, which are often insufficient to capture the intrinsic stochasticity of real-world environments. To address this challenge, we propose the first diffusion-based DFL approach, which trains a diffusion model to represent the distribution of uncertain parameters and optimizes the decision by solving a stochastic optimization with samples drawn from the diffusion model. Our contributions are twofold. First, we formulate diffusion DFL using the reparameterization trick, enabling end-to-end training through diffusion. While effective, it is memory and compute-intensive due to the need to differentiate through the diffusion sampling process. Second, we propose a lightweight score function estimator that uses only several forward diffusion passes and avoids backpropagation through the sampling. This follows from our results that backpropagating through stochastic optimization can be approximated by a weighted score function formulation. We empirically show that our diffusion DFL approach consistently outperforms strong baselines in decision quality. The source code for all experiments is available here .

## 1 Introduction

Many real-life decision-making tasks require selecting actions that minimize a cost function involving unknown, context-dependent parameters. These parameters must often be predicted from observed features. For example, in supply chain management, future product demand must be estimated before deciding how much inventory to order [31]. A common approach is the predict-then-optimize pipeline, where a predictive model is first trained using a loss function such as mean squared error (MSE), and the resulting predictions are then passed to an optimization solver to guide decisions. While simple and widely adopted, this two-stage method can be misaligned with the true objective: minimizing decision cost. In particular, lower prediction error does not always lead to higher-quality decisions [3, 6].

Decision-focused learning (DFL) addresses this misalignment by integrating the prediction and optimization stages into a single end-to-end framework [5, 36, 17]. Unlike the two-stage approach, DFL trains the prediction model specifically to improve decision outcomes, often resulting in solutions with lower regret. However, most existing DFL methods rely on point (deterministic) predictions as inputs to the optimization layer, despite the fact that in many real-world scenarios, the underlying parameters are inherently uncertain and may follow complex distributions. Ignoring this uncertainty can lead to overconfident models and degraded decision quality [13].

In this work, we introduce a novel DFL approach that leverages diffusion probabilistic models to capture the environment uncertainty in an end-to-end fashion. Here, we use a conditional diffusion model [32] to represent the distribution of uncertain parameters given contextual features. The advantage of integrating a diffusion model into DFL is that, unlike simple distribution predictions (e.g., Gaussian), diffusion models can capture multi-modal or complex distributions. However, the sequential sampling procedure of diffusion models introduces a challenge when training a diffusion model end-to-end for stochastic optimization. To address this, we develop two algorithms: reparameterization and score function. First, the reparameterization trick is a common approach that expresses a random sample as a deterministic function of the model parameters and some noise, and we can backpropagate through sampled prediction to solve the DFL problem.

However, this approach can be very costly in memory and computation because it requires differentiating (and therefore tracking gradients) through the diffusion sampling process. To address this, we introduce a lightweight score function estimator that avoids differentiating through the sampling process. Specifically, we use a score function surrogate to approximate the gradient of the diffusion predictor and plug it into the KKT (Karush-Kuhn-Tucker) implicit-differentiation approach to obtain the total derivative of the decision objective.
In addition, we further mitigate the high variance that arises from using only score functions for a few steps by employing a tailored importance sampling strategy.

We evaluate our proposed methods in various applications, including (synthetic) product allocation, energy scheduling, and stock portfolio optimization. Experimental results show that our diffusion DFL methods consistently outperform all baselines, with more improvements on larger problem sizes. Moreover, the score function estimator achieves decision quality comparable to that of the reparameterization method, while significantly reducing GPU memory usage from 60.75 GB to 0.13 GB. The contributions of this paper are the following:

- •
We introduce the first DFL method that uses diffusion models to capture the downstream uncertainty and employs the reparameterization trick for end-to-end gradient estimation.
- •
We propose a lightweight score function estimator that avoids backpropagating the reversing process in the reparameterization method, significantly reducing memory and computation cost.
- •
We evaluate our methods in three real-world optimization tasks and observe consistent improvements over strong baselines.

## 2 Related Works

#### Decision-focused learning

DFL is an emerging paradigm that
trains models end-to-end to directly optimize decision quality rather than minimizing prediction error [5, 36, 17].
Despite the success in aligning learning objectives with decision-making, a limitation of most existing DFL methods is that they typically rely on deterministic point predictions of uncertain parameters [36, 23]. By ignoring distributional uncertainty, deterministic point predictions cannot represent the full outcomes and may lead to lower decision quality [34]. Empirically, classic DFL was observed to struggle in high-dimensional and risk-sensitive real-world settings with significant uncertainty [16].

Therefore, the gap in uncertainty modeling motivates the need for more comprehensive DFL with stochastic predictions, where several works have started integrating uncertainty awareness into the DFL pipeline [26, 34, 25, 9].
For instance, [34] proposes a generative DFL approach (Gen-DFL) based on normalizing flow models as the predictor.
However, normalizing flows require a bijective network architecture, which restricts the expressiveness of the stochastic predictor.

In this paper, we propose using diffusion models [7] as a more expressive predictor. By leveraging diffusion models in the DFL paradigm, our approach extends DFL by predicting an accurate full distribution of the unknown parameters, which addresses the overconfidence of deterministic optimization and better aligns with downstream decision-making needs.

#### Diffusion model in optimization

Diffusion probabilistic models have achieved great success in modeling high-dimensional data distributions in recent years [27, 28, 4]. Originally popularized for image generation and related structured outputs, its ability to capture multi-modal and high-variety distributions has made it attractive beyond vision tasks, such as combinatorial optimization[30, 22], black-box optimization[15, 14].
To our best knowledge, however, no prior work has integrated diffusion models into a predict-then-optimize learning pipeline for decision tasks. This paper is the first to harness diffusion models in an end-to-end DFL framework. By using a conditional diffusion model, we can learn a rich distribution over the uncertain inputs and then propagate this uncertainty through to the downstream decision via gradient-based training (score function and reparameterization). This approach combines the strengths of expressive generative modeling and DFL to improve decision quality under uncertainty.

## 3 Problem Statement

### 3.1 Decision-focused learning

We consider a general predict-then-optimize setting [5, 6], where the goal is to make decisions under uncertainty about a key problem parameter. Given a feature vector $x\in\mathcal{X}$ and a prediction of an unknown parameter $y^{*}\in\mathcal{Y}$, the decision-maker selects $z\in\mathbb{R}^{d}$ to minimize a decision loss function $f:\mathcal{Y}\times\mathbb{R}^{d}\to\mathbb{R}$, which measures the cost of applying decision $z$ when the true parameter is $y^{*}$. We assume a joint distribution $\mathcal{D}$ over $(x,y^{*})$ pairs.

DFL integrates prediction and optimization into a unified framework. The goal is to learn a decision function $z_{\theta}^{*}:\mathcal{X}\to\mathbb{R}^{d}$, parameterized by $\theta$, that minimizes the expected decision loss,

$$ $\min_{\theta}F(\theta):=\operatorname{\mathbb{E}}_{(x,y^{*})\sim\mathcal{D}}[f(y^{*},z^{*}_{\theta}(x))].$ (1) $$

The decision $z^{*}_{\theta}(x)$ is typically obtained by solving an optimization problem involving a prediction of the uncertainty parameter. Most DFL methods [17] use a deterministic point prediction $y_{\theta}(x)$ of the uncertain parameter $y^{*}$:

$$ $z_{\theta}^{*}(x)=\operatorname*{arg\,min}_{z}f(y_{\theta}(x),z),\quad\text{s.t.}\ \ Gz\leq h,\ Az=b,$ (2) $$

where $G\in\mathbb{R}^{n\times d},h\in\mathbb{R}^{n},A\in\mathbb{R}^{p\times d},b\in\mathbb{R}^{p}$ are constraint problem coefficients(^1^11We consider affine constraints in our main paper for simplicity. The extension from affine constraints to general convex constraints $h(x,z)\leq 0$ follows a similar derivation as in the linear case.).

In contrast, we consider a probabilistic model $P_{\theta}(\cdot\mid x)$ for the uncertain parameter $y^{*}$ and let $z_{\theta}^{*}(x)$ be the solution to a stochastic optimization problem:

$$ $z_{\theta}^{*}(x)=\operatorname*{arg\,min}_{z}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)],\quad\text{s.t.}\ \ Gz\leq h,\ Az=b.$ (3) $$

We aim to learn the model parameter $\theta$ such that $z_{\theta}^{*}$ minimizes the expected decision loss $F(\theta)$. By the chain rule, the derivative of $F$ is

$$ $\frac{\mathrm{d}F(\theta)}{\mathrm{d}\theta}=\operatorname{\mathbb{E}}_{(x,y^{*})\sim\mathcal{D}}\left[\frac{\partial f(y^{*},z_{\theta}^{*}(x))}{\partial z}\frac{\mathrm{d}z_{\theta}^{*}(x)}{\mathrm{d}\theta}\right].$ $$

However, computing this gradient (specifically, the $\frac{\mathrm{d}z_{\theta}^{*}}{\mathrm{d}\theta}$ term) is challenging because $z_{\theta}^{*}$ is implicitly defined by a nested optimization problem.
A common solution is to differentiate the KKT system that implicitly defines $z^{*}_{\theta}$ w.r.t. $\theta$ [1].
Another crucial point is the selection of the stochastic predictor in DFL, which in the paper we choose to use diffusion models to represent $P_{\theta}$.

### 3.2 Diffusion probabilistic model

To generate complex multi-modal and high-dimensional distributions,
diffusion probabilistic models [7] are a promising way. It couples a fixed noising chain with a learned reverse denoising chain. Let $y_{0}\in\mathbb{R}^{d}$ denote a sample from the real data distribution $q(y_{0})$ and $\{\beta_{t}\in(0,1)\}_{t=1}^{T}$ denote the noise schedule. Define $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{i}$. The *forward process* $q$ adds Gaussian noise at each step $t$ to $y_{1}$ through $y_{T}$:

$$ $\displaystyle q(y_{t}\mid y_{t-1})=\mathcal{N}(y_{t};\sqrt{1-\beta_{t}}y_{t-1},\beta_{t}I),\quad t=1,\dots,T,$ (4) $$

which guarantees that $q(y_{T}\mid y_{0})$ becomes nearly standard normal distribution as $T\to\infty$ with common schedules ($\bar{\alpha}_{T}\to 0$). Note that $y_{t}$ can be equivalently sampled without iterating through intermediate time steps: $y_{t}=\sqrt{\bar{\alpha}_{t}}y_{0}-\sqrt{1-\bar{\alpha}_{t}}\epsilon$, where $\epsilon\sim N(0,I)$ is a Gaussian noise.

In the *reverse process* $p$, the diffusion model predicts the unknown added noise by

$$ $\displaystyle p_{\theta}(y_{t-1}\mid y_{t})=\mathcal{N}(y_{t-1};\ \mu_{\theta}(y_{t},t),\,\sigma_{t}^{2}I),$ (5) $$

whose mean $\mu_{\theta}(\cdot,t)$ is parameterized by a neural network predictor and variance is either fixed ($\sigma_{t}^{2}=\beta_{t}$) or learned.
The combination of $p$ and $q$ is equivalent to a hierarchical variational auto-encoder [33], and thus can be optimized by using the evidence lower bound (ELBO) as the loss function [8].

Conditional Diffusion Model. Throughout this paper, $x$ denotes contextual features, and every transition probability is conditioned on $x$ [32].
We use $P_{\theta}(\cdot|x)$ for the diffusion model’s conditional distribution for generated data and $p_{\theta}(y_{t-1}\mid y_{t},x)$ for its Markov transitions.

## 4 Stochastic Optimization and Reparameterization estimator

Real-world decision problems often face significant uncertainty in their parameters.
Optimizing with a stochastic predictor (e.g., diffusion model) yields better results than deterministic optimization,
by explicitly modeling the uncertainty and optimizing the expected cost.
[Figure 1](https://arxiv.org/html/2510.11590v1#S4.F1) illustrates a simple example: any deterministic solution ends up at an extreme decision with a higher expected cost, while the stochastic solution averages costs across likely outcomes and selects an interior decision with a lower expected cost.

Figure: Figure 1: A comparison of deterministic vs. stochastic optimization with cost function $\exp(-yz)$, as described in [Section 6.1](https://arxiv.org/html/2510.11590v1#S6.SS1). (a) Each curve represents a cost function given a sample $y$. For any fixed $y$, the deterministic optimization decision lies at one of the boundaries ($z^{*}=0$ or $z^{*}=C$). (b) When averaging the cost function over many samples of $y$, the stochastic optimization decision lies in the interior of the feasible region instead of on the boundary. Thus, any deterministic optimization decision is suboptimal. (c) A probabilistic (diffusion) model captures a distribution over $Y$ that closely resembles the true bimodal distribution.
Refer to caption: x1.png

#### Solving stochastic DFL.

Formally, in the stochastic case, the optimality condition for the decision problem must consider an expectation. The stationarity condition for decision problem Eq. [3](https://arxiv.org/html/2510.11590v1#S3.E3) becomes:

$$ $\displaystyle\nabla_{z}\mathcal{L}(\theta,z^{*},\lambda^{*},\nu^{*};x)=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{z}f(y,z)]+G^{\top}\lambda^{*}+A^{\top}\nu^{*}=0,$ (6) $$

where $\mathcal{L}$ denotes the Lagrangian. Note that the dependency on $\theta$ in the stationarity condition is in the distribution. Therefore, we need to handle this dependency carefully while differentiating the KKT system with respect to $\theta$:

$$ $\displaystyle\underbrace{\frac{\mathrm{d}}{\mathrm{d}\theta}(\nabla_{z}\mathcal{L}(\theta,z^{*},\lambda^{*},\nu^{*};x))}_{\text{distributional gradient}}=\frac{\mathrm{d}}{\mathrm{d}\theta}(\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\nabla_{z}f(y,z)]+G^{\top}\lambda^{*}+A^{\top}\nu^{*})=0.$ (7) $$

To resolve the dependence of both the predictive distribution $P_{\theta}(y|x)$ and the decision $z^{*}$ on $\theta$, we first adopt the reparameterization trick [11] for the diffusion model. From Section [3.2](https://arxiv.org/html/2510.11590v1#S3.SS2), recall that the diffusion sampling process introduces Gaussian noise at each step. Thus, we can reparameterize the reverse process by fixing all the random draws (Gaussian noises). Formally, a sample $y\sim P_{\theta}(y\mid x)$ can be expressed as a transformation $y=R(\epsilon,\theta\mid x)$ of a base Gaussian noise sample $\epsilon\sim P(\epsilon)$, where $R$ is differentiable in $\theta$. This makes the diffusion sampling a deterministic function of $\theta$. Then we have

$$ $\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}R(\epsilon,\theta|x))^{\top}\nabla_{y}f(y,z)].$ (8) $$

Next, we incorporate this into the optimization. Following Eq. [7](https://arxiv.org/html/2510.11590v1#S4.E7), we can formalize a KKT system that contains derivatives of $z^{*}_{\theta},\lambda^{*},\nu^{*}$ and $\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}$. Plugging the reparameterized gradient estimator into the KKT system, we can solve for $\frac{\mathrm{d}z_{\theta}^{*}}{\mathrm{d}\theta}$ and then obtain the total derivative of the final objective $F$ by multiplying $\frac{dF}{dz_{\theta}^{*}}$ [5] (proof can be found in [Section A.1](https://arxiv.org/html/2510.11590v1#A1.SS1)):

$$ $\displaystyle\frac{dF}{d\theta}=-\begin{bmatrix}\frac{dF}{dz_{\theta}^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&A^{\top}\\ D(\lambda^{*})G&D(Gz_{\theta}^{*}-h)&0\\ A&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}R(\epsilon,\theta|x))^{\top}\nabla^{2}_{zy}f(y,z_{\theta}^{*})]\\ 0\\ 0\end{bmatrix},$ (9) $$

where $H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\nabla_{zz}^{2}f(y,z_{\theta}^{*})]$ is the Hessian of the Lagrangian with respect to $z$, and $D(v)$ denotes a diagonal matrix with $v$ on its diagonal.
In practice, one can sample $\epsilon$ from a certain distribution (e.g., Gaussian) multiple times to estimate the expectation and then obtain the gradient. This gives us reparameterization-based diffusion DFL using Eq. [9](https://arxiv.org/html/2510.11590v1#S4.E9) to run stochastic DFL optimization.

## 5 Score function estimator

A major obstacle to implementing the total gradient (Eq. [9](https://arxiv.org/html/2510.11590v1#S4.E9)) is the need to backpropagate through the diffusion sampling process. In most cases, the diffusion model’s generative process is complex and multi-step (e.g., 1000 steps), which makes backpropagating through all those steps memory-intensive and prone to instability. To address this, we propose a score function(^2^22In this paper, score function refers to the statistical score $\nabla_{\theta}\log P_{\theta}(y|x)$ (gradient of log-likelihood w.r.t. model parameters), as opposed to Stein’s score $\nabla_{y}p(y_{t}|y_{t-1},x)$ often used in diffusion literature.) gradient estimator for the diffusion model, which circumvents explicit backpropagation through all sampling steps. The key idea is to rewrite the Jacobian $\nabla_{\theta}y$ in terms of the score $\nabla_{\theta}\log P_{\theta}(y\mid x)$, and then approximate the score with the diffusion model’s ELBO training loss.

### 5.1 Transform the Jacobian into Score Function

We begin by rewriting the gradient of expectation as an expectation of a score function using the *log-trick* [18]. Formally, if $y\sim P_{\theta}(\cdot|x)$ and $f(y)$ is any function not dependent on $\theta$, then by the log-trick we have

$$ $\displaystyle\nabla_{\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[f(y,z)]=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[f(y,z)\cdot\nabla_{\theta}\log P_{\theta}(y\mid x)].$ (10) $$

Intuitively, instead of differentiating the output $y$ through each diffusion step, we only need to compute the gradient for the final log-likelihood, which avoids the need to differentiate through the diffusion sampling process and yields an efficient estimator for the gradient.

Then, one remaining difficulty is that directly computing the exact $\nabla_{\theta}\log P_{\theta}(y|x)$ is complicated in practice because $P_{\theta}(y|x)$ is defined as the marginal probability of $y$ after integrating out the latent diffusion trajectory.
To obtain a computationally efficient estimator, we use the diffusion model’s training objective as a surrogate for the log-likelihood. Specifically, diffusion models are typically trained by maximizing an ELBO that lower-bounds the log-likelihood:

$$ $\displaystyle\log P_{\theta}(y_{0})$ $\displaystyle=\log\int p_{\theta}(y_{0}|y_{1})\,p_{\theta}(y_{1}|y_{2})\,\cdots p_{\theta}(y_{T-1}|y_{T})\,p_{\theta}(y_{T})\,dy_{1:T}$ $\displaystyle=\log\operatorname{\mathbb{E}}_{y_{t}\sim q(y_{t}|y_{t-1})\forall t\in[T]}\left[\prod_{t=1}^{T}\frac{p_{\theta}(y_{t-1}|y_{t})}{q(y_{t}|y_{t-1})}p_{\theta}(y_{T})\right]$ $\displaystyle\geq\operatorname{\mathbb{E}}_{y_{t}\sim q(y_{t}|y_{t-1})\forall t\in[T]}\left[\sum_{t=1}^{T}\log\frac{q(y_{t}|y_{t-1})}{p_{\theta}(y_{t-1}|y_{t})}+\log p_{\theta}(y_{T})\right]:=\operatorname{ELBO}(y_{0};\theta),$ $$

where the inequality is due to Jensen’s. To approximate $\nabla_{\theta}\log P_{\theta}(y_{0}|x)$ conditioned on $x$, we use the gradient of the conditional ELBO loss as a surrogate:

$$ $\displaystyle\nabla_{\theta}\log P_{\theta}(y_{0}|x)\approx\nabla_{\theta}\operatorname{ELBO}(y_{0}|x;\theta).$ (11) $$

Figure: Figure 2: Cosine similarity between the reparameterization and score function gradient across different dimensions.
Refer to caption: x2.png

In practice, we first sample a final output $y$ from the diffusion model given contextual features $x$. We then sample a subset of $k$ timesteps $\{t_{1},t_{2},\dots,t_{k}\}$ ($k\ll T$) and run forward noising process $q$ to generate the trajectory $\{y_{t_{1}},y_{t_{2}},\dots,y_{t_{k}}\}$. As in DDPM [7], we adopt the simplified form of $\operatorname{ELBO}\approx\operatorname{\mathbb{E}}_{t\sim[T],y_{0},\epsilon_{t}}[|\|\epsilon_{t}-\epsilon_{\theta}(y_{t},t)\|^{2}]$. We evaluate the $\operatorname{ELBO}$ on the sampled trajectories and compute its gradient w.r.t. $\theta$ as an estimation to the true score.

Empirical evidence suggests that the ELBO gradient closely tracks the true score, as shown in [Figure 2](https://arxiv.org/html/2510.11590v1#S5.F2), making Eq. [11](https://arxiv.org/html/2510.11590v1#S5.E11) a reliable proxy in practice.

### 5.2 Overall Gradient for Score Function

By plugging the ELBO gradient approximation from Eq. [11](https://arxiv.org/html/2510.11590v1#S5.E11) into Eq. [10](https://arxiv.org/html/2510.11590v1#S5.E10), we can express the KKT conditions without using reparameterization and thus obtain the score function-based derivative:

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle\approx-\begin{bmatrix}\frac{dF}{dz_{\theta}^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&A^{\top}\\ D(\lambda^{*})G&D(Gz_{\theta}^{*}-h)&0\\ A&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\nabla_{z}f(y,z_{\theta}^{*})(\nabla_{\theta}{\operatorname{ELBO}(y|x;\theta)}^{\top}]\\ 0\\ 0\end{bmatrix}.$ (12) $$

#### Practical algorithm – weighted ELBO gradient.

To compute the score surrogate in practice, we found it convenient to treat the total gradient as an importance-weighted form:

$$ $\displaystyle\frac{dF}{d\theta}\approx\frac{d}{d\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\underbrace{\operatorname{detach}[w_{\theta}(y)]}_{\text{importance weight, no grad in $\theta$}}\cdot\ \underbrace{{\operatorname{ELBO}(y|x,\theta)}}_{\text{1-step forward}}],$ (13) $$

where $w_{\theta}(y)$ is the importance weight simplified from Eq. [12](https://arxiv.org/html/2510.11590v1#S5.E12) (see [Section A.3](https://arxiv.org/html/2510.11590v1#A1.SS3) for complete form).
This yields a *weighted-ELBO* gradient estimator: we treat $w_{\theta}(y)$ as a stop-gradient weight and only differentiate the $\operatorname{ELBO}$ w.r.t. $\theta$, greatly reducing computations. We implement the entire gradient computation as a user-friendly PyTorch autograd module: the forward pass returns the optimal decision $z^{*}$ (and $\lambda^{*},\nu^{*}$), and the backward pass computes the gradient $\frac{dF}{d\theta}$ as derived above.

#### Variance-reduction strategy.

While the score-function estimator is effective, a naive implementation of the weighted ELBO loss in Eq. [13](https://arxiv.org/html/2510.11590v1#S5.E13) can suffer from high variance, leading to unstable training. In practice, we found that carefully designing the sampling strategy for the ELBO loss is crucial to obtaining low-variance and stable gradients. To reduce the variance, we utilize the method from Improved DDPM [19] for choosing diffusion steps. Specifically, instead of uniform sampling, we use *importance sampling* over timesteps with probability $p_{t}$ and weights $1/p_{t}$:

$$ $\displaystyle\nabla_{\theta}\text{ELBO}^{\text{IS}}=\operatorname{\mathbb{E}}_{t\sim p_{t}}\left[\frac{\nabla_{\theta}(\text{ELBO}_{t})}{p_{t}}\right],\ \text{where}\ p_{t}\propto\sqrt{\operatorname{\mathbb{E}}[\|\nabla_{\theta}(\text{ELBO}_{t})\|^{2}]}\ \text{and}\sum_{t}p_{t}=1.$ (14) $$

This method remains unbiased, but the variance is minimized.
In essence, this approach gives less weight to the early timesteps that have large gradients and more weight to later timesteps.

## 6 Experiments

We evaluate the performance of our diffusion-based DFL approaches on a variety of tasks, comparing against several baseline methods. Specifically, we consider:

- •
Two-stage predict-then-optimize baselines: a deterministic MLP, a Gaussian probabilistic model, and a diffusion model trained to minimize prediction error [6].
- •
Deterministic DFL: a deterministic MLP model with end-to-end DFL training [5].
- •
Gaussian DFL (both reparameterization and score function): a Gaussian probabilistic model with end-to-end stochastic DFL training [26]; see details in Appendix [A.5](https://arxiv.org/html/2510.11590v1#A1.SS5.SSS0.Px2).
- •
Diffusion DFL (ours): our diffusion model predictor, trained with either reparameterization or score-function gradient estimators.

### 6.1 Synthetic Example

In this example, we consider a factory that decides how much to manufacture for each of $d\in\mathbb{N}$ products. The parameter $Y\in\mathbb{R}^{d}$ represents the profit margin for each product, i.e., $Y_{i}$ is the profit per unit of product $i$; due to uncertainty in market conditions, $Y$ is uncertain. The factory’s decision $z\in[0,C]^{d}$ represents how much of each product to manufacture, where $C$ is the maximum capacity for each product. For simplicity, we do not consider any contextual features $x$ in this example. That means DFL learns a distribution that generates $y$ that can minimize the decision objective.

Suppose that the factory has a risk-averse cost function $f(y,z)=\exp(-y^{\top}z)$(^3^33Here, we have ignored the degenerate case $y=0$. To deal with the degenerate case, one could add a zero-centered bump function $c(y)$ to the objective $f(y,z)$.), which indicates that the factory wants to put a larger weight on the product with higher profit $Y_{i}$. Under uncertainty, the decision-maker seeks to minimize the expected cost by solving a stochastic optimization problem:

$$ $\displaystyle z_{\mathrm{sto}}^{*}$ $\displaystyle\in\operatorname*{arg\,min}_{z\in[0,C]^{d}}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\exp(-y^{\top}z)].$ (15) $$

In this stochastic case, the optimal investment $z_{\mathrm{sto}}^{*}$ typically lies in the interior of the feasible region, which balances the potential high reward of investing against the risk of losses.

Experimental setup.
We simulate the uncertain parameter $Y$ drawn from a mixture of Gaussians,

$$ $\displaystyle Y_{i}\ \stackrel{{\scriptstyle\mathrm{iid}}}{{\sim}}\ p\cdot\operatorname{\mathcal{N}}(a,\sigma^{2})+(1-p)\cdot\operatorname{\mathcal{N}}(-b,\sigma^{2}).$ (16) $$

Specifically, we set $p=0.8,a=1,b=3,\sigma=0.15,C=2$. We train each model (deterministic, Gaussian, diffusion) on this distribution in a decision-focused manner (for DFL methods) or on pure regression (for two-stage), and evaluate the expected cost achieved by the resulting decision $z_{\theta}^{*}(x)$. We present the results of one product ($d=1$) in [Figure 1](https://arxiv.org/html/2510.11590v1#S4.F1) and 10 products ($d=10$) in [Table 1](https://arxiv.org/html/2510.11590v1#S6.T1).

### 6.2 Power Schedule

In this experiment, we evaluate our method on a real-world energy scheduling problem from Donti et al. [5]. This task involves a 24-hour generation-scheduling problem in which the operator chooses $z\in\mathbb{R}^{24}$ (hourly generation). Given a realization $y$ of demand, the decision loss penalizes shortage and excess with asymmetric linear costs ($\gamma_{s}$ and $\gamma_{e}$) plus a quadratic tracking term; the decision must also satisfy a ramping bound $c_{r}$. Let $[v]_{+}:=\max(v,0)$. We have the decision loss as the quadratic problem:

$$ $\displaystyle\min_{z}\ \operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]=\sum_{i=1}^{24}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\gamma_{s}[y_{i}-z_{i}]_{+}+\gamma_{e}[z_{i}-y_{i}]_{+}+\frac{1}{2}(z_{i}-y_{i})^{2}],$ $\displaystyle\ \mathrm{s.t.}\ |z_{i}-z_{i-1}|\leq c_{r}\ \text{for all}\ i\in\{1,2,\dots,24\}.$ (17) $$

Experimental setup. We use more than 8 years of historical data from a regional power grid [20]. Feature $x$ includes the previous day’s hourly load, temperature, next-day temperature forecasts, non-linear transforms (lags and rolling statistics), calendar indicators, and yearly sinusoidal features. Given $x$, the prediction model $P_{\theta}(\cdot|x)$ outputs a distribution over $y\in\mathbb{R}^{24}$. We report the test decision cost in [Table 1](https://arxiv.org/html/2510.11590v1#S6.T1) and a held-out horizon in [Figure 7](https://arxiv.org/html/2510.11590v1#A1.F7).

### 6.3 Stock Market Portfolio Optimization

In this experiment, we apply our diffusion DFL approach to a financial portfolio optimization problem under uncertain stock returns. Here, the random vector $y\in\mathbb{R}^{n}$ represents the returns for the assets $n$ on the next day, and the decision $z\in\mathbb{R}^{n}$ represents the portfolio weights allocated to those assets. We consider a mean-variance trade-off decision loss: maximize expected return while keeping the risk (variance) low. This can be written as minimizing a loss that is a negative expected return plus a quadratic penalty on variance:

$$ $\displaystyle\min_{z}\ \operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}\left[\frac{\alpha}{2}z^{\top}yy^{\top}z-y^{\top}z\right],\quad\mathrm{s.t.}\quad z^{\top}\mathbf{1}=1,\ 0\leq z_{i}\leq 1,$ (18) $$

where $\alpha>0$ is a risk parameter and constraints enforce that $z$ is a valid portfolio. In practice, the deterministic solution may concentrate heavily on a few assets and yield a low average return, whereas a stochastic approach can achieve higher returns by accounting for variance.

Experimental setup. We have daily prices and volumes spanning 2004-2017 and evaluate on the S&P 500 index constituents [21]. The features $x\in\mathbb{R}^{28}$ include recent historical return, trading volume windows, and rolling averages. The immediate-return predictor $P_{\theta}(\cdot|x)$ is to predict the next day’s price.
We report the performance of different DFL baselines with 50 portfolios in [Table 1](https://arxiv.org/html/2510.11590v1#S6.T1) and other sizes of portfolios in [Section 7.2](https://arxiv.org/html/2510.11590v1#S7.SS2).

**Table 1: Results for different optimization tasks. Our two diffusion DFL methods achieve the best and second-best decision quality in all 3 tasks, significantly better than other baselines. Bolded values are the best in test task losses; underlined values are the 2nd-best. Mean $\pm$ standard error across 10 runs.**
|  | Synthetic Example | Power Schedule | Stock Portfolio |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Label / Method | RMSE$\downarrow$ | Task$\downarrow$ | RMSE$\downarrow$ | Task$\downarrow$ | RMSE$\downarrow$ | Task (%)$\uparrow$ |
| *Two-stage (TS)* |  |  |  |  |  |  |
| Deterministic TS | 0.639±0.00 | 1.987±0.00 | 0.120±0.00 | 41.239±3.18 | 0.027±0.00 | 0.04%±0.04 |
| Gaussian TS | 0.720±0.00 | 1.272±0.23 | 0.117±0.00 | 5.580±0.45 | 0.188±0.03 | 0.10%±0.04 |
| Diffusion TS | 0.905±0.00 | 0.393±0.00 | 0.147±0.00 | 7.901±0.76 | 0.455±0.00 | 0.13%±0.03 |
| *Decision-focused learning (DFL)* |  |  |  |  |  |  |
| Deterministic | 0.640±0.00 | 1.987±0.00 | 4.997±0.10 | 4.324±0.25 | 0.032±0.00 | 0.07%±0.00 |
| Gaussian Reparameterization | 0.707±0.00 | 1.169±0.03 | 4.525±0.12 | 3.724±0.05 | 0.189±0.03 | 0.08%±0.03 |
| Gaussian Score Function | 0.708±0.00 | 1.132±0.00 | 4.713±0.15 | 4.087±0.06 | 0.187±0.03 | 0.14%±0.05 |
| Diffusion Reparameterization | 0.852±0.01 | 0.365±0.00 | 3.141±0.06 | 3.152±0.03 | 0.063±0.00 | 4.17%±0.24 |
| Diffusion Score Function | 0.849±0.09 | 0.3620.362±0.00 | 2.893±0.03 | 3.171±0.02 | 0.067±0.00 | 3.98%±0.31 |

## 7 Discussion of Experimental Results and Ablation Study

### 7.1 Discussion of Results in Table 1

#### Two-stage vs DFL.

As shown in [Table 1](https://arxiv.org/html/2510.11590v1#S6.T1), across all three experiment tasks, we find that end-to-end DFL leads to better downstream decisions than the conventional two-stage approach. Conventional two-stage methods minimize RMSE during training, but this often leads to poor downstream decisions. In contrast, all variants of DFL directly minimize the decision cost during training and thus achieve lower decision costs.

#### Deterministic vs Stochastic Optimization.

Our results show that stochastic DFL methods outperform deterministic DFL in terms of decision quality on every task. By modeling uncertainty, stochastic predictors enable the decision optimization to account for risk and variability in outcomes. For instance, in the portfolio experiment, the deterministic DFL yields only $0.07\%$ return, whereas a Gaussian DFL modestly improves that, and our diffusion DFL achieves nearly $4\%$ average return. These gains come from the stochastic models’ ability to predict uncertainty: instead of committing to a point prediction of $y$, the stochastic DFL produces decisions for a range of possible outcomes.

#### Benefits of Diffusion DFL.

Among the stochastic approaches, including baselines using Gaussian models, our diffusion DFL method consistently delivers the best decision performance. In particular, the diffusion model’s strength is the capacity to capture complex, multi-modal outcome distributions that a simple parametric Gaussian cannot represent. The Gaussian DFL sometimes falls short of the optimal decision quality. The diffusion model, on the other hand, can represent more intricate distributions of $y$, leading to decisions that better reflect complex scenarios.

### 7.2 Ablation Study

Comparison Cost for Reparameterization and Score function.
A key finding from our ablation study is the computational advantage of score-function approach over the reparameterization. Here, we measure the trade-off between training cost and the final decision performance for different gradient estimators and sampling budgets.

In [Figure 4](https://arxiv.org/html/2510.11590v1#S7.F4) (a), we see that all variants reach similar final performance on the test set, indicating that even using as few as 50 samples is sufficient to optimize the decision quality accurately.
[Figure 4](https://arxiv.org/html/2510.11590v1#S7.F4) plots the GPU memory cost alongside the final test loss. The reparameterization method is very computationally expensive, requiring about 60 GB of GPU memory for backpropagating through all diffusion steps. In contrast, the score-function with 50 samples achieves virtually the same test loss as the reparameterization method (difference within 0.02) while using an order of magnitude less memory. Even with 10 samples, though slightly worse in loss, it still outperforms the deterministic baseline and uses a tiny fraction of the compute. These results validate that the score-function approach retains the decision-quality benefits of diffusion DFL while dramatically cutting computational requirements, making diffusion DFL practical even for complex problems.

Figure: Figure 3: Learning curves for (a) score function with 10 and 50 samples (sf 10 and sf 50) and reparameterization (rp), (b) score function and importance-weighted score function with 10 samples.
Refer to caption: x3.png

Gradient variance reduction.
As discussed in [Section 5.2](https://arxiv.org/html/2510.11590v1#S5.SS2), using the score function estimator allows us to avoid backpropagating through the entire diffusion sampling process by only sampling a limited number of diffusion timesteps per update. The reason behind this is that a naive implementation, sampling timesteps uniformly at random, would yield a very high variance in the gradient estimates, which then leads to unstable training. Intuitively, early diffusion steps (large noise levels) dominate the ELBO loss and its gradients, so if they happen to be sampled, they contribute disproportionately and noisily. With a small random subset of timesteps, the gradient estimate can thus be highly imbalanced and noisy, which causes training divergence in practice.

To address this, we adopt an importance sampling strategy for choosing diffusion timesteps. Empirically, as shown in [Figure 4](https://arxiv.org/html/2510.11590v1#S7.F4) (b), the learning curves with the importance-weighted sampler are much smoother and more stable than with the uniform sampler. The score-function DFL training no longer diverges; instead, it converges cleanly, indicating that our variance reduction strategy successfully stabilizes the training process for diffusion DFL.

Figure: Figure 5: Test regret vs. decision dimension $d$ in the stock portfolio task.
Refer to caption: x5.png

Comparison on different problem sizes.
A key challenge for DFL is scalability: as the decision dimension grows, many methods degrade significantly [17]. In this experiment, we investigate the performance of DFL methods under various decision dimensions in the stock portfolio. Specifically, we set the decision dimension range from 10 to 100 and report the test regrets. As summarized in [Figure 5](https://arxiv.org/html/2510.11590v1#S7.F5), the regret gap between DFL diffusion and Gaussian and deterministic methods increases with increasing dimension, which demonstrates that DFL diffusion scales effectively in more complex decision settings.

## 8 Conclusion

We propose the first diffusion-based DFL approach for stochastic optimization, which trains a diffusion model to capture complex uncertainty in problem parameters. We develop two end-to-end training techniques to integrate the diffusion model into decision-making: reparameterization and score function approximation. As demonstrated with empirical evidence, the score function method drastically reduces memory and computation cost while achieving similar performance to reparameterization and being easy to train. Empirically, diffusion DFL achieves state-of-the-art results on multiple benchmarks, consistently outperforming both traditional two-stage methods and prior DFL approaches.

## Reproducibility statement

We release an anonymized repository containing all code, configuration files, and scripts needed to reproduce our results, including data generation and figure plotting. All proofs for the main paper are stated in the appendix with explanations and proper assumptions.

## References

- Amos and Kolter [2017]
Brandon Amos and J. Zico Kolter.
OptNet: Differentiable Optimization as a Layer in Neural Networks.
In *Proceedings of the 34th International Conference on Machine Learning*, pages 136–145. PMLR, July 2017.
URL [https://proceedings.mlr.press/v70/amos17a.html](https://proceedings.mlr.press/v70/amos17a.html).
ISSN: 2640-3498.
- Arjevani et al. [2020]
Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan J. Foster, Ayush Sekhari, and Karthik Sridharan.
Second-Order Information in Non-Convex Stochastic Optimization: Power and Limitations.
In *Proceedings of Twenty Second Conference on Learning Theory*, pages 242–299. PMLR, July 2020.
URL [https://proceedings.mlr.press/v125/arjevani20a.html](https://proceedings.mlr.press/v125/arjevani20a.html).
ISSN: 2640-3498.
- Bertsimas and Kallus [2020]
Dimitris Bertsimas and Nathan Kallus.
From Predictive to Prescriptive Analytics.
*Management Science*, 66(3):1025–1044, March 2020.
ISSN 0025-1909, 1526-5501.
doi: 10.1287/mnsc.2018.3253.
URL [https://pubsonline.informs.org/doi/10.1287/mnsc.2018.3253](https://pubsonline.informs.org/doi/10.1287/mnsc.2018.3253).
- Dhariwal and Nichol [2021]
Prafulla Dhariwal and Alexander Nichol.
Diffusion Models Beat GANs on Image Synthesis.
In *Advances in Neural Information Processing Systems*, volume 34, pages 8780–8794. Curran Associates, Inc., 2021.
URL [https://proceedings.neurips.cc/paper_files/paper/2021/hash/49ad23d1ec9fa4bd8d77d02681df5cfa-Abstract.html](https://proceedings.neurips.cc/paper_files/paper/2021/hash/49ad23d1ec9fa4bd8d77d02681df5cfa-Abstract.html).
- Donti et al. [2017]
Priya L. Donti, Brandon Amos, and J. Zico Kolter.
Task-based End-to-end Model Learning in Stochastic Optimization.
In *Advances in Neural Information Processing Systems*, volume 30, Long Beach, CA, USA, December 2017. Curran Associates, Inc.
URL [http://papers.nips.cc/paper/7132-task-based-end-to-end-model-learning-in-stochastic-optimization](http://papers.nips.cc/paper/7132-task-based-end-to-end-model-learning-in-stochastic-optimization).
- Elmachtoub and Grigas [2022]
Adam N. Elmachtoub and Paul Grigas.
Smart “Predict, then Optimize”.
*Management Science*, 68(1):9–26, January 2022.
ISSN 0025-1909.
doi: 10.1287/mnsc.2020.3922.
URL [https://pubsonline.informs.org/doi/10.1287/mnsc.2020.3922](https://pubsonline.informs.org/doi/10.1287/mnsc.2020.3922).
Publisher: INFORMS.
- Ho et al. [2020]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising Diffusion Probabilistic Models.
In *Advances in Neural Information Processing Systems*, volume 33, pages 6840–6851. Curran Associates, Inc., 2020.
URL [https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html](https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html).
- Hoffman and Johnson [2016]
Matthew D Hoffman and Matthew J Johnson.
ELBO surgery: yet another way to carve up the variational evidence lower bound.
In *NIPS 2016 workshop*, 2016.
- Jeon et al. [2025]
Haeun Jeon, Hyunglip Bae, Minsu Park, Chanyeong Kim, and Woo Chang Kim.
Locally Convex Global Loss Network for Decision-Focused Learning.
*Proceedings of the AAAI Conference on Artificial Intelligence*, 39(25):26805–26812, April 2025.
ISSN 2374-3468.
doi: 10.1609/aaai.v39i25.34884.
URL [https://ojs.aaai.org/index.php/AAAI/article/view/34884](https://ojs.aaai.org/index.php/AAAI/article/view/34884).
- Kim et al. [2015]
Sujin Kim, Raghu Pasupathy, and Shane G. Henderson.
A Guide to Sample Average Approximation.
In Michael C Fu, editor, *Handbook of Simulation Optimization*, pages 207–243. Springer, New York, NY, 2015.
ISBN 978-1-4939-1384-8.
doi: 10.1007/978-1-4939-1384-8˙8.
URL [https://doi.org/10.1007/978-1-4939-1384-8_8](https://doi.org/10.1007/978-1-4939-1384-8_8).
- Kingma and Welling [2014]
Diederik P Kingma and Max Welling.
Auto-Encoding Variational Bayes.
In *International Conference on Learning Representations*, April 2014.
URL [http://arxiv.org/abs/1312.6114](http://arxiv.org/abs/1312.6114).
- Kleywegt et al. [2002]
Anton J. Kleywegt, Alexander Shapiro, and Tito Homem-de Mello.
The Sample Average Approximation Method for Stochastic Discrete Optimization.
*SIAM Journal on Optimization*, 12(2):479–502, January 2002.
ISSN 1052-6234, 1095-7189.
doi: 10.1137/S1052623499363220.
URL [http://epubs.siam.org/doi/10.1137/S1052623499363220](http://epubs.siam.org/doi/10.1137/S1052623499363220).
- Kochenderfer et al. [2015]
Mykel J. Kochenderfer, Christopher Amato, Girish Chowdhary, Jonathan P. How, Hayley J. Davison Reynolds, Jason R. Thornton, Pedro A. Torres-Carrasquillo, N. Kemal Üre, and John Vian.
*Decision Making Under Uncertainty: Theory and Application*.
The MIT Press, 1st edition, June 2015.
ISBN 978-0-262-02925-4.
- Kong et al. [2025]
Lingkai Kong, Yuanqi Du, Wenhao Mu, Kirill Neklyudov, Valentin De Bortoli, Dongxia Wu, Haorui Wang, Aaron M. Ferber, Yian Ma, Carla P. Gomes, and Chao Zhang.
Diffusion Models as Constrained Samplers for Optimization with Unknown Constraints.
In *Proceedings of The 28th International Conference on Artificial Intelligence and Statistics*, pages 4582–4590. PMLR, April 2025.
URL [https://proceedings.mlr.press/v258/kong25b.html](https://proceedings.mlr.press/v258/kong25b.html).
ISSN: 2640-3498.
- Krishnamoorthy et al. [2023]
Siddarth Krishnamoorthy, Satvik Mehul Mashkaria, and Aditya Grover.
Diffusion Models for Black-Box Optimization.
In *Proceedings of the 40th International Conference on Machine Learning*, pages 17842–17857. PMLR, July 2023.
URL [https://proceedings.mlr.press/v202/krishnamoorthy23a.html](https://proceedings.mlr.press/v202/krishnamoorthy23a.html).
ISSN: 2640-3498.
- Mandi et al. [2022]
Jayanta Mandi, Victor Bucarey, Maxime Mulamba Ke Tchomba, and Tias Guns.
Decision-Focused Learning: Through the Lens of Learning to Rank.
In *Proceedings of the 39th International Conference on Machine Learning*, pages 14935–14947. PMLR, June 2022.
URL [https://proceedings.mlr.press/v162/mandi22a.html](https://proceedings.mlr.press/v162/mandi22a.html).
ISSN: 2640-3498.
- Mandi et al. [2024]
Jayanta Mandi, James Kotary, Senne Berden, Maxime Mulamba, Victor Bucarey, Tias Guns, and Ferdinando Fioretto.
Decision-Focused Learning: Foundations, State of the Art, Benchmark and Future Opportunities.
*Journal of Artificial Intelligence Research*, 80:1623–1701, August 2024.
ISSN 1076-9757.
doi: 10.1613/jair.1.15320.
URL [https://www.jair.org/index.php/jair/article/view/15320](https://www.jair.org/index.php/jair/article/view/15320).
- Mohamed et al. [2020]
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih.
Monte Carlo Gradient Estimation in Machine Learning.
*Journal of Machine Learning Research*, 21(132):1–62, 2020.
ISSN 1533-7928.
URL [http://jmlr.org/papers/v21/19-346.html](http://jmlr.org/papers/v21/19-346.html).
- Nichol and Dhariwal [2021]
Alexander Quinn Nichol and Prafulla Dhariwal.
Improved Denoising Diffusion Probabilistic Models.
In *Proceedings of the 38th International Conference on Machine Learning*, pages 8162–8171. PMLR, July 2021.
URL [https://proceedings.mlr.press/v139/nichol21a.html](https://proceedings.mlr.press/v139/nichol21a.html).
ISSN: 2640-3498.
- PJM Interconnection [2025]
PJM Interconnection.
Data Miner, 2025.
URL [https://dataminer2.pjm.com/list](https://dataminer2.pjm.com/list).
- Quandl WIKI dataset [2025]
Quandl WIKI dataset.
Nasdaq Data Link, 2025.
URL [https://data.nasdaq.com](https://data.nasdaq.com).
- Sanokowski et al. [2025]
Sebastian Sanokowski, Sepp Hochreiter, and Sebastian Lehner.
A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization, August 2025.
URL [http://arxiv.org/abs/2406.01661](http://arxiv.org/abs/2406.01661).
arXiv:2406.01661 [cs].
- Shah et al. [2022]
Sanket Shah, Kai Wang, Bryan Wilder, Andrew Perrault, and Milind Tambe.
Decision-Focused Learning without Decision-Making: Learning Locally Optimized Decision Losses.
*Advances in Neural Information Processing Systems*, 35:1320–1332, December 2022.
URL [https://proceedings.neurips.cc/paper_files/paper/2022/hash/0904c7edde20d7134a77fc7f9cd86ea2-Abstract-Conference.html](https://proceedings.neurips.cc/paper_files/paper/2022/hash/0904c7edde20d7134a77fc7f9cd86ea2-Abstract-Conference.html).
- Shalev-Shwartz et al. [2009]
Shai Shalev-Shwartz, Ohad Shamir, Nathan Srebro, and Karthik Sridharan.
Stochastic Convex Optimization.
In *Proceedings of Thirty Third Conference on Learning Theory*. PMLR, 2009.
- Shariatmadar et al. [2025]
Keivan Shariatmadar, Neil Yorke-Smith, Ahmad Osman, Fabio Cuzzolin, Hans Hallez, and David Moens.
Generalized Decision Focused Learning under Imprecise Uncertainty–Theoretical Study, March 2025.
URL [http://arxiv.org/abs/2502.17984](http://arxiv.org/abs/2502.17984).
arXiv:2502.17984 [cs].
- Silvestri et al. [2023]
Mattia Silvestri, Senne Berden, Jayanta Mandi, Ali İrfan Mahmutoğulları, Maxime Mulamba, Allegra De Filippo, Tias Guns, and Michele Lombardi.
Score Function Gradient Estimation to Widen the Applicability of Decision-Focused Learning.
September 2023.
URL [https://openreview.net/forum?id=ty046JUllZ](https://openreview.net/forum?id=ty046JUllZ).
- Sohl-Dickstein et al. [2015]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep Unsupervised Learning using Nonequilibrium Thermodynamics.
In *Proceedings of the 32nd International Conference on Machine Learning*, pages 2256–2265. PMLR, June 2015.
URL [https://proceedings.mlr.press/v37/sohl-dickstein15.html](https://proceedings.mlr.press/v37/sohl-dickstein15.html).
ISSN: 1938-7228.
- Song and Ermon [2019]
Yang Song and Stefano Ermon.
Generative Modeling by Estimating Gradients of the Data Distribution.
In *Advances in Neural Information Processing Systems*, volume 32. Curran Associates, Inc., 2019.
URL [https://proceedings.neurips.cc/paper_files/paper/2019/hash/3001ef257407d5a371a96dcd947c7d93-Abstract.html](https://proceedings.neurips.cc/paper_files/paper/2019/hash/3001ef257407d5a371a96dcd947c7d93-Abstract.html).
- Song et al. [2020]
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-Based Generative Modeling through Stochastic Differential Equations.
October 2020.
URL [https://openreview.net/forum?id=PxTIG12RRHS&utm_campaign=NLP%20News&utm_medium=email&utm_source=Revue%20newsletter](https://openreview.net/forum?id=PxTIG12RRHS&utm_campaign=NLP%20News&utm_medium=email&utm_source=Revue%20newsletter).
- Sun and Yang [2023]
Zhiqing Sun and Yiming Yang.
DIFUSCO: Graph-based Diffusion Solvers for Combinatorial Optimization.
November 2023.
URL [https://openreview.net/forum?id=JV8Ff0lgVV](https://openreview.net/forum?id=JV8Ff0lgVV).
- Tang and Nurmaya Musa [2011]
Ou Tang and S. Nurmaya Musa.
Identifying risk issues and research advancements in supply chain risk management.
*International Journal of Production Economics*, 133(1):25–34, September 2011.
ISSN 0925-5273.
doi: 10.1016/j.ijpe.2010.06.013.
URL [https://www.sciencedirect.com/science/article/pii/S0925527310002215](https://www.sciencedirect.com/science/article/pii/S0925527310002215).
- Tashiro et al. [2021]
Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon.
CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation.
In *Advances in Neural Information Processing Systems*, volume 34, pages 24804–24816. Curran Associates, Inc., 2021.
URL [https://proceedings.neurips.cc/paper_files/paper/2021/hash/cfe8504bda37b575c70ee1a8276f3486-Abstract.html](https://proceedings.neurips.cc/paper_files/paper/2021/hash/cfe8504bda37b575c70ee1a8276f3486-Abstract.html).
- Vahdat and Kautz [2020]
Arash Vahdat and Jan Kautz.
NVAE: A Deep Hierarchical Variational Autoencoder.
In *Advances in Neural Information Processing Systems*, volume 33, pages 19667–19679. Curran Associates, Inc., 2020.
URL [https://proceedings.neurips.cc/paper/2020/hash/e3b21256183cf7c2c7a66be163579d37-Abstract.html](https://proceedings.neurips.cc/paper/2020/hash/e3b21256183cf7c2c7a66be163579d37-Abstract.html).
- Wang et al. [2025]
Prince Zizhuang Wang, Jinhao Liang, Shuyi Chen, Ferdinando Fioretto, and Shixiang Zhu.
Gen-DFL: Decision-Focused Generative Learning for Robust Decision Making, February 2025.
URL [http://arxiv.org/abs/2502.05468](http://arxiv.org/abs/2502.05468).
arXiv:2502.05468 [cs].
- Wang et al. [2024]
Yafei Wang, Bo Pan, Mei Li, Jianya Lu, Lingchen Kong, Bei Jiang, and Linglong Kong.
Sample Average Approximation for Conditional Stochastic Optimization with Dependent Data.
June 2024.
URL [https://openreview.net/forum?id=YuGnRORkJm](https://openreview.net/forum?id=YuGnRORkJm).
- Wilder et al. [2019]
Bryan Wilder, Bistra Dilkina, and Milind Tambe.
Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization.
*Proceedings of the AAAI Conference on Artificial Intelligence*, 33(01):1658–1665, July 2019.
ISSN 2374-3468.
doi: 10.1609/aaai.v33i01.33011658.
URL [https://ojs.aaai.org/index.php/AAAI/article/view/3982](https://ojs.aaai.org/index.php/AAAI/article/view/3982).

## Appendix A Appendix

#### Notation.

We let $\frac{\partial f}{\partial x}$ denote the Jacobian matrix where $\big(\frac{\partial f}{\partial x}\big)_{i,j}:=\frac{\partial f_{i}}{\partial x_{j}}$ and $\nabla_{x}f:=\big(\frac{\partial f}{\partial x}\big)^{\top}$ denote the gradient. For a vector $v$, $D(v)$ denotes a diagonal matrix with $v$ on its diagonal. Let $P(\cdot)$ denote a probability distribution and $p(\cdot)$ denote a probability density; in particular, for diffusion models we use $P_{\theta}$ for the model’s output distribution and $p_{\theta}$ for transition densities.

In this appendix, we derive the decision optimization problem with general convex constraints rather than merely linear constraints. Assume the optimization problem is

$$ $z_{\theta}^{*}(x)=\operatorname*{arg\,min}_{z}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)],\quad\text{s.t.}\ \ h(x,z)\leq 0,\ g(x,z)=0,$ (19) $$

where $h(x,z)\leq 0$ denotes the convex inequalities constraints and $g(x,z)=0$ denotes the equality constraints.

### A.1 Proofs for Section 4

###### Proposition A.1 (Reparameterization trick in diffusion models) .

Let $T\in\mathbb{N}^{+}$, and suppose the reverse diffusion model defines a Gaussian distribution in Eq. [5](https://arxiv.org/html/2510.11590v1#S3.E5) with fixed scalars $\sigma_{t}\geq 0$ and a standard normal prior $y_{T}\sim\mathcal{N}(0,I)$. Let $\{\epsilon_{t}\}_{t=0}^{T}$ be i.i.d. $\mathcal{N}(0,I)$. Then the model output $y$ can be expressed as a transformation $y=R(\epsilon_{0:T},\theta\mid x)$ of a base noise distribution $\epsilon\sim P(\epsilon)$, where $R$ is differentiable in $\theta$. Also assume $\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]$ is continuously differentiable. Then we have

$$ $\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}\left[\left(\sum_{s=1}^{T}\left(\prod_{u=1}^{s-1}J_{u}\right)A_{s}\right)^{\top}\nabla_{y}f(R(\epsilon,\theta\mid x),z)\right],$ (20) $$

where $A_{t}:=\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial\theta}$, $J_{t}:=\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial y_{t}}$, and we define $\prod_{u=1}^{0}J_{u}:=I$.

###### Proof.

The conditional diffusion reverse process is defined as

$$ $y_{t-1}=\mu_{\theta}(y_{t},t,x)+\sigma_{t}\epsilon_{t-1},\qquad y_{T}=\epsilon_{T},$ $$

where the noise term $\sigma_{t}\epsilon_{t-1}$ is $\theta$-independent. Differentiating both sides w.r.t. $\theta$ gives

$$ $\frac{\partial y_{t-1}}{\partial\theta}=\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial\theta}+\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial y_{t}}\frac{\partial y_{t}}{\partial\theta}.$ $$

Denote

$$ $A_{t}:=\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial\theta},\quad J_{t}:=\frac{\partial\mu_{\theta}(y_{t},t,x)}{\partial y_{t}},\quad G_{t}:=\frac{\partial y_{t}}{\partial\theta}.$ $$

Thus, we have

$$ $G_{t-1}=A_{t}+J_{t}G_{t},\quad G_{T}=0.$ $$

Our final goal is:

$$ $\displaystyle\nabla_{\theta}R(\epsilon_{0:T},\theta|x)$ $\displaystyle=\frac{\partial y_{0}}{\partial\theta}=G_{0}$ (21) $\displaystyle=A_{1}+J_{1}A_{2}+J_{1}J_{2}A_{3}+\cdots+J_{1}\cdots J_{T_{1}}A_{t}$ (22) $\displaystyle=\sum_{s=1}^{T}\left(\prod_{u=1}^{s-1}J_{u}\right)A_{s},$ (23) $$

where we define $\prod_{u=1}^{0}J_{u}:=I$.
Then, we have

$$ $\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}$ $\displaystyle=\nabla_{\theta}\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[f(R(\epsilon,\theta|x),z)]$ (24) $\displaystyle=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[\nabla_{\theta}f(R(\epsilon,\theta|x),z)]$ (25) $\displaystyle=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}\left[\nabla_{\theta}R(\epsilon,\theta\mid x)^{\top}\nabla_{y}f(R(\epsilon,\theta|x),z)\right]$ (26) $\displaystyle=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}\left[\left(\sum_{s=1}^{T}\left(\prod_{u=1}^{s-1}J_{u}\right)A_{s}\right)^{\top}\nabla_{y}f(R(\epsilon,\theta|x),z)\right]$ (27) $$

∎

###### Lemma A.2 (Gradient of Reparameterization method) .

Assume the model prediction $y$ can be expressed as a transformation $y=T(\epsilon,\theta\mid x)$, $\epsilon\sim P(\epsilon)$. The total derivative of the decision objective $F$ w.r.t. $\theta$ can be computed as

$$ $\displaystyle\frac{dF}{d\theta}=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}T(\epsilon,\theta|x))^{\top}\nabla^{2}_{zy}f(z^{*},y)]\\ 0\\ 0\end{bmatrix},$ (28) $$

where $H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(y,z^{*})]+\nabla_{zz}^{2}(\lambda^{*\top}h(x,z^{*}))$ is the Hessian of the Lagrangian with respect to $z$, $G=\nabla_{z}h(x,z^{*})$ is the Jacobian of the inequality constraints in $z^{*}$, and $Q=\nabla_{z}g(x,z^{*})$ is the Jacobian of the equality constraints in $z^{*}$.

###### Proof.

At the primal-dual optimal solution $(z_{\theta}^{*},\lambda_{\theta}^{*},\nu_{\theta}^{*})$ to Eq. [3](https://arxiv.org/html/2510.11590v1#S3.E3), the following KKT conditions must hold:

$$ $\displaystyle\nabla_{z}\mathcal{L}(\theta,z_{\theta}^{*},\lambda_{\theta},\nu_{\theta};x)$ $\displaystyle=0,$ $\displaystyle\lambda_{\theta}\odot h(x,z_{\theta}^{*})$ $\displaystyle=0,$ $\displaystyle g(x,z_{\theta}^{*})$ $\displaystyle=0$ $\displaystyle\lambda_{\theta}\geq 0,\nu_{\theta}$ $\displaystyle\geq 0,$ $\displaystyle h(x,z_{\theta}^{*})$ $\displaystyle\leq 0.$ $$

Since $h$ does not depend on $\theta$ here, we can apply Proposition [A.1](https://arxiv.org/html/2510.11590v1#A1.Thmtheorem1) to the KKT conditions to get

$$ $\displaystyle\frac{\partial\nabla_{z}\mathcal{L}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial z}\frac{\partial z^{*}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial\lambda^{*}}\frac{\partial\lambda^{*}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial\nu^{*}}\frac{\partial\nu^{*}}{\partial\theta}$ $\displaystyle=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}T(\epsilon,\theta|x))^{\top}\nabla_{y}(\nabla_{z}f(z^{*},y))]+\big(\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+\nabla_{zz}^{2}h(x,z^{*})\big)\frac{\partial z^{*}}{\partial\theta}$ $\displaystyle\qquad+\nabla_{z}h(x,z^{*})\frac{\partial\lambda^{*}}{\partial\theta}+\nabla_{z}g(x,z^{*})\frac{\partial\lambda^{*}}{\partial\theta}$ $\displaystyle=0.$ (29) $$

$$ $\displaystyle\frac{\partial\lambda^{*}\odot h(x,z^{*})}{\partial z^{*}}\frac{\partial z^{*}}{\partial\theta}+\frac{\partial\lambda^{*}\odot h(x,z^{*})}{\partial\lambda^{*}}\frac{\partial\lambda^{*}}{\partial\theta}=D(\lambda^{*})\nabla_{z}h(x,z^{*})\frac{\partial z^{*}}{\partial\theta}+D(h(x,z^{*}))\frac{\partial\lambda^{*}}{\partial\theta}=0.$ (30) $$

In matrix form, we have

$$ $\displaystyle\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\\ \frac{\partial\nu^{*}}{\partial\theta}\end{bmatrix}=-\begin{bmatrix}\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}T(\epsilon,\theta|x))^{\top}\nabla^{2}_{zy}f(z^{*},y)]\\ 0\\ 0\end{bmatrix},$ (31) $$

where $H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+\nabla_{zz}^{2}(\lambda^{*\top}h(x,z^{*}))+\nabla_{zz}^{2}(\nu^{*\top}g(x,z^{*}))$, $G=\nabla_{z}h(x,z^{*})$, and $Q=\nabla_{z}g(x,z^{*})$. Furthermore, if equalities and inequalities are affine (as in main paper), $H$ reduces to $\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(y,z^{*})]$ since $\nabla_{zz}^{2}h=\nabla_{zz}^{2}g=0$.

By chain rule, we have

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle=\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\end{bmatrix}$ $\displaystyle=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}T(\epsilon,\theta|x))^{\top}\nabla^{2}_{zy}f(z^{*},y)]\\ 0\\ 0\end{bmatrix}.$ $$

∎

###### Proposition A.1 (Reparameterization trick in diffusion models) .

###### Proof.

###### Lemma A.2 (Gradient of Reparameterization method) .

###### Proof.

### A.2 Proofs for Section 5

###### Proposition A.3 .

Let $f:\mathcal{Y}\times\mathbb{R}^{d}\to\mathbb{R}$ be any function that does not depend on $\theta$. If $y\sim P_{\theta}(\cdot\mid x)$, then

$$ $\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[f(y,z){\frac{d\log P_{\theta}(y\mid x)}{d\theta}}].$ (32) $$

###### Proof.

$$ $\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}$ $\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[f(y,z)]$ (33) $\displaystyle=\frac{\mathrm{d}}{\mathrm{d}\theta}\int f(y,z)P_{\theta}(y\mid x)\mathop{}\!\mathrm{d}y$ (34) $\displaystyle=\int P_{\theta}(y\mid x)\frac{\mathrm{d}}{\mathrm{d}\theta}f(y,z)+f(y,z)\frac{\mathrm{d}}{\mathrm{d}\theta}P_{\theta}(y\mid x)dy$ (35) $\displaystyle=\int P_{\theta}(y\mid x)\frac{\mathrm{d}}{\mathrm{d}\theta}f(y,z)+f(y,z)\frac{\mathrm{d}}{\mathrm{d}\theta}\log P_{\theta}(y\mid x)*P_{\theta}(y\mid x)dy$ (36) $\displaystyle=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\frac{\mathrm{d}}{\mathrm{d}\theta}f(y,z)]+\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[f(y,z){\frac{d\log P_{\theta}(y\mid x)}{d\theta}}].$ (37) $$

This immediately implies the results by noticing $f$ does not depend on $\theta$.
∎

###### Proposition A.4 .

Let $P_{\theta}(y\mid x)$ be a probability density parameterized by $\theta\in\Theta$, and let $f:\mathcal{Y}\times\mathbb{R}^{d}\to\mathbb{R}$ be a scalar-valued function that does not depend on $\theta$. Fix any $z\in\mathbb{R}^{d}$. Suppose that there exists some neighborhood $N(\theta_{0})\subseteq\Theta$ around $\theta_{0}\in\Theta$ such that the following 3 assumptions are satisfied:

- 1.
For all $\theta\in N(\theta_{0})$, the function $h(y):=P_{\theta}(y\mid x)\,f(y,z)$ is integrable;
- 2.
For all $\theta\in N(\theta_{0})$ and almost all $y\in\mathcal{Y}$, the gradient $\nabla_{\theta}P_{\theta}(y\mid x)$ exists; and
- 3.
There exists an integrable function $g:\mathcal{Y}\to\mathbb{R}$ that dominates $\nabla_{\theta}P_{\theta}(y\mid x)$. That is, for all $\theta\in N(\theta_{0})$ and almost all $y\in\mathcal{Y}$, $\left\lVert\nabla_{\theta}P_{\theta}(y\mid x)\right\rVert_{1}\leq\left\lvert g(y)\right\rvert$.

Then,

$$ $\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}=\operatorname{\mathbb{E}}_{y\sim P_{\theta_{0}}(\cdot\mid x)}[f(y,z)\cdot\nabla_{\theta}\log P_{\theta_{0}}(y\mid x)].$ $$

###### Proof.

We make use of the log-derivative trick:

$$ $P_{\theta_{0}}(y\mid x)\cdot\nabla_{\theta}\log P_{\theta_{0}}(y\mid x)=\frac{P_{\theta_{0}}(y\mid x)}{P_{\theta_{0}}(y\mid x)}\cdot\nabla_{\theta}P_{\theta_{0}}(y\mid x)=\nabla_{\theta}P_{\theta_{0}}(y\mid x).$ $$

Then

$$ $\displaystyle\nabla_{\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta_{0}}(\cdot\mid x)}[f(y,z)]$ $\displaystyle=\nabla_{\theta}\int_{\mathcal{Y}}f(y,z)P_{\theta_{0}}(y\mid x)\mathop{}\!\mathrm{d}y$ $\displaystyle=\int_{\mathcal{Y}}\nabla_{\theta}\left[f(y,z)\,P_{\theta_{0}}(y\mid x)\right]\mathop{}\!\mathrm{d}y$ Leibniz integral rule $\displaystyle=\int_{\mathcal{Y}}f(y,z)\,P_{\theta_{0}}(y\mid x)\,\nabla_{\theta}\log P_{\theta_{0}}(y\mid x)\mathop{}\!\mathrm{d}y$ log-derivative trick $\displaystyle=\operatorname{\mathbb{E}}_{y\sim P_{\theta_{0}}(\cdot\mid x)}\left[f(y,z)\nabla_{\theta}\log P_{\theta_{0}}(y\mid x)\right].$ $$

∎

###### Lemma A.5 (Gradient of Score Function) .

The total derivative of the decision objective $F$ w.r.t. $\theta$ can be computed as

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{z}f(z^{*},y)(\frac{dELBO}{d\theta})^{\top}]\\ 0\\ 0\end{bmatrix}.$ (38) $$

###### Proof.

Differentiate this KKT system w.r.t. $\theta$ and applying Proposition [A.4](https://arxiv.org/html/2510.11590v1#A1.Thmtheorem4) yields

$$ $\displaystyle\frac{\partial\nabla_{z}\mathcal{L}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial z}\frac{\partial z^{*}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial\lambda^{*}}\frac{\partial\lambda^{*}}{\partial\theta}+\frac{\partial\nabla_{z}\mathcal{L}}{\partial\nu^{*}}\frac{\partial\nu^{*}}{\partial\theta}$ $\displaystyle=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{z}f(z^{*},y)(\nabla_{\theta}\log P_{\theta}(y|x))^{\top}]+(\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+\nabla_{zz}^{2}(\lambda^{\star\top}h(x,z^{*})))\frac{\partial z^{*}}{\partial\theta}$ $\displaystyle\quad+\nabla_{z}h(x,z^{*})\frac{\partial\lambda^{*}}{\partial\theta}+\nabla_{z}g(x,z^{*})\frac{\partial\nu^{*}}{\partial\theta}$ $\displaystyle=0.$ (39) $$

$$ $\displaystyle\frac{\partial\lambda^{*}\odot h(x,z^{*})}{\partial z^{*}}\frac{\partial z^{*}}{\partial\theta}+\frac{\partial\lambda^{*}\odot h(x,z^{*})}{\partial\lambda^{*}}\frac{\partial\lambda^{*}}{\partial\theta}$ $\displaystyle=$ $\displaystyle D(\lambda^{*})\nabla_{z}h(x,z^{*})\frac{\partial z^{*}}{\partial\theta}+D(h(x,z^{*}))\frac{\partial\lambda^{*}}{\partial\theta}=0.$ (40) $$

In matrix form, this becomes

$$ $\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\\ \frac{\partial\nu^{*}}{\partial\theta}\end{bmatrix}=-\begin{bmatrix}\operatorname{\mathbb{E}}_{y}[\nabla_{z}f(z^{*},y)(\nabla_{\theta}\log P_{\theta}(y\mid x))^{\top}]\\ 0\\ 0\end{bmatrix},$ (41) $$

where $H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+\nabla_{zz}^{2}(\lambda^{*\top}h(x,z^{*}))$, $G=\nabla_{z}h(x,z^{*})$.

Applying the chain rule to $F$ now gives

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle=\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\\ \frac{\partial\nu^{*}}{\partial\theta}\end{bmatrix}$ (42) $\displaystyle=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{z}f(z^{*},y)(\nabla_{\theta}\log P_{\theta}(y\mid x))^{\top}]\\ 0\\ 0\end{bmatrix}.$ (43) $$

Then, we replace $\nabla_{\theta}\log P_{\theta}(y\mid x)$ with the gradient of ELBO score for sample $y$ and have

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{z}f(z^{*},y)(\frac{dELBO}{d\theta})^{\top}]\\ 0\\ 0\end{bmatrix}.$ (44) $$

∎

###### Remark A.6 (Why we cannot compute the gradient using score-matching) .

One may attempt to apply the chain rule $\nabla_{\theta}\log P_{\theta}(y|x)=\nabla_{\theta}y\nabla_{y}\log P_{\theta}(y|x)$, and then estimate $\nabla_{y}\log P_{\theta}(y_{t+1}|x)\approx\nabla_{y}\log P_{\theta}(y_{t+1}|y_{t},x)\approx s_{\theta}(y_{t},t,x)$ via score-matching [29] (using the learned score $s_{\theta}(y_{t},t,x)$ of the diffusion model). However, this approach is invalid in our setting: Under the log-trick, $y$ is treated as a free variable and $\theta$ enters only through $P_{\theta}(y|x)$, so the pathwise term $\nabla_{\theta}y$ does not exist (see [Section A.2](https://arxiv.org/html/2510.11590v1#A1.SS2) for derivation). Our ELBO-based surrogate (Eq. ([11](https://arxiv.org/html/2510.11590v1#S5.E11))) avoids this obstacle entirely.

###### Proposition A.3 .

###### Proof.

###### Proposition A.4 .

###### Proof.

###### Lemma A.5 (Gradient of Score Function) .

###### Proof.

###### Remark A.6 (Why we cannot compute the gradient using score-matching) .

### A.3 Proof for Eq. 13

Based on the results in [Lemma A.5](https://arxiv.org/html/2510.11590v1#A1.Thmtheorem5), we have

$$ $\displaystyle\frac{dF}{d\theta}$ $\displaystyle=-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\nabla_{z}f(z^{*},y)(\frac{d\operatorname{ELBO}}{d\theta})^{\top}]\\ 0\\ 0\end{bmatrix}$ (45) $\displaystyle=\underbrace{-\begin{bmatrix}\frac{dF}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}}_{:=u(\theta)^{\top}}\frac{d}{d\theta}\begin{bmatrix}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\nabla_{z}f(z^{*},y)\operatorname{ELBO}]\\ 0\\ 0\end{bmatrix}$ (46) $\displaystyle=\frac{d}{d\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\underbrace{u(\theta)^{\top}\begin{bmatrix}[\nabla_{z}f(z^{*},y)]\\ 0\\ 0\end{bmatrix}}_{:=w_{\theta}(y)}\operatorname{ELBO}]$ (47) $\displaystyle=\frac{d}{d\theta}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[{\mathrm{detach}[w_{\theta}(y)]}{\operatorname{ELBO}}]$ (48) $$

### A.4 Empirical evidence for ELBO gradient approximation

Assume our model is $\theta=(A,B,c)$, and the noise is predicted by $\epsilon_{\theta}(y_{t},t,x)=A_{t}y_{t}+B_{t}x+c_{t}$. True data $y\sim\mathcal{N}(Wx,I)$. In this way, there is a closed-form solution for true $\nabla_{\theta}\log p(y_{0}|x)$ since $y_{0}$ is a Gaussian and $y_{t}=\sqrt{\bar{\alpha}_{t}}y_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$.

Figure: Figure 6: Compare the cosine similarity between the true gradient and the estimated gradient using a linear model.
Refer to caption: x6.png

### A.5 Deterministic Optimization and Gaussian Model in Stochastic Optimization

#### Deterministic Optimization

Since deterministic can also be viewed as a reparameterization trick without any randomness $\epsilon$, we can reuse our derivation in Section [4](https://arxiv.org/html/2510.11590v1#S4) and compute the gradient $\nabla_{x}F(x)$ by Eq. [9](https://arxiv.org/html/2510.11590v1#S4.E9).

#### Gaussian Model in Stochastic Optimization

Since there are many recent papers that use the Gaussian model as a predictor for stochastic DFL, we also claim that it has the reparameterization and score function form.

- 1.
Reparameterization with Gaussian Model. Also using reparameterization trick:
$\displaystyle\nabla_{\theta}{\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]}=\operatorname{\mathbb{E}}_{\epsilon\sim P(\epsilon)}[(\nabla_{\theta}y\nabla_{y}f(y,z)].$
(49)
But with the predictor instantiated as a Gaussian model, i.e., $y$ is computed by $y=R(\epsilon,\theta|x))^{\top}$ and $R$ is a Gaussian model.
- 2.
Score Function with Gaussian Model. Recall that we need to approximate the term $\log P_{\theta}(y|x)$ for the diffusion model. However, this term has a closed-form for a Gaussian model. Assume our Gaussian model is $\mathcal{N}(\mu_{\theta},\Sigma_{\theta})$, then the negative log-likelihood is
$\displaystyle\log P_{\theta}(y\mid x)=-\frac{1}{2}[(y-\mu_{\theta})^{\top}\Sigma_{\theta}^{-1}(y-\mu_{\theta})+\log\det\Sigma_{\theta}+d\log(2\pi)].$
(50)
Then, the gradient for the score function can be calculated using $\nabla_{\theta}\log P_{\theta}(y\mid x)$.
Gaussian models are powerful tools for many DFL tasks. However, we want to claim that the diffusion model is more general and requires less model tuning and model assumptions.

### A.6 Training details

We summarize our model settings for the deterministic model, Gaussian model, and diffusion model in [Table 2](https://arxiv.org/html/2510.11590v1#A1.T2).

**Table 2: Comparison of deterministic, Gaussian, and diffusion model architectures.**
| Parameter | Deterministic MLP | Gaussian MLP | Diffusion |
| --- | --- | --- | --- |
| Trunk (layers $\times$ width) | $2\times 1024$ | $2\times 1024$ | $2\times 1024$ |
| Activation | ReLU / Swish (SiLU) | ReLU / Swish (SiLU) | Swish (SiLU) |
| Inputs to network | $\;x\in\mathbb{R}^{d_{x}}$ | $\;x\in\mathbb{R}^{d_{x}}$ | $\;[\,y_{t}\in\mathbb{R}^{d_{y}},\;t,\;x\,]$ |
| Time embedding | — | — | Sinusoidal $t$ (16-d) $\rightarrow$ 2 FC + SiLU |
| Output head | $\hat{y}\in\mathbb{R}^{d_{y}}$ | $\mu_{\theta}(x),\log\sigma_{\theta}^{2}(x)\in~\mathbb{R}^{d_{y}}$ | $\hat{\epsilon}=\epsilon_{\theta}(y_{t},t,x)\in\mathbb{R}^{d_{y}}$ |
| Uncertainty form | None (point) | Gaussian $\mathcal{N}(\mu,\mathrm{diag}(\sigma^{2}))$ | Non-parametric (learned $P_{\theta}(y\mid x)$) |
| Two-stage training loss | MSE on $y$ | Gaussian NLL | Weighted denoising MSE (ELBO-equiv.) |
| DFL gradient computation | Implicit differentiation through KKT | (1) Reparam: $y=\mu+\sigma\odot\epsilon$ + implicit diff; (2) Gaussian score-function | (1) Pathwise (backprop through sampler, $T$ steps) + implicit diff; (2) *Weighted-ELBO* score-function |
| Sample size $M$ | — | $10$ (synthetic), $25$ (power), $50$ (portfolio) | $10$ (synthetic), $25$ (power), $50$ (portfolio) |
| Learning rate (lr) | Task-specific, mostly $1\times 10^{-4}$ | Task-specific, mostly $1\times 10^{-4}$ | reparam: $1\times 10^{-5}$;  score-fn: $8\times 10^{-6}$ |
| Inference | Use one $\hat{y}$ in optimizer | Draw $M$ samples $y^{(m)}$ from Gaussian model | Reverse diffusion to sample $M$ samples $y^{(m)}$ |

For all experiments, we perform 10 random seeds to evaluate variability. We also find that adding a small regularizer during DFL training can help the model learning the data distribution and avoid some bad local minima, leading to a stable training process.

### A.7 Details of synthetic example

In this example, we consider a factory that decides how much to manufacture for each of $d\in\mathbb{N}$ products. The parameter $Y\in\mathbb{R}^{d}$ represents the profit margin for each product, i.e., $Y_{i}$ is the profit per unit of product $i$; due to uncertainty in market conditions, $Y$ is uncertain. The factory’s decision $z\in[0,C]^{d}$ represents how much of each product to manufacture, where $C$ is the maximum capacity for each product. For simplicity, we do not consider any contextual features $x$ in this example. That means DFL learns a distribution that generates $y$ that can minimize the decision objective.

Suppose that the factory has a risk-averse cost function $f(y,z)=\exp(-y^{\top}z)$, which indicates that the factory wants to put a larger weight on the product with higher profit $Y_{i}$. Intuitively, if the factory knew $Y$ exactly, then the optimal strategy would be all-or-nothing: set $z_{i}=C$ if $Y_{i}>0$, or $z_{i}=0$ if $Y_{i}<0$. Likewise, with respect to a point prediction of $Y$, the optimal deterministic decision $z_{\mathrm{det}}^{*}\in\{0,C\}^{d}$ is attained on the boundary of the feasible set.

Under uncertainty, the decision-maker seeks to minimize the expected cost by solving a stochastic optimization problem:

$$ $\displaystyle z_{\mathrm{sto}}^{*}$ $\displaystyle\in\operatorname*{arg\,min}_{z\in[0,C]^{d}}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\exp(-y^{\top}z)].$ (51) $$

In this stochastic case, the optimal investment $z_{\mathrm{sto}}^{*}$ typically lies in the interior of the feasible region, which balances the potential high reward of investing against the risk of losses.

Then, we compute the necessities for diffusion DFL:

$$ $\displaystyle H$ $\displaystyle=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}g(z^{*},y)]+(\lambda^{*})^{\top}\nabla_{zz}^{2}h(x,z^{*})=\exp(-y^{\top}z)yy^{\top}$ (52) $\displaystyle G$ $\displaystyle=\nabla_{z}h(x,z^{*})=-\exp(-y^{\top}z)y.$ (53) $$

For reparameterization, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\end{bmatrix}$ $\displaystyle=-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}(\nabla_{\theta}y_{i})^{\top}\nabla_{zy}^{2}g(z^{*},y_{i})\\ 0\end{bmatrix}$ $$

where $\nabla_{zy}^{2}g(z^{*},y)=\exp(-y^{\top}z)(yz^{\top}-I_{d})$ in this case.

For the score function, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y}[\nabla_{z}g(z^{*},y)(\frac{dELBO}{d\theta})^{\top}]\\ 0\end{bmatrix}$ $\displaystyle\approx-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}\nabla_{z}g(z^{*},y_{i})(\frac{dELBO_{i}}{d\theta})^{\top}\\ 0\end{bmatrix}.$ $$

### A.8 Details on power schedule task

This task involves a 24-hour electricity generation scheduling problem with uncertain demand. The decision $z\in\mathbb{R}^{24}$ represents the electricity output to schedule for each hour of the next day. The uncertainty $y\in\mathbb{R}^{24}$ represents the actual power demand for each of the 24 hours. The goal is to meet demand as closely as possible at minimum cost. We also consider a decision cost function that penalizes storage, excess generation, and ramping following [5]:

- 1.
Let $\gamma_{s}$ and $\gamma_{a}$ be the per-unit costs of shortage (not meeting demand) and excess (over-generation), respectively. We use $\gamma_{s}=50$ and $\gamma_{e}=0.5$ in our experiment.
- 2.
Let $c_{r}$ be a penalty on hour-to-hour changes in generation. We use $c_{r}=0.4$ in appropriate units.

Formally, if $z=(z_{1},\dots,z_{24})$ and $y=(y_{1},\dots,y_{24})$, the loss for a single day is

$$ $\displaystyle\min_{z}\ \operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]=\sum_{i=1}^{24}\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[\gamma_{s}[y_{i}-z_{i}]_{+}+\gamma_{e}[z_{i}-y_{i}]_{+}+\frac{1}{2}(z_{i}-y_{i})^{2}]$ $\displaystyle\ \mathrm{s.t.}\ |z_{i}-z_{i-1}|\leq c_{r}\ \text{for all}\ i\in\{1,2,\dots,24\}.$ (54) $$

Then, we compute the necessities for diffusion DFL:

$$ $\displaystyle H$ $\displaystyle=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}g(z^{*},y)]+(\lambda^{*})^{\top}\nabla_{zz}^{2}h(x,z^{*})=I_{n}$ (55) $\displaystyle G$ $\displaystyle=\nabla_{z}h(x,z^{*})=\begin{bmatrix}-1&1&0&\cdots&0\\ 0&-1&1&\cdots&0\\ \vdots&&\ddots&\ddots&\vdots\\ 0&\cdots&0&-1&1\end{bmatrix}.$ (56) $$

For reparameterization, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\end{bmatrix}$ $\displaystyle=-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}(\nabla_{\theta}y_{i})^{\top}\nabla_{zy}^{2}g(z^{*},y_{i})\\ 0\end{bmatrix}$ $$

where $\nabla_{zy}^{2}g(z^{*},y)=-I$ in this case.

For the score function, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y}[\nabla_{z}g(z^{*},y)(\frac{dELBO}{d\theta})^{\top}]\\ 0\end{bmatrix}$ $\displaystyle\approx-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}\nabla_{z}g(z^{*},y_{i})(\frac{dELBO_{i}}{d\theta})^{\top}\\ 0\end{bmatrix}.$ $$

#### Dataset.

We use real historical electricity load data from the PJM regional grid (a standard public dataset) [20]. The features $x$ for each day include: the previous day’s 24-hour load profile, the previous day’s temperature profile, calendar features, and seasonal sinusoidal features. In total, $d_{x}=28$ features for each day were constructed. We normalized all input features for training. The target label $y$ is the next day’s 24-hour load vector.

For completeness, we include an extended comparison of different sample sizes in [Figure 9](https://arxiv.org/html/2510.11590v1#A1.F9), which further highlights that additional samples yield diminishing returns in accuracy while linearly increasing compute cost.

Figure: Figure 7: Results on the 24-hour power grid scheduling task.
Refer to caption: x7.png

Figure: Figure 8: Score function vs. Reparameterization
Refer to caption: x8.png

Figure: Figure 9: Comparison between different sample sizes for score function and reparameterization.
Refer to caption: x9.png

### A.9 Details on stock portfolio task

We consider a mean-variance portfolio optimization problem with uncertain returns. Here $y\in\mathbb{R}^{n}$ represents the random next-day returns of $n$ assets (stocks), and $z\in\mathbb{R}^{n}$ are the portfolio weights we assign to each asset (the fraction of our capital invested in each stock).
Our goal is to maximize expected return while keeping the risk (variance) low. This can be written as minimizing a loss that is a negative expected return plus a quadratic penalty on variance:

$$ $\displaystyle\min_{z}\ \operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}[f(y,z)]=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot|x)}\left[\frac{\alpha}{2}z^{\top}yy^{\top}z-y^{\top}z\right],\quad\mathrm{s.t.}\quad z^{\top}\mathbf{1}=1,\ 0\leq z_{i}\leq 1.$ (57) $$

Then, we compute the necessities for diffusion DFL:

$$ $\displaystyle H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+(\lambda^{*})^{\top}\nabla_{zz}^{2}h(x,z^{*})+\nabla_{zz}^{2}(\nu^{*\top}g(x,z^{*}))=\alpha\operatorname{\mathbb{E}}_{y\sim P_{\theta}(y\mid x)}[yy^{\top}],$ (58) $\displaystyle G=\nabla_{z}h(x,z^{*})=\begin{bmatrix}I_{n},\\ -I_{n}\end{bmatrix}$ (59) $\displaystyle Q=\nabla_{z}g(x,z^{*})=\mathbf{1}^{\top}.$ (60) $$

For reparameterization, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}\frac{\partial z^{*}}{\partial\theta}\\ \frac{\partial\lambda^{*}}{\partial\theta}\\ \frac{\partial\nu^{*}}{\partial\theta}\end{bmatrix}$ $\displaystyle\approx-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}(\nabla_{\theta}y_{i})^{\top}\nabla_{zy}^{2}g(z^{*},y_{i})\\ 0\\ 0\end{bmatrix}$ $$

where $\nabla_{zy}^{2}f(z^{*},y)=\operatorname{\mathbb{E}}_{y}[2\alpha y^{\top}z-1]$ in this case.

For the score function, we have

$$ $\displaystyle(\frac{dloss}{d\theta})^{\top}$ $\displaystyle=\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\operatorname{\mathbb{E}}_{y}[\nabla_{z}g(z^{*},y)(\frac{dELBO}{d\theta})^{\top}]\\ 0\\ 0\end{bmatrix}$ $\displaystyle\approx-\begin{bmatrix}\frac{dloss}{dz^{*}}\\ 0\\ 0\end{bmatrix}^{\top}\begin{bmatrix}H&G^{\top}&Q^{\top}\\ D(\lambda^{*})G&D(h(x,z^{*}))&0\\ Q&0&0\end{bmatrix}^{-1}\begin{bmatrix}\frac{1}{M}\sum_{i=1}^{M}\nabla_{z}g(z^{*},y_{i})(\frac{dELBO_{i}}{d\theta})^{\top}\\ 0\\ 0\end{bmatrix},$ $$

where $\nabla_{z}g(z^{*},y_{i})=\alpha y_{i}y_{i}^{\top}z^{*}-y_{i}$.

#### Dataset.

We use daily stock prices from 2004–2017 for constituents of the S&P 500 index [21]. We obtained this data via Quandl’s API (specifically WIKI pricing data; the user will need a Quandl API key to replicate. We compute daily returns for each stock (percentage change). To construct features $x$, we use a rolling window of recent history for each asset. Specifically, for each day, a data point is for predicting next day’s returns: we include the past 5 days of returns for each of the $n$ assets, past 5 days of trading volume for each asset, plus some aggregate features. To avoid an explosion of dimension with large $n$, we also include PCA-compressed features: we take the top principal components of the last 5-day return matrix to summarize cross-asset trends. In the end, for $n=50$ assets, we ended up with $d_{x}=28$ features. All features are normalized and we use a time-series split: first 70% of days for training (2004–2013), next 10% for validation (2014), last 20% for test (2015–2017). We evaluate performance on the test set by simulating the portfolio selection every day and computing the average return achieved.

### A.10 (Additional task) Details on inventory stock problem

We also validate our approaches on a toy inventory control problem. In this task, the uncertain demand $y$ is drawn from a multi-modal distribution (a mixture of Gaussians), where we vary the number of mixture components $K$ to control the distribution complexity:

$$ $\displaystyle p(x)=\sum_{j=1}^{K}\pi_{j}\phi(x;\mu_{j},\Sigma_{j}),$ (61) $$

where $\pi_{j}$ is the probability of choosing component j and $\phi(x;\mu_{j},\Sigma_{j})$ is a multivariate Gaussian density with parameter $(\mu_{j},\Sigma_{j})$.
The cost function follows the standard newsvendor formulation with piecewise penalties for under-stock and over-stock:

$$ $\displaystyle f(y,z)=c_{0}z+\frac{1}{2}q_{0}z^{2}+c_{b}[y-z]_{+}+\frac{1}{3}r_{b}([y-z]_{+}^{3})+c_{h}[z-y]_{+}+\frac{1}{3}r_{h}([z-y]_{+}^{3}).$ (62) $$

Our learning objective is to minimize the expected cost over this stochastic demand, i.e., a stochastic optimization problem:

$$ $\displaystyle\min_{z}L(\theta)=\operatorname{\mathbb{E}}_{y\sim P(\cdot|x)}[f(y,z)]\quad\mathrm{s.t.}\ 0\leq z\leq z_{max}.$ (63) $$

Figure: Figure 10: Toy decision task comparing deterministic and diffusion DFL. *Left*: distribution of per-instance regret (lower is better). *Middle*: distribution of chosen decision $z$ in the lower-level; the stochastic method tracks the true distribution $z^{*}$ more closely. *Right*: pairwise win-rate on test set; a large fraction of costs from the stochastic method are lower than the deterministic one, indicating that modeling uncertainty yields better decisions.
Refer to caption: x10.png

Figure: Figure 11: Prediction distribution for inventory stock problem.
Refer to caption: x11.png

We compare a deterministic DFL model against our diffusion DFL model on this toy task. [Figure 10](https://arxiv.org/html/2510.11590v1#A1.F10) summarizes the results, where the diffusion model (stochastic DFL) achieves substantially lower regret on average than the deterministic model. Besides, we observe that the decision $z$ obtained by our diffusion method closely tracks the true optimal decisions $z^{*}$ by capturing the multi-modal demand uncertainty, whereas the deterministic predictor’s decisions deviate more. In [Figure 10](https://arxiv.org/html/2510.11590v1#A1.F10) (c), we directly compare the decision outcomes via a win-rate: the fraction of test instances where one method achieves lower cost than the other. The diffusion DFL method attains a win-rate of about 75% against the deterministic baseline, which confirms that modeling uncertainty leads to better downstream decisions

For data generation, we set $\mu=[-4,0,4],\Sigma=[0.15,0.25,0.15]^{\top}\mathbf{1}$ for $K=3$, $\mu=[-6.,-3.,0.,3.,6.],\Sigma=[0.15,0.25,0.35,0.25,0.15]^{\top}\mathbf{1}$ for $K=5$, and $\mu=[-8.0,-6.0,-4.0,-2.0,-1.0,0.0,1.2,2.8,4.5,7.5],\Sigma=[0.30,0.75,0.25,0.40,0.22,0.20,0.22,0.35,0.70,1.25]^{\top}\mathbf{1}$.

Following our previous derivation, we can compute the necessities for diffusion DFL by

$$ $\displaystyle H=\operatorname{\mathbb{E}}_{y\sim P_{\theta}(\cdot\mid x)}[\nabla_{zz}^{2}f(z^{*},y)]+(\lambda^{*})^{\top}\nabla_{zz}^{2}h(x,z^{*})=diag(q_{0})+q_{b}I_{(y>z)}+q_{h}I_{z>y}$ (64) $\displaystyle G=\nabla_{z}h(x,z^{*})=\begin{bmatrix}-I\\ I\end{bmatrix},$ (65) $\displaystyle D(h(x,z^{*}))=0.$ (66) $$

### A.11 Additional Related Works

#### Stochastic optimization

Making decisions under uncertainty is a classic topic in operations research and machine learning [24]. Stochastic optimization formulations explicitly consider uncertainty by optimizing the expected objective over a distribution of unknown parameters. A common approach is the Sample Average Approximation (SAA) [12, 2, 35], which draws many samples from the estimated distribution and solves an approximated deterministic problem minimizing the average cost. While SAA can handle arbitrary uncertainty distributions in theory, it becomes very computationally expensive and still does not consider the distribution during optimization [10]. It will lead to optimizing the *sample mean*, which may yield a decision that performs poorly if reality often falls into one of several distinct models far from the mean [10, 6].

## Peer Reviews

confidence: 5
contribution: 1
initial_rating: 0
presentation: 4
questions: Have I missed some critical contribution in the above points that should justify a different conclusion?
soundness: 3
strengths: I note two key strengths:
* **Writing attention to detail.** This paper is reasonably polished, and it is clear the authors know how to write a paper to an appropriate technical standard. 
* **Reasonably-careful experiments.** The authors not only demonstrate good performance on a key benchmark of interest compared to pure-deterministic methods, they also include various ablation studies to check individual parts of their conclusions.
 
In an ordinary review, I would offer many more details on these points. However, in light of a very serious issue which I will bring to light shortly, I will omit these in the interest of time.
summary: This paper studies a paradigm termed decision-focused learning, where the goal is to solve a decision-making problem that depends on uncertain parameters that must be learned. The authors develop a computational pipeline to propagate uncertainty from the learning procedure to the decision procedure using diffusion models, and show that incorporating uncertainty results in improved decision-making performance.
weaknesses: This paper has a critical weakness because of which I am confident it needs to be rejected: **a very significant portion of this work's findings are known and covered in standard textbooks, but are presented as novel.** In particular, the paper's actual novelty is limited to introducing diffusion models to well-known decision-making paradigms. If re-written completely in a manner that reflects this contribution, I do believe the authors have done sufficient development to warrant a paper, but the gap between the current manuscript and what would be appropriate is, to say it again, sufficient to warrant a complete rewrite. This is unfortunately far too much to address in the rebuttal phase.

Allow me to elaborate:
* **Decision-focused learning is presented as a novel emerging paradigm, but is not new, because the authors' formulation is equivalent to a (stochastic) contextual bandit.** Let me be completely precise: the authors define their setting via the optimization objective $\mathbb{E}f(y^\star, z^\star(x))$ where $(x,y^\star) \sim D$ and $z^\star$ is parameterized by some vector $\theta$ which is to be optimized. Let us define the notation $\phi^\star(z) = f(y^\star, z)$, so the objective becomes $\mathbb{E} \phi^\star(z^\star(x))$ where $(x,\phi^\star) \sim D$. With this notation change, it is clear that the authors' setting is a (stochastic) contextual bandit: we are given a random context vector x which is correlated with some random objective $\phi^\star$ (either in a known way in the Bayesian case, or unknown way in the non-Bayesian case), and our goal is to choose the optimal map from context to actions. With this notational link made, there is a substantial literature of prior work the authors could consider and compare with, which is not mentioned in this work. Note, in particular, that if D is assumed known, maximizing the objective is equivalent to minimizing regret, which is the usual performance objective considered because it behaves much better mathematically.
* **Using stochastic models inside dynamic programming is not new.** Indeed, this is covered in Chapter 1.5 of "Sequential Decision Analytics and Modeling" by Warren Powell, to name just one book that includes this topic.
* **In the contextual bandit literature, the need for uncertainty is well-understood.** This is the basis not only for classical upper confidence bound algorithms, but also methods based on approximate dynamic programming such as expected improvement algorithms from the Bayesian optimization literature.
* **Even if one is interested purely in practical aspects involving diffusion models, one should benchmark against standard algorithms.** For example, I see no reason why, in light of the above equivalence, any DFL algorithm in the authors' sense cannot directly be benchmarked against contextual kernel UCB.
* **Section 4 consists entirely of known ideas.** In particular, the application of the reparametrization trick to handle objectives defined as expectations of this kind is, at this point, completely standard and need not be introduced except in code. The discussion on differentiating through the solution of a parameterized optimization problem is also standard, though less than reparametrization tricks - this should be deferred to an appendix.
* **In Section 5, the key ideas involving score estimation are essentially the same as the standard REINFORCE trick from reinforcement learning. Moreover, the algorithmic justification for using ELBO as a surrogate to the score is not technically sound.** The authors' argument, essentially, is to state that $f(x) \leq g(x)$ on basis of an evidence lower bound, and argue that this implies $\nabla f(x) \approx \nabla g(x)$, which is not necessarily true even if $|f(x) - g(x)| \leq \epsilon$ for an arbitrary given $\epsilon$. Rather than relying on a theoretical explanation that does not follow, I believe the authors should instead lean on their empirical results in Figure 2, which shows these are close, as this avoids a misleading explanation.

Summarizing, I believe this work's only contribution which is novel is the exploration of diffusion models in the context of stochastic dynamic programming. This contribution is a reasonable fit for ICLR and could make for a good paper, but the paper needs to be re-written in a manner that reflects this.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
questions: * Can you confirm that the Gaussin DFL method is in fact a stochastic method, i.e. using a stochastic optimization problem with multiple samples downstream?
* How many samples were used in the empirical evaluation?
soundness: 3
strengths: The insight about the limited expressiveness (for lack of a better term) of deterministic DFL is correct and definitely useful, even if not entirely novel.

The empirical results on the stock portfolio problem provide evidence that a combination of stochastic prediction and a complex distribution can improve the effectiveness of the DFL approach very significantly, beyond what is possible for either technique. Remarkably, this is true even when the two-phase approach is rather expressive and should be capable of approximating well the ground truth parameter distribution (the PFL diffusion method).

The suggested differentiation method based on SFGE + ELBO appears to be rather good at improving the computational efficienty, with no major drop in terms of solution quality.

The SFGE/ELBO method also seems to be applicable, with a few minor modifications, to a _much_ more general setting than the one presented in eq. (2). Since the method only requires function evaluation, it seems it could easily account for solution repairs costs (similar to the other SFGE method by Silvestri et al.) and therefore, combinatorial problems, two-stage problems, or decision problems with uncertain parameters in the constraints.
summary: This paper presents a Decision Focused Learning method where the output of the estimator is stochastic, rather than being a point estimate, and the downstream problem is also a stochastic optimization one. As a main advantage, the approach can better address decision problems where no deterministic formulation can result in the correct optimal solution under uncertainty, e.g. because any individual realization leads to an either maximizing or minimizing a variables, while the best value under under uncertainty would be somewhat in the middle.

The authors use a diffusion process to generate their estimate and present two techniques to enable differentiation. The first -- technically more precise, but expensive in terms of computation time and memory -- is based on a reparameterization of each diffusion step and implicit differentiation. The second -- much lighter from a computational standpoint -- is based on Score Function Gradient Estimation and an Evidence Lower Bound. Using a diffusion process enables modeling complex distributions

The approach is tested on selected benchmarks from the literature, and at least in one realistic (but still far from true practical applicability) problem it achieved a very significant advantage in terms of solution quality.
weaknesses: It is my opinion that the paper strengths outweigh its weaknesses. That said, I can find some, of which a few are major and a few minor.

The first major weakness I can find is that the approach requires solving a stochastic optimization problem per example, and per epoch. Stochastic optimization problems can be very challenging to address and the corresponding solution methods are well known for their rather poor scalability. Combined with the already computational cost of DFL training, this significantly reduces the practical appeal of the method. That said, given the improvements reported in one benchmark, I think there are definitely practical settings with relatively lightweight optimization problems that can benefit from the proposed technique.

The second major weakness is again related to applicability. If enough data is available and a sufficiently expressive (distributional) ML model is used, the approximation error for the ground truth distribution will eventually converge to zero. Under these conditions, a Prediction Focused Learning (i.e. two phase) pipeline should be capable of reaching optimal or close to optimal results, similarly to what happens in a classical DFL setup with non-misspecified ML models. The empirical results corroborate this intuition, with the Diffusion TS method reaching high quality results in the synthetic benchmark. It would seem that the ideal application case for the proposed method is one where a deterministic decision problem is inadequate, a complex distribution is needed, and the problem cost is not well aligned with the likelihood function used in PFL. While situations like this certainly exists, the applicability of the method is also reduced.

Here is a list of some minor weaknesses:

* The fact that deterministic DFL may not be expressive enough for some optimization problems under uncertainty is also established in [1]. Given that this paper is very recent and that the expressiveness claim is not at the center of this work, this issue is definitely minor. That said, it might be worth including a citation
* The Gaussian method by Silvestri et al. is not stochastic, but deterministic: the Gaussian distribution is used just for smoothing, while the optimization problem works with a single parameter vector (both at training and inference time). By reading Appendix A.5, it seems the authors instead implemented an actual Gaussian stochastic DFL pipeline, based on their own techniques, but using a simpler distribution. This is again a minor issue, since the original method by Silvestri et al. would likely work even worse in this setting (though it would be much more scalable)
* This might be my fault, but I did not easily find the number of samples used in the empirical evaluation

[1] Schutte, Noah, et al. "Sufficient Decision Proxies for Decision-Focused Learning." arXiv preprint arXiv:2505.03953 (2025).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: The score-function estimator seems to rely on local smoothness between the diffusion model’s likelihood landscape and the downstream decision cost; can the authors discuss or empirically validate when this assumption might fail, e.g., in highly nonconvex or discontinuous decision problems? Have the authors tested Diffusion-DFL in nonconvex or combinatorial settings, where diffusion’s multimodality could be beneficial?
soundness: 3
strengths: The paper introduces the first integration of diffusion models into the decision-focused learning framework, allowing richer uncertainty modeling than deterministic approaches.

- The proposed score-function estimator is a practical contribution: it replaces costly pathwise differentiation with an ELBO-based surrogate, maintaining decision performance while reducing memory and compute by several orders of magnitude.
- The methodology is well-grounded theoretically (via implicit differentiation and the log-trick) and empirically validated on multiple domains, showing consistent improvements in decision quality.
- Clear ablation studies (memory vs. accuracy, variance reduction, decision dimension scaling) demonstrate that the approach is both effective and scalable.
summary: This paper introduces Diffusion-DFL, a first framework that integrates diffusion probabilistic models into decision-focused learning (DFL) — an approach that trains predictive models to directly optimize downstream decision quality rather than prediction accuracy.
Unlike prior DFL methods, which rely on deterministic or simple probabilistic predictors (e.g., Gaussian), Diffusion-DFL uses a conditional diffusion model to capture potentially complex and multimodal uncertainty in the parameters of a downstream stochastic optimization problem. The decision layer then optimizes the expected cost under this learned distribution.

To train the model end-to-end, the authors develop two gradient estimators:
1) Reparameterization estimator: backpropagates through the diffusion sampling process, providing an exact but memory-intensive gradient.
2) Score-function estimator: avoids backpropagating through diffusion by applying the log-trick, estimating gradients via the model’s ELBO  — drastically reducing memory and compute cost with minimal accuracy loss.

Experiments on synthetic manufacturing, power scheduling, and portfolio optimization tasks show that Diffusion-DFL consistently yields better decision outcomes than both two-stage and deterministic DFL baselines, while the score-function version achieves similar performance to reparameterization with over 400× lower memory usage.
weaknesses: - The score-function approximation relies on an implicit assumption of local smoothness between the model’s likelihood and decision-loss landscapes, which is not formally analyzed or guaranteed, although supported empirically. 
- Experimental settings are limited to convex or low-dimensional decision problems; performance on more complex or nonconvex tasks (where diffusion expressiveness would matter most) remains unclear.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - I'm not familiar with decision-focused learning. Are there any other benchmarks for evaluating the performance of decision-focused learning? It seems that the three optimization tasks provided in the paper are not complex enough.
soundness: 3
strengths: - Utilizing diffusion models in decision-focused learning is sound and novel
- Authors successfully mitigate issues raised when using diffusion models in decision-focused learning
- Experiments and further analysis are comprehensive
summary: The paper introduces a novel DFL approach that leverages diffusion models to represent the distribution of uncertain parameters. To mitigate the issue raised by using diffusion models, i.e., sequential sampling procedure, the authors develop two algorithms: reparametrization and score function. Experiment results demonstrate that using diffusion models outperforms prior baselines.
weaknesses: - The authors say one benefit of utilizing diffusion models is capturing a multi-modal distribution. Does the task provided benefit when we capture the multi-modal distribution? It would be nice to persuade readers why we should capture multi-modal distribution in decision-focused learning.


Decision: Accept (Poster)