Conformal Prediction Sets for Next-Token Prediction in Large Language Models: Balancing Coverage Guarantees with Set Efficiency

TL;DR

This paper introduces Vocabulary-Aware Conformal Prediction (VACP), a method that significantly reduces the size of prediction sets in large language models while maintaining coverage guarantees, improving uncertainty quantification in high-stakes applications.

Contribution

We propose VACP, a novel framework combining semantic masking and temperature adjustment to improve set efficiency without sacrificing coverage in large language models.

Findings

VACP achieves 89.7% empirical coverage at 90% target.

Reduces mean prediction set size from 847 to 4.3 tokens.

Demonstrates a 197x improvement in efficiency.

Abstract

Deploying large language models (LLMs) in high-stakes domains requires rigorous uncertainty quantification, yet standard softmax probabilities are often poorly calibrated. We present a systematic study of Adaptive Prediction Sets (APS) applied to next-token prediction in transformer-based models with large vocabularies (greater than 250,000 tokens). Our central contribution is the identification of a coverage-efficiency tradeoff: while naive conformal prediction achieves valid coverage, it produces prediction sets of hundreds of tokens, rendering them uninformative. We propose Vocabulary-Aware Conformal Prediction (VACP), a framework that leverages semantic masking and temperature-adjusted scoring to reduce the effective prediction space while provably maintaining marginal coverage. Experiments on Gemma-2B using SQUAD and WikiText benchmarks demonstrate that VACP achieves 89.7 percent empirical coverage (90 percent target) while reducing the mean prediction set size from 847 tokens to 4.3 tokens -- a 197x improvement in efficiency. We provide a theoretical analysis of vocabulary reduction and release our implementation for reproducibility.