Enhancing Conformal Prediction via Class Similarity

AI summary

Overview

Research area: machine learning, specifically conformal prediction for reliable classification. Technical level: Intermediate. Scope: This paper introduces a class-similarity penalty that can be added to any conformal prediction score function to produce smaller prediction sets that span fewer semantic groups.

What This Paper Is About

Conformal prediction (CP) generates prediction sets that contain the true label with a user-specified probability, but standard CP methods ignore whether the labels in a set are semantically related. This can produce sets mixing unrelated classes—for example, a diagnostic set containing both a benign condition and an unrelated serious disease—and can also be unnecessarily large. The paper's goal is to add a lightweight penalty based on class similarity that makes prediction sets both more semantically coherent and smaller, without sacrificing the coverage guarantee.

Key Contributions

1. A model-agnostic penalty, applied on top of any CP score function, that penalizes labels whose assigned group differs from the group of the model's top prediction.
2. A theoretical analysis proving the penalty cannot increase the number of distinct groups in a prediction set, and—counterintuitively—that for common class partitions it also reduces the average prediction set size for any score function.
3. A model-specific variant that derives class similarity from the classifier's own feature embeddings, removing the need for any human-defined class partition.
4. An extensive empirical study across multiple datasets, models, and prominent CP methods (LAC, RAPS, SAPS) showing consistent improvements in both the number of superclasses represented and the average set size.

Main Findings

• Coverage is preserved: Because the penalized score keeps calibration and test samples exchangeable, the standard marginal coverage guarantee of at least 1 − α still holds.
• Out-of-group labels cannot be added: Adding the penalty can only remove group-mismatched labels from the prediction set, never introduce new ones (Proposition 4.2), and the number of distinct groups represented cannot grow (Corollary 4.3).
• Average set size decreases in typical settings: Theorem 4.5 shows the sign of the change in expected set size at small penalty strength depends on the trade-off between p₁ · n̄₀ (probability of a cross-group error times the average group size) and p₀ · n̄₁ (probability of an in-group correct label times the average number of out-of-group classes). In practical settings the second term dominates, so a small penalty shrinks the average set.
• When the improvement can fail: The theorem also identifies extreme regimes—highly dominant groups combined with very weak classifiers—where the penalty could increase set size. The authors report never encountering these in their benchmarks.
• Model-agnostic variant works well: On CIFAR-100 with ResNet-50 and LAC, the model-agnostic version cut the average number of superclasses from 2.27 to 1.85 and the average set size from 3.68 to 3.17.
• Model-specific variant works better on size: Using cosine similarity between centered class means in the deep feature space, the model-specific variant further reduced average set size to 2.92 on the same setup, with comparable superclass counts.
• Consistent gains across CP methods: Similar improvements appear with RAPS (set size 3.83 → 3.17 for the model-specific variant on CIFAR-100/ResNet-50) and SAPS, and on datasets including CIFAR-100, ImageNet, and mini-ImageNet.
• A different trade-off from prior work: The compared hierarchical method (AIR) achieved the fewest superclasses but at a large cost in set size (6.80 versus 3.68 for standard LAC), whereas the proposed methods reduce both metrics simultaneously.

Methodology in Plain English

The starting point is the standard conformal prediction recipe: a model outputs scores for each candidate label, and a calibration set determines a threshold so that the true label is included with probability 1 − α. The authors modify this recipe by adding a term to every candidate label's score, equal to a penalty λ times a measure of how dissimilar that label is from the model's top-1 prediction.

Two ways of measuring dissimilarity are explored. The first uses a known partition of classes into groups—for instance, disease categories that share a treatment or animal families—and assigns a penalty of λ to any label outside the top prediction's group. The second abandons human-defined groups entirely: the researchers feed training data through the network, average the deep features for each class, center those averages, and compute pairwise cosine similarities between them. This yields a soft similarity score for every pair of classes, which plugs into the same penalty mechanism.

Theoretical work analyzes what happens to the expected prediction set size as the penalty strength grows from zero, expressing the answer in terms of how often the true class falls outside the predicted group, how large groups typically are, and how the score distributions behave at the conformal threshold. Experiments then apply both variants to LAC, RAPS, and SAPS across several dataset–model pairs, measuring average set size and average number of superclasses at a fixed coverage level.

Why This Matters

This work shows that post-processing any conformal predictor with a class-similarity penalty is a broadly applicable improvement, not a method-specific trick. It moves beyond the standard efficiency-versus-coverage framing by demonstrating that structural knowledge about labels—whether supplied by humans or extracted from the model itself—can shrink prediction sets rather than inflate them, which contrasts with prior structured, hierarchical, and group-conditional CP approaches.

Real-world applications:

• Medical diagnosis support, where a set of candidate conditions that share a treatment pathway is more actionable than one mixing unrelated diseases.
• Autonomous driving perception, where confusing two similar vehicle types is less consequential than confusing a vehicle with a pedestrian.
• Biodiversity and species classification, where listing several species from the same family is more useful than listing unrelated organisms.
• Content moderation or document tagging, where grouped categories make downstream human review faster.

Industry relevance: deploying CP in safety-critical pipelines typically requires balancing reliability with usability, and smaller, semantically coherent prediction sets reduce the human review burden. The model-specific variant is especially attractive in production because it requires no curated label taxonomy—only access to a trained model and, ideally, its training data.

Future Directions

• Extending the theoretical analysis of Theorem 4.5 from the binary group penalty to the continuous, embedding-based similarity penalty, which the authors explicitly flag as technically challenging.
• Optimizing how the class similarity matrix is constructed; the authors use centered class-mean cosine similarity but note that other choices might perform better.
• Characterizing the full range of penalty strengths λ over which average set size decreases, since the theorem only covers the local behavior near zero.
• Applying the framework to additional structured prediction settings, such as hierarchical label spaces or multi-label tasks, and to CP variants beyond the three evaluated.

Target Audience

Machine learning researchers working on conformal prediction, uncertainty quantification, and reliable classification, along with practitioners deploying CP in safety-critical systems who need prediction sets that are both small and semantically interpretable. Readers should have basic familiarity with classification models and the conformal prediction procedure; the theoretical sections assume comfort with probability and order statistics.

Authors’ abstract

Conformal Prediction (CP) has emerged as a powerful statistical framework for high-stakes classification applications. Instead of predicting a single class, CP generates a prediction set, guaranteed to include the true label with a pre-specified probability. The performance of different CP methods is typically assessed by their average prediction set size. In setups where the classes can be partitioned into semantic groups, e.g., diseases that require similar treatment, users can benefit from prediction sets that are not only small on average, but also contain a small number of semantically different groups. This paper begins by addressing this problem and ultimately offers a widely applicable tool for boosting any CP method on any dataset. First, given a class partition, we propose augmenting the CP score function with a term that penalizes predictions with out-of-group errors. We theoretically analyze this strategy and prove its advantages for group-related metrics. Surprisingly, we show mathematically that, for common class partitions, it can also reduce the average set size of any CP score function. Our analysis reveals the class-similarity factors behind this improvement and motivates a variant that can further reduce prediction set size by leveraging the model's embeddings, without requiring any human semantic partition. Finally, we present an extensive empirical study, encompassing prominent CP methods, multiple models, and several datasets, which demonstrates that our class-similarity-based approach consistently enhances CP methods.