Fair Conformal Classification via Learning Representation-Based Groups

Abstract: Conformal prediction methods provide statistically rigorous marginal coverage guarantees for machine learning models, but such guarantees fail to account for algorithmic biases, thereby undermining fairness and trust. This paper introduces a fair conformal inference framework for classification tasks. The proposed method constructs prediction sets that guarantee conditional coverage on adaptively identified subgroups, which can be implicitly defined through nonlinear feature combinations. By balancing effectiveness and efficiency in producing compact, informative prediction sets and ensuring adaptive equalized coverage across unfairly treated subgroups, our approach paves a practical pathway toward trustworthy machine learning. Extensive experiments on both synthetic and real-world datasets demonstrate the effectiveness of the framework.

• The paper introduces FaReG, a variational information bottleneck method that learns nonlinear, potentially unfair subgroups in latent space and calibrates conformal prediction sets for them.
• FaReG consistently achieves approximately 0.9 conditional coverage on synthetic, Nursery, and ACSIncome experiments, outperforming feature-based baselines while adding only modest prediction-set size.
• The paper proposes WSC⁺, a nonlinear worst-slab coverage metric that detects hidden coverage failures such as XOR and XNOR bias, although interpretability, small-group reliability, and distribution shift remain open challenges.

Motivation and problem setting

Conformal prediction (CP) provides distribution-free, finite-sample marginal coverage guarantees for classification, but marginal coverage can mask systematic under-coverage on specific subgroups — precisely where algorithmic bias manifests. Existing remedies face a fundamental tension: equalized coverage requires conditioning on groups of interest, yet the space of plausible groups is exponential in the number of features. Prior work addresses this only partially. The "Partial" method of Romano et al. conditions on single sensitive features, and AFCP (Zhou & Sesia) adaptively selects at most $k$ sensitive features via a greedy strategy built on Naïve Bayes scoring. Both are limited to groups defined by individual features or simple conjunctions; neither can capture nonlinear feature interactions such as XOR/XNOR relationships, where bias affects, e.g., "White Female or Black Male" while each single-feature slice appears balanced.

The paper proposes FaReG (Fair conformal prediction for Representation-based Groups), which learns unfair subgroups in a latent representation space rather than enumerating raw-feature conditions, together with a new evaluation metric $\text{WSC}^+$ for detecting nonlinearly defined coverage failures.

Method: learning groups via variational information bottleneck

FaReG encodes inputs into a latent $Z=f(X)$ using a stochastic encoder-decoder trained with a deep variational information bottleneck objective:

$$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$

where $S$ is a binary group-membership variable, $X$ is reconstructed through a decoder to preserve interpretability, and the negative term compresses instance-specific information. Variational lower bounds yield a loss combining three terms: an empirical conditional-coverage loss $\mathcal{L}_{CC}$, an MSE reconstruction loss $\mathcal{L}_{MSE}$, and a KL regularization term weighted by $\beta$. Group membership is predicted by a logistic head $q_\phi(s|z)$ over the latent code, and the group is selected by minimizing expected empirical conditional coverage subject to a minimum group-size constraint $\text{WSC}^+$0, solved by projected gradient descent with a closed-form $\text{WSC}^+$1 projection onto the feasible simplex-like set.

A finite-sample proposition bounds the gap between empirical and oracle group coverage by $\text{WSC}^+$2 when the membership classifier has VC-dimension at most $\text{WSC}^+$3. This motivates two design choices: keeping the decoder $\text{WSC}^+$4 low-complexity (logistic regression over latents) and enforcing a minimum group proportion $\text{WSC}^+$5. Both directions are empirically supported: sensitivity analysis shows coverage degrades as $\text{WSC}^+$6 shrinks, while results are insensitive to the sampling count $\text{WSC}^+$7 and the bottleneck weight $\text{WSC}^+$8.

At inference time, FaReG draws $\text{WSC}^+$9 Bernoulli membership vectors from the learned distribution, calibrates a conformal set on each induced subgroup, and outputs the union of these sets together with the standard full-data conformal set. A theorem establishes that this union satisfies adaptive equalized coverage under exchangeability. Notably, the proof argues that AFCP's guarantee implicitly requires its candidate group space to contain the oracle group — a condition violated for nonlinear groups since AFCP's selection rule is essentially a decision stump with VC-dimension 2, whereas FaReG's hypothesis class has VC-dimension $Z=f(X)$0 in model size. This is the paper's central expressiveness argument, though it should be read carefully: larger VC-dimension strengthens representational capacity but weakens the finite-sample estimation bound of Proposition 1, so the theory trades statistical tightness against expressiveness rather than dominating AFCP outright.

Computationally, FaReG runs in $Z=f(X)$1, i.e., linear in sample size, versus $Z=f(X)$2 for AFCP. Empirically, training dominates cost (161.3 s of 162.1 s total at $Z=f(X)$3), with inference overhead negligible.

The WSC⁺ metric

The paper extends the worst-slab coverage (WSC) metric of Cauchois et al. by replacing the linear slab map $Z=f(X)$4 with an arbitrary nonlinear function $Z=f(X)$5, instantiated as quadratic forms $Z=f(X)$6 sampled 1,000 times. A proposition gives a lower bound showing that if $Z=f(X)$7 achieves conditional coverage $Z=f(X)$8, then $Z=f(X)$9, with the $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$0 reflecting the VC-dimension of quadratic classifiers. On synthetic data with an XOR-defined protected group, $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$1 uncovers worse-covered slabs than linear WSC by up to 9.89% at $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$2, with the gap narrowing as $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$3 grows toward marginal coverage, as expected. This makes $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$4 a strictly more sensitive auditing tool for nonlinear bias, usable independently of FaReG itself.

Experimental results

Experiments cover a synthetic mental-illness-diagnosis dataset with an XNOR-defined disadvantaged subgroup (Color=Red ∧ Gender=Female, or Color=Blue ∧ Gender=Male), the UCI Nursery dataset with label noise injected into a conjunction-defined group, an eight-subgroup variant beyond XNOR, and ACSIncome from Folktables with injected bias on age–sex intersections.

Key findings:

• Coverage: On both synthetic and Nursery data, FaReG is the only method that consistently attains the target conditional coverage of 0.9 across all sample sizes, measured both by direct group coverage and by $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$5. Marginal and Partial fail badly on the XNOR group; AFCP1 cannot capture it at all.
• Efficiency trade-off: AFCP2, given the unrealistic prior knowledge that exactly two features matter, eventually exceeds 0.9 coverage but produces substantially larger prediction sets. FaReG reaches comparable or better coverage with near-baseline set sizes (e.g., average size 2.74 vs. 2.68 for AFCP2 on the eight-subgroup setup, with group coverage 0.901 vs. 0.849).
• Feature recovery: Perturbation-based attribution (Beta-VAE style) shows FaReG identifies the two target features more frequently than baselines, with the advantage growing with sample size; reconstruction visualizations confirm the encoder captures the XNOR interaction.
• Composition: Stacking FaReG's group discovery on top of CondCP (Gibbs et al.) raises group coverage from 0.819 to 0.903, indicating the contribution lies in adaptive group identification rather than the underlying calibration rule.
• Scalability: On ACSIncome (20,000 instances), FaReG achieves group coverage 0.898 vs. 0.839 for AFCP2, with modest set-size increase (3.31 vs. 3.17).

These results support the paper's claim that representation-based group discovery resolves the expressiveness–efficiency trade-off that limits feature-enumeration approaches.

Limitations and open questions

The paper concedes that representation-based groups partially sacrifice interpretability relative to groups defined on manifest features; it mitigates this through input reconstruction and perturbation analysis, but the recovered groups remain implicit functions of latents rather than human-readable rules. Several assumptions also bear scrutiny. The coverage theorem inherits AFCP's exchangeability framework and its oracle-group assumption, now satisfied by construction of a richer candidate space but untested under distribution shift. The minimum-group-size constraint $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$6 directly controls the estimation error bound, so very small disadvantaged subgroups may fall outside reliable detection — a scenario common in practice. Additionally, the quadratic instantiation of $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$7 provides only a lower bound on the true nonlinear worst-case coverage, and the choice of $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$8's function class shapes what biases can be audited. Whether the variational objective reliably discovers biased groups without the careful two-stage training schedule used here remains an open question, as does the behavior of the method when multiple, overlapping biased subgroups coexist.

Conclusion

FaReG reframes fair conformal classification as a representation-learning problem: a VIB-based encoder-decoder adaptively identifies unfairly treated subgroups, including those defined by nonlinear feature interactions that defeat single-feature methods, and conformal calibration on the discovered groups yields provable adaptive equalized coverage at linear computational cost. The companion $$\max_\theta\; I(Z,S;\theta) + I(Z,X;\theta) - \beta I(Z,i;\theta),$$9 metric offers a more expressive audit tool for conditional coverage. Experiments on synthetic, Nursery, and Folktables data consistently show valid 0.9 group coverage with minimal efficiency loss relative to baselines. The main open issues are interpretability of learned groups, robustness to small or shifted subpopulations, and the dependence of the guarantees on exchangeability.