Direct Prediction Set Minimization via Bilevel Conformal Classifier Training

TL;DR

This paper introduces DPSM, a bilevel optimization approach that directly minimizes prediction set sizes in conformal prediction, leading to more concise and practical uncertainty quantification in deep classifiers.

Contribution

The paper formulates conformal training as a bilevel optimization problem and proposes DPSM, a novel algorithm that improves prediction set size minimization with theoretical guarantees.

Findings

DPSM reduces prediction set size by over 20% compared to baselines.

DPSM has a learning bound of O(1/√n), better than prior methods.

Experiments validate DPSM's effectiveness across datasets and models.

Abstract

Conformal prediction (CP) is a promising uncertainty quantification framework which works as a wrapper around a black-box classifier to construct prediction sets (i.e., subset of candidate classes) with provable guarantees. However, standard calibration methods for CP tend to produce large prediction sets which makes them less useful in practice. This paper considers the problem of integrating conformal principles into the training process of deep classifiers to directly minimize the size of prediction sets. We formulate conformal training as a bilevel optimization problem and propose the {\em Direct Prediction Set Minimization (DPSM)} algorithm to solve it. The key insight behind DPSM is to minimize a measure of the prediction set size (upper level) that is conditioned on the learned quantile of conformity scores (lower level). We analyze that DPSM has a learning bound of $O(1/\sqrt{n})$ (with $n$ training samples), while prior conformal training methods based on stochastic approximation for the quantile has a bound of $\Omega(1/s)$ (with batch size $s$ and typically $s \ll \sqrt{n}$). Experiments on various benchmark datasets and deep models show that DPSM significantly outperforms the best prior conformal training baseline with $20.46\%\downarrow$ in the prediction set size and validates our theory.