Stochastic Grouping Conformal Prediction for Effective Subgroup Reliability
Organizations: University of Electronic Science and Technology of China · Kashi Institute of Electronics and Information Industry · Intelligent Digital Media Technology, Key Laboratory of Sichuan Province
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.
Figures & tables
| Method | SYN | Nursery | MIMIC-IV | BACH | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Marginal | 89.52 ±1.52 | 2.309 ±0.078 | 14.5 ±1.6 | 84.15 ±5.20 | 2.170 ±0.260 | 20.4 ±5.2 | 90.79 ±10.62 | 1.226 ±0.087 | 27.1 ±7.1 | 89.20 ±9.64 | 2.894 ±0.267 | 15.7 ±1.7 |
| Partial | 94.50 ±1.12 | 2.419 ±0.048 | 6.2 ±0.7 | 88.97 ±3.80 | 1.850 ±0.200 | 12.1 ±2.5 | 94.50 ±4.54 | 1.258 ±0.140 | 8.0 ±3.0 | 96.57 ±8.08 | 3.385 ±0.299 | 7.6 ±1.6 |
| Exhaustive | 95.85 ±0.91 | 3.342 ±0.053 | 6.0 ±0.6 | 96.80 ±3.10 | 2.850 ±0.220 | 15.5 ±3.2 | 96.45 ±4.22 | 2.800 ±0.142 | 19.7 ±3.2 | 97.41 ±6.89 | 3.867 ±0.118 | 11.6 ±1.5 |
| RLCP | 93.10 ±2.45 | 2.540 ±0.305 | 4.5 ±1.8 | 88.67 ±6.20 | 1.220 ±0.410 | 13.5 ±3.8 | 92.50 ±2.80 | 1.350 ±0.100 | 1.5 ±0.8 | 91.88 ±10.56 | 2.311 ±0.319 | 8.1 ±2.5 |
| CluCP | 89.80 ±1.78 | 2.426 ±0.050 | 4.9 ±0.8 | 85.08 ±1.60 | 1.121 ±0.064 | 12.8 ±1.3 | 92.22 ±1.28 | 1.302 ±0.059 | 3.5 ±0.8 | 95.91 ±3.59 | 3.257 ±0.220 | 10.7 ±2.6 |
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
| Dataset | Modality | #Classes | #Attrs | Sensitive / grouping attributes |
|---|---|---|---|---|
| SYN | tabular | 6 | 3 | gender, phenotype, age |
| Nursery | tabular | 4 | 5 | parents, children, finance, |
| social, health | ||||
| MIMIC-IV | tabular | 2/3 | 3 | minority, gender, public insurance |
| BACH | image | 4 | - | class label (for class-wise reliability analysis) |
| Hyperparameter | SYN | Nursery | M-ICU | M-Hosp | BACH | Description |
| (i) Latent calibration components | ||||||
| 3 | 6 | 3 | 3 | 6 | number of latent calibration components | |
| 16 | 8 | 8 | 8 | 64 | latent calibration dimension | |
| hidden dim | 64 | 64 | 32 | 32 | 64 | hidden dimension of grouping network |
| (ii) Score law learning | ||||||
| SGCP train epochs | 100 | 100 | 100 | 100 | 100 | number of SGCP training epochs |
| Method | SYN | Nursery | MIMIC | BACH | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Default SGCP | 90.13 ±0.39 | 2.225 ±0.035 | 4.2 ±0.7 | 91.65 ±0.70 | 1.089 ±0.035 | 11.8 ±0.7 | 90.13 ±0.07 | 1.107 ±0.001 | 0.8 ±0.0 | 91.90 ±2.70 | 2.224 ±0.215 | 6.2 ±1.1 |
| (a) w/o Stochastic | 89.82 ±1.48 | 2.348 ±0.081 | 7.35 ±1.18 | 91.38 ±1.15 | 1.123 ±0.051 | 13.42 ±1.76 | 89.94 ±0.96 | 1.053 ±0.013 | 2.36 ±0.48 | 91.76 ±4.42 | 2.447 ±0.308 | 8.84 ±2.35 |
| (b) Feature space | 90.96 ±1.96 | 2.476 ±0.127 | 7.18 ±1.46 | 90.92 ±1.43 | 1.248 ±0.079 | 14.16 ±2.05 | 90.48 ±1.47 | 1.179 ±0.044 | 4.08 ±1.08 | 89.61 ±5.74 | 3.094 ±0.446 | 12.41 ±3.16 |
| (c) w/o rank | 89.21 ±4.42 | 1.898 ±0.246 | 8.42 ±4.10 | 88.63 ±2.94 | 1.302 ±0.118 | 15.31 ±2.41 | 90.03 ±5.41 | 1.334 ±0.146 | 5.49 ±3.71 | 91.18 ±6.31 | 3.491 ±0.541 | 9.26 ±3.02 |
| (d) w/o NLL | 88.36 ±1.17 | 2.074 ±0.108 | 5.39 ±1.08 | 92.39 ±0.92 | 1.446 ±0.094 | 12.43 ±1.19 | 90.02 ±2.96 | 1.351 ±0.078 | 2.96 ±0.77 | 93.11 ±2.96 | 2.846 ±0.317 | 7.44 ±1.47 |
| Method | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Marginal | 90.15 ±6.50 | 2.150 ±0.420 | 32.5 ±6.5 | 90.20 ±4.20 | 1.950 ±0.400 | 21.5 ±6.2 | 90.12 ±4.10 | 1.820 ±0.380 | 19.7 ±6.0 | 89.52 ±1.52 | 2.309 ±0.078 | 14.5 ±1.6 |
| Partial | 95.50 ±3.20 | 4.520 ±0.300 | 9.5 ±2.8 | 94.20 ±3.00 | 2.950 ±0.280 | 7.9 ±2.5 | 93.50 ±2.80 | 2.820 ±0.250 | 7.5 ±2.2 | 94.50 ±1.12 | 2.419 ±0.048 | 6.2 ±0.7 |
| Exhaustive | 100.00 ±0.00 | 6.000 ±0.000 | 12.5 ±3.2 | 98.50 ±2.20 | 4.850 ±0.250 | 10.5 ±3.0 | 98.20 ±2.10 | 4.150 ±0.220 | 7.9 ±2.7 | 95.80 ±0.91 | 3.342 ±0.053 | 6.0 ±0.6 |
| RLCP | 95.80 ±3.80 | 2.910 ±0.520 | 9.2 ±8.2 | 95.10 ±3.20 | 2.850 ±0.480 | 8.5 ±6.8 | 94.60 ±2.90 | 2.850 ±0.460 | 6.4 ±5.5 | 93.15 ±2.45 | 2.540 ±0.305 | 4.5 ±1.8 |
| CluCP | 90.22 ±2.06 | 3.112 ±0.332 | 10.7 ±0.9 | 90.62 ±0.98 | 2.722 ±0.269 | 9.9 ±0.7 | 91.11 ±1.83 | 2.756 ±0.250 | 8.9 ±1.1 | 89.81 ±1.78 | 2.426 ±0.050 | 4.9 ±0.8 |
| Dataset | Source | Access note |
|---|---|---|
| Synthetic | Procedurally generated; see Appendix D.1 | Generated from the described simulation protocol |
| Nursery | https://archive.ics.uci.edu/dataset/76/nursery | Public UCI dataset |
| MIMIC-IV | https://physionet.org/content/mimiciv/3.1/ | Credentialed access via PhysioNet |
| BACH | https://zenodo.org/records/3632035 | Public Zenodo record |
| Method | Code source |
|---|---|
| SGCP | https://anonymous.4open.science/r/sgcp/ |
| Partial Equalized | https://github.com/yromano/cqr |
| RLCP | https://github.com/rohanhore/RLCP |
| CluCP | https://github.com/tiffanyding/class-conditional-conformal |
| AFCP | https://github.com/FionaZ3696/Adaptively-Fair-Conformal-Prediction |
| FaReG | https://github.com/Xusr1123/FaReG |
| Method | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Marginal | 48.25 ±8.50 | 1.120 ±0.350 | 29.5 ±6.5 | 57.59 ±8.20 | 1.080 ±0.330 | 29.1 ±5.2 | 59.72 ±7.80 | 1.050 ±0.310 | 26.2 ±5.8 | 84.15 ±5.20 | 2.170 ±0.260 | 20.4 ±5.2 |
| Partial | 84.61 ±4.80 | 2.850 ±0.310 | 17.7 ±3.8 | 86.07 ±4.50 | 2.620 ±0.280 | 16.8 ±3.5 | 87.86 ±4.20 | 2.450 ±0.250 | 15.2 ±3.2 | 88.97 ±3.80 | 1.850 ±0.200 | 12.1 ±2.5 |
| Exhaustive | 100.00 ±0.00 | 4.000 ±0.000 | 19.5 ±0.1 | 98.00 ±0.10 | 3.980 ±0.001 | 18.8 ±0.2 | 98.00 ±0.12 | 3.979 ±0.001 | 17.5 ±0.8 | 96.80 ±3.10 | 2.850 ±0.220 | 15.5 ±3.2 |
| RLCP | 82.45 ±14.20 | 2.250 ±0.580 | 18.5 ±9.5 | 83.36 ±13.80 | 1.720 ±0.550 | 17.8 ±6.2 | 83.31 ±11.50 | 1.480 ±0.510 | 16.5 ±5.8 | 88.67 ±6.20 | 1.220 ±0.410 | 13.5 ±3.8 |
| CluCP | 86.26 ±5.61 | 1.506 ±0.265 | 17.3 ±4.5 | 83.80 ±3.59 | 1.516 ±0.265 | 17.1 ±2.9 | 84.42 ±2.47 | 1.261 ±0.096 | 16.9 ±1.8 | 85.08 ±1.60 | 1.121 ±0.064 | 12.8 ±1.3 |
| Method | Full (All Samples) | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Marginal | 90.06 ±10.14 | 1.074 ±0.211 | 34.4 ±10.9 | 90.79 ±10.62 | 1.026 ±0.087 | 29.6 ±9.3 | 90.79 ±10.62 | 1.226 ±0.087 | 27.1 ±7.1 |
| Partial | 95.73 ±3.62 | 1.564 ±0.243 | 11.0 ±5.6 | 94.50 ±4.54 | 1.335 ±0.207 | 9.2 ±4.0 | 94.50 ±4.54 | 1.258 ±0.140 | 8.0 ±3.0 |
| Exhaustive | 100.00 ±0.00 | 3.000 ±0.000 | 18.7 ±8.0 | 98.86 ±5.21 | 2.989 ±0.004 | 21.4 ±5.3 | 96.45 ±4.22 | 2.800 ±0.142 | 19.7 ±3.2 |
| RLCP | 95.10 ±2.40 | 1.650 ±0.120 | 4.2 ±2.6 | 93.20 ±2.10 | 1.480 ±0.150 | 3.5 ±1.4 | 92.50 ±2.80 | 1.350 ±0.100 | 1.5 ±0.8 |
| CluCP | 95.11 ±1.92 | 1.654 ±0.192 | 4.5 ±1.5 | 93.19 ±1.28 | 1.422 ±0.059 | 3.5 ±1.0 | 92.22 ±1.28 | 1.302 ±0.059 | 3.5 ±0.8 |
| Method | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | Cov. | AvgSize | CovGap | |
| Marginal | 89.34 ±3.35 | 3.700 ±0.474 | 19.0 ±2.1 | 90.50 ±3.86 | 3.604 ±0.494 | 17.6 ±2.0 | 90.69 ±15.75 | 3.545 ±0.445 | 15.9 ±2.9 | 89.20 ±9.64 | 2.894 ±0.267 | 15.7 ±1.7 |
| Partial | 98.33 ±2.66 | 3.735 ±0.316 | 7.8 ±1.0 | 97.45 ±3.25 | 3.711 ±0.290 | 6.2 ±0.6 | 97.17 ±11.56 | 3.460 ±0.258 | 7.5 ±0.7 | 96.57 ±8.08 | 3.385 ±0.299 | 7.6 ±1.6 |
| Exhaustive | 100.00 ±0.00 | 4.000 ±0.105 | 11.7 ±0.7 | 100.00 ±0.00 | 4.000 ±0.115 | 9.2 ±0.5 | 98.73 ±0.63 | 3.837 ±0.233 | 11.0 ±0.1 | 97.41 ±6.89 | 3.867 ±0.118 | 11.6 ±1.5 |
| RLCP | 91.74 ±3.06 | 2.855 ±0.440 | 8.2 ±3.6 | 92.25 ±3.80 | 2.763 ±0.422 | 7.9 ±3.8 | 91.99 ±10.88 | 2.611 ±0.411 | 6.4 ±3.8 | 91.88 ±10.56 | 2.311 ±0.319 | 8.1 ±2.5 |
| CluCP | 97.80 ±3.14 | 3.936 ±0.199 | 8.7 ±1.5 | 96.60 ±3.29 | 3.882 ±0.180 | 8.1 ±1.9 | 98.34 ±0.50 | 3.985 ±0.110 | 10.0 ±2.6 | 95.91 ±3.59 | 3.257 ±0.220 | 10.7 ±2.6 |