Stochastic Grouping Conformal Prediction for Effective Subgroup Reliability

1 Introduction

• Population-level coverage can average away subgroup under-coverage, leaving clinically important populations with less reliable uncertainty estimates.Differences in prevalence, data quality, and feature distributions can induce distinct error patterns and nonconformity score distributions.Source: §1 · pp. 1–2
• Fine-grained or localized calibration can yield larger prediction sets and, for localization, greater instability across random seeds and data splits.The resulting sets may be less informative and impose excessive cognitive burden on decision makers.Source: §1 · p. 2
• Calibration-relevant heterogeneity may be latent or misaligned with predefined groups, while small subgroup samples can make fine-grained calibration conservative and unstable.The paper frames its goal as a middle ground between global calibration and approaches reliant on known groups or local neighborhoods.Source: §1 · p. 2
• SGCP learns a local score law through stochastic grouping, aiming to improve subgroup reliability without predefined sensitive groups while retaining finite-sample marginal validity.Its subgroup coverage gaps decrease when the learned local score law better captures meaningful subpopulations' score behavior.Source: §1 · p. 2
• Across synthetic and medical datasets, SGCP achieves a favorable tradeoff among overall coverage, subgroup reliability, and prediction-set efficiency.Its stochastic grouping lets samples draw calibration information from others with similar calibration behavior.Source: §1 · pp. 2–3

• Under exchangeability, multiclass conformal prediction uses nonconformity scores to construct prediction sets with marginal coverage at least 1−α.Source: §2 · p. 3
• Distribution-free conditional coverage would address disparities across heterogeneous groups, but finite samples cannot ensure tight, valid sets for arbitrary conditioning events.Source: §2 · p. 3
• Group-wise calibration [17] calibrates separately across groups, but relevant groups may be unknown and fine partitions can destabilize empirical quantiles.Source: §2 · p. 3
• AFCP [29] adaptively selects sensitive attributes or none, but remains tied to discrete attributes and costly, noisy leave-one-out calibration in sparse groups.Source: §2 · p. 3
• FAREG [27] discovers subgroups in latent representations, including nonlinear feature combinations, but is best suited to identifiable subgroups and may be conservative.Source: §2 · p. 3
• RLCP calibrates against neighborhoods containing each test point, avoiding explicit partitions, but worst-neighborhood coverage can inflate sets and feature proximity may miss error structure.Source: §2 · pp. 3–4

3.1 From subgroup reliability to local score laws

• SGCP approximates each sample’s oracle score law with a weighted mixture of latent calibration components, sharing evidence instead of estimating each distribution from scratch.Source: §3.1 · p. 4
• Locality means similarity in calibration-relevant score behavior, not geometric proximity in the input space.Source: §3.1 · p. 5

3.2 Learning a differentiable local score law

• The stochastic grouping map averages component memberships across latent-state posterior draws to produce stable sample-specific weights.Source: §3.2 · p. 5
• Training separates learning the predictor, latent-state posterior, and grouping map from fitting calibration components, then optimizes score uniformity, density fit, and informative assignments.Source: §3.2 · p. 6
• Under split-conformal exchangeability, SGCP retains finite-sample marginal validity regardless of how accurately its local score law is learned.Source: §3.2 · p. 6
• Subgroup coverage is controlled by the mismatch between subgroup and marginal transformed-score distributions, which vanishes in the oracle case.Source: §3.2 · p. 6

4 Numerical Experiments

• The four benchmarks span synthetic cardiopulmonary diagnosis, Nursery admissions, two MIMIC-IV clinical tasks, and BACH histopathology classification.The evaluation compares SGCP with seven baselines using the same backbone and nonconformity score; sensitive attributes are used only for fairness evaluation.Source: §4 · p. 7

4.1 Results and Discussion

• CovGap is 3.7 on SYN, 11.8 on Nursery, and 6.2 on BACH; on MIMIC-IV, SGCP’s 0.8 is close to FaReG’s 0.5.SGCP is the only method that consistently stays close to nominal coverage across benchmarks, with the smallest or nearly smallest AvgSize across datasets.Source: Table 1 · p. 8, §4.1 · p. 8
• Learned latent components have diverse nonconformity-score densities and serve as shared score-law building blocks, not demographic groups.SGCP mixes components into sample-dependent local score laws so samples with similar calibration behavior can share evidence across observed subgroup labels.Source: §4.1 · p. 8, Figure 3 · p. 8
• On BACH, SGCP stays near 0.9 coverage across classes, typically with smaller sets than competitors, and set sizes shrink as subgroup sample size grows.The plots show that gains occur across individual groups, not only in aggregate.Source: §4.1 · p. 8, Figure 4 · p. 9
• Across repeated runs and varying sample sparsity, SGCP is consistently among the methods with the lowest median CovGap and relatively small variation.Its advantage is most pronounced in sparse-data settings; as sample size grows, the gap narrows for most methods while SGCP remains competitive.Source: §4.1 · p. 9, Figure 5 · p. 9

5 Conclusion

• SGCP improves subgroup reliability while maintaining efficient prediction sets, and its theoretical analysis preserves the standard split-conformal marginal coverage guarantee.The conclusion reports experiments across synthetic and real-world benchmarks and identifies personalized medicine as future work.Source: §5 · p. 9

A.1 Proof of Proposition 1

• With the oracle local score law, the percentile transform makes scores uniform within any measurable subgroup, eliminating subgroup reliability inconsistency.This explains why the method models bF(s | x) rather than operating directly on raw scores.Source: §A.1 · p. 12

A.2 Proof of Proposition 2

• Conditional on the training split, exchangeable transformed calibration and test scores yield split-conformal coverage of at least 1 − α.The proof also allows auxiliary randomness when it is independent across calibration and test samples.Source: §A.2 · pp. 12–13

A.3 Proof of Proposition 3

• Accurate local score-law learning provides a sufficient condition for small subgroup coverage error, through its control of residual transformed-score mismatch.When a subgroup’s transformed-score distribution is close to the global distribution, the global conformal cutoff keeps coverage close to nominal.Source: §A.3 · p. 14

A.4 Empirical Subgroup Coverage

• Corollary 1 says that when the learned transform reduces ΔG and subgroup test sample size nG is sufficiently large, empirical subgroup coverage concentrates near nominal 1−α, up to conformal correction 1/(m+1).Source: §A.4 · p. 15

B Model Details

• SGCP learns sample-specific component memberships and combines component score transforms into a local score transform used for conformal calibration.The algorithm applies the learned transform to calibration scores, computes a conformal cutoff, and evaluates transformed scores for candidate labels.Source: B Model Details · p. 16
• SGCP’s latent calibration components share calibration evidence among samples with similar score behavior, approximating local score laws without assuming exact local exchangeability.The approach addresses subgroup heterogeneity while retaining the standard marginal validity guarantee.Source: B Model Details · p. 16

B.1 Training Objective

• Component score laws are trained to map routed scores to comparable percentile scales while fitting the observed true-label nonconformity scores.The objective uses rank-uniformization and local density-fitting terms.Source: §B.1 · p. 17
• Information maximization discourages both component collapse and nearly uniform assignments that carry little sample-specific information.The marginal entropy term encourages component use across the batch, while conditional entropy encourages sharper assignments per sample.Source: §B.1 · pp. 17–18
• The full objective combines rank-uniformization with weighted local density fitting and stochastic-grouping information maximization.λ controls density-fitting strength, and β controls information-maximization strength.Source: §B.1 · pp. 17–18

C Complexity Analysis

• Training overhead is linear in N, up to an optional O(N log N) sorting step for empirical component CDFs.This excludes backbone training and treats epochs, posterior samples, components, latent dimension, classes, and batch size as constants.Source: C Complexity Analysis · p. 18
• Calibration costs O(Ncal log Ncal) with sorting or O(Ncal) with linear-time quantile selection, while test-time cost is O(Ntest).The test-time calculation has constant overhead per sample.Source: C Complexity Analysis · p. 18
• Overall additional complexity is O(N + Ncal log Ncal + Ntest), or O(N + Ncal + Ntest) with linear-time quantile selection.Source: C Complexity Analysis · p. 18

D.1 Additional Benchmark Details

• The four benchmarks cover synthetic cardiopulmonary diagnosis, Nursery, MIMIC-IV clinical prediction, and BACH histopathology, with dataset-specific grouping attributes.Sensitive attributes are used for subgroup construction and fairness evaluation, not as model inputs.Source: §D.1 · pp. 19–20
• Synthetic and Nursery evaluations construct disadvantaged groups through subgroup-dependent uncertainty or noise and downsampling.The synthetic setting uses an intersectional sex–phenotype subgroup; Nursery downsamples a parents-attribute subgroup to 10% and perturbs its labels.Source: §D.1 · p. 19
• MIMIC-IV evaluates in-hospital mortality and ICU-delay prediction across race, gender, and insurance groups, while BACH tests subgroup reliability across disease classes with visual heterogeneity.BACH uses class labels as a proxy grouping variable because demographic sensitive attributes are unavailable.Source: §D.1 · pp. 20–21
• Across synthetic sample sizes, SGCP stays close to nominal coverage while achieving the smallest subgroup coverage gap and comparatively small prediction sets.Figure A6 reports coverage across classes and notes smaller sets than more conservative baselines.Source: Table A5 · p. 21, Figure A6 · p. 22

D.2 Implementation Details

• Experiments use a workstation with an Intel i5-14600KF CPU, 32 GB RAM, and one NVIDIA GeForce RTX 4090 D GPU, with multiple random seeds.The code and external dataset access information are provided in the supplementary material and cited tables.Source: §D.2 · p. 21

D.2.1 Baseline implementation details

• The six baselines include group-conditional calibration, clustered conformal prediction, adaptive attribute selection, representation-based group discovery, and weighted localization.Partial Equalized calibrates one sensitive attribute at a time, whereas Exhaustive Equalized uses intersectional groups; RLCP weights calibration samples by similarity to each test point.Source: §D.2.1 · pp. 21–22

E.1 Ablation Study

• The default model achieves nominal marginal coverage on all datasets and best or tied-best CovGap on three of four, with smallest or near-smallest prediction sets.The full objective therefore balances validity, subgroup reliability, and efficiency.Source: §E.1 · p. 22
• Removing posterior stochasticity increases CovGap on every dataset, while feature-space grouping raises CovGap and set sizes, especially on MIMIC and BACH.The comparison supports using stochastic averaging and calibration similarity based on score behavior rather than raw features.Source: §E.1 · pp. 22–23
• On Nursery, SGCP stays near nominal subgroup coverage and among the most efficient methods; its subgroup coverage gap is smallest across sparse settings.Figure A7 shows subgroup coverage versus average set size at sample size 4000 and set size versus subgroup sample size; Table A8 summarizes sample-sparsity performance.Source: Figure A7 · p. 23, Table A8 · p. 23
• Removing rank-uniformization makes coverage less stable and CovGap larger; removing density fitting worsens coverage–efficiency tradeoffs, and removing mutual information raises CovGap and run-to-run variance.The ablations associate these losses with score normalization, fitting the observed score distribution, and informative stochastic components, respectively.Source: §E.1 · p. 23

E.2.1 Effect of sample sparsity

• Across datasets, SGCP remains near the nominal target with competitive prediction-set sizes and consistently low subgroup disparity, especially in sparse regimes where competing methods become more conservative or variable.On MIMIC-IV (target 1 −α = 0.90), SGCP achieves near-nominal coverage and the lowest CovGap at N = 2000 and 5000. On BACH, under extreme sparsity, SGCP recovers fastest as sample size increases, with near-nominal coverage, a smaller prediction set, and lower subpopulation disparity.Source: §E.2.1 · p. 24, Table A9 · p. 24, Table A10 · p. 25

E.2.2 Effect of the number of latent components

• On SYN and MIMIC-IV, moderate K captures dominant score-law patterns, while larger K adds overlapping or weakly differentiated components.On MIMIC-IV, increasing K refines the low-score region and tail behavior; on SYN, it refines distinct low- and high-score regimes.Source: §E.2.2 · pp. 24–25

E.2.3 Hyperparameter sensitivity

• Across SYN and MIMIC-IV, SGCP is generally insensitive to tested hyperparameters; allocating more data to score-law learning has a more visible effect on reliability and efficiency.The sensitivity results report coverage near nominal across tested settings and reduced CovGap with more score-law learning data.Source: §E.2.3 · p. 25

E.3 Wall-Clock Time

• SGCP’s runtime is moderate and scales smoothly with calibration sample size; it is substantially less costly than AFCP while remaining practical across tested sizes.Marginal CP and Partial CP are fastest; SGCP and RLCP incur additional sample-dependent calibration cost.Source: §E.3 · p. 25

E.4 Limitations

• Experiments use retrospective benchmarks, so prospective validation is needed to assess practical impact in clinical workflows; the framework currently covers multiclass classification, not regression or other settings.Evaluation under deployment-level constraints is also left to future work.Source: §E.4 · p. 25

F Supplement to Related Work

• Fairness-aware methods target coverage parity across specified or adaptively selected groups; SAGCCI shares calibration across clusters of protected groups, unlike SGCP’s sample-dependent latent mixtures.FAREG [27] learns representation-based subgroups, while AFCP [29] selects groups adaptively.Source: F Supplement to Related Work · pp. 26, 28

Supplementary analyses

• SGCP remains among the most efficient methods across datasets and sample sizes, has moderate runtime, and is faster than AFCP.Figure A10 compares prediction-set sizes; Figure A15 compares wall-clock time with a log-scaled horizontal axis.Source: Figure A10 · p. 27, Figure A15 · p. 29
• On SYN and MIMIC-IV, increasing K refines learned score-law components but can create overlapping or redundant densities.Figure A11 shows the SYN pattern; Figure A12 shows refinement of MIMIC-IV’s low-score region and tail behavior.Source: Figure A11 · p. 27, Figure A12 · p. 28

F Supplement to Related Work

• Clustered conformal prediction [8] calibrates label clusters, while Kandinsky Conformal Prediction and learned homogeneous groups extend calibration structures beyond predefined sensitive groups.SGCP instead learns a local score law through stochastic mixing and applies split conformal calibration to transformed scores.Source: F Supplement to Related Work · pp. 27–28

Supplementary analyses

• SGCP is largely insensitive to tested hyperparameters, while more score-law learning data reduces CovGap on SYN; on MIMIC-IV, coverage stays near nominal.Figures A13 and A14 report sensitivity across hyperparameters and D1 split ratios.Source: Figure A13 · p. 28, Figure A14 · p. 29