Decision-aware Training of Spatiotemporal Forecasting Models to Select a Top-K Subset of Sites for Intervention
Authors: Kyle Heuton
Frederick Muench
Shikhar Shrestha
Thomas J. Stopka
Michael C Hughes
Venue: ICML 2025 poster (2025)
Keywords: spatiotemporal forecasting
top K
decision theory
probabilistic models
Primary area: probabilistic_methods
TL;DR: We develop methods for selecting exactly K sites out of many with probabilistic spatiotemporal forecasting models.

Abstract: Optimal allocation of scarce resources is a common problem for decision makers faced with choosing a limited number of locations for intervention. Spatiotemporal prediction models could make such decisions data-driven.
A recent performance metric called fraction of best
possible reach (BPR) measures the impact of using a model’s recommended size K subset of sites compared to the best possible top-K in hindsight. We tackle two open problems related to BPR. First, we explore *how to rank* all sites numerically given a probabilistic model that predicts event counts jointly across sites. Ranking via the per-site mean is suboptimal for BPR. Instead, we offer a better ranking for BPR backed by decision theory. Second, we explore*how to train* a probabilistic model's parameters to maximize BPR. 
Discrete selection of K sites implies all-zero parameter gradients which prevent standard gradient training. We overcome this barrier via advances in perturbed optimizers. We further suggest a training objective that combines likelihood with a BPR constraint to deliver high-quality top-K rankings as well as good forecasts for all sites. We demonstrate our approach on two where-to-intervene applications: mitigating opioid-related fatal overdoses for public health and monitoring endangered wildlife.

Decision-aware Training of Spatiotemporal Forecasting Models to Select a Top-K Subset of Sites for Intervention
Kyle Heuton
F. Samuel Muench
Shikhar Shrestha 2 Thomas J. Stopka? Michael C. Hughes
Abstract
from mitigating overdoses in public health (Marks et al., 2021b) to forest fire management (Cheng & Wang, 2008) to wildlife monitoring (Golden et al., 2022; Hefley et al., 2017). Across these domains, there a pressing need for predictive models that can make accurate predictions of near-term future events fine spatiotemporal resolutions. Such models can enable inform data-driven decisions about how to allocate limited resources maximize utility.
Optimal allocation of scarce resources is a com- mon problem for decision makers faced with choosing limited number of locations for in- tervention. Spatiotemporal prediction models could make such decisions data-driven. A recent performance metric called fraction of best possi- ble reach (BPR) measures the impact of using a model's I recommended size K subset of sites com- pared to the best possible top-K in hindsight. We tackle two open problems related to BPR. First, explore how rank all sites numerically given probabilistic model that predicts event counts jointly across sites. Ranking via the per-site mean is suboptimal for BPR. Instead, we offer a bet- ter ranking for BPR backed by decision theory. Second, we explore how train probabilistic model's parameters to maximize BPR. Discrete selection of K sites implies all-zero parameter gradients which prevent standard gradient train- ing. We overcome this barrier via advances in perturbed optimizers. We further suggest a train- ing objective that combines likelihood with a BPR constraint to deliver high-quality top-K rankings as well as good forecasts for all sites. We demon- strate approach two where-to-intervene applications: mitigating opioid-related fatal over- doses for public health and monitoring endan- gered wildlife.
In this work, we seek to help decision pecision-makers select where to intervene. Given historical data for a fixed set of S candi- date spatial sites, we develop models that can recommend a specific subset of given size K for some action or interven- tion. We think of hyperparameter К as setting the budget for interventions. In an ideal world, ddecision-maker could afford interventions in all S sites. However, when resource constraints allow only K sites to receive interventions, de- cision selecting a specific K-of-S subset is required. While such decisions may often be heuristic in current practice, we hope to offer data-driven solutions.
With this goal in mind, choosing sensible performance metric is critical assessing which models have real-world utility. Common metrics such as squared error or absolute error are not well matched to where-to-intervenn decisions because they treat all sites equally. Recent work on over- dose forecasting has suggested a metric termed the fraction of best possible reach, or BPR (Heuton al., 2022; 2024). BPR measures a ratio of event counts. The numerator sums over the model's recommended K sites, while the denomi- nator sums over the best possible K selected in hindsight. This type of evaluation has been used in preregistered trial (Marshall et al., 2022) for assessing forecasts of opioid overdoses in Rhode Island, as well as a follow-up feasabil- ity study (Allen et al., 2023). BPR is applicable to many tinerett-intervene problems beyond public health.
1. Introduction
Statistical machine learning methods for spatiotemporal forecasting can play a vital role in high-stakes applications
While some publications have reported BPR in evaluations, we suggest that natural goal would be for this perfor- mance metric to inform two other key parts of data-driven decision-making: model-based ranking and model training. By ranking, we mean that given fixed model parameters, the knowledge of BPR as the metric of interest should im-
'Department of Computer Science, Tufts University, Med- ford, Massachusetts, United States Department of Public Health and Community Medicine, Tufts University School of Medicine, Boston, Massachusetts, United States. Correspondence to: Kyle Heuton <kyle.heuton@ee eduu.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Code: http:://giihub..oouustt--ss- wtt.///eciiiinwaaareeeeo

---

Decision-Aware Training of Spatiotemporal Forecasting Models
pact the numerical score assigned to each site to determine the top K. By training, we mean how to update model pa- rameters to achieve high BPR. This paper contributes new methods for solving both ranking and training problems when BPR is the preferred metric.
strate these contributions first on synthetic data, where we reveal how off-the-shelff methods without our innovations can be suboptimal for decision-making We further eval- uate against alternatives on two applications: mitigating opioid-related fatal overdoses and monitoring endangered birds. We hope our contributions spark interest in where- to-intervene problems in the methodological community and also lead effective deployments of data-driven top-K decision-making in public health and beyond.
Our work overcomes several technical barriers. The first barrier in ranking. Given a fixed probabilistic model, determining how to compute per-site score, which will later be sorted to find the top K, is not obvious. It may be tempting to use the model's per-site mean, but decision the- ory suggests not all loss functions recommend the per-site mean as the best estimator (Berger, 2013; Murphy, 2022). As a relatively new metric, the problem of how to assign a numerical ranking to sites for optimal BPR decision-making currently open. We contribute a tractable ranking method that is provably best for reasonable bound on BPR. We further show how the score function estimator (Kleijnen & Rubinstein, 1996; Mohamed et al., 2020) can be used to cal- culate gradients of this ranking with respect to parameters.
Related work. The application of machine learning to decision-making problems is widely studied in operations research literature (Bertsimas & Kallus, 2020; Sadana et al., 2025), including the problem of how to train a model for downstream decision-making (Mandi et al., 2024). We re- view several model training approaches later in Sec. 4.
Other researchers have used decision-aware objectives to solve limited resource allocation problems. Chung et al. (2022) study how to allocate essential medicines in Sierra Leone across hospital sites. Gupta et al. (2024) pursue a mnlrre-to-intevven task in urban planning, selecting where to build speed humps to reduce pedestrian injuries. Our work differs in its focus on the BPR metrie, our hybrid decision aware objective that preserves likelihood, and eval- uations that forecast the future given the recent past.
The second barrier prevents training parameters to improve BPR. First-order gradient descent is a common, effective algorithm we would like to use. However, gradients of BPR with respect to parameters are problematic. While small changes to parameters induce some changes to per- site scores, only changes large enough to move a site into or out of the top-K ranked sites will adjust BPR. Thus, gradi- ents of BPR with respect to parameters will be zero almost everywhere, preventing gradient methods from ever moving beyond subpar initial parameters. To fix this, we leverage recent advances in perturbed optimizers (Abernethy al., 2016; Berthet et al., 2020) to yield effective and efficient gradient estimation for BPR. Our team explored this idea for optimizing BPR alone in earlier non-archival work (Heuton et al., 2023); this paper offers an expanded treatment with more accessible presentation, while addressing two more barriers.
Throughout this paper, a recurring takeaway is that conven- tional training based on maximizing likelihood can yield suboptimal decision meriiiion-mmkinggg for BPR, especially when fore- casting models are misspecified. In this vein, our work shares similar goals direct loss minimization (Wei et al., 2021) and loss-calibrated methods (Lacoste-Julien et al., 2011). We are inspired by the way these works combine task- specific losses and decision theory to improve probabilistic models. Others have extended loss-calibration to neural nets (Cobb et al., 2018) and to continuous actions (Kuśmier- czyk et al., 2019). Yet a tractable and scalable recipe that prioritizes top- there-to-intervent decision-makinn is not immediate next step from this work.
The final barrier is designing training objective achieve applied goals. We find that optimizing BPR alone can lead to predictions with far lower likelihood than conventional training. This raises concerns about overall model quality and generalization. To address this, we pose a constrained optimization problem to maximize likelihood subject to a BPR quality constraint. This combined objective delivers quality top-K K recommendations and good forecasts for all sites. Our objective is reminiscent of past additive combi- nations of regression loss and decision loss (Kao et al., 2009). Unlike that work, ours pursues non-convex losses and directly enforces decision quality via constraints.
2. Technical Background and Problem Setup
Notation. Mathematically, some of our notation follows Sander et al. (2023). Function TOPKM ASK takes as input a vector of length S and an integer K. It produces a binary vector of length S with exactly K entries equal to I. Each I entry corresponds to value in the top K largest entries of the input vector. The s-th entry of the output is
{ {
if RANK(()) S K otherwise
TOPKMASK(r, K), =
(1)
We ultimately contribute methods for how-to-rank and how- to-train when making where-to-interver decisions. Using these tools, a variety of models can be directly optimized to make effective top K site recommendations. We demon-
Here, function RANK provides a numerical ranking (largest- to to-smallest) from to S for each entry of an S-dimensional input vector. The function ToPK IIDI acts on the same input
2

---

Decision-Aware Training of Spatiotemporal Forecasting Models
Decision P Preto F Frontier
Where to place K cameras to observe the most crones next year?
B R: fraction of Best Poss ible Reach
(ight
High Good Decatrn
BPR. Radia al ovens in recommer and kis aduual event in porfect cradica top sitas Between 0.0 and 1.0, higher better.
Likelihood
Good loz Dediions
Data-driven top-K
BPR
ranking scores
K=2 decision
How to Rank
How to Train
2014
2014
2024
for 2025 Q1
Q1
Q2
Q4
Pol Pol Maa....
as ainary mask
Neiver per-site mean
BPR
Lik.
0
0.01
0
(Ø) = Epg Uv
max Lik (ML)
X
14
1
0
0.02
max BPR
0
3
2
0.65
1
Ours: Rank by bound BPR
r (6) = Epp
3/a
Decision Aware ML
0.03
1
3/
2
2
0
0.44
(DAML)
Figure Visual overview of our approach and contributions to the how to rank and how to train open problems.
as TOPKMASK, but return the size K set of integer indices corresponding to top K values.
of K sites among all S sites. For the rest of this section, let y E TR:S be the vector of observations at the target time / 1-4- (we skip time subscripts on y to keep notation simple). The vector y is not available when the decision of R is made. Following Heuton et al. (2022), we define BPR as:
Problem definition. We wish to probabilistically model events that occur across S distinct spatial sites over time. At each site, indexed by s, we can observe non-negative scalar count or value ys 2 O. We assume that larger ys corresponds to greater value in intervention at site In our public health applications, Vs represents counts of fatal opioid-related related overdoses. In our later wildlife monitoring case study, V/ counts how often a rare animal appears. At each time f, we stack all observations into a vector y/: e R$S We assume that the true data-geus gata-aeeng distribution for each vector z/: given past history does not change over time, including between the training and test periods.
EVER Us
BPR(R, y) - HETOOIDDSs Us
Both terms in this fraction can only be evaluated in hind- sight, after the vector is realized at time t + 1. The numer- ator counts how many events the model's 's recommendation would reach. The denominator counts how many events a perfect oracle with knowledge of the future could reach on the same budget of K sites. Overall, we interpret BPR the fraction of events of interest the current model's selection R would reach compared to perfect knowledge of the future. Higher BPR indicates a better model for choosing where to intervene. BPR's ' best value is 1.0, its worst is 0.0.
We denote our joint model for this vector Po (y:), where model parameter vector & defines the density over r.v. yt- We assume that explicitly evaluating this pdf and sampling values yt are both feasible. Given these assumptions, our framework is quite flexible: the vector & could represent the weights of neural network or the coefficients of a logistic regression Bayesian hierarchical model.
Ranking sites. Given a fixed model é and target time, the ranking problem is how to assign numerical values to all S sites that if the top K sites reap high BPR scores. We need to define a ranking vector E R of numerical scalar scores for all S sites. Higher r. values indicate greater priority for site s.
assigned R, we
Given training set of T times, our model family in general factorizes p(y1.r) = II (M 11: To ease notation throughout, we omit conditioning on past history or other exogenous features. So (y.) below should be read as equal p (3 for models with such dependencies.
Suppose we have a loss function L(r, y) (not necessarily related to BPR) that produces a scalar value indicating the overall quality of taking an "action" r and then realizing outcome y. Lower values of T. indicate better decisions. A natural framework for making decisions about actions (Mur- phy, 2022) is to minimize the expected loss:
Our goal is to use this probabilistic model to solve where- to-intervene decision making problem. We primarily intend to use the model to numerically rank all S sites, then select the top K sites in this ranking for fear-future intervention.
argmin Ey-pe [T.(r, y)
(3)
Definition of BPR. At current time t, we evaluate model o's ability to select top-K K subset of sites for intervention before time ++ Let R be the model's recommended subset
As a simplistic example, if the loss is defined as the sum of squared errors, L(r, y) E (r+ - -ya)² the optimal ac-
3

---

Decision-Aware Training of Spatiotemporal Forecasting Models
tion is provably the per-site mean: r. = Ep [ys]. Similarly, for the loss that sums up absolute errors, the optimal action is the per-site median (Schwertman al., 1990; Balkus, 2024). We tackle defining an optimal ranking for BPR.
all S terms instead of the top K: E ETTTOrIDD(()) Vs S Σ -==1 ys 1 y. This bound tighter when the sum of all entries is less than K U, which is typically true of sparse y vectors in our applications. Our how-to-rank problem then becomes
To pose our ranking problem formally, we need to convert the higher-is-better BPR metric into a lower-is-better loss that depends on r. Define negative BPR loss as
argmin r
Ey Byypp
ToPKMASK(r, K)
(6)
y
TP (r,y) y) y ToPKMAsK(r, K)
(4)
<BPR(r.))
y TOPKMASK K)
Again, we solve via any vector r" whose Kop-K binary vector TOPKMAsK(r", K) equals TOPKMASK(Ep.) 2 1 K).
This way of writing the loss with dot products of top-K binary vectors is equivalent to -PR((TTPKIDS(T y)
One valid solution is the expected ratio of vector y to its sum, which we nickname the ratio estimator
Connection to 0-1 knapsack. Given a fixed y vector, the problem of selecting K sites to minimize LBP can reduce to the canonical 0-1 knapsack problem (Dantzig, 1957) where each site would have value ys and weight 1, and the budget constraint allows just K of all S sites. Our how-to-rank contribution solves a more general problem: how to set r when y is not given but must be forecasted by our model.
1
M
y (m)
r' (0) = E Eyrps
(7)
1.y.
M m=
1 yy(m
This ranking is distincr from the per-site mean: there exist fixed models & where the ratio estimator and the per-site mean would select different subsets of the same S sites. See App. B for concrete cases where the ratio estimator earns BPR 2.5x to 5x higher than the per-site mean, even when all estimators know the true data- gentraannn model.
In Sec. below, we show how analysis of tractable bounds of the loss in Eq. (4) suggests high-quality ranking function (o) for BPR. This ranking is usable across different model families, as long the model Pe allows generating many samples of events y. Later in Sec. 4, we show how to train parameters Φ with gradient descent to yield better top-K decisions as judged by BPR.
When exact computation of this expectation is not easy, we recommend a Monte Carlo approximation using M samples { (m)(m) Jm= drawn iid from ps, as in Eq. (7). This a stochastic estimator; rankings can differ across repeat trials if M is not large enough.
3. Methods for Ranking
4. Methods for Training
Loose bound justifies per-site mean ranking. A natural first guess for ranking is the per-site mean: F = Ep, [y]. We can show this is justifiable way to minimize expected loss on a simplistic upper bound on BPR. Assume there exists an upper limit U such that for all 8, we can guarantee U 2 ys- The sum over any entries of vector y in the denominator of BPR is then bounded by K U. Plugging this bound into the minimize expected loss problem and simplifying with linearity of expectations yields
We now consider various ways to train the parameters of our probabilistic model on a training set that covers T distinct time periods indexed by
4.1. Maximum likelihood (ML) estimation
Conventional training would maximize the likelihood, or cquivalently minimize negative log likelihood (NLL):
J NL (d)= - - (3)))
(8)
Eps [v] K TOOKMAAK(r,, K)
r argmin
(5)
We assume Ps is differentiable, so solving for point es- timate à is possible via gradient descent. If we add an optional prior term log p(0) to enforce an inductive bias or control fver-fitting, this is known as MAP estimation.
<BPR(r.y)
Many solutions exist: any vector r' that satisfies ToPKMASK(P",, K) (E,, [y), K) can an argmin. One valid solution here the per-site mean F(d) Ep, yy]. However, this solution is optimal for potentially quite loose bound on BPR that approximates the denominator with the constant K U.
If the model is well-specified and training set size T is large enough, this is reliable strategy to estimate Φ. We could then use the ranking methods from Sec. 3 for where-to- intervene decisions. However, popular wisdom reminds us that "all models are wrong" in some way for real-world data. As we will show in later experiments, fitting a misspeci- fied model via ML estimation can produce 6 that deliver suboptimal BPR, even using the optimal ranking for that Φb
Tighter bound suggests the ratio estimator for rank- ing. Instead of bounding with constant K U, we can up- per bound the denominator in Eq. (4) by summing over

---

Decision-Aware Training of Spatiotemporal Forecasting Models
4.2. Direct loss minimization for BPR
3, ~ ~N(, Then, we compute
Inspired by the broad goal of direct loss minimization (Wei et al., 2021), another approach would be to find parameters that minimize our BPR-specific decision making loss L BPR In this strategy, we seek & values that minimize
8 , br)), by(r) = (
V = 10 Ei (U))
(11)
JBPR(() = Σ L BPR (r) (o),
(9)
The Jacobian v, is an x matrix, where entry j. k gives the scalar derivative Dru. The conceptual justification for the Jacobian estimator comes from Abernethy et al. (2016) and Berthet et al. (2020). Noise level o > is a hyperparameter that sets the strength of stochastic smoothing. It must be carefully selected in practice to add enough noise so that indicators bj change for different samples 2j. but not too much noise so the b; preserve the signal in r.
_b:
Here, for r" (δ) we use the ratio estimator in Eq. (7).
To train with modern gradient methods, we'd need to com- pute the gradient VoJ E. Var, : Vr, Lt. However, technical difficulties arise with each term in this chain rule expansion. Below, we propose practical estimators for each term that overcome these difficulties.
Gradient V.r:.. The difficulty here is differentiating through the expectation in Eq. (7), especially when y is a discrete random variable (integer counts in our later over- dose or wildlife applications). We use the score function trick (Mohamed et al., 2020), popularized by Ranganath et al. (2014) yet dating back decades (Kleijnen & Rubin- stein, 1996), sometimes also called REINFORCE (Williams, 1992). We can draw M samples y
Putting our ccore-function trick and perturbed optimizer estimators together, we compute the overall gradient of loss at index I as a product of individual estimators: Volr = V.r, b,,, Lz. We compute the last term V Le via automatic differentiation.
Armed with this gradient estimator, we can pursue direct minimization of 7BPR via stochastic gradient descent meth- ods. Stochasticity here comes from both M score function samples and J perturbation samples. For convenience and reliability, we use all T records in the training set in every estimate, avoiding minibatching over time.
(m))
Pd. then compute
M
(m)
y
Vor r
V (B)))
(10)
M
(mm)
Yt
We assume the true data- gettagenerating distribution is un- changed across train and test time periods. If this does not hold, objectives that just average over as in Eq. (9) may have disadvantages. Instead we could upweight later 4, or minimize out-of-sample error as in Gupta et al. (2024).
This estimator can reuse the M i.ddd samples already used to evaluate in a forward pass. This is easy to implement for any model Pe where sampling and evaluating the pdf is feasible, as we have assumed. We use automatic differentia- tion to compute V log.o.g (3))).
Other methods for decision-aware training. Our approach to direct BPR optimization here is example of decision- aware decision-focused training. In the taxonomy of Mandi et al. (2024), our approach is in the family of differ- entiable perturbed optimizers. Other work instead pursues surrogate losses. The SPO+ method (Flmachtoub & Gri- gas, 2022) finds a convex surrogate for the "smart predict then optimize" optimization problem. The perturbed gradi- ent (PG) method (Huang & Gupta, 2024) develops a more sophisticated surrogate, with theory and experiments sug- gesting utility even with misspecified models. In both cases, surrogate bounds make SGD-based learning tractable for a wide set of optimization tasks, including our knapsack-like BPR problem but also other tasks like shortest path finding.
A downside of this estimator is high variance. We miti- gate this with large M values, though future work may use control variates (Ranganath al., 2014; Mohamed et al., 2020) or try other estimators for gradients of discrete expec- tations (Maddison et al., 2017; Dimitriev & Zhou, 2021).
Gradient Vr,Le. For losses L defined in terms of TOP- KMASK binary vectors, like BPR, it is difficult to compute useful gradients because this loss is flat almost everywhere with respect to the input rankings r. To overcome this bar- rier, we leverage recent advances in perturbed optimization (Berthet et al., 2020), also referred to as stochastic smooth- ing (Abernethy et al., 2016). A recent computer vision method (Cordonnier et al., 2021) shows how these ideas en- able selecting top-K set of patches from a high-resolution image for downstream prediction. We adapt this top-K ap- proach to spatiotemporal forecasting for intervention.
Downsides of only fitting BPR. When models are misspec- ified, directly estimating @ to minimize JBPR should yield better BPR than some o' fit via conventional loss JN.IL, and better BPR means better decisions. However, probabilistic forecasts of near-future outcomes yt++1 produced by BPR- trained & have questionable utility. Nothing in the JBPR objective makes Ps accurately reconstruct even the train set yı.r: only relative ranking of sites matters. Even with
Concretely, Cordonnier et al. (2021) obtain tractable J- sample Monte Carlo estimates of both the top-K indica- tor vector b = TAPK.M K) and the Jacobian V,b needed for backpropagation. First, we draw J indepen- dent samples of a standard Gaussian noise vector of size S:
5

---

Decision-Aware Training of Spatiotemporal Forecasting Models
high-quality decisions, model which produces unlikely forecasts may be difficult to interpret or verify.
recommend setting M and J high as affordable. We found setting both to 100 worked for tasks in Sec. 5.
Proper selection of a, the standard deviation of the Gaussian noise in the perturbed estimator, is also vital. If too small, estimated gradients will be zero; too large will swamp out any data-driven signal for learning. In practice, we found that o values bit smaller than the largest elements of rating r' (Ø) worked well. Because our ratio estimator produces values between 0 and 1, we set between 10 3 and 10-!
4.3. Decision-aware maximum likelihood
To jointly achieve the goals of good top-K decisions and accurate forecasts across all sites, propose to find param- eters that solve constrained optimization problem:
argmin Σ log (y))) s.t. (( S 1402,
(12)
5. Experiments and Results
where (((0) = (. (.).
Decision aware training can benefit a variety of models in diverse problem domains. The subsections below cover toy and real applications. On each task, we fit models using all three training approaches from Sec. 4: ML estimation, BPR optimization, and DAML. We wish to verify common hypotheses throughout: (i) ML training can yield suboptimal top-K decisions; (ii) direct optimization of BPR improves this at the expense of likelihood; (iii) our DAML allow navigating tradeoff's between likelihood and BPR.
Here, is a desired lower bound on a tolerable BPR for the decision task. Function go (Ø) checks if the constraint is satisfied at time I, returning a non-positive value when BPR, 2 and a positive value otherwise. A feasible so- lution & must deliver BPR as good or better than on the provided training set. Practitioners can set e to achieve a de- sired minimum value for BPR. For example, if BPR below 60% was unworkable stakeholders, set 0.6 enforce 0.6 < BPR, recalling by definition BPR = -,
Common setup. In each task below, for fixed task-specific K we first train models for MI. and BPR. Using their final BPR values as guidelines, we select suitable range of values for DAML and fit each one, intending to explore in- termediate points on the two-ohjective Pareto frontier (Costa & Lourenço, 2015). When training each objective, we run many random initializations to convergence across a range of learning rates and other hyperparameters (see App. C.2 for details). We keep the model & from one run that best achieved its objective as needed. A model's ultimate top-K rankings can have some stochasticity. Thus, later figures show estimated dis- tributions of BPR across 1000 trials of a M=1000 sample Monte Carlo estimate of r r()) after training is completed.
We call this combined objective decision-aware ML estima- tion, or DAML. If the model is well-specified and training data are plentiful, DAML should deliver the same param- eters as ML estimation when E is low enough. However, when the model is misspecified and is higher than the BPR delivered by ML-estimated (, we argue DAML's constraint will produce better top-K decisions than MI. alone, trading lower likelihood for higher BPR. Compared to direct mini- mization of JBPR DAML can deliver similar BPR but more accurate forecasts of y for all sites. Additionally including the ML objective offers the ability to include an optional prior term log p()) to incorporate any prior knowledge.
a validation set, early stopping
To solve in practice, we use the penalty method (Chong & Zak, 2013) to convert to an unconstrained loss:
5.1. Synthetic Data
JDDML (Ø)= (b)=Etn À (d), 0)- log))( (13)
We begin with an illustrative synthetic case study chosen to highlight the trade-off's that decision-aware training allows a modeler to make when working with misspecified models.
This DAML formulation makes estimating & via gradient descent possible. Here, À > 0 is a nuisance hyperparameter that must be tuned. When the constraint is not satisfied, larger À values force gradient updates to move parameters further in directions that might satisfy the constraint. In practice, we set À such that both components of the loss are of similar magnitude during early training.
Task. We create a synthetic dataset of S 7 sites over T = 500 times. Integer data Yta is generated i.i.d over time from a quantized Gaussian with site-specific mean and small constant variance. The first six sites have means evenly spaced between 10 and 60; the last site has a large mean of 100. Fig. (top left) shows the training data.
Implementation. Pseudocode for DAML training is pro- vided in the supplement (Alg. A. There are several key hyperparameters. First, M and J are the number of Monte Carlo samples used during training to estimate gradients via the score function trick and perturbed optimizer method. Setting M and J larger produces lower variance estimates, but at the cost of runtime and memory. Consequently, we
Model. To show the benefits of our framework, we fo- cus a misspecified model. In particular, we model each site with I-component positive Gaussian mixture model, with global mean and variance parameters and site-specific
6

---

Decision-Aware Training of Spatiotemporal Forecasting Models
raw data
densiyy density densiyy dendry
0.12
0.08-
Training Frontier
0.04
-35
0.0
Opt NLL only
8-11
0.12
πI
67%
0.08
0,04
0.0
-55
DAML == 0.94
8:88
0.12
8
π
0.08
0.04
-60
Optimisaiion
0.12
0.0
π
Opt BPR only
8.05
8:05
In To tap
108%
8%
-10
D4P1, robi
0.08
B
108%
0.04
109%
-50
0.0
0.85
0.88
y-0
20 40
BP4
Figure 2: Synthetic ID data: learned models and Pareto frontier. Left Row 1: histograms of ye values by site (circled numbers). Sites 3-7 should be the top K=S under true model. Left Rows 2-4: Learned Gaussian components, with site-specific weights na marked as horizontal position between pure green and blue. Text provides ranking r with how often that site is in top K 5 over 200 trials. Right: Likelihood vs. BPR tradcoff frontier for final models delivered by different training objectives.
component frequencies:
5.2. Opioid-related Overdose Forecasting
Py())) Eter The No of)
(14)
Motivation. The ongoing opioid overdose epidemic in the United States has incurred over 500,000 deaths in the past decade, with more than 80,000 fatal opioid-related overdoses in 2023 alone (Ahmad et al., 2025). Possible evidence-basee interventions to mitigate overdose fatali- ties include overdose education and nalaxone distribution. Scarce resources require local decision makers to allocate these interventions to small areas that are high-risk (Allen et al., 2024), with co-incident education and support for proper follow-through, Public health agencies could use forecasting to help allocate limited resources towards the goal of harm reduction.
Here Ø = {HI:L, TL:L, }- We reparameterize to unconstrained real values to make gradient-based learning possible: see App. D for details.
With L=7 components, the model would be well-specified and could recover the true data- tatar-gene process. How- ever, we focus on misspecification, so we fit with L=2 components. This will hurt likelihood performance, as sites with distinct true means will need to use common means. However, we wish to show that solid top-K decision-making can still happen even with such severe misspecification.
Several efforts have developed small-area forecasting mod- els of opioid-related events (Marks et al., 2021b; Neill & Herlands, 2018; 8; Bauer et al., 2023). A preregistered trial for overdose reduction in Rhode Island (Marshall al., 2022) used BPR-like metrics to evaluate the top- K predictions of conventionally-trained models. This past work does not rank or train to improve BPR, as we do.
Experiment setup. We use BPR with K = 5 for this task. For reproducible details, including all hyperparameters, see App. D. We only report training set metrics here for simplic- ity: later tasks assess generalization to test data.
Results and analysis. From results in Fig. 2. we draw several conclusions. First, training to optimize NLL alone delivers subpar BPR for this task. Second, optimizing for BPR alone yields much better BPR values, suggesting that even this mispecified model can deliver much better top-K. K decision making than the off-the-shelf ML solution. How- ever, BPR alone yields nonsensical likelihood values, as nothing in the objective forces & to be good at modeling the outcomes y, only at relative ranking of the sites. In Fig. 2, we see how our DAML hybrid objective allows user to traverse the Pareto frontier of likelihood and BPR by enforcing desired threshold on minimum BPR. As the desired minimum BPR threshold e increases, we can sweep the tradeoff between likelihood and BPR. Ultimately, our DAML yields the best high- lige-iiiel, high-BP BPR solutions in the tight corner of the Pareto plot.
We study the capabilities of different methods on two datasets of historical opioid-related overdose mortality. Our IRB provided Not Human Subjects Research determi- nation for analysis of this decedent data. The first dataset, MA Fatal Overdoses, covers opioid-related overdose deaths in the state of Massachusetts from 2001-2022 This dataset is publicly available upon request from the MA Registry of Vital Records and Statistics. The second dataset, Cook County IL Fatal Overdoses, tracks appioid-involve overdose deaths in the greater Chicago area between 2015 and 2022. This is an open dataset obtained via the public website of the Cook County Medical Examiner Case Archive (Cook County, IL, 2014 2014-present.
Datasets.
7

---

Decision-Aware Training of Spatiotemporal Forecasting Models
MA Fatal Overdoses 2020-2021
Cook Countyy IL Fatal Overdoses 2021-2022
ANWR TX Cranes, 2009-2010
Test Log LiNeincdd
Trast Log Lhelthoo
Thsst Log LiNelnood
BAR Only
SFR Only
BPR Only
DANL 5070
JAML (000.0))
DANL 2:75
-20
JAML (5-0.5)
DOML 20w0
DAML (0-0.01
(ow0.751
(0-0.251
005
Only
-29
-10
Test
BPR
Test BPR
Figure 3: Pareto frontier of best possible reach (BPR, x-axis) and log likelihood (LL, y-axis) for real-world tasks, Higher is better on hoth axes. Each panel how the final models estimated by different training methods score the test of forecasting task defined in Sec. 5. capture the stochasticity of BPR due to our sampling-hasee -aased ranking estimator, for each model we show an estimated density for BPR over 1000 trials. Uncertainty in this plot only corresponds to uncertainty in BPR, log likelihood is a point estimate. In all three tasks, our proposed decision- aware ML (DAML) delivers better top- K decisions measured by BPR than ML estimation. DAML also delivers likelihood comparable ML methods and much better than directly optimizing BPR. In the Cranes dataset, the DAML objective surprisingly offers better BPR than the BPR-only objective, although the magnitude of this difference is small and perhaps due to the small-scale and sparsity of this dataset.
We follow previous evaluations of these datasets in Heuton et al. (2024). Each dataset was processed to common for- mat of fatal overdose counts per time and spatial unit. For the spatial units, we chose census tracts. Each tract by de- sign contains mean population of 4000 people, a scale that allows capturing variation in overdoses at the neighborhood level. For temporal binning, we picked calendar years to reflect the frequency at which health agencies might enact policy changes. Summary facts are in Tab. C.I.
ear function of feature vector Xxt as well as a site-specific intercept bos and site-specific dis weight on time t. The parameter g is probability of success shared by all sites.
Feature vector as includes tract s's overdose gravity, a recent average of opioid-related overdose deaths in tracts spatially near to s, as well as measures of sociological vul- nerability. See App. C.2.1 for details.
The overall parameters to estimate during training are So, 3, b0.1:s,b1,1:5) To make constrained param-
é = {q, eters amenable to gradient descent, we employ suitable one-to-one transforms to unconstrained spaces. Random effect weights b are regularized via prior that assumes a zero-mean Normal distribution with learnable covariance parameters. Full details are provided in App. E.
Task. Our forecasting task is to predict the next year's count of opioid-related fatal overdoses in each census tract. For S=120 tracts, we train on data from 2011-2018, 18, tune hyperparameters on validation data from 2019, and test on 2020 and 2021. For Cook County IL's S=1328 tracts, we train on data from 2015-2019, tune on 2020, and test on 2021 and 2022. These splits follow Heuton et al. (2024).
Competitor methods. We compare to two other decision- aware methods discussed above: Perturbed Gradient (PG) (Huang & Gupta, 2024) and SPO+ (Elmachtoub & Grigas, 2022), as implemented in PyEPO software (Tang & Khalil, 2024). To be fair, each uses the same model pa. the ratio estimator to rank sites, and the gradient of this estimator in Eq. (10). We conduct a hyperparameter search over learning rate (and perturbation noise for PG), selecting the model with the best loss value on validation data.
Given trained model, we assess heldout likelihood over all sites as well as BPR with where-to-intervene ranking. A 1=100 budget was selected to reflect realistic public health budgets, and is similar to values used in other studies (Marshall et al., 2022).
to measure the model's 's
Model. A recent benchmark (Heuton et al., 2024) compared many models designed for fatal overdose forecasting, includ- ing neural architectures with attention (Ertugrul et al., 2019) and more classical statistical models. A lop-performing model is negative binomial mixed- effects regression (Marks et al., 2021a). The generative model can be expressed as
Setup. We followed the common setup described above. For reproducible details specific to overdose tasks, see App. C.2.
Results. Results on test data for both MA and Cook County IL tasks are shown in Pareto frontier plots in Fig. 3. For both datasets, we see that ML training yields suboptimal top-K decisions. Direct optimization of BPR can improve BPR, though gains on test data over MI. vary (+0.04 on Cook County; less than +0,01 on MA). Direct BPR and surrogate
Yst ~ NegBin(JJst, 4);
(15)
log(Use) = 30 + 8"'x Xxt + bos + byst.
Here, count yoe is modeled as a Negative Binomial, where the log of the number of successes to stop at pst a lin-
8

---

Decision-Aware Training of Spatiotemporal Forecasting Models
loss (SPO+, PG) solutions can yield much worse likelihood, as expected. The surrogate loss methods provide lower quality decisions than the DAML and direct BPR approach. Our DAML approach provides better top-K K decisions than the MI. approach, with no visible decay in likelihood.
Results. From results in Fig. 3 (right panel), we find that optimizing NLL or BPR alone will yield poor results in the other metric. In this case, direct optimization of BPR results in the best top-K K decisions, but dramatically worse likelihood. Our DAML. models improve BPR by up to 0.05 over conventional ML training with only modest decay in likelihood. Our DAML and direct BPR objectives also de- liver better BPR than previous decision-award aware methods (PG, SPO+). This is even after providing smarter initializations to these methods, as we found their training had trouble im- proving on our common random initialization of o, perhaps due to this task's 'sss much sparser y values.
5.3. Endangered Bird Forecasting
Motivation. In the 1940s, whooping cranes were almost completely extinct in the U.S., with only 20 existing in the wild (Cannon, 1996). Thanks to efforts over the years to preserve their population, there are roughly 650 wild cranes today. This key species is still listed as endangered in 2025.
A major flock of cranes, known the Aransas- pop- ulation, winters at the Aransas National Wildlife Refuge (ANWR) along the Gulf Coast of Texas (Vartanian, 2023), while spending summers north in Canada. Ecologists wish to actively monitor this population. Regular acrial sur- veys of the Texas wintering region have been conducted for decades (Taylor et al., 2015). Using binned spatiotem- poral data of sighting counts over time from these surveys, we wish to offer data-driven forecasts of where cranes may be found. This could help conservationists decide where to send future human monitors or where to place K fixed- location cameras to efficiently track population health.
6. Discussion
We have addressed two open challenges related to top-K re- source allocation problems guided by the best possible reach (BPR) performance metric. We provided ranking strategy that can outperform simple per-site means. We posed a train- ing objective that strives for high likelihood across all sites while ensuring top-K decisions meet a bukeholder-speeifffee quality level. Our experimental evaluations suggest our ap- proach can better manage tradeoff's in likelihood and BPR than conventional training methods or previous decision- aware methods.
Data. We use raw data from Taylor et al. (2015), which digitized decades of aerial surveys of ANWR that marked individual crane sightings on paper maps. We processed this ANWR TX Cranes data into a common format of sighting counts over time and space, selecting bin sizes to support our goal of using the top-K sites to improve monitoring. Quality cameras or binoculars that could be used to track cranes can reasonably capture a 1.5 meter tall whooping crane in a 250-meter radius. Therefore, for spatial units, we divided the ANWR into 1338 boxes, each 500 meters per side. For temporal binning, we use 2 month periods. This choice catches seasonal variation, but avoids how finer scales might burden staff to move cameras too often.
There are several limitations to this study. We focused on showing the tradeoffs between likelihood and BPR for a fixed model family in each task without comparing a wide variety of possible models. Only one type of model misspec- ification is considered each in the synthetic and real-world experiments; perhaps decision-aware objectives are more less different to maximum likelihood estimates depending on the kind of model misspecification. Our evaluations use fixed K values and do not explore sensitivity to K or other hyperparameters. Practitioners may need multiple metrics to assess overall utility; focusing myopically on BPR may not always be wise. Additionally, our framework assumes any site with large outcome ys is better candidate for interven- tion. Future work could explore data-driven site selection that considers how site-specific attributes might make some interventions more or less effective.
Task. The experiment on these data was trained on years 2002-2006, validated on 2007-2008, and tested on 2009- 2010. In this context, "year" represents one wintering season; winter 2010 means the winter that began in October 2010 and ended in April 2011. We choose K = 50 for this task. This is an estimate of how many sites scientists might reasonably monitor every two months on a limited budget.
Looking forward, we hope to see applications of these ideas inform public health, wildlife conservation, and other where-to-intervene decision-making problems across the private and public sectors. We also hope that future method- ological work could scale up to much larger problems (S > 2000, T > 20) as we found that our current experi- ments tested the limits of commodity GPUs.
Models. The same negative binomial mixed-effect model was used as in Sec. 5.2. Features E.t for site s at time include historical bird counts at s, the latitude and longitude of the site's 's centroid, time of year, and the overall time.
Acknowledgments
Setup. We followed the common setup described above. See reproducible details in App. C.3.
Authors KH, FSM, and MCH are supported in part by the U.S. National Science Foundation (NSF) via grant IIS #
9

---

Decision-Aware Training of Spatiotemporal Forecasting Models
2338962. Author TJS was supported by the U.S. National Institute on Drug Abuse via grant # R01DA054267. Our team also acknowledges support from the American Public Health Association for data science demonstration project. We are thankful for computing infrastructure support pro- vided by Tufts University, with hardware funded in part by NSF award OAC CC 2018149.
C., Hallowell, B. D., Goedel, W. C., Basta, M., Yedinak, J. L.,, Li, Y., Cartus, A. R., Marshall, B. D. L.., Cerdá, M., Ahern, J., and Neill, D. B. PROVIDENT: Development and Validation of a Machine Learning Model to Predict Neighborhoood-level Overdose Risk in Rhode Island. Epidemiology, 35(2):232, March 2024. ISSN 1044-3983. URL https:
doi: 10. //journals 1ww. wwtp://ww///p/cc//c///..cco ftweeemmfulllt///ooeeo 2024 00224/C3000/provvddee___d_deede development and_ ind__aalidaaionnof_eeeemm
We are grateful to several anonymous reviewers, who pointed to related work in operations research literature and helped us connect our BPR metric to work on knapsack problems.
13. aspx.
Balkus, S. Proof #471: The median minimizes the mean absolute error. In The Book of Statistical Proofs, 2024. doi: 10.528//2emod0044359499 URL https: //statproofboodt git.ggtthhh.... io/P /med-mae.
Impact Statement
Recent work on where-to-intervenne decision making has emphasized the importance of incorporating additional con-
Bauer, C., Zhang, K., Li, W., Bernson, D., Dammann, O., LaRochelle, M. R., and Stopka, J. Small Area Forecast- ing of Opioid-Related Mortality: Bayesian Spatiotempo- ral Dynamic Modeling Approach. JMIR Public Health and Surveillance, 9(1 )ee4450, February 2023. doi: 10. 2196/41450, http:: / htpp/publchheeeee th. jmir. o org/2023/1/e41455
straints into the K site selection budget to meet the needs of all constituents. For example, Marshall et al. (2022) sought a balance of rural and urban sites for intervention in overdose mitigation efforts in the state of Rhode Island. We believe exciting methodological work could extend the ideas in this paper to handle such constraints.
Berger, J. O. Statistical decision theory and Bayesian anal- ysis. Springer Science & Business Media, 2013.
We also acknowledge that our decision-awars methods are currently technically more complex than existing off-the- shelf solutions, presenting barriers to adoption in many downstream applications where practitioners have limited hardware and limited expertise. We hope to work with stakeholders to build open-source packages that would allow DAML models to be created with off- thel- ease on new dataset of interest.
Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.-P., and Bach, F. Learning with Differentiable Perturbed Op- timizers. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Bertsimas, D. and Kallus, N. From Predictive to Prescrip- tive Analytics. Management Science, 6623):1025-1044, March 2020. ISSN 0025-1909, 3253.
doi: (0.1128//msc:201
References
Aberntthy, J., Lee, C.,, and Tewari, A. Perturbation Tech- niques in Online Learning and Optimization. In Hazan, T., Papandreou, G., and Tarlow, D. (eds.), Perturbations, Optimization, and Statistics, pp. 233-264. The MIT Press, 2016.
Cannon, J. R. Whooping crane recovery: case study in public and private cooperation in the conservation of endangered species. Conservation Biology, 10(3):813- 821, 1996.
CDC ATSDR. CDCIATSDR Social Vulnerabil- ity Index Massachusetts, 2022. URL https:
Ahmad, F,, Cisewski, J., Rossen, L., and Sut- ton, P. Provisional drug overdose death counts, 2025. URL https: // /ww.. edc. wwwgggoo/////ngoeeee parr/drug
/ atsdr. edc. gov/pl wtwp..go/opaacaanneaaaalt.. aceandhealthh htt.:/occcceaaaaioooodoooo...d.d html.
htm.
Cheng, T. and Wang, Integrated spatio-temporal data mining for forest fire prediction. Transactions in GIS, 12 (5), 2008
Allen, B., Neill, D. B., Schell, R. C., Ahern, J.. Hallow- ell, B. D., Krieger, M., Jent, V. A., Goedel, W. C., Car- tus, A. R., Yedinak, J. L., Pratty, C., Marshall, B. D. L., and Cerdá, M. Translating predictive analytics for pub- lic health practice: A case study of overdose preven- tion in Rhode Island. American Journal of Epidemi- ology, 2023. URL https:/ //o.: orgg10.1193) aje/kwadlls
Chong, F., K. P. and Zak, S. H. An Introduction Op- timization. John Wiley & Sons, January 2013. ISBN 978-1-118-27901-4
Chung, T.-H., Rostami, V., Bastani, H., and Bastani, O. Decision-Aware Learning for Optimizing Health Supply Chains, November 2022. URL http: : /////aaaa /aaaaaivvvo//// org/ abs 08507, arXiv:2211 100507 [cs].
Allen, B., Schell, R. C., Jent, V. A., Krieger, M., Pratty,
10

---

Decision-Aware Training of Spatiotemporal Forecasting Models
Cobb, A. D., Roberts, S. J., and Gal, Y. Loss-Calibrated Ap- proximate Inference in Bayesian Neural Networks, 2018. URL http: / aawii.iiiii.///ii///a.a/aaaaaiva org/ aaaaabss/1110 03901.
Heuton, K., Shrestha, S., Stopka, T. J., Pustz, Liu, L.-P. and Hughes, M. C. Predicting spatiotemporal counts of opioid-related fatal overdoses via zero-inflated gaussian processes. In The 2022 NeurlPS Workshop on Gaus- sian Processex, Spatiotemporal Modeling, and Decision- Making Systems., 2022.
Cook County, IL. Medical Examiner Case Archive, 2014-present. URL https::/datacatall/ cookcountyil. spv/Public-Safety madicaa--Hxe Deaexammnne--Cae-Arrhhh cjeg-bst bs86.
Heuton, K., Shrestha, S., Stopka, T., and Hughes, M. C. Learning where to intervene with differentiable top-k operator: Towards data-driven strategies to prevent fa- tal opioid overdoses. In ICML 3rd Workshop on Inter- pretable Machine Leurning in Healthcure (IMLH), 2023.
Cordonnier, J.-B., Mahendran, A., Dosovitskiy, A., Weis- senborn, D., Uszkoreit, J., and Unterthiner, T. Dif- ferentiable Patch Selection for Image Recognition. In IEEE Conf. on Computer Vision and Pattern Recognition (CVPR). arXiv, 2021. URL http:/ htap:///axxv.ox. abs/2104 03059.
Heuton, K., Kapoor, J., Shrestha, S., Stopka, T. J., and Hughes, M. C. Spatiotemporal forecasting of opioid- related fatal overdoses: Towards best practices for model- ing and evaluation. American Journal of Epidemiology. pp. kwae343, 2024.
Costa, N. R. and Lourenço, J. A. Exploring pareto frontiers in the response surface methodology. Transactions on Engineering Technologies: World Congress on Engineer- ing 2014, pp. 399-412. Springer, 2015.
Huang, M. and Gupta, V. Decision-Focused Learning with Directional Gradients, October 2024. URL http:// arxiv wwxiv.org/aba........o.........a wttp/bs/g/aaaaeeeeeeeeeeeeeeg 2402. 03256. arXiv:2402.03256 [cs].
Dantzig, G. B. Discrete- Variable Extremum Problems. Op- erations Research, 5(2):266-277, 1957. ISSN 0030- 364X.
Kao, Y.-h., Roy, B., and Yan, X. Directed Re- gression. In Advances in Neural Information Processing Systems, volume 22. Curran Asso- ciates, Inc., 2009. URL https://papers: nige::cc/apee...lle////rrr/////2 0c11511180024222222455333333333 html.
Dimitriev, A. and Zhou, M. Carms: Categorical- antithetic-reinforce multi-sample gradient estimator. In Advances in Neural Information Processing Systems, 2021. URL https::/proceeddgggs..
oc cc/paper hwtpppper_filee//eeeeeeeeeeeeeee wtt.ss//aaaeet//2///1eee
6e1 8e11655afeee1de77331117777fffff Paper. pdf.
Kleijnen, J. P. and Rubinstein, R. Y. Optimization and sensitivity analysis of computer simulation models by the score function method. European Journal of Operational Research, 88(3):413-427, 1996. ISSN 0377-2217. doi: http://2bbaa///..000665222111110 URI. http:s://ww..sceeeeeddieee ccciane//rticee//////77727700000
Elmachtoub, A. N. and Grigas, P. Smart "Predict, then Optimize". Management Science, 68(1):9-26, January 2022. ISSN 0025-1909. doi: 10. (0.11887/mss.2020.333
Ertugrul, A. M., Lin, Y.-R., and Taskaya-Temizel, T. CAST- Net: Community Atttnnitve Temppral Networks for Opioid Overdose Forecasting. In Machine Learning and Knowledge Discovery in Databases: European Con- ference (ECML. (KDD)) 2019. URL http: p /tt////xxii/ arxiv. org/ org/abs// abs/ 1905 04714, arXiv: 1905.04714.
Kušmierczyk, T., Sakaya, J., and Klami, A. Variational Bayesian Decision-makiene making for Continuous Utilities. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Golden, K. E., Hemingway, B. L., Frazier. A. E., Scholtz, R., Harrell, W., Davis, C. A., and Fuhlendorf, S. D. Spa- tial and temporal predictions of whooping crane (grus americana) habitat along the us gulf coast. Conservation Science and Practice, 4(6), 2022.
Lacoste-Julien, S., Huszár, F., and Ghahramani, Z. Ap- proximate inference for the loss-calibrated Bayesian. In Artificial Intelligence and Statistics, 2011. URL http:
/ //proceedings. mlr htt.://ww...............eeeemee. wwtps/sssssssllssttttttoo pdd...
Gupta, V., Huang, M., and Rusmevichientong, P. Decision- Aware Denoising, February 2024. URL https:// papers, ssrn. con/ abstract=4714305
Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Rep- resentations (ICLR), 2017. URL https: /arxiv. org/pdf/1611. 00722.
Hefley, T. J., Hooten, M. B., Hanks, E. M., Russell, R. E., and Walsh, D. Dynamic spatio-temporal models for spatial data. Spatial statistics, 20:206-2200 2017.

---

Decision-Aware Training of Spatiotemporal Forecasting Models
Mandi, J., Kotary, J., Berden, S., Mulamba, M., Bucarey, V., Guns, T., and Fioretto, F. 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.1532) URL htwph: jwi.pp::///w..ai...... org/ index.
Sadana, U.. Chenreddy, A., Delage, E., Forel, A., Frejinger, E., and Vidal, T. A survey of con- textual optimization methods for decision-making under uncertainty. European Journal of Oper- ational Research, 320(2)2211-29,, January 2025. ISSN 0377-2217. doi: 10.1016/f.jjo:22244.3. 020. URL htt.:/tpss/ww..cceeeeeeeeeeeeee com/ ccteee//watiaee///////7777222222
Marks, C., Abramovitz, D., Donnelly, C. A., Carrasco- Escobar, G., Hemáádez- R., Ciccarone, D., Gonzalez-Izquierdo, A., Martin, N. K., Strathdee, S. A., Smith, D. M., and Bórquez, A. Identifying counties at risk of high overdose mortality burden during the emerging fentanyl epidemic in the USA: predictive statistical modelling study. The Lancer Public Health, 6(10)::700=<728, October 2021a. ISSN 2468-2667. doi: httpp:://ww.tthelareet
Sander, M. E., Puigcerver, J., Djolonga, J., Peyré, G., and Blondel, M. Fast, differentiable and sparse top-k: a convex analysis perspective. In International Conference on Machine Learning, pp. 19919-299366 PMLR, 2023.
Schwertman, N. C., Gilks, A. J., and Cameron, J. A simple noncalculus proof that the median minimizes the sum the absolute deviations. The American Statistician, (1), 1990. doi: 10475690.
URL Larputa/aartielle/r11s4--------- (21) 00080-3/fulltexe Publisher: Elsevier.
Tang, B. and Khalil, E. B. Pyepo: a pytorch- based end- to-end
Marks, C., Carrasco-Escobal G., Harnasco- R., Johnson, D., Ciccarone, D., Strathdee, S. A., Smith, D., and Bórquez, A. Methodological approaches for the prediction of opioid use-related epidemics in the United States: A narrative review and cross-disciplinary call to action. Translational research the jour- nal of laboratory and clinical medicine, 234:88-113, 2021b. URL https::tt///ooc TtRLhhts::///wtt //www http://wwww..i.ii.cc//./......tt ncbi n mlm. nih. gov/ pmc/ pec/articn wttttiiccceee/MMM22117
library for linear and inte- ger programming. Mathematical Programming Com- putation, July 2024. ISSN 1867-2957. doi: 10.1007/ s $12532-024-00255-x
Taylor, L. N., Ketzler, L. P., D., R., Strobel, B. N., Met- zger, I.., and Butler, M. J. Observations of whooping cranes during winter aerial surveys: 1950-2011. 1. Techni- cal report, Aransas National Wildlife Refuge, U.S. Fish and Wildlife Service, Austwell, Texas, USA, 2015. URL http: /dxx/xxx/////////// doi orr.orrgg///.oo 10. 700477944//44
Marshall, B. D. L., Alexander-Scott, N., Yedinak, J. L., Hal- lowell, B. D., Goedel, W. C., Allen, B., Schell, R. C.,, Li, Y., Krieger, M. S., Pratty, C., Ahern, J., Neill, B., and Cerdá, M. Preventing Overdose Using Information and Data from the Environment (PROVIDENT): Protocol for randomized, population-based, community intervention trial. Addiction (Abingdon, England), 1 117(4): 1162, 2022. URL ht http:///wwww//tbc ncbi n m. nim.iihhh.. nov. pmc/ http:/arrcccee//c//0///
Vartanian, J. Report on whooping crane recovery ac- tivities. 2023. URL https:// fws.
Heport_Whoopinge htt.://wwooppp......eeeeeeeeeee. Aransas.pdf. Wei, Y., Sheth, R., and Khardon, R. Direct loss minimiza- tion for sparse gaussian processes. In International Con- ference on Artificial Intelligence and Statistics. PMLR, 2021.
Mohamed, S., Rosca, M., Figurnov, M.. and Mnih, A. Monte carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132), 2020.
Williams, R. J. Simple statistical folloingg algo- rithms for connectionist reinforcement learning. Machine learning, 8, 1992.
Murphy, K. S. Probabilistic Machine Learning: An Intro- duction, chapter 5.1: Bayesian Decision Theory. MIT Press, 2022.
Neill. D. B. and Herlands, W. Machine Learning for Drug Overdose Surveillance. Journal of Technology in Hu- man Services, 36(1):8-14, 2018. URL https: //dpi..//d.... org/10. 10w.:000000000 1ww11100111111113 2017 1416511.
Ranganath, R., Gerrish, S., and Blei, D. Black box varia- tional inference. In Artificial intelligence and statistics, pp. 822. PMLR, 2014.
12

---

Decision-Aware Training of Spatiotemporal Forecasting Models
Appendix
This appendix includes additional experimental results and information for understanding and reproducing experiments. Reproducible code used for experiments included an open-source repository:
com.tuufts https: https://github. wtw/tuftttteeeeeoo cesftttsmmmmeciiiioooiii aware topk/
Contents
Pseudocode for Training
14
Ranking Demo for BPR
14
B.1 Justification for ratio estimator
14
B.2 Demo experiment
15
Experimental Details and Results from Real-World World Data
16
C.I Results on alternative models and metrics
16
C.2 Opioid-related Overdose Forecasting Results
16
C,3 Endangered Bird Forecasting Results
17
Experimental Details and Results from Synthetic Data
17
Model for Negative Binomial
18
13

---

Decision-Aware Training of Spatiotemporal Forecasting Models
A. Pseudocode for Training
Algorithm A.1 Decision-aware ML training
Input:
(yr)em train data, each E I 20
Φ E RR:; parameter vector for model
M: int num MC samples for score func estimator
J: int num MC samples for stochastic smoothing estimator
* > 0: float stddev of stochastic smoothing estimator
c E 10, Desired minimum BPR value. Will try to enforce constraint BPR 2 c.
À > 0: Strength multiplier when constraint is violated.
Output: Trained model parameter Φ
Procedure:
1: while not converged do
2: VoJ +0
I| P X vector to store grad wrt params
3: for time t E (1,2...2 .
do
4: (y M
m=1 ~
I/ Monte Carlo (MC) samples
5:
Fs Σ m 1.
I| S ranking vector via MC
6:
7:
b, +- TOPKMASK(T:)
8:
LAPR t- y1 MU-SSK) (Ye b.)
// Scalar loss, -BPR
9:
g. BPR +- E + LBPR
// Scalar. Negative if BPR 2 is satisfied.
J.HPR + max(1g BPR 0)
// Scalar ultimate loss for BPR
11:
12:
V Tr M Em V log log V.rh 1
11 P S matrix, score func. est. of V
13:
(2, }ー ~ Is)
14:
Vr be < 1 Σj= OUTER OUTERUTTTRNNNNO (TOPKMASK(T; + o2;),
II S x 5 matrix, perturbed estimate of V I| S X vector, nonzero if constraint unsatisfied.
15:
V J.APR +- Ui TOPKMASKE))) Vr 1 0]
16:
V J.PRR + (V bi) (BP.)
|/ P X 1 vector
17:
18:
JALL +- - log (v:))
|/ Scalar ultimate loss for NLL
JoJ.NI NLL V log (B:)
20:
21:
VoJ T NLL J.B.
22: end for
23: ④ + GRADDESCENTUPPATEIOO VOJ)
// Use steepest descent or Adam or
24:
end while
25:
return
B. Ranking Demo for BPR
B.1. Justification for ratio estimator
In the text, we propose the ratio estimator for ranking locations: = 24, and justify its use by demonstrating that it is a better bound on our decision loss BPR than the mean estimator. However, a natural question is, why not use a ranking estimator which more closely resembles BPR, such as r = El TopKMask(y.K) yy. In practice, this estimator is equivalent to the ratio estimator in expectation, and will select the same top-K locations. Accordingly, we use the ratio estimator for its improved computational speed and simplicity.
14

---

Decision-Aware Training of Spatiotemporal Forecasting Models
B.2. Demo experiment
To gain understanding about the problem of ranking optimize the fraction of best possible reach (BPR) performance metric, here present detailed analysis of toy problem with S 9 sites. We'll Wllll assume the true data-ge gatterenng process for each site is completely known throughout and that each site can be modeled independently of other sites.
9
p
II p(v))
(16)
For each site, we select from 3 possible site-specific model archetypes, nicknamed type A, type B, and type C.
sites #1, #2, and #3 each iid with type A PMF
sites #4, #5, and #6 are each iid with type B PMF
sites #7, #8, and #9 are each iid with type C PMF
Each type's PMF function over the non-negative integers defined in the table below.
Type A
Type B
Type C
0.0
if ys = 0
0.35
y, 0
0.90
f. =
p(y))=
1.0
if y/s
p(y=)
0.65
if y/ 10
p(v.)
0.10
if y, S0
7.0 0.0
otherwise
0.0
otherwise
0.0
otherwise
Mean
6.5
8.0
Median
7.0
10.0
0.0
have now defined the joint PMF p(y) over the sites.
We can compare two possible ways to compute numerical ranking of the S sites:
Mean estimator, which computes the per-site mean: r EEyy
Ratio estimator, which computes the expectation of normalized by its sum: EE 1.y
For each possible ranking strategy, we repeated BPR calculations over 10000 trials. To ensure accuracy of Monte Carlo estimates, we average over M 50000 samples to estimate the expectation defining each r.
Results are provided in the table below. We have two key findings. First, our proposed ratio estimator can select very different sites than the pir-site mean, even when both estimators have access to the true data- aata--gnerating model. Using K 3, our estimator would select all type-A sites as the top 3; contrast the per-site mean would select all type-C sites (#7-9). Second, this can produce very different BPR values. Even in this simple example, BPR of over 0.4 between the different rankings at K = 1 and over 0.39 at K 3, which is huge shift for metric that is bounded between 0.0 and 1.0.
absolute difference in
BPR
fraction of trials each site in top K-3 of r
1=1
K=3
K=6
A:#1
A:#2
A:#3
B:#4
B:#5
B:#6
C##7
C##8
C:#9
mean r y
0.107
0.231
0.636
0.0
0.0
0.0
0.0
0.0
0.0
1.0
1.0
1.0
ratio
r 1
0.538
0.625
0.810
1.0
1.0
1.0
0.0
0.0
0.0
0.0
0.0
0.0
These results can be replicated via scripts provided in the code repository.
15

---

Decision-Aware Training of Spatiotemporal Forecasting Models
# Spatial Sites
Temporal Scale
Outcomes
Features
Count of opioid-related fatal overdoses
1620 Census Tracts
20 years, 2001-2021
SVI, Past Deaths, Location, Time
MA Fatal Overdoses
Count of opioid-related
Cook County IL Fatal Overdoses
1328 Census Tracts
years,
SVI, Past Deaths, Location, Time
2015-2022
fatal overdoses
years, 1952-2011 (last 10
Count of bird spotted in aerial survey
Past observations, Location, Time, Month
Aransas TX Whooping Cranes
1338 boxes, each
500mx 500m
years used)
Table C.I: Comparison of Real datascts, in terms of number of spatial sites S, temporal scales, outcomes and features.
ANWR TX Cranes
MAE RMSE
BPR-50
Cook County IL Overdose
MAE
RMSE BPR-100
Last timestep
0.24
1.04
0.27
Last timestep
1.07
1.66
0.76
Avg over 10
0.25
0.78
0.38
Avg over 5
0.99
1.58
0.80
Chance decision, 1 = 0
0.15
0.79
0.05
Chance decision, y =
1.37
2.46
0.20
EpiGNN
0.38
0.70
0,06
EpiGNN
1.33
2.03
0.32
PG
1.03
4.85
0.35
PG
5.65
17.03
0.80
SPO+
0.15
0.83
0.18
SPO+
1.25
2.38
0.74
NLL Only
0.35
1.14
0.38
NLL Only
1.03
1.58
0.78
DAML
9.08
35.89
0.39
DAML
1.12
1.84
0.80
BPR Only
40495
40495
0.41
BPR Only
70.31
152.72
0.82
(a) ANWR TX Cranes 2009-2010.
(b) Cook County IL. 2021-2022.
Table C.2: Performance comparison between different model families for error-based and decision metries. The top rows represent simple historical baselines: using either previous timestep alone. or an average over a larger amount (10 bi-months for crane dataset, and years Cook County), Next, the Chance decision model chooses spatial locations by random chance, and uses all 0's for predicted in MAE and RMSE calculations. Despite the simplicity, this is a competitive model on the error- tror-bhsed metrics due to sparsity. It outperforms all others MAE the Cranes dataset. Next is EpiGNN, spatiotemporal forecasting model based graph neural networks, trained to minimize RMSE. Although this model has the best performing RMSE the crane dataset, makes poor top-K decisions both. Finally, the three rows show different ways to train the negative binomial mixed effects regression model 14. NLL Only trains to maximize likelihood alone (8), BPR Only is our direct-loss minimization trained to maximize BPR alone (9), and DAML (13) is our hybrid loss that allows user to explore the pareto frontier between likelihood and BPR-K K.
C. Experimental Details and Results from Real- World Data
C.1. Results on alternative models and metrics
C.2. Opioid-related Overdose Forecasting Results
C.2.1. FEATURES
The feature vector Sst for site and time f includes:
Latitude and Longitude of the centroid of the census tract
The current timestep
Overdose gravity, an weighted average of the prior year's overdose deaths in all contiguous tracts. We construct the feature in the same way (Marks et al., 2021a). Note that the original paper describes weighted average over all regions within radius, but the provided code uses only immediately contiguous locations. We follow the implementation from the code.
Covariates from the Social Vulnerability Index (SVI) (CDC ATSDR, 2022). These covariates are available as 5-year estimates for every census trast in the United States and are updated every 2 years. The SVI measures report the
16

---

Decision-Aware Training of Spatiotemporal Forecasting Models
percentile ranking of every census tract according themes: Socioeconomic, Household Composition & Disability, Minority Status & Language, and Housing Type & Transportation. We use 5 variables: each tract's ranking in each of the four themes as well as its composite ranking.
We include 5 temporal lags: the number of fatal opioid overdoses in tract in each of the past 5 years.
C.2.2. HYPERPARAMETER RANGES
Hyperparameter ranges explored include:
Perturbation noise: 0.1, 0.01, and 0.001.
Adam step size: 0.1, 0.01, and 0.001.
Number of samples for score function trick estimator: 100
Number of samples for perturbation estimator: 100
BPR constraint e: 5 possible values of the penalty threshold: 1.0, for a threshold that always encourages better BPR, as well values selected be around the best BPR obtained on the training data.
Multiplier on the penalty for DAML: 30. This value was chosen so that the BPR and likelihood components were roughly the same magnitude after 100 epochs of training.
For each location, training objective (likelihood, direct loss minimization, and DAML), as well as for each threshold for the hybrid model, we selected the model based on validation dataset performance. For the maximum likelihood model and direct BPR models we picked the model that best maximized their respective objective. For the DAML models, we found that given the larger number of hyperparameter configurations, small amount of validation data, and BPR's sensitivty, selected model based on BPR lead to overfitting. To ameliorate this, we chose the model with the highest likelihood provided that the BPR was greater than given threshold. Because models failed meet the target c, chose the BPR of the maximum likelihood model on the validation dataset the maximum BPR model. For the We did this by evaluating the validation performance every 10 epochs, and saving checkpoint for the model with the lowest loss on validation data.
threshold. If no models met this threshold, selected
C.3. Endangered Bird Forecasting Results
1. FEATURES
The feature vector v., for site and time t includes the following variables: the past 5 count values at site latitude & longitude of site s's centroid, an enumerated timestep indicating time passed at since the starting timestamp of the dataset, and a monthly indicator variable. The dataset includes three bimonthly periods per wintering season: Oct 20-Dec 24, Dec 25-Feb 27, and Feb 28 28-Apr 30. The monthly indicator received value of I, 2, or 3, respectively, depending on which set of months an observation spanned.
C.3.2. HYPERPARAMETER RANGES
The same hyperparameter and model selection was performed Whooping Crane data as the opioid experiment in Section C.2.
D. Experimental Details and Results from Synthetic Data
Model details. As explained in the main paper, we use a mixture of Gaussian distributions where each is truncated the positive reals. We give each of the S locations their mixture weights T., where Σ T.. 1 and म. 0. Our model for individual location is then:
plya)= Σ naf No (ysluu,
(17)
17

---

Decision-Aware Training of Spatiotemporal Forecasting Models
Our parameter vector then consists of the of all Pt of and T..-- We have that both j4 and of should be positive, as are modeling positive counts and standard deviation is defined positive. To accomplish this, we transform these variables using the sofiplus function when performing gradient based learning. Additionally, we constrain or 2 0.2 to avoid degeneracy. Finally, to make a valid pdf, we have that the mixture weights must sum to one: Σ T..i = 1. To enforce this, we transform unconstrained variables using the softmax transform. This creates an subtle issue with model identifiability, as there are many unconstrained values that will lead to similar weights, but we do not find that this impacts our ability to train models effectively to achieve good likelihood or good BPR with the appropriate objective.
Training Procedures. We seek train 7 models: one that optimizes only for model log likelihood, one that optimizes only BPR-5, and 5 models penalized to achieve certain threshold values of BPR. Each model is given 20 random initializations of the model parameters. We use learning step-size of 1. For models with BPR in the objective, we try a perturbation noise of both 0.01 and 0.05, with 500 samples. In the hybrid models, we use À = 30, selected so that the likelihood term and the BPR penalty term are on the same order of magnitude after several epochs of training.
Tradeoffs between likelihood and BPR when 1.=2. If this model had I. = 7 components, would be well-specified and recover the true data generating process. However, with only 2 components, it is forced to group locations with distant mean values, which will come cost to likelihood and predictive capability as assessed by BPR.
For this example, we will consider BPR-5 as our decision making metric. A model that correctly ranks the top-5 locations will achieve perfect BPR. Our misspecified model is capable of this by learning 2 distinct mean values (4k. one higher than the other. As long as the mixture weights for the top-5 locations T. assign all probability to the high component, and the mixture weights for the bottom 2 locations assign all probability to the low component, the model will have perfect BPR.
However, this is not what the model with the best possible likelihood looks like. To maximize likelihood, a model will assign the 6 low locations to one component, and the one high location to another, as in Fig. 2.
Our hybrid objective DAML can explore the Pareto frontier between maximizing for likelihood and BPR-5. By including more locations into the high-valued component, BPR-5 will increase as log likelihood slightly decreases. Our hybrid objective formulation allows us to control this tradeoff by specifying the threshold at which the penalty term takes effect.
Results. Results are shown in Figure in the main paper. Here we see the ability of the hybrid Decision-awand object to traverse the Pareto frontier between the best possible likelihood and BPR. The lowest threshold of 0.5 is trivially satisfied by the maximum likelihood model. The highest threshold of 1.0 is only satisfied by perfect BPR, while the 3 intermediate thresholds were chosen to explore the solution frontier that is possible by including excluding a particular component from one of the learned mixtures. We see that by using the decision-aware training objective, we can maximize BPR while still obtaining a highly likely model.
E. Model for Negative Binomial Mixed-Effects
Here we provide further detail on the model described in Eq. (15).
For some elements of the parameters Φ, we have a prior that informs point estimation in MAP fashion. The random effects are assumed to follow multivariate normal prior
bos b1o)
N (()2z)
(18)
where the covariance matrix Σ is parameterized as:
Σ
of poo01
panar of
(19)
Here, dg is the standard deviation of the random intercepts, 01 is the standard deviation of the random slopes, and p is the correlation between the random intercepts and slopes, where -1 S p S 1. We pack these hyperparameters into separate vector 7 = (oo, 01,P). p).
Transforming to unconstrained parameters. To ensure parameter constraints are satisfied throughout gradient descent optimization, we employ invertible transformations that map constrained domains to unconstrained spaces. We reparam- eterize the correlation coefficient p E (-1,1) using u +- arctanh(p), allowing u to be optimized over IK while ensuring
18

---

Decision-Aware Training of Spatiotemporal Forecasting Models
la)h(u) remains within its valid bounds. For the strictly positive parameters oü 01 > 0, optimize unconstrained parameters Ea-E1 E TR and apply the softplus transformation d; +- In (65), which guarantees positivity while maintaining smoothness. Similarly, for the probability parameter (0, 1), we optimize unconstrained parameter Ç E IR and apply the sigmoid transformation q - 1/(1 + t During gradient descent, we optimize these unconstrained parameters, and transformed to the constrained version during forward sampling or pdf evaluation the model.
MAP estimation. Together, the parameters & and hyperparameters 7) comprise this model. Both are point estimated to maximize the MAP objective. Thus, what is marked in the paper as NLL optimization is really best viewed MAP or penalized NLL estimation, where the loss is
J(o,n) n)) logp(y1:10)) log p(o/n)
(20)
Similarly, when fit this model with DAML, the above loss what minimized subject to the BPR constraint.
19


Reviews:
[{"review_id": "M7IJ3KQuLz", "reviewer": "Reviewer_7cWK", "summary": "The paper tackles two main issues related with a metric called Best Possible Reach (BPR): (1) the ranking problem, basically how to rank sites numerically to select the top K for intervention, based on a probabilistic method, to solve this, the paper works on a tighter bound on BPR and utilizes the ratio estimator. (2). the training problem, basically how to optimize the model's parameter to maximize BPR performance. The paper tackles the difficult problem of training models to directly optimize the BPR metric, which involves a discrete top-K selection that leads to zero gradients. It uses the perturbed optimizers and the score function trick to estimate gradients and enable training. Experiments are done on synthetic data and two real world datasets: opioid overdose mitigation and endangered bird monitoring.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful feedback. We are glad that they thought our DAML method “provides a nice way to balance the goals of achieving high BPR for decision-making and maintaining good likelihood for overall forecast accuracy.”\n> the approach is evaluated only on …Gaussian mixtures and negative binomial mixed effects, which might seem too simple for realistic use cases? Are these still the most common used models?\n\nFor our real applications, we selected the negative binomial mixed effects regression with spatially-lagged features (Sec. 5.2), specifically because it had competitive heldout performance in a recent evaluation of methods for opioid-overdose forecasting published in 2024 by Heuton et al. [A]. In that study, the negative binomial model with conventional training (not the decision-aware training in our submission) tied or beat an attention-based neural network as well as a boosted ensembles of trees model and a Gaussian process model in terms of test-set BPR on two different datasets.\n\nThese spatiotemporal forecasting tasks are overall quite difficult; there is not often a very strong signal for predicting observed counts of overdose deaths or bird sightings given the limited available features. In such cases, somewhat simpler models tend to work well, even when substantial effort is put into tuning hyperparameters to avoid overfitting with complex tree ensembles or neural nets.\n\n**New Results with GNNs**. We have also added a [New Table](https://anonymous.4open.science/r/DecisionAwareMaximumLikelihood/NewTable.pdf) as a standalone PDF in our anonymous code repo. This compares our methods to a recent graph neural networks for spatiotemporal forecasting (from Xie et al. ‘24 [B]) on the Crane and Cook County tasks, in terms of BPR, RMSE, and MAE. This GNN is trained to minimize squared error, and thus can do best on RMSE but is inferior to our DAML for top-K decisions as measured by BPR.\n\nWe do hope to investigate a broader class of probabilistic models directly trained by our DAML in future work.\n\n[A] Heuton et al. Spatiotemporal forecasting of opioid-related fatal overdoses. Amer. J. of Epidemiology, 2024.\n\n[B] Xie et al. EpiGNN: Exploring Spatial Transmission with Graph Neural Network for Regional Epidemic Forecasting. ECML PKDD ‘22. \n\n> concerned about the scalability of the proposed method, it seems evaluated on a small scale of parameters T, S, and K. . Could the author comment more on the possible approximate algorithms \n\nWe agree that scalability is an important practical consideration. Our current code can scale to number of timesteps T in the dozens, budget K in the hundreds, number of sites S in the thousands, and number of parameters P in the thousands. We have found this sufficient for the practical problems in our paper.\n\nOur algorithm can further be scaled up to larger problems by processing one time record (out of T) at a time, and also by computing gradients with respect to one parameter (out of P) at a time. We have prototyped code to do this already in the last few weeks.\n\nWith further implementation effort, our methods could be easily adapted to process minibatches of timesteps, to scale even further.\n\n> The approach comes with a pre-defined K. Typically how people find this value? … does the method perform robustly across a wide range?\n\nTypically, the practical intervention budget of the stakeholders determines K. For example, in the whooping crane monitoring application shown in Fig. 1, if the monitoring agency can only afford 10 new cameras, then we would set K=10. \n\nWe will try to do some further experiments on different K values for one of the applied datasets in the next few weeks. If accepted, we promise to include such experiments in the appendix for the camera-ready deadline.\n\n> Both sigma and lambda are key hyper-parameters needs to be tuned. Could the authors provide some rule of thumb? \n\nWe will revise the paper to provide this information. To complete our experiments in Fig. 3, we select sigma from a grid of values between 0.001 and 0.1 based on validation-set loss. We observed the experiments to be insensitive to lambda, and it was fixed at 30 to make the BPR and likelihood components of the loss similar in magnitude during early training.\n\n> the synthetic experiment uses a relatively simple form of misspecification. It would be interesting to see how methods perform under more severe or diverse types of misspecification.\n\nWe agree that exploring more types of misspecification would be interesting. However, with limited time in this rebuttal period, we elect to leave this to future work. We will revise to acknowledge this limitation in our Discussion section.", "claims_and_evidence": "Some claims made by the paper:\n- Ranking via the per-site mean is suboptimal for BPR. Section 3.1 shows the per-site mean minimize expected loss on a simplistic upper bound on BPR. \n- The effectiveness of the proposed ratio estimator. Intuitively, it works under the sparse vector setting, and does provide a tighter upper bound. This is also empirically proved in Fig 2.\n- DAML helps navigating the pareto frontier, both experiments on synthetic data and real-world datasets consistently verify this empirically.", "methods_and_evaluation_criteria": "- In the synthetic experiments, the approach is evaluated only on the very specific models: Gaussian mixtures and negative binomial mixed effects, which might seem too simple for realistic uses cases? Are these still the most common used models? Can the authors comment more on the model's performance on a more complex model families?\n- I am concerned about the scalability of the proposed method, it seems evaluated on a small scale of parameters $T, S$ and $K$. Could the author comment more on the possible approximate algorithms people might be able to use when face more real-world large scale problems?\n- The approach comes with a pre-defined $K$. Typically how people find this $K$ value? How does the method perform robustly across a wide range of $K$'s? Can we add more empirical results on this?\n - Both the hyper-parameter $\\sigma$ and $\\lambda$ are key hyper-parameters needs to be tuned. Could the authors provide some rule of thumb (or automatic selection methods) in selecting them under different use cases?\n- Though the paper explicitly discusses the benefits of their approach under model misspecification, the synthetic experiment uses a relatively simple form of misspecification. It would be interesting to see how the methods perform under more severe or diverse types of misspecification.", "theoretical_claims": "I roughly looked at the ones in the main paper, and does not identify any obvious issue.", "experimental_designs_or_analyses": "See Methods And Evaluation Criteria.", "relation_to_broader_scientific_literature": "Might be related with spatial-temporal forecasting topics; method-wise, might related with BPR, direct loss minimization, and score function estimator.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strength:\n- The paper is well-written and solves some key problems with BPR metrics, in terms of ranking and optimization.\n- The authors propose a \"ratio estimator\" for ranking sites based on a tighter bound of the BPR metric. It moves beyond the suboptimal per-site mean ranking and is theoretically justified. Empirically, this method significantly outperforms the per-site mean.\n- The proposal of DAML provides a nice way to balance the goals of achieving high BPR for decision-making and maintaining good likelihood for overall forecast accuracy. \n\nWeakness:\n- The scalability of the method.\n- The effectiveness of the method moving beyond the two families of the models examined in the paper.\n- The robustness regarding to the pre-determined K value.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "See Methods And Evaluation Criteria.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vg13uxBeQD", "reviewer": "Reviewer_ZTv9", "summary": "The paper proposes an approach for learning a potentially misspecified spatialtemporal probabilistic model for decision-making settings. They specifically focus on the decision problem optimizing best possible reach (BPR) which closely corresponds to the ranking top-K items problem. Their approach consists of first proposing a metric similar to BPR in order to rank the items. They then show how to compute the metric given a probabilistic model. Given the metrics they then show how to compute the BPR loss and how to optimize the loss relative to the probabilistic model. Putting it all together, they propose a loss that blends decision loss (BPR loss) with traditional ML loss (negative log likelihood). This approach allows users to trade off interpretability (ML loss) with decision performance (BPR loss) via a tuning parameter. \n Finally, they test their approach on three real-world datasets to highlight the benefits of their approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank ZTv9 for their thoughtful review, especially in introducing related work and PyEPO. We offer brief replies to key points below. We will revise to address all points raised.\n\n## Essential References\n\nThanks for several useful references. Please see “Essential References (Common Issue)” in Response to Reviewer His4 for our revision plan. \n## W1 & Q2: Does decision loss optimize for time T+1?\n> Optimizing over average historic BPR instead of the BPR of time period T+1 may reduce the decision quality\n\nEmpirically, we find decisions at heldout times are good (Fig 3). Revision will discuss justification below, and mention out-of-sample bounds by Gupta et al (SSRN ‘24). \n\nWe assume the true data-generating distribution of outcomes at time t given the recent past, $p(y_t | y_{t-W:t-1})$, is unchanged across train and test time periods. We will make this explicit in revision. Many forecasting methods use this assumption and use average loss to extrapolate to the future. We could weight recent timesteps more.\n\nAs further protection, we already use a “leave future out” experimental design: e.g. for MA overdose task, we train on 2011-18, tune/validate on 2019, and test on 2020-21.\n\n## W2 & Q1: 0-1 Knapsack for BPR / PyEPO baseline\n\nThanks for suggesting the 0-1 knapsack formulation. The problem of “how to rank” for BPR given outcome y can reduce to 0-1 knapsack where each site has weight 1 and the budget constraint allows K of S sites.\n\nAs suggested, we tried both SPO+ loss and PG loss from PyEPO on the 3 datasets in our Fig. 3. Both losses are direct competitors to our BPR-only (Sec 4.2), as they try only to improve top-K decisions and do not account for likelihood. \n\nResults are in [revised Fig 3](https://anonymous.4open.science/r/DecisionAwareMaximumLikelihood/Figure3Update.pdf) in our anonymous code repo. \n\nWe find SPO+ and PG do not advance any panel’s Pareto frontier compared to our methods. In terms of top-K decisions, each method can sensibly beat NLL-only but not our DAML or BPR-only. PG loss gets higher BPR than NLL-only on Cook, SPO+ does the same for MA. However, these objectives ignore likelihood, so they naturally produce models with low likelihood. \n\n## W3: Novelty\n\nOur DAML approach in Sec 4.3 is not wholly new in combining classic objectives with decision-aware losses. We will revise to cite Kao et al. (NeurIPS 2009)’s convex-combination of these two loss types, albeit for directed regression rather than our top-K where-to-intervene tasks. Compared to that work, DAML is distinct in its *constrained* approach: as soon as the desired BPR is achieved, training focuses only on likelihood. \n\nAnother novel aspect of our work is the ratio estimator for the how to rank problem (Sec 3). While previous works used BPR as a metric, they defaulted to the per-site mean, unaware that it may be suboptimal (see App. B for demo of suboptimality).\n\nA final novel aspect is our analysis of 3 real where-to-intervene datasets, covering whooping-crane and opioid-overdose forecasts, each with 1000+ sites. These tasks are not yet in any decision-aware literature to our knowledge, and use models much bigger than in Gupta et al. (SSRN ‘24)\n\n## Q3: Why does using only BPR loss provide worse Test BPR for the MA overdose forecasting? \n\nYes, the left panel of Fig 3 shows a different relative ranking of BPR-only vs. our DAML method. Optimizing BPR is prone to local optima and difficult loss landscapes. The extra NLL loss in DAML appears to find better BPR solutions sometimes.\n\n## Q4: Can you control BPR with epsilon?\n\nYes, see Pareto frontier in our Fig 2. On this toy task, as we raise epsilon from low to 0.86 to 1.0, the BPR on test data from the true generating process follows the expected trajectory. On real data (Fig. 3), the epsilon at training has a looser relationship with BPR on a limited test set due to mismatched assumptions.\n\n## Q5: Can you use a ranking with top-K in denominator?\n\nThe proposed estimator is:\n$$\nr(\\phi) = \\mathbb{E}[ \\frac{ y }{ y \\cdot {TopKMask}(y, K) } ]\n$$\nThis differs from our ratio estimator (Eq 7) by including the top-K binary mask in the denominator, instead of an all-ones vector.\n\nIn expectation, our ratio estimator will have the same ranking of the S sites as this proposal. Thus, this estimator will produce the same top-K decisions as our ratio estimator. We prefer our ratio estimator because it is simpler and faster (avoids top-K for each sample y). \n\nWe will update the Appendix to discuss our reasons for preferring our ratio estimator.\n\n## Other \n\n> better justify the choices in Sec. 3\n\nSee App. B for detailed examples where given the same parameters, our ratio estimator can deliver BPR 2x higher than the per-site mean. We will revise to clarify further.\n\n> consideration of the spatial or temporal aspects\n\nAll models in Fig. 3 use as a feature for site s the gravity of its spatial neighborhood, that is, a recent average of events in sites spatially near to site s.", "claims_and_evidence": "In general, this paper provides examples and empirical evidence to support their proposed approach. The paper does not explicitly provide any theoretical justification for their approach, but do cite other papers when deriving key quantities like the gradient for the ratio estimator. The paper does not have any obvious problematic claims.", "methods_and_evaluation_criteria": "The high level idea makes sense as the method provides users a heuristic approach for trading off decision quality and prediction quality. It makes sense to construct a surrogate for decision-loss in order to incorporate it into learning the underlying probabilistic model. The datasets provided by the authors also fit the setting of optimizing best possible reach and seem to have been used in previous papers for the same application.", "theoretical_claims": "The paper makes no theoretical claims.", "experimental_designs_or_analyses": "I reviewed the synthetic data experiment, the opioid-related overdose forecasting experiment, and the endangered bird forecasting experiment. The synthetic data experiment provided a simple example highlighting the potential pitfalls of only using traditional ML loss and only using BPR loss. The real-world datasets seem to closely follow the experimental set-up of previous works.", "relation_to_broader_scientific_literature": "The key contribution of the paper is showing how to solve the best possible reach (BPR) problem in a decision-aware way. This builds off existing decision-focused and end-to-end learning literature [1][2].  They use compelling real-world data sets and show that decision-aware methods can be effective. The paper outlines the challenges of solving the decision-aware problem and highlights solutions that leverage the score function trick [3] and perturbed optimization [4]. \n\nIn the decision-aware literature, this work most closely resembles [5] which also focuses on applying decision-aware approaches to concrete real world problem. \n\n-----\n[1] Mandi, Jayanta, et al. \"Decision-focused learning: Foundations, state of the art, benchmark and future opportunities.\" Journal of Artificial Intelligence Research 80 (2024): 1623-1701.\n[2] Tang, Bo, and Elias B. Khalil. \"Pyepo: A pytorch-based end-to-end predict-then-optimize library for linear and integer programming.\" Mathematical Programming Computation 16.3 (2024): 297-335.\n[3] Mohamed, S., Rosca, M., Figurnov, M., and Mnih, A. Monte carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132), 2020.\n[4] Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.-P., and Bach, F. Learning with Differentiable Perturbed Op- timizers. In Advances in Neural Information Processing Systems (NeurIPS), 2020.\n[5] Chung, Tsai-Hsuan, et al. \"Decision-aware learning for optimizing health supply chains.\" arXiv preprint arXiv:2211.08507 (2022).", "essential_references_not_discussed": "The paper does not really cite decision-aware literature which can be is well summarized in [1],[2], and [3]. The last paper especially provides computational approaches in a python package for computationally solving the decision-aware BPR problem. \n\nThey also do not cite applications of decision-aware approaches such as [4].\n\nFinally, the authors may also benefit from comparing their approach to [5] which also proposes a decision-aware approach and a case study that closely resembles the BPR problem. \n\n[1] Mandi, Jayanta, et al. \"Decision-focused learning: Foundations, state of the art, benchmark and future opportunities.\" Journal of Artificial Intelligence Research 80 (2024): 1623-1701.\n[2] Tang, Bo, and Elias B. Khalil. \"Pyepo: A pytorch-based end-to-end predict-then-optimize library for linear and integer programming.\" Mathematical Programming Computation 16.3 (2024): 297-335.\n[3] Sadana, Utsav, et al. \"A survey of contextual optimization methods for decision-making under uncertainty.\" European Journal of Operational Research 320.2 (2025): 271-289.\n[4] Chung, Tsai-Hsuan, et al. \"Decision-aware learning for optimizing health supply chains.\" arXiv preprint arXiv:2211.08507 (2022).\n[5] Gupta, Vishal, Michael Huang, and Paat Rusmevichientong. \"Decision-aware denoising.\" Available at SSRN 4714305 (2024).", "other_strengths_and_weaknesses": "Strengths\n1. The paper provides compelling real-world datasets and experiments to highlight the benefits their approach. \n2. The paper provides good motivation for blending traditional ML loss with decision-aware loss.\n\nWeaknesses\n1) The paper provides limited theoretical or practical justification for their proposed BPR loss. My main concern is that the loss proposed by the authors seem to equally weight the BPR loss of each time period $t$. Thus, they are maximizing the average historic BPR. However, since they consider temporal elements, a more sensible decision loss to optimize would be the BPR loss of time period T+1. Optimizing over average historic BPR instead of the BPR of time period T+1 may reduce the decision quality of the learned probabilistic models. [1] highlights how to estimate the decision loss at time period T+1 which may be an interesting benchmark to consider. \n\n2) The paper does not consider other more popular decision-aware approaches. It can be shown that the BPR problem could be formulated as a 0-1 knapsack problem since the denominator term $\\bf{y} \\cdot \\text{TopKMask}(\\bf{y},K)$ is a constant (does not depend on $r$) and $\\text{TopKMask}(\\bf{r},K)$ selects the $K$ largest elements in $\\bf{r}$. This should allow the authors to optimize the BPR loss with decision-aware approaches found in the python package PyEPO [2]. The approaches in the package should be compatible with the BPR problem since the package only requires users providing the gradients of $r^{*}(\\phi)$ and the 0-1 knapsack formulation. As shown in [3], the choice of the approach can affect the resulting decision quality of the learned prediction model. \n\n3) The paper's novelty seems limited. The paper's contributions can be grouped into i) computation and ii) formulation. From a computational perspective, the challenges can be addressed by existing work as discussed in the previous point 2) and the authors do not compare their approach to these existing approaches. From a formulation perspective, the paper's constrained optimization problem reduces to an unconstrained optimization problem that takes a weighted combination of the ML loss and decision loss. This has been previously proposed in [4]. Moreover, the construction of the decision loss is not well justified as highlighted in point 1). \n_____\n[1] Gupta, Vishal, Michael Huang, and Paat Rusmevichientong. \"Decision-aware denoising.\" Available at SSRN 4714305 (2024).\n[2] Tang, Bo, and Elias B. Khalil. \"Pyepo: A pytorch-based end-to-end predict-then-optimize library for linear and integer programming.\" Mathematical Programming Computation 16.3 (2024): 297-335.\n[3] Huang, Michael, and Vishal Gupta. \"Decision-focused learning with directional gradients.\" Advances in Neural Information Processing Systems 37 (2024): 79194-79220.\n[4] Kao Yh, Roy B, Yan X (2009) Directed regression. In: Bengio Y, Schuurmans D, Lafferty J, Williams C, Culotta A (eds) Advances in Neural Information Processing Systems, Curran Associates, Inc., vol 22", "other_comments_or_suggestions": "1. It may help the paper to better justify the choices in Section 3. The paper seems to present two methods of ranking, but in the experiments only present the results of the second approach. Section 3 seems to be the most novel component of the decision-aware approach, so explaining the modeling choices would help highlight the differences between existing decision-aware learning approaches.\n\n2. The paper's title mentions spatialtemporal decision-making settings, however, aside from Eq. (15) there seems to be little consideration of the spatial or temporal aspects of the data. Highlighting components of the paper that leverage the structure of space and time would help differentiate it from existing work.", "questions_for_authors": "Below summarizes some of the questions mentioned in the rest of the review:\n1. Is top K not just a 0-1 knapsack problem of size K? In that case why can you not reformulate the problem as a linear program? You have the gradient of the estimator in (10). Methods found in the PyEPO package (https://github.com/khalil-research/PyEPO) like SPO+, DBB, and PG Loss are compatible. SPO+ is a convex surrogate while PG Loss works well for misspecified models. \n\n2. Why is your proposed loss (9) a good measure of decision loss? Is it some unbiased estimate of the decision loss of time T+1? \n\n3. Why does using only BPR loss provide worse Test BPR for the MA opioid-related overdose forecasting? This seems different than the other two Pareto frontiers shown in Figure 3.\n\n4. In your results, can you show that you can control the BPR with $\\epsilon$?\n\n5. Can you apply the score function trick to the quantity $\\frac{y_i}{\\bf{y} \\cdot \\text{TopKMask}(\\bf{y},K)}$? Using this quantity instead of the ratio estimator would be a more direct way to rank the items.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DSGr8fJGfV", "reviewer": "Reviewer_wWV7", "summary": "This paper studies measures of best possible reach (BPR) to select the best subset of interventions. They analyze different measures based on a probabilistic model, as well as different way to train these probabilistic models from historical data. They also propose new methods for training, including a decision-aware maximum likelihood solution which empirically achieves the best trade-off between BPR and test log-likelihood. They evaluated their methods on synthetic data and real overdose forecasting data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and helpful feedback. We are glad to hear the overall story of our manuscript made sense to you. \n\nWe offer a few responses to the questions you raised below:\n\n> One suggestion I could make is to expand on the minimization of JBPR when the model is misspecified (lines 255-260), and why the forecasts [for BPR only] have questionable utility. \n\nThanks for this idea, we will revise to improve the clarity of this point.\n\n\nEssentially, the issue is that when only training to maximize BPR, all that matters is the ranking of each site in the r vector. Nothing about the implied distribution over y values is forced to match the true distribution of y, beyond just getting the ranking order correct so the predicted set of top K sites aligns with the true top K\nConsider the seven sites in the toy example in Fig 2. Imagine two possible r vectors, each defined as a per-site mean of a distinct model for $p(y_1, … y_7)$\n\n```\nsite  1  2  3  4  5  6   7\n rA: 10 20 30 40 50 60 100\n rB:  1  2  3  4  5  6   7\n```\nBoth rA and rB would have the same BPR, reaching a perfect 1.0, because they rank the top 5 sites (#3 - #7) correctly. However, only model A is at all near the true per-site means indicated in the true-distribution histograms in Fig. 2. Model B has very questionable utility for forecasting y values, because it says the per-site mean of site 7 is 7, which far lower than the actual mean of site 7 (around 100).\n\n> Another suggestion is to highlight the fact that DAML can also integrate prior knowledge if available (which I imagine might be the case in practice?). This is similar to how we can move from MLE to MAP (as mentioned by the authors), but is crucially different from the BPR loss where there is a priori no obvious way to incorporate this type of knowledge.\n\nThanks, we will revise to highlight the ability to integrate prior knowledge via the model. Indeed, in the real applications, we do incorporate a prior on the random effect coefficients (see App. E), as recommended by past work on this negative binomial regression model in conventional settings.", "claims_and_evidence": "The claims made in the submission are clear, and convincingly supported by experiments.", "methods_and_evaluation_criteria": "The methods (ranking and training) make sense for the application at hand.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experimental designs and analyses are sound.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Despite not being familiar with this literature, this paper seems to be a strong practical contribution for ICML. The pedagogical approach of the different methods for ranking and training is greatly appreciated. The rigor shown in the experimental results is valuable, and the results convincing on the advantages of DAML as a trade-off between test log-likelihood and BPR.\n\nThe authors are also very explicit in the limitations of their work, and possible (realistic) directions for future work.", "other_comments_or_suggestions": "One suggestion I could make is to expand on the minimization of $J^{BPR}$ when the model is misspecified (lines 255-260), and why the forecasts have questionable utility. This is seems to be related to the observation in Figure 2 (left), but this is not obvious to a reader less familiar with this field.\n\nAnother suggestion is to highlight the fact that DAML can also integrate prior knowledge if available (which I imagine might be the case in practice?). This is similar to how we can move from MLE to MAP (as mentioned by the authors), but is crucially different from the BPR loss where there is a priori no obvious way to incorporate this type of knowledge.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2t7cBokTwD", "reviewer": "Reviewer_His4", "summary": "In a spatiotemporal forecasting setting, the authors consider traning prodiction models adapted to the task of selecting the top-K sites for intervention. The authors consider BPR as the desired metric and develop an algorithm for training models using a gradient-based approach. The main difficulty of this approach lies in the uninformative gradients that the BPR objective, combined with the top-K selection provides. To differentiate through the sampling of discrete variables $r$, the authors adapt a REINFORCE-like approach, while deriving the Loss (which includes the top-K ranking) is done using a stochastic smoothing approach. In a last step, a combination of DPR and log-likelihood is used to propose DAML, which is intended to train decision aware models which provide good likelihood estimates.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive comments about our work. We try to address key points below.\n\n> RE Question about “the title of Algorithm A.1”\n\nAlgorithm A.1 is meant to summarize the decision-aware ML training (DAML) approach described in Sec. 4.3. Current title is “Decision-aware ML training for top-K tasks judged by BPR”. We will revise to just say “Decision-aware ML training” in case the latter half of the title was confusing\n\n## Essential References (Common Issue)\n\n> I think it would be valuable to discuss the connection to other works in the literature related to decision-aware learning, end-to-end learning and contextual optimization\n\nWe agree there’s a broader set of related work, especially from the OR community. We plan to cite and discuss the survey by Sadana et al. (2024), as well as valuable works pointed out by ZTv9.\n \nSpecifically, we plan revisions that will\n* cite and discuss the application of “decision-aware denoising” to the problem of where to place speed humps to calm traffic (Gupta et al., SSRN preprint from 2024)\n* cite and discuss the health supply chain task of Chung et al. (ML4H 2022)\n* cite and discuss the PyEPO package (Tang & Khalil '24), especially how it can support solving problems like ours\n* cite the survey on “decision-focused learning” by Mandi et al. (2024) and the survey by Sadana et al. (2025), and relevant references therein\n\n> uncertainty plots in Figure 3 are quite confusing. Do they provide uncertainty only in the BPR direction? \n\nYes, the uncertainty in this plot is only in the BPR direction. We will revise the caption to clarify. \n\nIn Figure 3, we intend to show how each method’s estimated parameter performs across two metrics: log likelihood and BPR. Log likelihood is well-summarized by one deterministic number (vertical position of the marker). In contrast, BPR is not deterministic, as it depends on calculating the ratio estimator for r in Eq 7 which requires an average over M *samples* of y. Each method in Fig 3 has BPR visualized as a histogram showing the distribution across 1000 trials, with each trial using M =1000 samples. \n\n> 1. … How would you propose to sample from the joint model p if it was parametrized by a multilayer neural network?\n\nThere’s a substantial literature on deep probabilistic models, which combine neural networks for flexible parameterization with classic statistical distributions with well-known sampling routines. \nFor example, suppose our joint model for vector y was multivariate Gaussian with some mean and some covariance. The mean vector could be parameterized as the output of a neural network. So could the covariance matrix, assuming symmetry and positive definite constraints were enforced. Sampling from a Gaussian given its mean and covariance is then straightforward.\n\n> 2. In Figure 2, left panel, “OPT BPR only”. I think this result is quite interesting, … one would expect this to carry over to any top-K setting: No matter the complexity of the task, a mixture of 2 gaussians should be enough … for identifying the selected IDs (being identified with the gaussian with higher mean) and non-selected IDs … Is this intuition true?\n\nYes we agree with this intuition. Deliberate construction of component weights $\\pi$ for each site to assign the top K sites to a “high mean” component and other sites to a “low mean” component should be enough to get good BPR on training data, regardless of the task.\n\nNaturally, whether this can generalize to test data depends on the model’s (mis)match to the true data-generating process.\n\n> 3. Can we view DAML as “finding the global optima” (in terms of BPR) that “has the highest log likelihood”?\n\nYes, we agree with this view. There can be different models $\\phi$ that have equivalent rankings r and thus equivalent BPR value. We can view DAML with a high BPR constraint as seeking $\\phi$ that have both high BPR and high log likelihood.\n\n> Do the results remain similar if minibatching is used? I understand that for the datasets used minibatching may not be necessary, but for larger datasets one would certainly like to use it.\n\nMinibatch gradient descent to minimize the BPR only loss in Eq 9 or the DAML objective in Eq 13 appears straightforward in both cases. In each case, the overall “whole dataset” loss is a sum over the loss at each timestep indexed by t. Taking a random subset of timesteps as a minibatch and computing the gradient of that minibatch should be an unbiased estimate of the gradient of the whole-dataset loss.\n\nWe have not tested minibatching in our experiments, as it wasn’t necessary and would introduce extra noise to our gradient-based learning that is already noisy (due to score function trick and perturbations). Tuning learning rates and batch sizes carefully would be important, but if done well, minibatching seems feasible.", "claims_and_evidence": "The claims are well supported.", "methods_and_evaluation_criteria": "The evaluation methods are suitable.", "theoretical_claims": "I followed the mathematical derivation in the main text, which looks correct to me.", "experimental_designs_or_analyses": "As far as I can tell, the experimental designs make sense.", "relation_to_broader_scientific_literature": "The authors use a decision aware objective for training a model adapted to the top-K selection problem. I am not aware of any other works that tackle decision aware training for the top-K problem. Decision aware training is however a growing trend in the literature, see Sadana et al. (2024).", "essential_references_not_discussed": "I think it would be valuable to discuss the connection to other works in the literature related to decision-aware learning, end-to-end learning and contextual optimization, see for instance Sadana et al. (2024) and references therein.\n\nSadana, Utsav, et al. \"A survey of contextual optimization methods for decision-making under uncertainty.\" *European Journal of Operational Research* (2024).", "other_strengths_and_weaknesses": "The authors tackle an interesting problem and propose a new training scheme which is well adapted to the the problem and improves performance compared to the baselines. The paper is well written and explained, I specifically appreciated the interesting, and practically relevant experiments section.", "other_comments_or_suggestions": "The uncertainty plots in Figure 3 are quite confusing. Do they provide uncertainty only in the BPR direction? If not, why does uncertainty increase for lower log-likelihood values?", "questions_for_authors": "1. On line 91ff, right side, you state that the framework is flexible and can be applied to neural networks. How would you propose to sample from the joint model $p_\\phi$ if it was parametrized by a multilayer neural network?\n2. In Figure 2, left panel, “OPT BPR only”. I think this result is quite intersting, especially looking at the weights $\\pi$. The model identifies the IDs to select (high $\\pi$) and the ones not to select (low $\\pi$). Intuitively, one would expect this to carry over to any top-K setting: No matter the complexity of the task, a mixture of 2 gaussians should be enough as it is a proxy for identifying the selected IDs (being identified with the gaussian with higher mean) and non-selected IDs (being identified with the gaussian with higher mean). Is this intuition true?\n3. In Figure 2, on the right panel: DAML attains the same BPR as directly training on the BPR loss, but with better log-likelyhood. The structure of the BPR loss in equation (4) suggests that minimizing (9) could have multiple global optima. I think one view DAML (with appropriately chosen $\\varepsilon$ as “finding the global optima of $\\min(9)$ that has the highest log likelihood.\n4. On line 252ff, left side: “For convenience and reliability, we use all T records in the training set in every estimate, avoiding minibatching over time.” Do the results remain similar if minibatching is used? I understand that for the datasets used minibatching may not be neccessary, but for larger datasets one would certainly like to use it.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: I concur with the overall positive sentiment from the reviewers and agree the submission would make a nice addition to the conference. The reviewers gave some great feedback that should be carefully incorporated into the final camera ready. The most important is making very clear the connections to contextual stochastic/linear optimization (using the reduction explained by reviewer ZTv9) and referencing the relevant papers in that literature (for reference a comprehensive list would be these arxiv ids 1402.5481 1710.08005 1905.11488 2008.07473 2011.03030 2402.03256 2405.16564) as well as discussing the other decision-aware learning papers mentioned.
Decision: Accept (poster)