Learning decision-focused uncertainty representations

Machine learning can significantly improve performance for decision-making under uncertainty in a wide range of domains. However, ensuring robustness guarantees and satisfaction of risk constraints requires well-calibrated uncertainty estimates, yet there may be many valid uncertainty estimates, each with their own performance profile—i.e., not all uncertainty is equally valuable for downstream decision-making. To address this problem, we developed an end-to-end framework to learn uncertainty representations for stochastic, robust, and risk-constrained optimization in a way that is informed by the downstream decision-making loss. For stochastic optimization, we train a conditional diffusion model from which we sample scenarios, using a novel approximation to backpropagate gradients through the diffusion model sampling procedure, achieving 60x GPU memory savings over standard techniques. In robust optimization, we approximate arbitrary convex uncertainty sets with sublevel sets of partially input-convex neural networks. Finally, for risk-constrained optimization, we devise a general, differentiable risk control procedure that generalizes conformal risk control. Our approach consistently improves upon two-stage estimate-then-optimize baselines on concrete applications across energy systems and finance.