Vocabulary-Aware Conformal Prediction (VACP)

• VACP is a framework that enhances next-token prediction by restricting the vocabulary and sharpening probability distributions for efficient uncertainty quantification.
• It employs structural filtering and temperature scaling on a reduced token set, achieving prediction sets up to 200× smaller than standard APS methods.
• Empirical results on benchmarks like SQuAD and WikiText-103 validate its near-nominal coverage and significant efficiency improvements over full-vocabulary approaches.

Vocabulary-Aware Conformal Prediction (VACP) is a framework designed for efficient and reliable uncertainty quantification in next-token prediction with LLMs that possess extremely large vocabularies (e.g., |V| ≈ 256,000). While Adaptive Prediction Sets (APS) provide provable marginal coverage guarantees for classification and next-token tasks, their naive application to LLMs results in prediction sets containing hundreds of tokens, which are inefficient and uninformative. VACP addresses this by systematically reducing the effective prediction space and sharpening uncertainty quantification, while maintaining marginal coverage guarantees comparable to standard APS (Kotla et al., 27 Dec 2025).

1. Formal Description of VACP

Let $V$ denote the full vocabulary of the LLM. Standard APS constructs the prediction set for a new context $x$ as

$$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$

where

$$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$

is the cumulative softmax mass up to token $y$, with $p_V(x)_y$ the softmax probability for token $y$ and $\tau_{1-\alpha}$ the $(1-\alpha)(1+1/n)$ quantile of calibration scores.

VACP modifies this in three principal steps:

1. Vocabulary restriction: Select a subset $V^* \subset V$ such that $x$0.
2. Temperature-adjusted scoring: Apply temperature scaling to sharpen the prediction distribution over $x$1, using $x$2 for model logits $x$3 and temperature $x$4.
3. APS on $x$5: Run APS on $x$6 with scores $x$7 and threshold $x$8.

The resulting prediction set

$x$9

achieves the same marginal coverage as APS on $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$0 provided $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$1.

2. Algorithmic Workflow

VACP is structured as follows:

Step A: Effective Vocabulary Construction

• Structural filtering: Remove control tokens (e.g., <pad>, <eos>, <unused...>), non-printables.
• Empirical filtering: For each $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$2 (remaining), compute $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$3. Remove all $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$4 with $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$5 ($$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$6).
• The surviving set defines $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$7 (e.g., $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$8 for Gemma-2B).

Step B: Temperature Selection

• For each $$C_V(x) = \{ y \in V : s_V(x, y) \leq \tau_{1-\alpha} \}$$9 in a grid $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$0:
  • For $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$1:
  • Mask $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$2 for $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$3.
  • Compute $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$4.
  • Compute $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$5.
  • Compute threshold $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$6 as above; select $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$7 giving coverage $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$8 with smallest mean set size.

Step C: Prediction

• Given new $$s_V(x, y) = \sum_{j: p_V(x)_{(j)} \geq p_V(x)_y} p_V(x)_{(j)}$$9, mask logits for $y$0.
• Compute $y$1.
• Form $y$2 as above.

3. Mathematical Guarantees and Core Formulas

With $y$3, temperature $y$4, and nonconformity score $y$5 defined as above, VACP sets use the threshold

$y$6

Marginal coverage satisfies

$y$7

provided $y$8. If $y$9, coverage is at least $p_V(x)_y$0.

To bound $p_V(x)_y$1, Proposition 2 gives $p_V(x)_y$2, so choosing $p_V(x)_y$3 yields $p_V(x)_y$4; structural filtering yields $p_V(x)_y$5.

4. Theoretical Analysis

Theorem 1 (Marginal Coverage under Vocabulary Restriction):

If $p_V(x)_y$6 satisfies $p_V(x)_y$7, then $p_V(x)_y$8 constructed via APS restricted to $p_V(x)_y$9 ensures $y$0. The proof follows by noting that calibration and prediction are both restricted to $y$1, and thus exchangeability holds under the true distribution, preserving the conformal validity argument.

Proposition 2 asserts that empirical filtering controls the size of $y$2, with the $y$3 bound being worst-case; actual LLM output distributions are even sparser (≈3,200 tokens above $y$4 on average). Exact marginal coverage in practice depends on $y$5, a property observed empirically on benchmarks such as SQuAD and WikiText.

5. Empirical Results and Efficiency

Experiments on the Gemma-2B model with SQuAD and WikiText-103 benchmarks compared standard APS and VACP for next-token prediction at $y$6 (nominal 90% coverage). Key outcomes are summarized below:

Method	Empirical Coverage	Mean $y$7 (tokens)	Efficiency Ratio
Standard APS (full $y$8)	91.2% ± 1.3%	847	1×
VACP (mask + $y$9)	89.7% ± 1.4%	4.3	197×

• Median prediction set size: standard APS median $\tau_{1-\alpha}$0, VACP median $\tau_{1-\alpha}$1.
• Ablations: APS + mask only (coverage 90.8%, mean $\tau_{1-\alpha}$2), APS + temperature only ($\tau_{1-\alpha}$3; coverage 76.4%, mean $\tau_{1-\alpha}$4).
• On WikiText-103 (no recalibration): coverage 88.4%, mean $\tau_{1-\alpha}$5.
• After recalibration: coverage 90.1%, mean $\tau_{1-\alpha}$6.
• By true-token confidence:
  • $\tau_{1-\alpha}$7: coverage 91.3%, mean $\tau_{1-\alpha}$8
  • $\tau_{1-\alpha}$9: coverage 89.4%, mean $(1-\alpha)(1+1/n)$0
  • $(1-\alpha)(1+1/n)$1: coverage 88.1%, mean $(1-\alpha)(1+1/n)$2

These results show VACP achieves near-nominal coverage rates with prediction sets that are approximately 200× smaller than those from uninformed APS.

6. Significance and Robustness

VACP demonstrates that conformal prediction methods can be made practical for LLM next-token prediction, with substantial improvements in efficiency while maintaining formal statistical guarantees under strong conditions ($(1-\alpha)(1+1/n)$3). Empirical filtering and structural masking eliminate tokens with negligible probability under practical workloads, ensuring practical set sizes without loss of relevant support for $(1-\alpha)(1+1/n)$4 in evaluation. Furthermore, VACP transfers effectively to new benchmarks without requiring extensive reparameterization or set recalibration.

The design achieves a balance between computational tractability, interpretability (reasonably small prediction sets), and theoretical rigor in coverage. These aspects support its use in high-stakes LLM deployments requiring calibrated set-based uncertainty quantification (Kotla et al., 27 Dec 2025).

7. Context and Relation to Prior Work

VACP builds on the Adaptive Prediction Sets (APS) framework of Romano et al. (2020), whose direct application in large-vocabulary regimes results in unwieldy prediction sets. By integrating semantic masking and temperature-based distribution sharpening, VACP adapts the conformal prediction toolkit to the unique challenges of LLM architectures. The framework is primarily motivated by deficiencies in softmax calibration for LLMs and the need for uncertainty quantification in open-vocabulary, high-entropy predictive tasks. VACP offers theoretical and algorithmic foundations to extend conformal validity to applications previously impeded by efficiency constraints (Kotla et al., 27 Dec 2025).