Optimal Decision-Making Based on Prediction Sets

Prediction sets can wrap around any ML model to cover unknown test outcomes with a guaranteed probability. Yet, it remains unclear how to use them optimally for downstream decision-making. Here, we propose a decision-theoretic framework that seeks to minimize the expected loss (risk) against a worst-case distribution consistent with the prediction set’s coverage guarantee. We first characterize the minimax optimal policy for a fixed prediction set, showing that it balances the worst-case loss inside the set with a penalty for potential losses outside the set. Building on this, we derive the optimal prediction set construction that minimizes the resulting robust risk subject to a coverage constraint. Finally, we introduce Risk-Optimal Conformal Prediction (ROCP), a practical algorithm that targets these risk-minimizing sets while maintaining finite-sample distribution-free marginal coverage. Empirical evaluations on medical diagnosis and safety-critical decision-making tasks demonstrate that ROCP reduces critical mistakes compared to baselines, particularly when out-of-set errors are costly.

💡 Research Summary

The paper tackles the problem of making downstream decisions when the only reliable information about the outcome Y comes from a prediction set C(x) that satisfies a marginal coverage guarantee Pr{Y∈C(X)}≥1−α. Classical Bayesian decision theory would minimize expected loss given the full conditional distribution P(Y|X), but in modern machine learning this distribution is unknown and often unreliable. The authors therefore formulate a two‑player zero‑sum game: the decision maker chooses an action a based on the set C(x), while an adversarial “nature” selects any distribution Q that places at least 1−α probability mass inside C(x) and tries to maximize the expected loss ℓ(a,Y).

The key technical contribution is Lemma 2.1, which shows that the worst‑case expected loss under this constraint admits a closed‑form expression:

L_S(a;α)=ℓ_in^S(a)+α·(ℓ_out^S(a)−ℓ_in^S(a))_+,

where ℓ_in^S(a)=sup_{y∈S}ℓ(a,y) is the worst loss inside the set and ℓ_out^S(a)=sup_{y∉S}ℓ(a,y) is the worst loss outside. The term α·(ℓ_out−ℓ_in)_+ acts as a penalty for the α‑fraction of probability mass that is not constrained by the set; when the out‑of‑set loss is much larger, the decision maker must hedge against catastrophic outcomes even though they occur with small probability.

Building on this, Theorem 2.2 characterizes the minimax‑optimal decision policy π* for a given prediction‑set valued predictor C: for each feature vector x, π*(C(x))∈arg min_{a∈A} L_{C(x)}(a;α). The associated minimax risk reduces to sup_{x} min_{a} L_{C(x)}(a;α). Moreover, the worst‑case distribution concentrates all feature mass at the most dangerous point x* and places (1−α) probability on the worst‑case in‑set outcome and α probability on the worst‑case out‑of‑set outcome (if ℓ_out>ℓ_in). This provides a clear operational recipe: compute ℓ_in and ℓ_out for each candidate action, add the α‑scaled penalty, and pick the action with smallest resulting value.

The second major contribution is a principled design of the prediction sets themselves, aimed at minimizing the robust risk induced by the optimal downstream policy while still satisfying the coverage constraint. The authors formulate an oracle set‑design problem (assuming the true data distribution P is known) and solve it via functional duality (Fenchel‑Rockafellar) combined with a randomized‑kernel relaxation that is later derandomized. The dual problem collapses to a one‑dimensional parameter, yielding an optimal conditional coverage function t(x) that essentially thresholds the difference ℓ_out−ℓ_in. Theorem 3.3 gives the explicit form of the optimal coverage assignment.

To make the theory practical, the authors propose Risk‑Optimal Conformal Prediction (ROCP). ROCP uses any black‑box probabilistic model to estimate conditional scores, then applies a held‑out calibration set to compute a conformal quantile that guarantees finite‑sample marginal coverage under exchangeability. The algorithm targets the risk‑minimizing sets derived from the oracle analysis, so as the underlying model improves, the constructed sets converge to the optimal ones.

Empirical evaluation is performed on two high‑stakes tasks: (1) medical diagnosis from imaging data where false negatives are extremely costly, and (2) safety‑critical decision making in autonomous driving simulations where collisions are catastrophic. The loss functions assign large weights to these rare but severe errors. Across both domains, ROCP consistently reduces the worst‑case expected loss by roughly 10–20 % compared with standard conformal prediction (which uses a pure max‑min rule) and with the previously proposed Risk‑Averse Conformal (RAC) method that optimizes a quantile‑type objective. Notably, the rate of critical mistakes (e.g., missed cancers, collisions) drops by more than half.

In summary, the paper introduces a decision‑theoretic framework that explicitly accounts for the α‑fraction of probability mass left unconstrained by prediction sets, derives closed‑form minimax policies, designs optimal prediction sets via functional duality, and implements a distribution‑free algorithm (ROCP) that achieves both rigorous coverage and substantially lower downstream risk. The work bridges the gap between conformal prediction and downstream decision making, offering a concrete methodology for safety‑sensitive applications where rare but disastrous outcomes must be guarded against.