Stochastic Grouping Conformal Prediction for Effective Subgroup Reliability

Meihui Zhong; Wenxin Tai; Ting Zhong; Fan Zhou

cs.LGOct 8, 2026

Stochastic Grouping Conformal Prediction for Effective Subgroup Reliability

Authors: Meihui Zhong, Wenxin Tai, Ting Zhong, Fan Zhou

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

Figure 1: Tradeoff between subgroup coverage and prediction set size on MIMIC-IV Johnson et al. (2024) (

1-\alpha=0.9

). Results are shown for RLCP ( Hore and Barber, 2024 ) , CluCP ( Ding et al., 2023 ) , AFCP ( Zhou and Sesia, 2024 ) , and FaReG ( Xu et al., 2026 ) . The red dashed line marks half of the total number of classes. Several fairness-aware CP methods improve sensitive-group coverage, but often do so with much larger prediction sets, making them less informative in practice.

Figure 2: Overview of SGCP. SGCP learns a local score law by mapping $x$ to a stochastic grouping map $w(x)$ over latent calibration components $\{F_{k}\}_{k=1}^{K}$ . Candidate labels are transformed by this local law and kept when their calibrated scores fall below the conformal threshold.

Method	SYN			Nursery			MIMIC-IV			BACH
Method	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$
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

Table 1: Each method is run 20 times, and we report the mean and standard deviation. All results are reported for the largest sample size. Results with nominal coverage are shown in blue , optimal values are boldfaced , and AvgSize exceeding half of the total number of categories are highlighted in red .

Panel of Figure 3 — Figure 3: Learned latent component densities on four datasets. The components reveal distinct calibration regimes in the nonconformity-score space.

Figure 4: Group coverage on BACH dataset. SGCP remains close to the target coverage across classes while using smaller prediction sets than more conservative baselines.

Figure 5: Distribution of CovGap over repeated runs under varying sample sparsity.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Modality	#Classes	#Attrs	Sensitive / grouping attributes
SYN	tabular	6	3	$A_{1}=$ gender, $A_{2}=$ phenotype, $A_{3}=$ age
Nursery	tabular	4	5	$A_{1}=$ parents, $A_{2}=$ children, $A_{3}=$ finance,
				$A_{4}=$ social, $A_{5}=$ health
MIMIC-IV	tabular	2/3	3	$A_{1}=$ minority, $A_{2}=$ gender, $A_{3}=$ public insurance
BACH	image	4	-	class label (for class-wise reliability analysis)

Appendix

Table A2: Summary of datasets and grouping attributes used for evaluation.

Hyperparameter	SYN	Nursery	M-ICU	M-Hosp	BACH	Description
(i) Latent calibration components
$K$	3	6	3	3	6	number of latent calibration components
$d_{z}$	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

Appendix

Table A3: Dataset specific hyperparameter settings for SGCP.

Method	SYN			Nursery			MIMIC			BACH
Method	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

Appendix

Table A4: Ablation study of SGCP over 20 runs. Results with nominal coverage are shown in blue , optimal values are boldfaced , and AvgSize exceeding half of the total number of categories is highlighted in red . CovGap is reported in percent.

Method	$\mathbf{N=500}$			$\mathbf{N=1000}$			$\mathbf{N=2000}$			$\mathbf{N=4000}$
Method	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$
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

Appendix

Table A5: Performance comparison w.r.t. sample sparsity on the Synthetic Dataset (Target

1-\alpha=0.90

). Takeaway: SGCP remains consistently close to nominal coverage across all sample sizes while achieving the smallest subgroup coverage gap and comparatively small prediction sets. And the performance of each specific sensitive group is shown in the Fig. A6 .

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

Appendix

Table A6: Dataset access and preprocessing sources.

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

Appendix

Table A7: Code availability and baseline implementations.

Figure A6: Group coverage on SYN dataset. SGCP remains close to the target coverage across classes while using smaller prediction sets than more conservative baselines.

Figure A7: Group coverage on Nursery. Top: subgroup coverage vs. average prediction set size at a sample size of 4000, with the dashed line indicating nominal coverage of

0.9

. Bottom: prediction set size varies with subgroup sample size. SGCP stays close to the nominal target while remaining among the most efficient methods across subgroup definitions.

Method	$\mathbf{N=500}$			$\mathbf{N=1000}$			$\mathbf{N=2000}$			$\mathbf{N=4000}$
Method	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$
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

Appendix

Table A8: Performance comparison w.r.t. sample sparsity on the Nursery Dataset (Target

1-\alpha=0.90

). Takeaway: SGCP remains near nominal coverage even in the sparse regime and consistently achieves the smallest subgroup coverage gap, while competing methods improve subgroup reliability only with larger prediction sets or at the cost of under- or over-coverage. And the performance of each specific sensitive group is shown in the Fig. A7 .

Method	$\mathbf{N=2000}$			$\mathbf{N=5000}$			Full (All Samples)
Method	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$
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

Appendix

Table A9: Performance comparison w.r.t. sample sparsity on MIMIC-IV (target

1-\alpha=0.90

). Each method is run 20 times. Takeaway: On the MIMIC-IV dataset, SGCP consistently achieves near-nominal coverage while maintaining the lowest subpopulation disparity (CovGap) in sparse settings (

N=2000,5000

). Even with the full dataset, SGCP provides a highly efficient prediction set with excellent fairness metrics compared to conventional baselines.

Figure A8: Group coverage on MIMIC-IVdataset. SGCP remains close to the target coverage across classes while using smaller prediction sets than more conservative baselines.

Method	$\mathbf{N=50}$			$\mathbf{N=150}$			$\mathbf{N=200}$			$\mathbf{N=300}$
Method	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$	Cov. $\uparrow$	AvgSize $\downarrow$	CovGap $\downarrow$
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

Appendix

Table A10: Performance comparison w.r.t. sample sparsity on BACH (target

1-\alpha=0.90

). Each method is run 20 times. Takeaway: Under extreme sparsity, most methods become conservative, but with the increase in sample size, SGCP recovers the fastest, with near-nominal coverage, a smaller prediction set, and lower subpopulation disparity, outperforming other baseline methods. The performance of each individual subgroup is further shown in Fig. 4 .

Figure A9: Distribution of overall coverage over repeated runs under varying sample sparsity. The dashed line marks the nominal target coverage of

0.9

. Takeaway: SGCP remains consistently close to the target across datasets and sample sizes, whereas global calibration may under-cover in sparse regimes and stronger equalization baselines often become overly conservative.

Figure A10: Distribution of average prediction set size over repeated runs under varying sample sparsity. Smaller values indicate better efficiency. Takeaway: SGCP remains among the most efficient methods across datasets and sample sizes, often matching or improving upon the smallest prediction sets while avoiding the severe conservatism of explicit equalization and localized baselines.

Panel of Figure 21 — Figure A11: Effect of $K$ on learned latent component densities for SYN. Increasing $K$ refines the score law components, but large $K$ introduces overlapping or redundant densities.

Panel of Figure 22 — Figure A12: Effect of $K$ on learned latent component densities for MIMIC-IV. Larger $K$ refines the dominant low score region and tail behavior, but also introduces increasingly overlapping components.

Figure A13: Hyperparameter sensitivity of SGCP on SYN. We vary

\lambda_{\mathrm{nll}}

\lambda_{\mathrm{bal}}

K

, and

M_{\mathrm{eval}}

, and report coverage, average set size, and subgroup coverage gap under different

D_{1}

split ratios. Shaded bands show one standard deviation. Takeaway: SGCP is largely insensitive to these hyperparameters, while larger samples for score-law learning reduce covgap.

Figure A14: Hyperparameter sensitivity of SGCP on MIMIC-IV (ICU delay). Takeaway: SGCP is largely insensitive to these hyperparameters. Coverage remains close to the nominal level, and most of the variation in efficiency and subgroup reliability is driven by the

D_{1}

split ratio rather.

Figure A15: Wall clock time comparison w.r. calibration sample. We report the wall-clock time for each method, averaged over repeated runs; error bars indicate one standard deviation. The horizontal axis is shown on a log scale. Takeaway: SGCP has moderate computational cost and is faster than AFCP, while remaining competitive with other subgroup aware baselines across all sample sizes.

Explore similar work

Self-Organized Conformal Prediction: Reducing Regional Coverage Gaps with Unsupervised Group Discovery
Jun 28, 2026Louis Berthier, Ahmed Shokry, Maxime Moreaud +2Group-Conditional Conformal Prediction Conformal Prediction
Centre de Mathématiques Appliquées, Ecole Polytechnique, Palaiseau, France · Manufacture Française des Pneumatiques Michelin, Clermont-Ferrand, France
Efficient Federated Conformal Prediction with Group-Conditional Guarantee
Mar 15, 2026Haifeng Wen, Osvaldo Simeone, Hong XingGroup-Conditional Conformal Prediction Conformal Prediction
1IoT Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Institute for Intelligent Networked Systems (INSI), Northeastern University London, London, UK · Department of ECE, The Hong Kong University of Science and Technology, HK SAR
Cluster Frequency Conformal Prediction for Local Coverage
May 24, 2026Tomer Lavi, Bracha Shapira, Nadav RappoportConformal Prediction Multiclass Classification
Institute of Applied AI Research (AAIR) Faculty of Computer and Information Science Ben-Gurion University of the Negev, Beer-Sheva 84105, ISRAEL · Institute for Interdisciplinary Computational Science (ICS) Faculty of Computer and Information Science Ben-Gurion University of the Negev, Beer-Sheva 84105, ISRAEL

Jun 28, 2026, stat.MLJ/K move · Enter open · S save

Self-Organized Conformal Prediction: Reducing Regional Coverage Gaps with Unsupervised Group Discovery

Louis Berthier, Ahmed Shokry, Maxime Moreaud, Guillaume Ramelet+1

Centre de Mathématiques Appliquées, Ecole Polytechnique, Palaiseau, France · Manufacture Française des Pneumatiques Michelin, Clermont-Ferrand, France

Group-Conditional Conformal Prediction Conformal Prediction Self-Organizing Maps

Open paper AskPDF arXiv

Conformal prediction guarantees marginal coverage, but a pooled calibration quantile can hide systematic undercoverage across heterogeneous regions of the feature space. We introduce Self-Organized Conformal Prediction (SOCP), a calibration scheme that discovers input-space groups with an unsupervised Self-Organizing Map (SOM) trained without calibration labels. At prediction time, the query's best-matching unit (BMU) draws a calibration buffer from one cell, a fixed grid neighborhood, or a prototype-based enlargement. When fixed neighborhoods are too sparse, Regime 3 adds cells by prototype distance, using a global budget selected from training-cell occupancies and the planned calibration size before any calibration score is observed. The predictor and nonconformity score remain unchanged. Cell-only retrieval has exact cell-conditional validity, and each fixed union of cells has exact retrieved-set validity. Interpreting a neighborhood threshold at its central cell incurs an explicit Kolmogorov-Smirnov (KS) bias term. Across ten regression and classification benchmarks, SOCP reduces the weighted coverage gap relative to pooled split conformal prediction on nine datasets. The mean relative change is

-14.3\%

, at a mean output-size change of

+3.0\%

. Under fixed-neighborhood retrieval, SO composition lowers the ten-seed mean WCovGap in

43

of the

50

dataset-score comparisons, while SO-SCP lowers it on average over paired seeds for every dataset at all three tested external partition granularities. These results provide a concise route to group-local calibration without supervised partitions or predictor retraining with a diagnostic toolkit, while keeping the cost and limits of locality explicit.

Abstract

Explore similar work

Self-Organized Conformal Prediction: Reducing Regional Coverage Gaps with Unsupervised Group Discovery

Efficient Federated Conformal Prediction with Group-Conditional Guarantee

Cluster Frequency Conformal Prediction for Local Coverage