Source-linked AI summary

Conformal Language Modeling

Victor Quach, Adam Fisch, Tal Schuster, Adam Yala, Jae Ho Sohn, Tommi S. Jaakkola, Regina Barzilay

arXiv:2306.10193v2cs.CLcs.LG

TL;DR

Language models generate from an enormous output space, making uncertainty quantification and standard conformal prediction difficult. The paper calibrates sampling, stopping, and rejection to construct statistically valid candidate sets and identifies correct response components, demonstrating risk control with efficient sets across diverse tasks. Its guarantees assume access to a generative model that can be sampled and impose practical limits when the sampling budget is finite.

  • Problem

    Applying conformal prediction to language models is challenging because their output space is unbounded and candidate responses are difficult to explore tractably.

  • Method

    The method calibrates a stopping rule for sampling LM outputs, a rejection rule for removing poor candidates, and component-level conformal prediction using model likelihoods and admission functions.

  • Results

    The approach demonstrates valid risk control and meaningful, efficient output sets across multiple tasks and language models, while also identifying independently correct generation components.

  • Takeaways & Limitations

    Conformal sampling extends statistical coverage guarantees to generative LM outputs and can help users assess both whole responses and correct subcomponents.

  • Takeaways & Limitations

    The procedure assumes a generative model that can produce diverse samples with evaluated likelihoods, and finite kmax restricts the achievable error tolerance range.

Abstract

from arXiv · show

We propose a novel approach to conformal prediction for generative language models (LMs). Standard conformal prediction produces prediction sets -- in place of single predictions -- that have rigorous, statistical performance guarantees. LM responses are typically sampled from the model's predicted distribution over the large, combinatorial output space of natural language. Translating this process to conformal prediction, we calibrate a stopping rule for sampling different outputs from the LM that get added to a growing set of candidates until we are confident that the output set is sufficient. Since some samples may be low-quality, we also simultaneously calibrate and apply a rejection rule for removing candidates from the output set to reduce noise. Similar to conformal prediction, we prove that the sampled set returned by our procedure contains at least one acceptable answer with high probability, while still being empirically precise (i.e., small) on average. Furthermore, within this set of candidate responses, we show that we can also accurately identify subsets of individual components -- such as phrases or sentences -- that are each independently correct (e.g., that are not "hallucinations"), again with statistical guarantees. We demonstrate the promise of our approach on multiple tasks in open-domain question answering, text summarization, and radiology report generation using different LM variants.

1 INTRODUCTION

The paper adapts conformal prediction to generative language models by calibrating sampled output sets rather than enumerating an intractable text space. It also controls rejection and identifies independently correct components, demonstrating valid and precise sets across diverse tasks.

  • LM uncertainty is difficult to quantify because outputs can be incorrect despite often being coherent and contextually relevant.
  • Standard conformal prediction is hard to apply to LMs because their output space is unbounded and candidate responses are only approximately explored by sampling.
  • The method calibrates a stopping rule that samples diverse LM outputs into a growing candidate set until an acceptable response is likely present.
  • A calibrated rejection rule removes redundant, incoherent, or lower-confidence samples while preserving coverage and making output sets smaller.
  • The framework extends conformal prediction to confident components of long generations, such as independently correct report findings.
  • Across diverse tasks and language models, the approach demonstrates valid risk control while retaining efficient, precise output sets compared with baselines.

2 RELATED WORK

The related work situates the paper at the intersection of uncertainty estimation, reliable generation, and conformal risk control for language systems. Existing approaches address confidence, filtering, verification, and reranking, while this work targets systematic control at both generation and component levels.

  • Conformal and risk-control methods have been extended to language tasks including finite-class classification and token-level prediction.
  • Prior uncertainty-estimation work finds that language-model logits can be overconfident even when outputs are wrong.
  • Reliable-generation systems commonly apply post-hoc classifiers, source-grounding checks, or reranking to decoded outputs.
  • This work provides systematic filtering or flagging of poor-quality outputs at both full-generation and component levels, incorporating auxiliary signals such as NLI models.

3 BACKGROUND

The background reviews conformal prediction and Learn Then Test as risk-control frameworks, then identifies exhaustive candidate enumeration as infeasible for language models. The paper therefore casts calibration parameters as settings of a sampling algorithm.

  • Standard conformal prediction uses nonconformity scores to accept or reject candidate labels and construct prediction sets for new inputs.
  • Split conformal prediction provides marginal coverage guarantees when calibration and test examples are exchangeable and the score function is fixed.
  • Learn Then Test selects parameter configurations by testing risk-control hypotheses with super-uniform p-values and family-wise error control.
  • For language models, iterating over every candidate sequence is intractable, so the paper treats λ as a configuration of a sampling algorithm rather than a filter over the output space.

4 CONFORMAL LANGUAGE MODELING

The method generates conformal prediction sets for generative LMs by sampling responses, rejecting low-quality or redundant candidates, and calibrating when to stop. It also calibrates selection of individually confident components within responses, with statistical guarantees for both levels.

  • Conformal sampling with rejection: The procedure samples candidate responses from an LM, rejects candidates that are insufficiently confident or too similar, and stops when set confidence reaches a calibrated threshold.The configuration λ=(λ1,λ2,λ3) controls similarity rejection, quality rejection, and the stopping criterion.
  • Conformal sampling with rejection: The method assumes a sampleable generative model, an admission function for response acceptability, and functions measuring set confidence, similarity, and response quality.The quality function may use model likelihood, while similarity functions help preserve diversity.
  • Calibration with learn then test: Learn-then-test calibration identifies valid threshold configurations and selects one balancing final set size with excess samples while retaining risk control.The selected configuration empirically minimizes a weighted combination of output size and excess sampling.
  • Guarantees: Theorem 4.2 guarantees that the sampled prediction set satisfies the target coverage condition, provided the calibrated configuration is used.The guarantee applies to the prediction set returned by Algorithm 1.
  • Conformal selection of individual components: Component selection splits responses into components and retains those exceeding a calibrated confidence threshold, separately controlling correctness at the component level.The component procedure is calibrated independently and Proposition 4.4 provides its guarantee.
  • Guarantees: The set-level and component-level guarantees hold simultaneously with probability 1−2δ by the union bound.

5 EXPERIMENTAL SETUP

The experiments evaluate conformal sampling across radiology report generation, news summarization, and open-domain question answering, using task-specific language models and admission criteria. They compare likelihood-based and rejection-aware scoring strategies while measuring validity, sampling cost, and final set size.

  • Experimental design: The experiments use kmax=20 across tasks and evaluate both prediction-set validity and efficiency.Efficiency includes excess samples and final output size, reflecting computation and usability.
  • Radiology report generation: Radiology experiments use MIMIC-CXR with a ViT encoder and GPT2-small decoder, judging admission through Clinical Efficacy labels.
  • News summarization: Summarization experiments use CNN/DM with T5-XL and define admissibility using ROUGE-L thresholds of 0.35 for generations and 0.4 for components.
  • Open-domain question answering: Question-answering experiments sample few-shot responses from LLaMA-13B on TriviaQA and require exact normalized matches to annotated answers.
  • Scoring functions: The study uses length-normalized LM likelihood as Q, ROUGE-L as S, and compares FIRST-K, FIRST-K+REJECT, MAX, and SUM set scoring functions.FIRST-K does not use rejection, whereas FIRST-K+REJECT adds duplicate rejection.

6 EXPERIMENTAL RESULTS

The experiments show that conformal sampling achieves the intended validity while likelihood-based scoring improves prediction efficiency across tasks. Rejection and component-level scoring provide additional ways to control output quality and identify reliable response parts.

  • Validity of conformal sampling with rejection: Average set loss often matches but never exceeds the target risk level, confirming empirical validity of conformal sampling.Logit-aware methods remain below the diagonal, indicating validity without conservatism.
  • Prediction efficiency: Likelihood-based approaches outperform the uniform FIRST-K baseline across all three tasks in prediction efficiency.
  • Prediction efficiency: In QA, MAX and SUM achieve expected-set-size AUCs below half the FIRST-K AUC.The comparison concerns Figure 2c and expected set size.
  • Prediction efficiency: For relative excess samples, MAX largely outperforms SUM and FIRST-K, while FIRST-K+REJECT has similar size efficiency but lacks sampling efficiency.
  • Individual components: Component evaluation compares SPAN-LOGITS and an application-specific CLASSIFIER against a RANDOM baseline for selecting confident components.SPAN-LOGITS may underestimate a correct component that follows an incorrect one because it conditions on previous context.

7 CONCLUSION

The paper presents conformal prediction for generative LMs as a way to obtain statistically guaranteed prediction sets despite combinatorial output spaces. It also identifies confident subcomponents and empirically obtains efficient sets in size and total sampling cost.

  • Conclusion: The method calibrates iterative sampling with rejection to produce LM prediction sets with statistical guarantees in combinatorial output spaces.
  • Conclusion: The approach separately identifies confident answer subcomponents, helping assess long responses that may mix correct and incorrect content.
  • Conclusion: Experiments demonstrate efficient prediction sets in both final size and total required samples.

REPRODUCIBILITY STATEMENT

The code release provides implementations, preprocessing, metric computation, plots, tables, and detailed experiment specifications.

  • The codebase includes implementations of Algorithms 1 and 2, preprocessing code, metric functions, and plotting and table-generation utilities.
  • Sections 5 and Appendix F document the datasets, language models, scoring and admission functions, and hyper-parameters used.

A ASSUMPTIONS

The method assumes i.i.d. calibration and test prompts, sampling access to the LM, and an admission function that meaningfully measures response quality. Its applicability also depends on good responses being expressible and sampled often enough, while finite sampling budgets and probabilistic guarantees constrain achievable settings.

  • Data and LM assumptions: Calibration and test prompts are assumed to be independent and identically distributed, excluding dependent multi-turn dialogue and distribution shift.
  • Data and LM assumptions: The LM must support sampling outputs from its conditional distribution, with no other assumptions imposed on the model or sampling process.
  • Data and LM assumptions: The method requires an expressible good response and sufficient probability mass on good responses for tractable sampling.
  • Limitations: Finite kmax can make some error targets unattainable, causing the algorithm to fail to return a risk-controlling configuration.
  • Admission-function assumptions: The admission function must be a good proxy for generation quality, and may be manually designed, user-based, or automatically constructed.
  • Admission-function assumptions: Conservative admission functions can preserve the validity of the selected configuration relative to the true admission function.
  • Limitations: The guarantees rely on probabilistic calibration assumptions and may require careful choices of admission functions, set sizes, computational budgets, and δ and ϵ.

D.2 PROOF OF THEOREM 4.2

The supplied passages describe Pareto Testing as a two-stage procedure that searches configuration trade-offs and then validates promising configurations with fixed-sequence testing. The experimental materials specify datasets, preprocessing, models, and admission functions for radiology report generation, while the TriviaQA prompt uses 32 question-answer pairs.

  • Stage 1: Constructing the Pareto frontier: Stage 1 constructs an approximate Pareto frontier of configurations with trade-offs among constrained objectives and objectives to optimize.
  • Stage 1: Constructing the Pareto frontier: The frontier is ordered by estimated p-values, prioritizing configurations likely to satisfy the coverage constraint.
  • Stage 2: Fixed sequence testing: Stage 2 applies Fixed Sequence Testing and stops at the first configuration satisfying the maximum-p-value criterion; the evaluated non-rejected configurations form Λvalid.
  • Stage 2: Fixed sequence testing: The procedure controls the family-wise error rate at level δ and can identify many valid configurations through the constructed sequence.
  • F.1 Radiology report generation: Radiology experiments use preprocessed MIMIC-CXR data, with training and validation for encoder-decoder training, dev for conformal prediction, and test unused.
  • F.1 Radiology report generation: Candidate radiology reports are generated with a ViT encoder and GPT2-small decoder, then admitted when all 14 CheXbert labels match the reference report.
  • TriviaQA prompt: The TriviaQA dev prompt used for answer generation contains 32 question-answer pairs.

F.4 DATASET DETAILS

The experiments use standard dataset splits with reserved unseen data for calibration, plus additional calibration splits. Evaluation reports validity and efficacy, while component analyses track selected-component counts and recall across α.

  • Standard dataset splits reserve unseen data for calibration experiments, with MIMIC-CXR instead using part of the official training set because validation and test sets are too small.
  • Final evaluation uses a calibration test set across 100 trials, with the remaining 70% measuring validity and efficacy.
  • Component-selection analyses report the number of components in Cinnerγ and its AUC over α.
  • All experiments apply length-normalization to the model logits.

G ADDITIONAL RESULTS

Additional results characterize how effectively the method selects components. Scoring functions based on auxiliary classifiers outperform uncertainty measures based only on model-provided span logits.

  • Expected recall measures the fraction of reference sentences that almost match at least one selected component, averaged across examples.
  • The component-selection analysis reports expected recall in Figure G.2.
  • Auxiliary-classifier scoring functions produce component sets that outperform measures based solely on span logits.

H QUALITATIVE RESULTS

Qualitative examples show the method sampling until candidate sets are acceptable, selecting confident components while rejecting or withholding low-confidence content. They also expose limits from coarse admission labels, collective stopping scores, and model confidence failures.

  • Radiology report generation: In one radiology example, the method stops after three samples and selects several correct sentences while avoiding a low-confidence right-apical-scarring finding.
  • Radiology report generation: CheXbert’s label granularity cannot distinguish left from right, allowing a generation with the wrong side attribution to be treated as matching the label.
  • Radiology report generation: In a harder radiology example, samples capture cardiomegaly and edema but are rejected because they also include effusion, so no components meet the confidence threshold.
  • News summarization: For CNN/DM, component selection highlights main ideas and excludes lower-confidence minor details, while a fully correct sample may require the 19th draw.
  • News summarization: Sampling can stop when the collective score exceeds λ3 = 1.02 rather than when the latest sample has the highest individual score.
  • News summarization: Higher individual scores do not necessarily indicate more acceptable generations, and high confidence can accompany summaries missing main reference ideas.
Loading 2306.10193v2…