Stochastic grouping improves prediction reliability across subgroups
Stochastic Grouping Conformal Prediction for Effective Subgroup Reliability
Machine LearningArtificial Intelligence
Summary
Predicting outcomes reliably for everyone is hard because some groups get less accurate results. The authors found that treating every group separately can make overall predictions worse or harder to use. They designed a new method that mixes data from similar cases to improve fairness and accuracy for all groups without making predictions too broad. Their approach keeps the strong guarantees of standard methods but improves reliability across different subgroups.
What this means in practice
- •For clinical data teams: Create more reliable prediction models that ensure fairness across patient groups without increasing prediction complexity.
- •For financial risk analysts: Improve uncertainty estimates across diverse client segments, reducing bias in risk predictions without inflating risk margins.
Authors
Meihui Zhong, Wenxin Tai, Ting Zhong, Fan Zhou
Abstract
Conformal prediction offers a distribution-free coverage guarantee, making it especially attractive for clinical applications. Standard conformal prediction, however, provides such guarantees only at the population level, and its prediction sets can exhibit coverage disparities across clinically important subgroups. A natural remedy is to calibrate within predefined groups. However, this can require access to sensitive subgroup attributes and is prone to a worst-group bottleneck: protecting the most difficult subgroup can inflate prediction sets for all, increasing cognitive burden on decision makers. To this end, we propose Stochastic Grouping Conformal Prediction (SGCP), a conformal framework for subgroup-reliable uncertainty quantification. It learns a stochastic grouping map that allows each sample to draw calibration information from others with similar calibration behavior, yielding a local score law that boosts reliability across subpopulations. We prove that SGCP retains the standard coverage guarantee. Experiments on synthetic and real-world benchmarks show that it consistently reduces subgroup coverage gaps while achieving smaller or comparable prediction set sizes relative to existing baselines.