Two-stage Risk Control with Application to Ranked Retrieval

1 Introduction

As machine learning models become more integrated into our daily lives, the need for transparency and reliability in their predictions is crucial. Moving beyond black-box approaches, the ability to understand and quantify uncertainty in these models is increasingly important to ensure their effectiveness in real-world applications. Conformal prediction, a distribution-free and statistically valid approach that is straightforward to integrate with existing models, has emerged as a promising solution for quantifying uncertainty in machine learning [Vovk et al., 2005]. Its primary objective is to generate uncertainty sets for model predictions while ensuring a specified coverage level. Recently, the conformal risk control framework [Angelopoulos et al., 2024] has expanded upon traditional miscoverage control by enabling control over the expected value of any loss function. This extension greatly enhances its applicability across a wider range of contexts.

Most existing research on conformal prediction focuses on a single-stage process, where the machine learning system processes the input and generates the prediction in a single step. However, this assumption does not hold for many realworld systems, which often involve two or more concatenated stages. One notable example is ranked retrieval systems, such as search engines, where the task is to retrieve and rank documents based on their relevance to a user's query. These systems typically involve two sequential stages: (1) the retrieval stage, which identifies a set of candidate documents from a large repository, and (2) the ranking stage, which refines and orders these candidates to produce the final ranked list presented to the user [Yin and et al, 2016; Khattab et al., 2020]. This two-stage approach is necessary because the massive volume of documents often exceeds the capacity of a single-stage ranking model, particularly when employing computationally intensive methods. In such twostage problems, each stage is designed with distinct optimization objectives, and errors from one stage can propagate to the next. Consequently, the two-stage process introduces additional complexity, making it more challenging to accurately quantify and control uncertainty.

To address these challenges, we propose two-stage conformal prediction methods to quantify and control the uncertainty inherent in such problems. Specifically, we apply the learn-then-test framework [Angelopoulos et al., 2021a] and extend the recently developed single-stage conformal risk control framework [Angelopoulos et al., 2024] to a two-stage setup, where each stage has its own distinct risk control requirement. Risk control is achieved by identifying parameters that jointly satisfy the risk constraints for both stages. Furthermore, to address the specific purpose of the two stages in ranked retrieval problems, we introduce the retrieval risk and the ranking risk, respectively, and then apply our proposed two-stage risk control methods to derive their corresponding prediction sets while controlling both risks at the prespecified levels. Our proposed methods are model-agnostic

systems.

1.1 Related work

Conformal prediction Conformal prediction, originally developed by Vovk and collaborators, has recently emerged as a prominent method for uncertainty quantification in statistical machine learning [Vovk et al., 1999; Papadopoulos et al., 2002; Vovk et al., 2005; Lei et al., 2015]. A recent survey by Angelopoulos and Bates outlines the significance and wide applications of the topic [Angelopoulos and Bates, 2021]. Our work builds upon the learn-then-test framework [Angelopoulos et al., 2021a] and the recently developed conformal risk control (CRC) framework [Angelopoulos et al., 2024]. The application discussed in our work shares similarities with [Angelopoulos et al., 2023], which uses the LTT technique [Angelopoulos et al., 2021a] to control the false discovery rate in recommender systems and optimize recommendation diversity. However, despite addressing the same ranked retrieval challenges, our work and that of [Angelopoulos et al., 2023] differ in both objectives and methodologies.

Ranked retrieval Ranked retrieval has been extensively studied, with models evolving from traditional IR approaches like BM25 [Baeza-Yates and Ribeiro-Neto, 1999; Stephen and K., 1976] to modern learning-to-rank algorithms [Liu, 2009]. Recent advances in deep learning have also been successfully applied to ranked retrieval [Severyn and Moschitti, 2015; Guo et al., 2016], with methods broadly categorized into pointwise [Crammer and Singer, 2001; Chu and Ghahramani, 2005], pairwise [Burges et al., 2005; Freund et al., 2003], and listwise approaches [Burges et al., 2006; Cao et al., 2007], based on their loss functions. A recent study [Wang and Joachims, 2023] also examines a similar two-stage problem in recommender systems but focuses on group fairness in the first stage, differing from our objective. Another work [Guo et al., 2023] introduces stochastic ranking at inference to ensure utility or fairness in learning-torank models. In contrast, our work addresses a distinct focus.

2 Problem Setup

Formally, in the first stage, consider an i.i.d collection of non-increasing, right-continuous random functions L(1) i : Λ → [0, 1], i = 1, . . . , n + 1, representing the associated losses. We denote by λ the tuning parameter in this stage. In the second stage, we consider another i.i.d collection of random functions, L(2) i : Λ × Γ → [0, 1], i = 1, . . . , n + 1, to represent the associated losses in the second stage, incorporating an additional tuning parameter γ. Here, L(2) i (λ, γ) is assumed to be non-increasing and right continuous in each coordinate. Furthermore, the following conditions hold: L(1) i (0) = 1, L(1) i (1) = 0, L(2) i (0, 0) = 1, and L(2) i (1, 1) = 0. We use R(1) (λ) and R(2) (λ, γ) to denote the expected risk functions at each stage with fixed tuning parameter λ and γ, i.e., R(1) (λ) = EL(1) n+1 (λ) and R(2) (λ, γ) = EL(2) n+1 (λ, γ), and use ̂R(1) n (λ) and ̂R(2) n (λ, γ) to denote the empirical risk functions, i.e., ̂R(1) n (λ) = 1n ∑ni=1 L(1) i (λ) and ̂R(2) n (λ, γ) = 1n ∑ni=1 L(2) i (λ, γ). Without loss of generality, we assume that the parameter λ is chosen from a finite set Λ = {λi : i ∈ [m]}, and the parameter γ is chosen from a finite set Γ = {γi : i ∈ [m]}. We assume the values are ordered such that 0 ≤ λ1 < λ2 < · · · < λm ≤ 1 and 0 ≤ γ1 < γ2 < · · · < γm ≤ 1.

For pre-specified risk levels α1, α2, we aim to utilize random functions {L(1) i (λ)}ni=1 and {L(2) i (λ, γ)}ni=1 to identify data-dependent tuning parameter pairs (λ, γ) that satisfies the expected risk control guarantee

When there exists a set of feasible tuning parameter pairs, R, that satisfy (1), we are also interested in a uniform expected risk control guarantee

as this would allow us to select tuning parameter pair fromR based on certain objective function with valid expected risk control guarantee. Note that the selection of λ and γ encodes the prediction set sizes. Throughout the paper, we assume α1, α2 are fixed to be in the interval [0, 1], which implies feasible risk control.

With regard to two-stage risk control, one might wonder: if we can manage risk effectively in the second stage, why invest effort in controlling it in the first stage? The reason is that controlling risk in the first stage is foundational to the entire process. In the case of the ranked retrieval problem, retrieving all documents in the first stage imposes a significant computational burden on the ranking stage. Conversely, retrieving too few relevant documents in the first stage undermines the feasibility of second-stage risk control and compromises ranking quality. Thus, this paper focuses on simultaneously controlling risks at both stages.

2.1 Data Structure in Ranked Retrieval Problem

Before discussing how to achieve risk control as specified in equations (1) and (2), we first outline the structure of the data for the ranked retrieval problem. Consider a set of i.i.d calibration data points {(Xi, Yi,Zi)}ni=1 , with (Xi, Yi,Zi) ∈ X × Y × Z, where Yi and Zi are the labeled responses of feature vector Xi corresponding to the two stages respectively. In the ranked retrieval problem, Xi represents a query along with its associated candidate documents. We let Xi = {qi, {di,j}Ni j=1 }, where qi denotes the user query and {di,j}Ni j=1 denotes its associated documents with size Ni. Each pair of qi and di,j is associated with a ground truth relevance score ri,j ∈ {0, 1, . . . ,R}, where a higher value of ri.j indicates a higher relevance of di,j to qi. This score is only observable for the training data and is hidden for the test data. Here we denote the relevant documents with a relevance score great than 0 in the retrieval results of qi by Yi = {di,j : ri,j > 0}. In the ranking stage, with a focus on the ranking quality of documents with a ground truth relevance level r0 ∈ [R] or above, the set of r0-relevant documents for qi: {di,j : ri,j ≥ r0} is considered. Then, Zi denotes the ordered set of the r0-relevant documents, sorted in descending order based on the ground truth relevance scores with ties broken arbitrarily: Zi = {di,(1), di,(2), . . .}.

Without loss of generality, we assume that each stage is associated with a model learned on the training data for all queries, denoted by Mretrieval for the retrieval model and by Mrank for the ranking model, respectively. Note that the form of the retrieval model is flexible; it can range from a simple Okapi BM25 model, which counts word occurrences, to a more complex large language model that generates embeddings for embedding-based retrieval. Typically, the model used in the retrieval stage is more efficient but less powerful than the one in the ranking stage. For both stages, by leveraging the pre-trained model and calibration data, we can construct prediction sets ˆ C (1) (x; λ) for the unknown response y ∈ Y and ˆ C (2) (x; λ, γ) for the unknown response z ∈ Z, given a test data point (x, y, z) ∈ X ×Y ×Z. By employing the loss functions l(1) and l(2) for the first and second stages, respectively, we obtain:

We will specify the choice of ˆ C (1) (·; λ), ˆ C (2) (·; λ, γ), l(1) , and l(2) for ranked retrieval problem in Section 4.

3 Two-stage Risk Control

In this section, we present two approaches for achieving expected risk control. The first is a direct application of the LTT framework [Angelopoulos et al., 2021b], included for completeness, as high-probability risk control offers a stronger guarantee than expected risk control. In the second approach, we extend the conformal risk control framework [Angelopoulos et al., 2024] to accommodate the two-stage setting.

3.1 LTT Framework

In the first stage, the value of λ is determined by evaluating its risk function R(1) (λ) through the following hypothesis tests for each λi ∈ Λ:

Given λ = λi ∈ Λ, the value of γ is selected in the second stage by using the corresponding risk function R(2) (λi, γ). For each γj ∈ Γ, the following hypothesis tests are performed:

Let F (1) = {H(1) i : i = 1, . . . ,m} denote the collection of all hypotheses tested in the first stage. Given λ = λi ∈ Λ, let F (2) i = {H(2) i,j : j = 1, . . . ,m} denote the collection of all hypotheses tested in the second stage for a fixed λi. Finally, let F (2) = ⋃mi=1 F (2) i represent the complete collection of all hypotheses tested in the second stage. To conduct hypothesis testing across both stages, we aim to control the global family-wise error rate (FWER), defined as the probability of making at least one Type I error across the two families F (1) and F (2) . For each individual hypothesis H(1) i ∈ F (1) in the first stage and H(2) i,j ∈ F (2) in the second stage, we use the Hoeffding-Bentkus inequality to compute p-values p(1) i and p(2) i,j , as introduced in [Angelopoulos et al., 2021a]. These p-values are valid under their respective null hypotheses and are defined as follows:

Here the function h(a, b) is given by:

To sequentially test the families F (1) and F (2) i for each i ∈ [1,m], one can employ the FWER controlling algorithm to obtain the tuning parameter pairs [Angelopoulos et al., 2021a]. Given that the loss function in the first stage, L(1) i (λ), is non-increasing in λ, and the loss function in the second stage, L(2) i (λ, γ), is non-increasing in both λ and γ, the fixedsequence procedure tests the hypotheses in these families in reverse order. Specifically, for F (1) , the m hypotheses are tested sequentially in the order H(1) m, H(1) m-1 , . . . ,H(1) 1 . Similarly, for F (2) i , the hypotheses are tested sequentially in the order H(2) i,m, H(2) i,m-1 , . . . ,H(2) i,1 . Building on these considerations, we propose the following procedure for simultaneously testing F (1) ∪F (2) , ensuring control of the global FWER at the pre-specified level δ:

Procedure:

1. Test F (1) using the Bonferroni procedure:

Apply the Bonferroni procedure to the p-values p(1) i to simultaneously test F (1) at level δ. Let R(1) denote the index set of the rejected hypotheses:

2. Test F (2)

For each i ∈ R(1) , use the fixed-sequence procedure on the p-values p(2) i,j to simultaneously test F (2) i at level δ/m. Let R(2) i denote the index set of the rejected hypotheses:

• i using the fixed-sequence procedure:

1. Determine the final set of rejected hypothesis pairs: The set R of tuning parameter pairs that correspond to the rejected pairs of hypotheses (H(1) i ,H(2) i,j ) is:

The procedure described above can be viewed as a special case of the sequential graphical approach for multiple testing [Bretz et al., 2009; Angelopoulos et al., 2021a]. Consequently, it strongly controls the global FWER at the prespecified level δ. We note that some other alternative global FWER-controlling procedures can be applied; a few examples are provided in the appendix. However, determining optimal procedure for a specific setting remains an open problem. In the following, we present the expected risk control guarantee for general FWER-controlling procedure.

Theorem 1. Let R denote the collection of tuning parameter pairs returned from a FWER controlling algorithm testing F (1) ∪F (2) at level δ. Then, we have

Therefore, we have

A few comments are in order regarding the comparison with multiple risks in [Angelopoulos et al., 2021a]. The authors therein discuss a similar setting involving multiple risk functions. Specifically, consider the case where the sets of tuning parameters have size m, i.e., |Λ| = |Γ| = m. Their procedure uses the set of all parameter pairs {(λ, γ) : λ ∈ Λ, γ ∈ Γ}, as input for running FWER controlling algorithm. In contrast, our approach leverages the sequential nature of the problem and the monotonicity of the risk functions, which reduces the computational burden and enhances overall effectiveness.

3.2 Expected Risk Control

In this section, we describe methods that leverage conformal risk control framework. To proceed, we consider α1, α2 ∈ ( 1 n+1 , 1] and introduce a few additional notations. For the purpose of risk control in the first stage, define:

Similarly, to ensure the feasibility of risk control in the second stage, define:

Lastly, for any fixed λ ∈ Λ, we define:

When the set in equation (3) is empty, we set ˆ γ (2) 0 (λ) = 1. To ensure risk control in both stages, we set the value of λ as a linear combination of ˆλ(1) 0 ∨ ˆλ(2) 0 and 1:

Here ˆλ(1) (t) is defined to be the ceiling of t(ˆλ(1) 0 ∨ˆλ(2) 0 )+(1-t), restricted to values in Λ, and t ∈ [0, 1] is a tuning parameter. This formulation ensures that ˆλ(1) (t) ≥ ˆλ(1) 0 ∨ ˆλ(2) 0 for any t. Using ˆλ(1) (t), we determine the value of γ by defining

With the above definitions, we are ready to state our results for expected risk control.

Theorem 2. For any t ∈ [0, 1], tuning parameter pair (ˆλ(1) (t), ˆ γ (2) (t)) achieves first-stage risk control at level α1 with a finite-sample guarantee and achieves second-stage risk control at level α2 asymptotically, i.e.,

Since Λ and Γ are finite sets, we further have the following corollary.

Corollary 1. Uniform risk control for both stages can be achieved in the set

i.e.,

Finite-sample second-stage risk control

To ensure finite-sample guarantee for the second stage, a datasplitting approach can be employed. Let calibration data be divided into two non-overlapping parts with index sets I1 and I2. Let α1, α2 ∈ (1/(1 + |I1|), 1]. Define

To proceed, we need to impose the following additional assumption to ensure finite-sample risk control in the second stage:

Assumption 1. There exists a known constant λ0 ≤ 1 such that L(2) i (λ0, 1) ≤ α2 for i ∈ [n+1].

The above Assumption 1 enables feasibility of finite-sample second-stage risk control. Next, for λ ∈ Λ ∩ [λ0, 1] define

We define ˜ γ (2) 0 (λ) = 1 when the set is empty. For t ∈ [0, 1], we then define:

Theorem 3. For any t ∈ [0, 1], tuning parameter pair (˜λ(1) (t), ˜ γ (2) (t)) achieves finite-sample first-stage and secondstage risk control at level α1, α2, respectively, i.e.,

Note that for any t ∈ [0, 1], we have ER(2) (˜λ(1) (t), ˜ γ (2) (t)) ≤ α2. However, there is no guarantee that

as no monotonic relationship exists. Therefore, to obtain a finite-sample uniform risk control in the second stage, we define

Corollary 2. Uniform risk control for both stages can be achieved in the set

i.e.,

We remark that our data-splitting approach, which ensures a uniform finite-sample risk control guarantee, can be regarded as a special case of the method proposed in Section 4.3 of [Angelopoulos et al., 2024].

4 Application to Ranked Retrieval

In this section, we apply our proposed methods to ranked retrieval problem. We begin by introducing the loss functions defined for each stage.

4.1 Loss Function for Retrieval Stage

As defined in the Section 2, Yi is the set of relevant documents with respect to the query qi, i.e., the set of documents with ground truth relevance scores greater than 0. To represent the set of documents fetched byMretrieval used in the retrieval stage, we define the retrieved document set ˆ C (1) (Xi; λ) for the retrieval stage as:

where λ ∈ [0, 1] and Mretrieval(qi, di,j) denotes the model score for the query-document pair (qi, di,j), as computed by the retrieval model Mretrieval. Note that a good retrieved document set ˆ C (1) (Xi; λ) associated with the query qi should aim to cover as many relevant documents in Yi as possible. To measure the miscoverage of Yi by ˆ C (1) (Xi; λ), we define retrieval loss as:

Note that the loss function defined in equation (5) is nonincreasing, right-continuous in λ, and bounded within [0, 1]. This form of loss function quantifies the missed fraction of relevant documents retrieved by model Mretrieval. By choosing an appropriate λ, we aim to control the risk at level α1.

4.2 Loss Function for Ranking Stage

Similarly, in the ranking stage, we first define the prediction set ˆ C (2) (Xi; λ, γ) as:

In contexts like information retrieval, search engines, and recommendation systems, nDCG [J¨arvelin and Kek¨al¨ainen, 2000] is widely used to evaluate the ranking quality of algorithms or systems. Motivated by the intuition from nDCG that a ranked list can be evaluated by rewarding relevance while considering ranking positions, we slightly modify the definition of nDCG to suit our setting, and define:

Notably, when ˆ C (2) (Xi; λ, γ) contains all the r0-relevant documents, DCGmod attains its maximum value. To normalize the DCGmod, we define the modified Ideal Discounted Cumulative Gain (iDCGmod) for query qi as:

Correspondingly, we define the modified Normalized Discounted Cumulative Gain (nDCGmod) as

The loss function for the ranking stage is then defined as

The proposed loss function in equation (7) addresses the ranking order within the ground truth set Zi. When selecting an appropriate γ that satisfies risk control guarantee, greater weight is assigned to documents in Zi with higher ranking positions, ensuring that the most relevant documents are prioritized for inclusion in the prediction set. Note that, given λ specified in the first stage, the ranking loss function is nonincreasing, right-continuous in γ, and bounded within [0, 1]. Our goal is to determine γ to control the ranking loss at specified level α2.

4.3 Parameter Pair Selection via Empirical Set Size Minimization

Given a collection of tuning parameter pairs that achieve risk control at both stages, we determine the tuning parameter pair (ˆλ, ˆ γ) through optimization with a objective function L:

In many applications (recommendation in mobile devices etc.), it is desirable to produce a smaller prediction set in the second stage. Therefore, in this paper, we consider objective function

We comment that, the choice of the loss function should be driven by the specific problem at hand. For instance, in scenarios where the objective is to reduce computational cost during the ranking stage, a smaller prediction set size in the retrieval stage reduces the number of documents evaluated by the computationally intensive ranking model. To address this, the term 1n ∑ni=1 | ˆ C (1) (Xi; λ)| can be incorporated into the loss function L, thereby systematically mitigating the computational burden.

4.4 Experiments

We address the ranked retrieval problem using three approaches: the learn-then-test framework (LTT), two-stage conformal risk control (tCRC), and two-stage conformal risk control with data splitting (tCRC-s 1 ). The tuning parameter pairs (ˆλ, ˆ γ) for these methods are determined by equation (8). Accordingly, these selected tuning parameter pairs ensure expected risk control guarantees, as established in Theorem 1, Corollary 1, and Corollary 2, respectively. Our experiments are conducted on two datasets: MSLR-Web, and Yahoo LTRC. For each dataset and experiment, we split the data into a calibration set and a test set, and then apply different methods to the data. The calibration set is used to determine the tuning parameter pair, while the test set is used for evaluation. Using the test data and the tuning parameter pair computed from the calibration data, we compute the following quantities for evaluation: empirical risks in two stages (Risk (1) represents empirical risk in the first stage, Risk (2) represents empirical risk in the second stage), average prediction set size in the second stage, average recall for documents with a relevance level greater than 2, average recall for documents with a relevance level equal to 1, and average precision for documents with relevance level greater than 1. Then, we replicate each experiment 10 times and report their averages.

In Table 1, we summarize results for the dataset MSLR. From the table, we observe that both Recall (≥ 2) and Recall (1) exhibit relatively high values, indicating that the prediction sets effectively cover a substantial proportion of relevant documents. Additionally, Recall (≥ 2) consistently achieves higher values than Recall (1). This confirms that our proposed ranking loss effectively prioritizes documents with higher relevance levels, resulting in prediction sets that are more likely to include highly relevant documents. In the table, we vary the risk levels and consider (α1, α2) ∈ {(0.1, 0.1), (0.01, 0.1), (0.1, 0.2)}. Notably, LTT exhibits the lowest Risk (1) and Risk (2) values, indicating that LTT is the most conservative approach. This can also be explained by the fact that the size of the feasible set R for LTT is, on average, the smallest. Conversely, we observe that tCRC achieves the smallest prediction set size after tuning parameter selection via equation (8), attributed to its larger feasible set R. Notably, compared with tCRC-s, tCRC maintains a larger set size in the second stage but achieves a smaller Recall (≥ 2). While this seems counterintuitive, it can be explained by that fact that tCRC-s has a larger risk in the first stage. This is achieved by using a smaller ˆλ, which results in fewer documents with higher relevance levels being retrieved in the first stage, ultimately impacting the performance in the second stage. When α1 is reduced from 0.1 to 0.01, the metrics of tCRC and LTT exhibit limited variation, whereas tCRCs demonstrates greater sensitivity to this change. This increased sensitivity can be attributed to the construction of the feasible set by tCRC-s. Additionally, when α2 is increased from 0.1 to 0.2, we observe that Risk (2) slightly exceeds the target risk level of 0.2. This deviation can be explained by approximation error and the asymptotic validity of tCRC, as established in Corollary 1.

1 While the required information λ0 for tCRC-s is unknown in practice, we use the calibration data with index set I1 to obtain an estimate.

Figure 2 illustrates (ˆλ, ˆ γ) and the corresponding risks over two stages, with α1 = 0.1 and α2 varying in the set 0.05, 0.075, 0.1, 0.125, 0.15, 0.175, 0.2. Further visualizations are included in the appendix. All three methods effectively manage risk within the desired levels. Notably, as α2 varies, tCRC achieves the smallest prediction set size in the second stage, LTT attains the highest recall rate, and both tCRC and LTT demonstrate superior precision. In practice, we recommend using LTT or tCRC. When the sample size is small, tCRC-s may suffer from lower sample efficiency, compounded by the challenge of the unknown parameter λ0. Lastly, we present similar results for the Yahoo dataset in Table 2, with the corresponding figures deferred to the appendix.

5 Conclusion &Discussion

In this paper, we study expected risk control under a twostage setup. We propose methods aimed at simultaneously controlling risk, and establish theoretical guarantees for these methods. The effectiveness of our approach is validated through experiments on two large-scale, widely used datasets: MSLR-Web and Yahoo LTRC. Below, we discuss several extensions of the current framework.

Non-monotone risk function When risk function is nonmonotone, the following monotonization procedure can be applied. At the first stage, the loss function is modified as follows:

At the second stage, we define the modified loss function as

With these modifications to the loss functions, the results from the previous sections can be directly applied.

Multiple-stages We extend the problem to the scenario where there are K > 2 stages. Let θ ∈ RK denote the parameter of interest. At the k-th stage, the loss function for sample i is denoted as L(k) i (θ1:k), where θ1:k is shorthand for (θ1, . . . , θk). The sequential nature of the problem determines that loss function at stage k only involves parameters θ1:k. The goal is to determine parameter set θ that satisfies

or alternatively,

When loss functions at different stages are monotonic in each parameter θi with i ∈ [K], one can develop analogous approaches to achieve risk control with high probability and expected risk control.

Contribution Statement

Yunpeng Xu and Mufang Ying contributed equally in this work.