Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive Uncertainty Quantification for Generative AI

Jungeum Kim, Sean O'Hagan, Veronika Rockova

TL;DR

The paper addresses predictive uncertainty for black-box models trained on unavailable data by developing Conformal Tree, a locally adaptive conformal prediction framework. It combines self-grouping of conformity scores via robust, add-one-in dyadic regression trees with leaf-wise conformal calibration to produce adaptive prediction bands while preserving a near-nominal coverage. Theoretical results quantify the trade-off via a small loss term $\delta(n,m)$ and show conditional coverage $\ge 1-\alpha-\delta(n,m)$ and marginal coverage near $1-\alpha$, with practical extensions to full conformal and forest variants. Empirically, Conformal Tree yields substantial local tightening on simulated data and two LLM-based tasks (legislator state prediction and skin-disease diagnoses) while maintaining coverage, illustrating its practical impact for reliable, locally aware uncertainty quantification in generative AI contexts.

Abstract

This work is concerned with conformal prediction in contemporary applications (including generative AI) where a black-box model has been trained on data that are not accessible to the user. Mirroring split-conformal inference, we design a wrapper around a black-box algorithm which calibrates conformity scores. This calibration is local and proceeds in two stages by first adaptively partitioning the predictor space into groups and then calibrating sectionally group by group. Adaptive partitioning (self-grouping) is achieved by fitting a robust regression tree to the conformity scores on the calibration set. This new tree variant is designed in such a way that adding a single new observation does not change the tree fit with overwhelmingly large probability. This add-one-in robustness property allows us to conclude a finite sample group-conditional coverage guarantee, a refinement of the marginal guarantee. In addition, unlike traditional split-conformal inference, adaptive splitting and within-group calibration yields adaptive bands which can stretch and shrink locally. We demonstrate benefits of local tightening on several simulated as well as real examples using non-parametric regression. Finally, we consider two contemporary classification applications for obtaining uncertainty quantification around GPT-4o predictions. We conformalize skin disease diagnoses based on self-reported symptoms as well as predicted states of U.S. legislators based on summaries of their ideology. We demonstrate substantial local tightening of the uncertainty sets while attaining similar marginal coverage.

Adaptive Uncertainty Quantification for Generative AI

TL;DR

The paper addresses predictive uncertainty for black-box models trained on unavailable data by developing Conformal Tree, a locally adaptive conformal prediction framework. It combines self-grouping of conformity scores via robust, add-one-in dyadic regression trees with leaf-wise conformal calibration to produce adaptive prediction bands while preserving a near-nominal coverage. Theoretical results quantify the trade-off via a small loss term and show conditional coverage and marginal coverage near , with practical extensions to full conformal and forest variants. Empirically, Conformal Tree yields substantial local tightening on simulated data and two LLM-based tasks (legislator state prediction and skin-disease diagnoses) while maintaining coverage, illustrating its practical impact for reliable, locally aware uncertainty quantification in generative AI contexts.

Abstract

This work is concerned with conformal prediction in contemporary applications (including generative AI) where a black-box model has been trained on data that are not accessible to the user. Mirroring split-conformal inference, we design a wrapper around a black-box algorithm which calibrates conformity scores. This calibration is local and proceeds in two stages by first adaptively partitioning the predictor space into groups and then calibrating sectionally group by group. Adaptive partitioning (self-grouping) is achieved by fitting a robust regression tree to the conformity scores on the calibration set. This new tree variant is designed in such a way that adding a single new observation does not change the tree fit with overwhelmingly large probability. This add-one-in robustness property allows us to conclude a finite sample group-conditional coverage guarantee, a refinement of the marginal guarantee. In addition, unlike traditional split-conformal inference, adaptive splitting and within-group calibration yields adaptive bands which can stretch and shrink locally. We demonstrate benefits of local tightening on several simulated as well as real examples using non-parametric regression. Finally, we consider two contemporary classification applications for obtaining uncertainty quantification around GPT-4o predictions. We conformalize skin disease diagnoses based on self-reported symptoms as well as predicted states of U.S. legislators based on summaries of their ideology. We demonstrate substantial local tightening of the uncertainty sets while attaining similar marginal coverage.
Paper Structure (38 sections, 10 theorems, 57 equations, 12 figures, 6 tables, 3 algorithms)

This paper contains 38 sections, 10 theorems, 57 equations, 12 figures, 6 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Notations
  3. Conformal Tree
  4. Calibration by Self-grouping
  5. Robust Regression Trees
  6. Full-Conformal with Robust Trees
  7. Theoretical Underpinnings
  8. Simulation and Benchmark Studies
  9. Conformal Tree is Adaptive
  10. Benchmark Datasets
  11. Large Language Model Uncertainty Quantification
  12. Classifying States of Legislators using an LLM
  13. Naive Uncertainty Quantification with ChatGPT
  14. Conformalizing ChatGPT Diagnoses of Skin Diseases
  15. Concluding Remarks
  16. ...and 23 more sections

Key Result

Lemma 1

Let $(X_1,Y_1),\ldots,(X_{n+1},Y_{n+1})$ be i.i.d. samples from a joint probability measure $\mathbb{P}$. Let $\widehat{\mathfrak{X}}(\mathcal{D}_{1:j})$ denote the partition resulting from Algorithm alg:robust_tree on data $(X_1,S(X_1,Y_1)),\ldots,$$(X_{j},S(X_{j},Y_{j})))$ for $j\in\{n,n+1\}$. Den Then If we fix $m=c\cdot (n+1)$ for $c\in(0,1)$ such that $c\cdot(n+1)\in\mathbb{N}$, we have a mo

Figures (12)

  • Figure 1: Conformal Tree can adaptively reflect locality without an auxiliary model of quantiles or conformity scores fit on the training dataset. We display the average width, empirical coverage on held-out test data, and proportion better. The interval corresponds to a single random trial among the 10 random trials.
  • Figure 2: Example of conformal sets for a test point with DW-NOMINATE coordinates $(0.647,0.033)$, corresponding to a representative from Indiana. Conformal tree is able to shrink the size of the conformal set for test points in this region of DW-NOMINATE space.
  • Figure 3: (Left) Histogram of conformal set sizes for U.S. state classification from ideology using GPT-4o using Conformal Tree, standard split conformal prediction, and naive UQ for $\alpha=0.2$. (Right) Empirical test coverage of split-conformal prediction and Conformal Tree on heldout test data, for varying $\alpha$ levels and $m=200$. Coverage is on heldout data averaged over ten folds.
  • Figure 4: State-wise average prediction set sizes for each U.S. state. Lower is better. Conformal tree reduces the median state-wise prediction set size by 4.0 when compared to standard split-conformal prediction.
  • Figure 5: Obtaining diagnoses by asking a large language model. Exact list of covariates and prompt used are available in Section \ref{['sec:full_skin_detail']}.
  • ...and 7 more figures

Theorems & Definitions (28)

  • Remark 1
  • Definition 2.1
  • Definition 2.2: Candidate node
  • Definition 2.3
  • Remark 2
  • Remark 3
  • Lemma 1
  • Theorem 3.1
  • Theorem 3.2
  • Remark 4
  • ...and 18 more