How to Teach Label to Understand Decisions: A Decision-aware Label Distribution Learning Framework

ICLR 2026 Conference Submission4760 Authors

Keywords: Contextual Stochastic Optimization, Label Distribution Learning, Uncertainty modeling

TL;DR: We introduce a Decision-aware Label Distribution Learning framework that embeds decision knowledge into labels, enabling them to better guide robust decision-making.

Abstract: Contextual Stochastic Optimization (CSO) aims to predict uncertain, context-dependent parameters to inform downstream decisions. A central challenge is that high predictive accuracy does not necessarily translate into optimal decisions. Existing approaches typically rely on custom loss functions, but these often suffer from non-differentiability, discontinuity, and limited modularity. To address these limitations, we propose a decision-aware Label Distribution Learning (LDL) framework that retains standard loss functions to avoid computational issues, while encoding decision knowledge entirely at the level of data representation. Our approach models uncertainty as full label distributions and reshapes them during the label enhancement stage to reduce predictive mass in high-risk regions. Scalar targets are transformed into individualized mixture distributions using decision-aware similarity matrices, and a dual-branch neural network is trained to learn decision-aware label distributions. Extensive experiments on synthetic benchmarks (e.g., newsvendor, network flow) and real-world datasets demonstrate consistent regret reduction across different sample sizes, with particularly strong improvements in low-data regimes. These results highlight LDL as a promising new pathway for achieving robust and principled decision-making under complex cost structures.

Primary Area: optimization

Submission Number: 4760