Active, anytime-valid risk controlling prediction sets

Additional comments on problem setting. For example, we may wish to calibrate the aggressiveness of a robot's behavior policy that is parameterized by $\beta$ (e.g., [5]), where the goal is to reach a destination while avoiding obstacles. The action $A$ may be the control policy of the robot based on the environment, and the risk would be the distance to the nearest obstacle over the entire trajectory of the robot as it travels to its destination. We may want this risk to be controlled on average over the distribution of environments. Our framework applies to such as setup, and hence we term the output based on the covariate $X$ an action $A = f(X)$ rather than limiting it to only being a prediction set.

We also include our top level comment on related work to active learning here:

Comparison to active learning. We briefly summarize the differences and similarities between active learning and our problem. Our problem objective and the methods use to prove their validity are different from typical active learning methods. Active learning (stream-based and pool-based) aims to minimize label queries and still learn the best possible machine learning predictor. Guarantees in this area usually are model-based or learning theoretic, i.e., they propose a model update/selection procedure and query procedure that minimizes the true risk of the model over a known or arbitrary function class, and derive results using a notion of class complexity (when one is developing a procedure that agnostic to the exact function class), or the methods are evaluated empirically for specific models, without guarantees. In contrast, our procedure tunes a calibration parameter that can be wrapped around any black-box model to provide a statistically rigorous risk guarantee. As a result, it means that other types of querying strategies that are deterministic (e.g., disagreement based, diversity based, etc.) cannot be directly imported to our problem setting, since the statistical guarantees we derive require that our queries are probabilistic. Further, we do not think existing active learning necessarily tackle the same objective, since they focus primarily on optimizing the performance of a classifier, rather than guaranteeing risk control while calibrating a parameter. Further development of how to leverage active learning methods in our setting is a fruitful direction for future work.

References

[1] I. Waudby-Smith and A. Ramdas. Estimating means of bounded random variables by betting. Journal of the Royal Statistical Society Series B (Statistical Methodology), 2023.

[2] I. Waudby-Smith, L. Wu, A. Ramdas, N. Karampatziakis, and P. Mineiro. Anytime-valid off-policy inference for contextual bandits. ACM / IMS Journal of Data Science, 2024.

[3] P. Casgrain, M. Larsson, and J. Ziegel. Sequential testing for elicitable functionals via supermartingales. Bernoulli, 2024.

[4] I. Waudby-Smith, P. B. Stark, and A. Ramdas. Rilacs: Risk limiting audits via confidence sequences. International Joint Conference on Electronic Voting, 2021.

[5] J. Lekeufack, A. N. Angelopoulos, A. Bajcsy, M. I. Jordan, and J. Malik. Conformal Decision Theory: Safe Autonomous Decisions from Imperfect Predictions. arXiv:2310.05921, 2024.

[6] D. Cacciarelli and M. Kulahci. Active learning for data streams: a survey. Machine Learning, 2024.

Our responses to the highlighted weaknesses are as follows.

• Novelty of e-processes. Thank you for pointing out the error in the proof of Theorem 1 --- we will correct as you have described. While we agree that betting e-process was introduced in [1] and an inverse propensity weighted form was proposed in [2], we disagree that is simply an application of existing work. Our methods are novel in several ways.
  1. The betting processes of [1, 2] are used to derive confidence sequences for the mean of a bounded (or lower bounded with bounded mean) random variable. In our setting, each betting process does not assume there is a different mean of the random variable tested under the null --- it assumes the same lower bound on the mean, that is $E[r(X, Y, \beta)] \geq \theta$, for each value of $\beta$. Here, our goal is to estimate $\beta^*$ --- we only use the fact that under the null, we can derive a random variable that has a bounded mean, and has a lower bound. Our methods are more similar to those in Lemmas 3.1 and 3.2 in [3]. In that work, however, the authors are only interested in designing confidence sequences for estimation, rather than outputting a parameter with risk control.
  2. The anytime-valid risk control guarantee in Definition 1 is not the same as a merely providing a confidence sequence (CS) for $\beta^\*$. If one only provided a confidence sequence, it would not be clear which $\beta$ one should choose in the CS to ensure that $\rho(\beta) \leq \theta$ --- one of them is $\beta^*$ with high probability, but no guarantees would be made about any of the risk of any other $\beta$, we would not be able to guarantee we can always output a $\hat{\beta}_t$ that has risk control at each time step. Thus, key to showing the risk control guarantee in Definition 1 is also the assumption that $\rho$ is monotonic in $\beta$. The choice of $\beta$ that are safe and provide risk control are the ones that have been eliminated from the CS. In fact, the null of our e-process is the opposite of the risk control we want to achieve, since our resulting CS captures the set of $\beta$ that still might have risk that is at least $\beta$. This is a relatively novel formulation. Creating a CS to capture the opposite of the desired result one would want probabilistic guarantees for (in some sense) has been used in auditing election results with risk-limiting audits [4], but the guarantees and use in that application are quite different than ours.
  3. In [2], there is an assumption that one always has access to an outcome $Y$ for each time step(i.e., the reward), since the application of interest is off policy evaluation, and some action is taken at each time step, and the corresponding reward is obtained. In our setting, we do not receive the outcome $Y$ at all if we did not query a label --- however, the weighted estimator is an unbiased estimator of the actual $r(X, Y, \beta)$ incurred at each time step, so our e-process is still correct.
• Relate work on active learning. See our point on comparing with active learning in the top-level rebuttal.
• Problem setting seems confusing. In line 50, we say that our action is a prediction set for $Y$ in many applications, b/c the risk we wish to control miscoverage. However, the action doesn't need to be a prediction set (e.g., could be the behavior policy of a safe robot [5]). Our usage of $X$ and $Y$ as covariate and label is standard from previous papers on RCPS, but we will make clear what $A$ means and refers to in our final version. We will also clarify the relationship between $\rho$ and $f$ explicitly. The relationship is as follows: we define $\rho(\beta) = E[r(X, Y, \beta)] = E[\ell(Y, f(X, \beta))]$.

Our responses to the questions are as follows.

1. Assumption of monotonicity on $\rho$. The assumption on $\rho$ is slightly more general than assuming $r$ is monotonic directly. We appreciate the suggestion and will clarify that monotonicity of $r$ implies monotonicity of $\rho$ and this is a useful special case of our assumption that is distribution free.
2. How the predictor reduces variance. The e-process in section 2.2 has $\bar{R}(\beta)$ inversely weighted by the probability of labeling when a label is queried. We can see that without the predictor set to always output 0, the variance of $\hat{r}(X, \beta) + L / q(X) * \bar{R}(\beta)$ is larger than $r(X, Y, \beta)$. However, in the extreme, if we had a perfect predictor, i.e., $r^*(X, \beta) = r(X, Y, \beta)$, then we recover exactly $r(X, Y, \beta)$. This has lower variance, and would result in our growth rate being identical to the e-process in Section 2 that queries every label. For more formal analysis, Theorem 3 encapsulates how the accuracy of the predictor results in changes in the growth rate of our e-processes.
3. Relationship to variance maximization in active learning. Our optimal policy aims to sample covariates $X$ with probability proportional to square root of the expected conditional risk squared --- when the optimal predictor is used, this quantity is the same as the conditional standard deviation of the risk. This is similar in spirit to existing active learning algorithms that aim to select data points that have the highest variance of predictions (e.g., [6]). There is a subtle point here that variance in predictions is not precisely the same as the variance in risk (i.e., in a prediction set setting, if the different predictions are covered by the prediction set anyways, then there is no variance in risk). In addition, many active learning algorithms are deterministic, i.e., they pick the points that best fit a criterion(e.g., conditional variance, disagreement, diversity, etc.) rather than sampling them with some probability.
4. Comparison to standard active learning policies. See previous points on related work in active learning and the relationship to variance maximization.

Our response to each highlighted weakness and question are as follows.

• Notational usage and clearer introduction of concepts. We will clean up our notational usage and introduce more unfamiliar concepts more comprehensively (e.g., risk controlling prediction sets, e-processes, anytime-validity, etc.) in our final draft.
• Symbol usage. We will simplify this as well to make it more accessible, including incorporating the suggestions put forth by reviewer XpAy about the notation for the problem setting.
• Goal of method is easy to miss. The goal of our method is specified in Definition 1 and (3) in our paper --- we will highlight this definition and make it more apparent in the final version.
• Experimental section requires reading theory section. We will separate out the methods we are employing from the theoretical section so they can be read in a standalone fashion (i.e., in an algorithm environment). We do have meaningful conclusions concerning our experiments --- the "pretrain" and "learned" label policy/predictor combinations that both utilize the machine learning model being calibrated outperform (in terms of label querying efficiency) the naive baseline labeling policies/predictor combos. This shows that, in practice, the models we are calibrating can be used to estimate the conditional risk and variance to a sufficient degree of accuracy such that it significantly improves the label efficiency over the baseline approaches. We will clarify this emphasis in our final version.
• Limited experiments: We will aim to add more experiments from different machine learning methods and risk functions to illustrate the efficacy of our procedure (i.e., QA task for LLMs, and image labeling using the COCO-MS dataset).
• Practicality of methods. We will provide more concrete examples of applications, but we think our framework is quite widely applicable. Here are some examples of applications:
  • Reduce query cost in medical imaging. A medical imaging system that outputs scores for each pixel of image that determines whether there is a lesion or not would want to utilize labels given by medical experts for unlabeled images from new patients. Since the cost of asking experts to label these images is quite high, one would want to query experts efficiently, and only on data that would be most helpful for reducing the number of highlighted pixels.
  • Domain adaptation for behavior prediction. One reason we would want online calibration in a production setting is that we may have much different distribution of data that we do not have access to before deployment. For example, during a navigation task for a robot, we may want to predict the actions of other agents and avoid colliding into them when travelling between two points [1]. Since agents may behave differently in every environment, it makes sense to collect the behavior data in the test environment and update the behavior prediction in an online fashion to get accurate predictions calibrated for specifically the test environment.
  • Safe outputs for large language models (LLMs). One of the goals with large language models is to ensure their responses are not harmful in some fashion (e.g., factually wrong, toxic, etc.). One can view this as outputting a prediction set for the binary label set of $Y \in \\{`harmful`,\ `not harmful`\\}$. Many pipelines for modern LLMs include some form of a safety classifier, which scores the risk level of an output, and determines whether it should be output to the user or not [2, 3], or a default backup response should be used instead. One would want to label production data acquired from user interaction with the LLM and used to calibrate cutoff for the scores that are considered low enough for the response to be allowed through.
• Choice of $\theta$ and $\alpha$. The choice of $\theta$ and $\alpha$ will depend on the application the user has in mind. Though, since $\alpha$ is determines the probability the bound holds, reasonable default is to choose it $\alpha=0.05$. We would like to reiterate that a particular utility of our method is that the probability bound holds uniformly over every time step, that is the probability that there is a single $\hat{\beta}_t$ for all $t \in \mathbb{N}$ that has risk greater than $\theta$ is less than $\alpha = 0.05$. $\theta$ should be chosen based on the risk metric being controlled, and a reasonable default choice could also to be choose $\theta=0.05$. In the context of image classification, this would mean that the prediction set of possible classes will not cover the true label only 5% of the time on average. We will elucidate this point more in the final version.

References

[1] J. Lekeufack, A. N. Angelopoulos, A. Bajcsy, M. I. Jordan, and J. Malik. Conformal Decision Theory: Safe Autonomous Decisions from Imperfect Predictions. arXiv:2310.05921, 2024.

[2] T. Markov, C. Zhang, S. Agarwal, F. E. Nekoul, T. Lee, S. Adler, A. Jiang, and L. Weng. A Holistic Approach to Undesired Content Detection in the Real World. AAAI, 2023

[3] L. Hanu and Unitary team. Detoxify. Github. https://github.com/unitaryai/detoxify, 2020.

Here are our point-by-point responses to each of the highlighted weaknesses.

1. More examples of the method being used in practice, along with examples of choice of estimators for the labeling policy and control variate regression. We agree that providing more examples would illustrate the utility of our method better, and will add more concrete examples to our final draft. In particular, we note that often machine learning models have some estimate $\hat{P}(X)$ of the conditional distribution of $Y \mid X$ (e.g, class probabilities, conditional diffusion models, LLMs, etc.). Thus, for any realized covariate $x$, we can derive use $\mathbb{E}_{Y \sim \hat{P}(x)}[r(X, Y, \beta) \mid X = x]$ from the machine learning model as our choice of $\widehat{r}(x)$. This expectation can either be calculated analytically (as we do in or classification examples in our experiments) or derived using Monte Carlo (for generative models/LLMs where one can sample from the conditional distribution). In essence, we are already getting a predictor (for free in some sense) from the very model we are calibrating.
2. More extensive experiments and empirical analysis of regret bound. We will aim to add more experiments from different machine learning methods and risk functions to illustrate the efficacy of our procedure (i.e., QA task for LLMs, and image labeling using the COCO-MS dataset), as well as an empirical analysis of our regret bound in our experiments.
3. Experiments with estimators of varying rates of convergence. We agree that more experiments in easier or harder settings could illustrate how different rates of convergence affect our method. However, in practice, one would not train a model from scratch on the labeled data, but use the labeled data to fine-tune existing pretrained models. In fact, our framework is able to learn a labeling policy or predictor using outputs of the pretrained machine learning model we are calibrating, making the learning task much easier. The regret bound elucidates how the error of the label policy and predictor propagates into the growth rate. We do agree that the efficacy of these pretrained models would be different for different tasks, and we will be sure to include empirical analysis of how the accuracy pretrained (or learned) label policies and predictors affect the accuracy of our $\hat{\beta}$ estimate.

Our responses to your questions are as follows.

1. Does the convergence rate of the estimation terms largely dominate the $O(\log(T))$ rate from the $G^\beta$-contribution? In terms of asymptotic rates, we agree with your point that the $O(\log(T))$ regret incurred by online Newton step (or other online learning algorithms) will be dominated by a typical estimation rate that scales with $O(1 / \sqrt{T})$ if we are training the label policy and predictor completely from scratch. When our pretrained machine learning models are already accurate, however, we can have fast rates. An example of this is if we assume that the optimal policy/predictor lies in a class of functions that is formed by a linear combination of multiple existing pretrained models we have access to (or we are only interested in comparing against the best linear combination). In that case, the estimation error is the regret of an online linear regression algorithm, and it will have $O(\log(T) / T)$ convergence (although we would be directly using the $G^\beta$ regret result, rather than the estimation error result). So our convergence rate does depend on the assumption we make about the accuracy of our pretrained models.
2. Obtaining an adaptive rate. Our result is already instance-adaptive in the sense that the optimal $G^\beta$ is adaptive to the $X$-conditional variance of the risk, i.e., it is lower bounded by a term that increases as $\mathbb{E}[\sigma_\beta(X)]$ decreases. It is possible that when $\mathbb{E}[\sigma_\beta(X)]$ is low, the optimal predictor is estimated more quickly, and we believe exploration of adaptive convergence rates of different estimators would be an interesting problem to solve in future work. Our focus in this work is to show that the $G^\beta$ regret depends on estimator error of the policy and predictor in a direct way, and hence having pretrained models that can approximate those well will also improve the efficiency of estimating $\beta^*$.