Conformal Risk Training: End-to-End Optimization of Conformal Risk Control

AI summary

Overview

Research area: machine learning — conformal prediction, risk control, and end-to-end training of risk-aware models.

Technical level: Advanced. The paper is written for readers comfortable with conformal risk control, convex risk measures (optimized certainty equivalents, conditional value-at-risk), and differentiating through optimization problems.

Scope: The paper extends conformal risk control (CRC) from expected-loss control to the broader class of optimized certainty-equivalent (OCE) risks, and introduces "conformal risk training," an end-to-end method that differentiates through conformal OCE risk control during model training or fine-tuning.

What This Paper Is About

Conformal risk control is a distribution-free, finite-sample technique that picks a decision parameter so that the expected value of a bounded monotone loss stays below a chosen level α, and it is normally applied after training, to a frozen model. The authors address two limits of that approach: it only controls expected loss rather than tail risks like CVaR (an open question posed in the original CRC paper), and because it gives no feedback to the model, it can badly degrade average-case performance. Their goal is a method that controls general OCE risks while jointly training the model and the risk-control procedure.

Key Contributions

1. Risk control for the OCE family. The authors develop a risk-control method for optimized certainty-equivalent risks, a broad class of risk measures that includes expected loss (recovering the original CRC method) and conditional value-at-risk as special cases. The key insight is that any OCE risk can be bounded by a monotone transformation of the loss, preserving the monotonicity that CRC-style methods require. The CVaR formulation partially answers the open question posed by the original CRC paper.

2. Conformal risk training. They propose an end-to-end method that trains a model together with the conformal risk control procedure, substantially generalizing prior conformal training of uncertainty sets to the conformal risk control setting. The trained model becomes "risk aware," learning predictions that maximize performance while minimizing downstream risk.

3. Gradient computation. They show how to compute the derivative of the risk-controlling parameter with respect to model parameters, giving closed-form expressions under two sets of conditions — one covering piecewise-constant losses with strictly decreasing costs, and one covering convex losses with strictly convex or strictly monotone costs. They show the gradient used in prior conformal training follows as a special case.

4. Empirical demonstration. They show that fine-tuning with conformal risk training improves average-case performance while guaranteeing the risk constraint, on controlling false negative rate in tumor image segmentation and controlling tail risk of losses in battery storage operation.

Main Findings

• CVaR and OCE risks can be conformally controlled. Theorem 1 shows that for any OCE risk measure with disutility function φ, transforming each loss as t + φ(L_i(λ) − t) and applying the CRC-style bound yields the guarantee R[L_{N+1}(λ)] ≤ α. Setting φ(x) = x recovers the original CRC result exactly.
• CVaR needs only monotonicity, not monotone increasing losses. Theorem 2 relaxes the assumption that losses are nondecreasing, replacing it with the milder requirement that all losses and the upper bound are monotonic (some may be nonincreasing), and still guarantees CVaR_δ[L_{N+1}(λ)] ≤ α for t in [B(λ_min), α]. The authors state this can be extended to some, but not all, other OCE risk measures.
• A new hyperparameter t is required, and must not depend on calibration data. Both new theorems introduce t; the authors recommend choosing it using an additional held-out set of losses, picking the value that yields the largest risk-control parameter.
• Post-hoc CRC can be very costly in false positives. On colonoscopy gut polyp image segmentation, the post-hoc CRC and cross-entropy baselines reached as high as 80% false positive rate when the target false negative rate was 1%.
• Conformal risk training reduces that cost substantially. Across the tested α levels (0.01, 0.05, 0.1), conformal risk training reduced FPR on average by 23-42% versus the baselines, while all three methods kept FNR controlled at the target level α.
• The trained models are less conservative. The method generally selected a higher classification threshold λ than the baselines, suggesting it maintains the risk guarantee with less conservatism.
• Battery storage experiments test tail risk. A second experiment controls the CVaR of losses in grid-scale battery storage operation while maximizing average profit. The truncated paper content describes the setup but does not report the numeric results for this task.

Methodology in Plain English

The starting point is standard conformal risk control: given a set of losses computed on calibration data, you find the largest value of a decision threshold λ whose average loss (plus an upper-bound term) stays under the target α, found by bisection search because the quantity is monotone in λ.

To move beyond expected loss, the authors use a standard representation of OCE risk measures as an infimum over a scalar t of t plus the expected value of a disutility function applied to the loss minus t. This turns an OCE risk bound into an ordinary expectation bound on a transformed loss, so the same conformal machinery applies. That gives their conformal OCE risk control algorithm, which searches for λ by bisection on the transformed quantity, with a special case for CVaR.

For training, they pose a bilevel problem: an outer objective minimizes a differentiable cost (for example, false positives), while an inner problem picks the risk-controlling λ subject to the conformal constraint. To optimize the outer objective by gradient descent, they need the derivative of λ with respect to the model parameters, which requires differentiating through the inner optimization. They give closed-form expressions for this derivative in two practical cases. The training loop splits each minibatch into a pseudo-calibration half (used to compute λ) and a prediction half (used to compute the cost and its gradient), mirroring prior conformal training. After training, a fresh calibration set is used to compute λ for test inputs.

Why This Matters

Impact on research. This is described as the first conformal-style approach that gives a certain (rather than high-probability) bound on risk measures beyond expected loss for monotone loss functions, and it partially resolves an open question from the original CRC paper about CVaR. It also connects conformal risk control to decision-focused uncertainty quantification and provides a framework for differentiating through risk-control procedures.

Real-world applications (as discussed in the paper):

• Medical image segmentation, where false negative pixels (missed tumors) must be provably bounded while keeping false positives low — the paper's own tumor segmentation experiment.
• Tumor classification, where the health risk of misclassifying a tumor as benign must be prioritized with a provably bounded false negative rate.
• Energy and finance, where a battery operator needs optimal expected profit while controlling the tail risk (CVaR) of losses from poor electricity price forecasts.
• Predictive uncertainty sets targeting a coverage level, and factuality of large language model outputs, both named as examples of risk-control problems in the introduction.

Industry relevance. The two demonstration domains — clinical imaging and grid-scale battery operation — are high-stakes settings where practitioners need guarantees, not just good average accuracy. Because the guarantees are distribution-free and finite-sample, and because the method plugs into fine-tuning of existing models such as a pre-trained PraNet, the approach is directly relevant to deployment pipelines that already use conformal methods post-hoc.

Future Directions

• Extending the relaxed monotonicity guarantee from Theorem 2 to OCE risk measures beyond CVaR. The authors note the result generalizes to some, but not all, OCE risks (Section C.3).
• Computing the gradient of the risk-controlling parameter in settings outside the two cases covered by Theorem 3, since the inner problem may in general be nonconvex.
• Better strategies for choosing the hyperparameter t, which must be selected without using the calibration data used to pick λ; the paper offers a held-out-set heuristic but no formal procedure.
• Applying conformal risk training to the other risk-control problems named in the introduction, such as large language model factuality and predictive uncertainty sets with coverage targets, which are not tested here.

Target Audience

Researchers and practitioners in trustworthy machine learning, conformal prediction, and decision-focused learning who already understand conformal risk control and convex risk measures. It is also relevant to applied scientists in medical imaging and energy systems who need provable risk guarantees alongside competitive average-case performance, and to readers interested in differentiating through optimization layers during training.

Authors’ abstract

While deep learning models often achieve high predictive accuracy, their predictions typically do not come with any provable guarantees on risk or reliability, which are critical for deployment in high-stakes applications. The framework of conformal risk control (CRC) provides a distribution-free, finite-sample method for controlling the expected value of any bounded monotone loss function and can be conveniently applied post-hoc to any pre-trained deep learning model. However, many real-world applications are sensitive to tail risks, as opposed to just expected loss. In this work, we develop a method for controlling the general class of Optimized Certainty-Equivalent (OCE) risks, a broad class of risk measures which includes as special cases the expected loss (generalizing the original CRC method) and common tail risks like the conditional value-at-risk (CVaR). Furthermore, standard post-hoc CRC can degrade average-case performance due to its lack of feedback to the model. To address this, we introduce "conformal risk training," an end-to-end approach that differentiates through conformal OCE risk control during model training or fine-tuning. Our method achieves provable risk guarantees while demonstrating significantly improved average-case performance over post-hoc approaches on applications to controlling classifiers' false negative rate and controlling financial risk in battery storage operation.