Enhancing Conformal Prediction via Class Similarity

Conformal Prediction (CP) has emerged as a powerful statistical framework for reliable classification, which generates a prediction set, guaranteed to include the true label with a pre-specified probability. The performance of CP methods is typically assessed by their average prediction set size. In setups where the classes can be partitioned into semantic groups, e.g., based on shared downstream actions or more interpretable coarse labels, users can benefit from prediction sets that are not only small but also contain a limited number of 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.

When an AI classifies something (like identifying an animal in a photo), instead of giving one answer, Conformal Prediction (CP) gives a set of possible answers that is statistically guaranteed to include the correct one. Smaller sets are better — they're more useful and informative. This paper introduces a smarter way to build those sets by taking advantage of the fact that many categories naturally cluster into groups (e.g., "dogs" and "cats" both belong to "mammals"). We show that penalizing predictions that mix unrelated groups leads to tidier, more coherent prediction sets — and, perhaps surprisingly, also smaller ones on average. Even better, we propose a model-specific variant based on automatic detection of class similarities from the AI model itself, without needing humans to define them, which further reduces the size of the sets. Tests across many datasets and models confirm that the approach consistently improves existing CP methods. In short: smarter grouping of classes helps CP methods give more focused and compact sets of predictions.