Kandinsky Conformal Prediction: Beyond Class- and Covariate-Conditional Coverage

We thank you for your review. Below, we address your questions:

Q1. Most conformal prediction algorithms share a common structure, involving quantile regression followed by constructing prediction sets via comparison with a quantile threshold. However, significant differences emerge in the precise choice of the quantile function class and the selection of samples used for regression—these subtle yet critical choices underpin our algorithmic contributions.

Weight functions jointly dependent on covariates and labels: Although inspired by Gibbs et al. (2023), we generalize their linear quantile regression approach from weight functions defined solely on covariates to those defined jointly on covariates and labels. This generalization is key to achieving broader and practically relevant coverage guarantees (e.g., overlapping and latent groups).

Efficient quantile regression without monotonicity assumptions: While our Alg. 3 directly extends Gibbs et al. (2023), ensuring computational efficiency is non-trivial due to the loss of monotonicity structure utilized in their method. This monotonicity structure no longer holds when we consider groups jointly defined by $X$ and $Y$. To address this, our main method, Alg. 1, extends the computationally efficient quantile regression technique from Jung et al. (2023)—originally limited to covariate-conditional guarantees—to handle weighted coverage guarantees. This algorithm avoids the need to solve a quantile regression problem for all candidates y at test time.

Generalization of Mondrian prediction sets: Alg. 2 further adapts the classic Mondrian conformal prediction approach to our generalized weighted setting.

Q2. The primary challenge arises because both the conformity score and the quantile $\hat{q}$ depend on both $X$ and $Y$. For example, in the analysis of quantile regression we want to bound the number of calibration samples $(X_i, Y_i)$ for which the score equals the quantile. In general, when this number is small, you can get a more accurate quantile estimator. As in prior work, we do this by conditioning on the random variables that fix the quantile’s value. However, in our case, if the score function is deterministic, then fixing the values of $X$ and $Y$ fully determines the score (as opposed to Gibbs et al. (2023) where they only need to condition on $X$). To address this, we show that it suffices to condition on $\Phi (Χ, Y)$, which preserves uncertainty in the score. For non-continuous $S|\Phi$, we can consider any randomized score function $\tilde{S}$ that satisfies the continuity assumption.

Another challenge arises when given a quantile function $\hat{q}$ and the features of a test point $X$, we want to compute the prediction set. The issue is that both the score function and the quantile depend on the value of $X$ and $y$. Inspired by Mondrian conformal prediction, we obtain the prediction sets by returning all the labels y where the score $S(X,y) \leq \hat{q}(X,y)$.

Q3. Areces et al. (2024) prove uniform convergence of the weighted coverage for the entire function class. We sharpen the dependence on $d$ by focusing on the basis weight functions, since the function class is a vector space.

Q4. The alternative test-time algorithm in Sec. 4 achieves a different type of coverage guarantee based on exchangeability arguments. Specifically, this guarantee incorporates randomness over both the calibration dataset and the test point. In contrast, Alg. 1 and 2 provide a training-conditional guarantee, conditioning only on the fixed realization of the calibration data. Although these two types of guarantees are not directly comparable, exchangeability-based guarantees generally yield smaller coverage error bounds, as in our results in Sec. 4.

“the paper [1] considers … weight function.”

Thank you for bringing [1] to our attention. Both this paper and ours address potential distribution shifts between the calibration and the test data, where the distribution over the features and the label changes. The main difference between the two papers is that we get coverage guarantees for a class of distribution shifts, while they consider a single distribution shift that they can estimate using samples from the test distribution.

“particularly odd issues … discussion of the score function.”

In Sec. 3 we mention that “As is common in conformal prediction, our method ensures the desirable coverage for any score function $S$ [that satisfies the continuity assumption]”. For deterministic scores violating the continuity assumption, we mention that “we allow for the use of a randomized score function to break potential ties in quantile regression while keeping the assumptions about distribution $D$ minimal”. In our experiments we use the deterministic Conformalized Quantile Regression score function (Romano et al., 2019), and the randomized Adaptive Prediction Sets (APS) score function (Romano et al., 2019; Ding et al.,2023).

We would like to thank you for your review. We address your question and your comment below.

1. I would like to see more details regarding the implementation of the algorithm (e.g., expanding section A.3), especially as the code was not shared in the current submission.

This links to the code: https://doi.org/10.5281/zenodo.15104271. We will include the code in the publication and we are ready to share further implementation details.

2. Is there any reason why comparison with Gibbs et al., 2023 and Ding et al., 2023 (class conditional conformal prediction) is not included in the experiments? I believe this would help understand the performance improvements offered by the proposed method in the very specific cases introduced in the paper.

We have compared Kandinsky conformal prediction with class conditional conformal prediction with the CivilComments dataset. However, the class conditional method is not implementable for the regression task in ACSIncome. The method of Gibbs et al., 2023 does not scale for the image tasks we consider because it trains the image classifier from scratch to learn the weight functions. This violates the post-hoc nature of conformal prediction. Kandinsky trains the group classifier using the two-dimensional input of logits and labels, which is scalable.

Thank you for your constructive feedback. We address your concerns below.

1. “The paper provides the same flavor of bounds as GCC'24, but for weight functions that depend on both X and Y. Some initial work in this direction was completed by Blot et al (https://arxiv.org/abs/2406.17819), but this paper takes it yet another step further by proving a larger set of results including two-sided bounds.”

Thank you for pointing us to the paper by Blot et al. In our paper, we provide results that are stronger than Blot et al. in two ways. First, Algorithm 3 and the bound in Theorem 4.1 in our paper have a similar flavor to the results in Section 2.2 of Blot et al. However, their result is a one-sided bound, whereas we derive a two-sided bound assuming that the data are i.i.d. and that the distribution of $S(X,Y)|\Phi(X,Y)$ is continuous. Second, in Section 3 we obtain high probability bounds (that are in general stronger than expectation bounds) using Algorithms 1 and 2, under the same assumptions. We will incorporate this citation and discussion into our paper.

2. All the arguments in GCC'24 already hold for randomized score functions. There is no novelty in using a randomized score in the quantile regression.

The use of randomized score functions is not the key contribution of our work. The main reasons why randomized score functions are used in Gibbs et al. (2023) and our paper though are somewhat different. In Gibbs et al. (2023) the authors use randomized score functions to obtain exact coverage. The difficulty when extending quantile regression to work for weight functions of the form $w(X,Y)$ is that conditioning on $X$ and $Y$, the value of a deterministic score function $S$ is fully defined (as opposed to the setting in Gibbs et al. (2023)). Therefore, we add randomization to the score function as one of the ways to overcome the obstacle of $S(X,Y)|X,Y$ not being continuous for a deterministic $S$.

3. Unclear what randomness wCD is integrating over. if C is an argument, then it can be randomized (as in the case of a random score). Then in what sense does the inequality in Theorem 3.1 hold? Almost surely?

In Section 2 we define wCD as an expectation over the test point $(X,Y)$ and the internal randomness of $\mathcal{C}$ (which exists when as you mention we use a randomized score of the test point). Since wCD is by definition integrating over the randomness of $\mathcal{C}$, in Theorem 3.1 the dependence on the randomness of the score of the test point is accounted for in wCD. To explain things further, the guarantee in Theorem 3.1 holds with high probability over the randomness of the calibration dataset and the noise $\{\varepsilon_i\}_{i\in [n]}$ that is used to randomize the scores of the calibration data points.

We would like to thank you for your review. We answer your question below.

1. “... how should we design the basis (\phi) in practice? … how do we systematically figure out what kind of conditioning event we have to include or more generally what kind of distribution shifts we have to consider … How can we design a meaningful basis as a function of X and Y?”

This is an interesting research direction. Towards this direction, we analyze theoretically and evaluate empirically the application of fractional group membership where we obtain the desired coverage for groups defined by unobserved attributes. For this application in Section 3.1 we define $\Phi(X,Y)$ as the probability that the test point $(X,Y)$ belongs to a group $G$ based on the unobserved attributes conditioned on the value of $(X,Y)$. Our theoretical analysis is for the case where we know the true probabilities that define $\Phi(Χ,Υ)$. In practice, we can use the calibration data that includes the value of the protected attribute to estimate the probabilities for the basis $\Phi$. We implement this approach in our experiments in Section 5.

We would like to thank you for your valuable comments. We address your questions and concerns below.

1. “It is unclear why the proposed method performs better…Kandinsky beats other methods.” “Your test examples seem to show… alternatives allround?”

We provide two concrete example scenarios where KCP's generality directly improves coverage:

• Overlapping Groups: We consider overlapping subgroups jointly defined by covariates $X$ and labels $Y$. For example, in our experiments with the ACSIncome dataset (income prediction task), groups are naturally overlapping and jointly defined by subsets of demographic attributes and income levels. In such scenarios, standard Mondrian conformal prediction does not apply, as it requires disjoint groups. Furthermore, the conservative baseline method introduced by Barber et al. (2021)—which assigns the most conservative prediction sets to test points based on all overlapping groups to which they belong—tends to produce excessively large and thus impractical prediction sets. Our experiments also validate this result.

• Latent or Implicit Groups: KCP effectively handles scenarios involving latent groups or implicit subpopulations defined by attributes unavailable at test time. Previous methods that rely exclusively on explicit group annotations are inadequate for such cases. For example, class-conditional conformal prediction uses only labels $Y$ as explicit group identifiers, making it insufficient for capturing latent groups. In contrast, the Kandinsky framework jointly leverages covariates $X$ and labels $Y$ to infer implicit, probabilistic group memberships. Our experiments with the CivilComments dataset precisely demonstrate this scenario: group annotations are unavailable at test time, and our results show that KCP successfully maintains the desired coverage guarantees, whereas baseline approaches like class-conditional CP fail to do so.

Additionally, we want to emphasize that KCP unifies existing methods (class-conditional, covariate-group-conditional, Mondrian) without increasing complication. It reduces general group-conditional coverage to a single quantile regression over an appropriately chosen function class. Our main contribution rigorously shows that this approach achieves stronger conditional coverage guarantees through a practical and simple algorithm.

2. “I would be interested to know what degree of overlap between classes would make Kandinsky better than the alternative. “

In the CivilComments test dataset, the number of groups that a sample belongs to follows the distribution below:

3. Are there applicative situations where Kandinsky CP "breaks down"?

We conducted experiments to systematically analyze scenarios where Kandinsky Conformal Prediction (KCP) provides an advantage. Specifically, we explored the effects of varying two critical parameters: the number of groups and the sample size per group. In the ACSIncome dataset, we fixed the number of samples per group while progressively increasing the total number of groups. In the CivilComments dataset, we varied the total sample size while keeping the number of groups constant.

Our results show that KCP scales robustly with an increasing number of groups, whereas baseline methods deteriorate significantly (Fig 1a). Additionally, KCP’s performance consistently improves with larger sample sizes. However, we observed a scenario where KCP underperforms compared to class-conditional conformal prediction when the sample size per group falls below approximately 250 in the CivilComments dataset (Fig 1d). This occurs because class-conditional prediction assumes binary groups defined solely by labels Y, giving it a larger "effective" sample size per group—up to eight times more than KCP, which handles 16 overlapping groups jointly defined by X and Y. Consequently, class-conditional prediction may initially appear competitive at small sample sizes, but its inherent limitations become clear as sample sizes grow. In summary, KCP demonstrates substantial practical advantages, particularly in settings with many groups and moderate to large sample sizes per group.

We would like to thank you for your feedback. We address your comments below.

1. “Add explicit definition of C(X_{n+1}) in Algorithm 3.”

The definition of $C(X_{n+1})$ is given in the equation of the second line of Algorithm 3, where it includes all $y$ such that the score function on the test point is smaller than a threshold $\hat q_y$. The threshold $\hat q_y$ is defined in the next equation, which is computed by quantile regression on both calibration samples and test samples.

2. “Provide examples or experiments for Fractional group memberships.”

The ACSIncome experiment in the paper is designed with fractional group memberships. The dataset is derived from US Census data collected from different states, which are considered as groups. The state attribute is omitted from the covariates and test point labels, leading to fractional group membership: as defined in L220, we consider fractional group membership with groups defined by unobserved attributes, such that the membership can only be inferred by a probabilistic function over $X$ and $Y$. Given the ground-truth group, the membership is always deterministic. But the membership is fractional for the predictor which observes $X$ but does not observe the group. We propose another example in the introduction (L74). Due to regulatory restrictions, sensitive attributes are sometimes prevented from being used for prediction. Grouping by unobserved sensitive attributes causes fractional membership.

3. “Discuss the technical reason for improving the convergence rate over Roth (2022); Jung et al. (2023), and the sharper dependence on $|G|$ compared with Areces et al. (2024).”

Compared to Roth (2022); Jung et al. (2023):

We establish a connection between the ​subgradient of the empirical quantile loss and the ​empirical coverage guarantee, then prove concentration for the empirical coverage. Roth (2022) and Jung et al. (2023) connect the ​subgradient of the expected quantile loss to the ​expected coverage, so they must prove the concentration for the quantile loss itself. The empirical coverage yields tighter concentration because it is the sum of binary-valued functions. Roth (2022) and Jung et al. (2023) only prove concentration for a ​regularized quantile loss, and the regularization worsens sample complexity.

Compared to Areces et al. (2024):

Areces et al. (2024) prove uniform convergence of the weighted coverage for the whole function class, but we prove uniform convergence only for the basis weight functions since the function class is a vector space. This sharpens the dependence on |G|.

4. “Lack of comparison results with Romano et al. (2020a).”

Romano et al. (2020a) consider a different setting from ours. They assume that the group label for the test point is known, while we address the more general case where the group label is latent. Therefore, their method cannot be directly applied in our experiments.