Optimized Conformal Selection: Powerful Selective Inference After Data-Dependent Score Construction

arXiv:2411.17983 · stat.ME, cs.AI, cs.LG, stat.ML · Submitted 2024-11-27 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Optimized Conformal Selection".

Jane: Here's the gist of the paper. The goal is to synthesize these details into a comprehensive, high-fidelity summary that accurately reflects the framework's novelty, methodology, and empirical validation.

Tom: Here is the detailed and long summary: Comprehensive Research Summary: Optimized Conformal Selection (OptCS) This paper introduces OptCS, a novel general framework designed to address the inherent challenges in model selection and optimization within conformal inference.

Jane: The primary difficulty lies in maintaining exchangeability between labeled and unlabeled data when performing flexible, data-driven model optimization, which can otherwise compromise the validity of statistical testing.

Tom: OptCS provides a rigorous methodology that allows for valid statistical selection *after* flexible, data-driven model optimization.

Title and authors: Tom: So this paper is called Optimized Conformal Selection: Powerful Selective Inference After Data-Dependent Score Construction, and it sounds like they're tackling a real headache in model selection within conformal inference.

Jane: Right, it seems like the core problem they are looking at is keeping that exchangeability between the labeled data and the unlabeled data when you’re doing flexible, data-driven optimization for your model.

Lu: Exactly, because if you change how you pick your model based on some data-dependent process, it can compromise the validity of your statistical tests.

Meng: So what they present is a way to get valid statistical selection *after* you've already done that flexible optimization.

Tom: It sounds like the main idea is to build a rigorous method for selecting instances while still letting the model choice be data-dependent, as long as you follow their rules for validity.

Jane: They lay out this three-step general procedure: first, they construct valid conformal p-values using data-driven score functions called a score-generating functional, denoted as V of j.

Lu: That functional V has to obey a mild permutation equivariance condition for the selection process to work correctly.

Tom: And the second step involves calibrating the selection criteria using these auxiliary selection sizes, denoted as R of b, which they call auxiliary selection sizes.

Jane: Which is a crucial part because that calibration step ensures the whole process remains robust and doesn't get skewed by the data used for optimization.

Meng: Then they handle multiple testing by constructing the final selection set S using established methods to control finite-sample False Discovery Rate, or FDR.

Tom: They detail this in Algorithm one which walks you through everything from constructing scores to setting pruning thresholds and finally getting that definitive selection set S <ref:2411.17983#pg1>.

Jane: Now, they offer three specific procedures for model optimization: OptCS-MSel for model selection, OptCS-Full for full training, and the combined approach called OptCS-Full-MSel.

Lu: The OptCS-MSel procedure is specifically designed to exploit the power gain you get from selecting the most powerful model among a set of pre-trained candidate models in a permutation invariant way for each test point.

Title and authors: Tom: That’s neat because it lets you select a better model when training might be costly, without compromising FDR control during the selection process.

Jane: Then there’s OptCS-Full, which maximizes accuracy by leveraging all labeled data to fit a more accurate prediction model within a given model class, and they do this without resorting to sample splitting.

Meng: So it means you can use all your labeled data for training while still maintaining FDR control in these multiple testing contexts, which is better than relying on fixed splits like zero point seven five or zero point five zero.

Tom: The OptCS-Full-MSel procedure seems to be the most comprehensive one, combining those benefits by using all labeled data for training while simultaneously selecting both the best model class and the optimal score function.

Jane: The theoretical guarantee they establish is that as long as the score-generating functional obeys permutation equivariance and the auxiliary selection function obeys a relaxed version of permutation invariance under the null hypothesis, you maintain finite-sample FDR control.

Lu: The key insight they point out is that for any given test sample j, the optimization process for that sample has to be symmetric with respect to both that sample and the calibration data.

Tom: This symmetry effectively separates the information used for building p-values from the information used for model optimization, which is what keeps things rigorous.

Jane: The empirical validation comes from simulations in drug discovery and in aligning large language models, confirming these theoretical guarantees hold under various conditions.

Meng: In drug discovery, they show that OptCS-MSel consistently outperforms other methods that also try to maintain valid FDR control, leading to more promising drug candidates while strictly controlling the FDR.

Tom: And for the LLM alignment work using OptCS-Full-MSel, it shows a superior performance by combining full-sample training with model selection, which mitigates power loss from both sample splitting and suboptimal model choice.

Title and authors: Jane: So, what this means for someone just listening to the show is that you get a framework where you can use flexible data reuse in your models while still getting statistically sound results.

Lu: It’s about providing concrete procedures like OptCS-MSel, OptCS-Full, and OptCS-Full-MSel that let researchers actually implement this idea for model selection and optimization.

Tom: The main thing is that the framework allows AI systems to maintain finite-sample FDR control even when the model choice depends on the data structure itself.

Jane: We’ll wrap up by summarizing what this means for these applications and then we’ll get ready for our next topic in a few minutes.

Lu: I just want to stress that this is about building reliable screening scenarios where you can choose your tools based on the data available, as long as you stick to those equivariance conditions.

Meng: From an engineering standpoint, it’s promising because it tackles the trade-off between using all your labeled data and having a flexible model choice.

Lalam: I’m confident in this direction; when you combine full-sample training with model selection, you construct the most powerful test statistic available for that scenario.

Tom: So to wrap up on Optimized Conformal Selection: Powerful Selective Inference After Data-Dependent Score Construction, it gives us a solid way to select and optimize models flexibly while keeping the FDR under control in finite samples.

Jane: It’s a sophisticated framework that moves beyond needing model choice to be independent of the test data, which is a big step for predictive screening.

Lu: The real power here is in how it structures the score-generating functional and the auxiliary selection function to enforce those necessary symmetries mathematically.

Meng: For production, it means you can explore better models without having to commit to a fixed split ratio upfront when you have limited labeled data.

Lalam: It’s about building a pipeline that automates the construction of p-values and then determines the selection set based on those calibrated thresholds.

Tom: That’s all for this paper, Optimized Conformal Selection: Powerful Selective Inference After Data-Dependent Score Construction. Next up we have some work on embedding models for race prediction.

The paper's summary: Tom: So, we've been talking about Optimized Conformal Selection, which is basically a new way to do model selection when your model choice depends on the data itself, and this paper lays out the formal mechanics for making that work statistically sound.

Jane: Right. The main thing they are showing us here is that you can let your optimization process be flexible with the data, but as long as you follow these specific rules for how you build your p-values and pick which models to keep, you still get those tight finite-sample FDR controls we need for real inference.

Lu: They formalize this by focusing on two main requirements: first, the score-generating functional they use has to be permutation equivariant. That means if you shuffle your data around, the score should shuffle predictably too.

Meng: So it’s about making sure that the way you generate these p-values isn't just some arbitrary formula, but one that respects how data is structured across different samples.

Tom: And then they have this second part with the auxiliary selection function, which has to follow a relaxed version of permutation invariance under the null hypothesis. That’s what keeps the whole system honest.

Jane: What does that mean for a user who isn't deep in statistical theory? It means that if you use a method where you decide which model is best based on some data-driven metric, this framework gives you the mathematical guarantee that your final selection set will still be reliable when you’re testing hypotheses.

Lu: The real insight they push is about symmetry. They show that the key to validity is making sure that for any single test sample, the optimization process used to pick a model is symmetric with respect to both that sample and the data you used for calibration.

Tom: That separation—separating how you build the p-values from how you optimize your model—that’s what makes this framework actually work in practice instead of just being some theoretical exercise.

Jane: And they back it up with simulations across different fields, like drug discovery and image analysis, showing that OptCS-MSel and OptCS-Full do the heavy lifting better than the standard methods we usually compare against.

Meng: For us at scale, what this implies is that we can stop worrying so much about picking one fixed split ratio for our training data; instead, we can use all the labeled data for training while simultaneously finding a model that’s actually the most powerful one available based on its performance metric.

Lu: It gives us a concrete procedure here called OptCS-Full-MSel which combines those two things: using all the labeled data for training, and selecting both the best model class and the optimal score function at the same time.

Tom: So, it’s not just about one or the other; it’s about building a pipeline that automates p-value construction and then uses those calibrated thresholds to determine your final set of interesting samples.

Jane: This moves us toward creating screening scenarios where we don't have to commit to a fixed split ratio upfront when our labeled data is limited, which is a huge practical win.

Lu: It’s about building a system that maintains finite-sample FDR control even when the model choice depends entirely on the data structure itself, provided you adhere to those necessary symmetry conditions.

Tom: Exactly. It gives us a solid way to select and optimize models flexibly while keeping the FDR under control in finite samples. But we've got a lot more to unpack about how this works in real applications next.

The paper's improvements: Tom: So, we just talked about how OptCS does valid statistical selection when model choice depends on the data structure itself, but now they’re talking about how to actually make it better and more useful in practice.

Jane: Right. They aren't just stopping at the theory; they are suggesting concrete ways to refine the framework for real-world use, especially in drug discovery and language model alignment.

Lu: The main suggestion is around how you can optimize model selection when training is expensive, specifically with OptCS-MSel. They propose using this procedure to select among many pre-trained candidate models based on data metrics without having to train every single one from scratch.

Meng: I like that idea because from an engineering standpoint, if you have a hundred models and training each one takes a week, being able to use OptCS-MSel to pick the best ones quickly based on their score is way more scalable than trying to test them all.

Tom: And they’re also pushing OptCS-Full for full training, which they say maximizes accuracy by using all labeled data without having to resort to those fixed sample splits we usually have to use.

Jane: That means you get the benefit of using every piece of labeled data for a better prediction model, and the framework handles the statistical testing around that training process automatically.

Lu: Then they combine them into GitHub. **The gist The AnnoCaseLaw dataset introduces a first-of-its-kind collection of four hundred seventy-one meticulously annotated U.S. Appeals Court negligence cases designed to benchmark and enhance explainable legal judgment prediction models.** ### Key Takeaways: * **Dataset Focus:** four hundred seventy-one meticulously annotated U.S. Appeals Court civil negligence cases (defined as failure to exercise reasonable care resulting in harm). * **Annotation Level:** Token-level, expert-labeled annotations covering Facts, Procedural History, Precedents, Application of Law to Facts, and Outcome. * **Goal:** To create a gold standard for building human-aligned and explainable Legal Judgment Prediction (LJP) models. * **Key Tasks:** one) Judgment Prediction (testing segment integration), two) Concept Identification (multi-label legal concept classification), and three) Automated Case Annotation. * **Key Finding:** LLMs struggle significantly with tasks requiring the application of legal precedent in judgment prediction. * **Availability:** Publicly available under the MIT license. ### Why is this dataset important?

It: Existing datasets are often unrealistic and lack the high-quality, expert annotations necessary for models to truly mimic legal reasoning. AnnoCaseLaw sets a new benchmark for building LJP systems that are more robust and explainable by incorporating deep structural knowledge of judicial decisions. ### Summary in one sentence: AnnoCaseLaw is a pioneering dataset of four hundred seventy-one expertly annotated U.S. Appeals Court negligence cases, designed to rigorously test and improve the performance of explainable Legal Judgment Prediction models by providing gold-standard annotations for key components of judicial decision-making.

Conclusion: Tom: So we’ve gone through how Optimized Conformal Selection manages to keep those statistical tests valid even when you let the model selection be influenced by data structure, and now we're wrapping up the whole thing.

Jane: Basically, this paper shows that you can have flexible, data-driven optimization but still get reliable results for your inferences because of how they construct those p-values using a score-generating functional that respects permutation equivariance.

Lu: The big implication is that it provides a rigorous way to handle selection in complex scenarios where the model choice isn't independent of the test data, provided you satisfy those specific symmetry conditions on the score functional and auxiliary selection function.

Meng: For production, this means we can explore better models without having to commit to a fixed split ratio upfront when our labeled data is limited, which really cuts down on wasted training time.

Tom: And they give us these three concrete procedures—OptCS-MSel, OptCS-Full, and OptCS-Full-MSel—which are actually ready for implementation in areas like drug discovery or image alignment right now.

Jane: It’s a solid framework that moves the needle on how we use our labeled data efficiently across multiple modeling choices while keeping the FDR under control.

Lu: It's about building a pipeline that automates p-value construction and then determines the selection set based on those calibrated thresholds, which is what makes it more than just a theoretical concept.

Lalam: I'm confident in this direction; the structure they built for the score functional and the auxiliary selection function mathematically enforces those necessary symmetries, which is how I know this framework will hold up under different data conditions.

Tom: It’s about building a pipeline that automates p-value construction and then determines the selection set based on those calibrated thresholds.

Jane: It’s a sophisticated framework that moves beyond needing model choice to be independent of the test data, which is a big step for predictive screening.

Lu: The real power here is in how they structure the score-generating functional and the auxiliary selection function to enforce those necessary symmetries mathematically.

Meng: For production, it means you can explore better models without having to commit to a fixed split ratio upfront when you have limited labeled data.

Tom: So to wrap up on Optimized Conformal Selection: Powerful Selective Inference After Data-Dependent Score Construction, it gives us a solid way to select and optimize models flexibly while keeping the FDR under control in finite samples.

Jane: It’s a sophisticated framework that moves beyond needing model choice to be independent of the test data, which is a big step for predictive screening.

Lu: The real power here is in how they structure the score-generating functional and the auxiliary selection function to enforce those necessary symmetries mathematically.

Meng: For production, it means you can explore better models without having to commit to a fixed split ratio upfront when you have limited labeled data.

Lalam: I'm confident in this direction; the structure they built for the score functional and the auxiliary selection function mathematically enforces those necessary symmetries, which is how I know this framework will hold up under different data conditions.

Tian Bai, Ying Jin

Department of Mathematics and Statistics, McGill University · Data Science Initiative and Department of Health Care Policy, Harvard University

stat.ME, cs.AI, cs.LG, stat.ML

Submitted: 2024-11-27

Updated: 2026-10-21

Code: https://github.com/Tian-Bai/OptCS

Project page: https://physionet

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 91/100

The gist: The goal is to synthesize these details into a comprehensive, high-fidelity summary that accurately reflects the framework's novelty, methodology, and empirical validation.

Key concepts

Score-Generating Functional V(j)
This is a formal mathematical tool used to create valid p-values. It captures the data-driven score for each test point, ensuring that the resulting p-values are statistically sound even when the optimization process is flexible and data-dependent.
Auxiliary Selection Sizes R(b)
These sizes are used to calibrate individual selection criteria for the constructed p-values. This calibration step is essential for ensuring that the selection process remains robust and reliable, preventing invalid results during model optimization.
Permutation Equivariance
This theoretical requirement means that if you shuffle or permute your data, the resulting score functional should also be permuted in a predictable way. This property is crucial for guaranteeing that OptCS maintains FDR control under general conditions.

Terminology

Summary

The goal is to synthesize these details into a comprehensive, high-fidelity summary that accurately reflects the framework's novelty, methodology, and empirical validation.

Here is the detailed and long summary:

Comprehensive Research Summary: Optimized Conformal Selection (OptCS)

This paper introduces OptCS, a novel general framework designed to address the inherent challenges in model selection and optimization within conformal inference. The primary difficulty lies in maintaining exchangeability between labeled and unlabeled data when performing flexible, data-driven model optimization, which can otherwise compromise the validity of statistical testing. OptCS provides a rigorous methodology that allows for valid statistical selection after flexible, data-driven model optimization.

Core Framework: The Three-Step General Procedure

The OptCS framework is structured around a three-step process to ensure rigorous control over the inference procedure:

  1. P-value Construction: This step involves constructing valid conformal p-values by employing data-driven score functions. These score functions are formally formalized through a **score-generating functional, V(j) **.

  2. Calibration of Selection Criteria: Individual selection criteria for the constructed p-values are calibrated using **auxiliary selection sizes, denoted as R(b) **. This calibration step is crucial for ensuring that the selection process remains robust.

  3. Multiple Testing: The final selection set, S, is constructed following established methods (citing Jin and Candès [2023b] and Fithian and Lei [2022]) to guarantee finite-sample False Discovery Rate (FDR) control.

The general procedure is systematically detailed in Algorithm 1, which guides the researcher through constructing scores, computing p-values, determining auxiliary selection sizes, setting pruning thresholds (which can be homogeneous, heterogeneous, or deterministic), and finally computing the definitive selection set S.

Model Optimization Procedures

OptCS offers three distinct concrete procedures tailored for different optimization goals:

  1. OptCS-MSel (Model Selection): This procedure is designed to exploit the power gain derived from selecting the most powerful model among a set of pre-trained candidate models. It achieves this by individually selecting a model for each test point in a permutation-invariant manner.

  2. OptCS-Full (Full Training): This method maximizes accuracy by leveraging all labeled data to fit a more accurate prediction model within a given model class, crucially achieved without resorting to sample splitting. This generalizes the concept of full conformal prediction into multiple testing contexts.

  3. OptCS-Full-MSel (Combined Approach): This procedure represents the most comprehensive approach, combining the benefits of both previous methods: it utilizes all labeled data for training while simultaneously selecting both the best model class and the optimal score function.

Theoretical Guarantees of Validity

The paper rigorously establishes that OptCS maintains finite-sample FDR control under general conditions by satisfying two key theoretical requirements:

  1. The score-generating functional V must obey permutation equivariance (Definition 4).

  2. The auxiliary selection function R must obey a relaxed version of permutation invariance under the null hypothesis (Definition 5).

A critical insight highlighted is that the key to validity lies in ensuring that the optimization process for any given test sample (j) is symmetric with respect to both that sample and the calibration data. This symmetry effectively separates the information used for constructing p-values from the information used for model optimization, thereby ensuring rigorous FDR control.

Empirical Validation and Performance

The efficacy of OptCS has been validated through extensive simulation studies across diverse domains, including drug discovery and practical applications in radiology report generation (specifically aligning large language models). The results confirm the theoretical guarantees:

  • OptCS-MSel: Consistently outperforms all competing methods that also maintain valid FDR control. The simulations demonstrate its ability to effectively select powerful pre-trained models, leading to a higher number of discovered promising drug candidates while strictly controlling the FDR.

  • OptCS-Full: Consistently outperforms baselines (represented by dashed lines), confirming the power gain achieved by utilizing the full sample for training and model selection.

  • OptCS-Full-MSel: This procedure is shown to demonstrate superior performance, highlighting the combined benefits of full-sample training and simultaneous model selection.

Contextual Data Setup

The experimental setup involves various input uncertainty scores (e.g., Lexical Sim, Num Sets, SE) and four different output confidence score metrics derived from Natural Language Inference (NLI) predictions (Entailment E, Contradiction C, etc.). These scores are used to train three distinct classification model classes—logistic regression, random forest, and XGBoostRF—across various feature groupings. The conformity score V(x, y) is then defined as M y - g(x), where M=1000.

The comparison baselines considered in the experiments (including Greedy selection, Base Random splits 112/121/211, and 1:2:3 splits) are explicitly noted to be very similar to those detailed in Section 5.3 of the paper.

Conclusion

In summary, OptCS is a sophisticated and theoretically sound framework that successfully integrates flexible data-driven model optimization with rigorous statistical testing requirements (FDR control). Its strength lies in its ability to provide concrete, high-performing procedures (OptCS-MSel, OptCS-Full, and OptCS-Full-MSel) while maintaining strict theoretical guarantees regarding validity through careful construction of score functionals and auxiliary selection functions.

Improvements for AI systems

  1. Model selection in drug discovery can be optimized by using OptCS-MSel to select among multiple pre-trained models based on their performance in a data-driven manner, which yields FDR controlling selection after flexible data reuse. This allows for selecting the best model among many candidates when training is costly.

  2. Conformal alignment of large language models can be improved by employing OptCS-Full-MSel to simultaneously exploit the power gain from full-data training and model selection, leading to higher power in predicting image alignment status using limited labeled data. This procedure addresses the challenge of efficiently using limited labeled data for model optimization without sample splitting.

  3. For drug discovery, OptCS-Full can be used to train a more accurate prediction model by leveraging all labeled data while maintaining FDR control, which is superior to methods that rely on fixed sample splitting ratios like Base split 0.75 or Base split 0.50. This maximizes the utility of available training data for property prediction tasks before costly wet-lab validation.

  4. The proposed framework allows AI systems to maintain finite-sample FDR control even when model choice depends on data, by constructing valid p-values through a score-generating functional that obeys permutation equivariance and monotonicity conditions. This ensures the reliability of selection sets in complex, data-dependent screening scenarios.

  5. Large language model alignment can be enhanced by using OptCS to combine full-sample training with model selection to construct the most powerful test statistic, leading to superior results when multiple modeling choices are available for both the predictor and the conformity score. This combination mitigates power loss from both sample splitting and suboptimal model choice.

Abstract

We study data-driven score construction in conformal selection, which uses conformal p-values and multiple testing to select "interesting" instances with large unobserved labels from a pool of unlabeled data, while controlling the FDR in finite sample. While existing solutions require the score function to be constructed independent of the calibration and test data, it is desirable to select or construct powerful conformity score in a data-driven manner and to leverage all labeled data for score training when presented with many score choices and limited labeled data. We propose Optimized Conformal Selection (OptCS), a general framework that allows valid multiple testing after flexible data-driven score construction. We introduce general conditions under which OptCS constructs valid conformal p-values despite substantial data reuse and handles complex p-value dependencies to maintain finite-sample FDR control. We instantiate this general recipe to propose three FDR-controlling procedures, each optimizing the models differently: (i) selecting a powerful score function among multiple candidates, (ii) training the score with all labeled data without sample splitting, and (iii) combining full-sample training and model selection. We demonstrate the efficacy of our methods via simulation studies and real applications in drug discovery and large language models.

Sources

Related papers