[Paper] Beyond Surface Statistics: Robust Conformal Prediction for LLMs via Internal Representations

Source: arXiv - 2604.16217v1

Overview

Large language models (LLMs) are being used for high‑stakes tasks such as question answering, but the usual confidence signals (token probabilities, entropy, self‑consistency) often break down when the model is deployed on data that differs from its training set. This paper introduces a new way to apply conformal prediction—a statistical technique that guarantees a user‑specified error rate—by tapping into the model’s internal hidden states instead of its surface‑level outputs. The result is a more reliable “confidence interval” for LLM answers, especially under domain shift.

Key Contributions

• Layer‑Wise Information (LI) scores: a novel non‑conformity metric that quantifies how much the model’s internal entropy changes across layers when conditioned on a given input.
• Conformal prediction pipeline built on LI: integrates the LI scores into a standard split‑conformal framework, preserving finite‑sample validity under exchangeability.
• Empirical validation on QA benchmarks: demonstrates superior validity‑efficiency trade‑offs on both closed‑ended (multiple‑choice) and open‑domain question answering tasks, with the biggest gains when test data comes from a different domain than training data.
• Insight into representation‑level uncertainty: shows that hidden‑layer dynamics can be more stable than surface statistics, offering a new angle for robustness research in LLMs.

Methodology

1. Collect internal activations – For each input question, the authors extract hidden representations from every transformer layer of a pre‑trained LLM.
2. Compute layer‑wise entropy – At each layer they treat the representation as a distribution over the vocabulary (via a softmax over the next‑token logits) and calculate the predictive entropy.
3. Derive the LI score – The LI score is the difference between the entropy of the unconditioned model (no input) and the entropy after conditioning on the actual question, aggregated across layers. Intuitively, a large drop means the model’s internal knowledge aligns strongly with the input, indicating higher confidence.
4. Split‑conformal calibration – A held‑out calibration set is used to turn LI scores into quantile thresholds that define prediction sets (e.g., a set of answer candidates) with a user‑specified risk level (e.g., 10 % error).
5. Inference – At test time, the same LI score is computed for a new question, compared against the calibrated threshold, and the corresponding answer set is returned. If the set contains the correct answer, the method is considered valid for that instance.

The pipeline requires no changes to the LLM’s training objective; it only adds a lightweight post‑processing step that reads hidden states.

Results & Findings

• Validity (the proportion of times the true answer lies inside the prediction set) consistently meets the nominal risk level, confirming the conformal guarantee.
• Efficiency (average size of the prediction set) improves markedly, especially under domain shift, meaning developers get tighter confidence bounds without sacrificing reliability.
• Ablation studies show that aggregating entropy across all layers outperforms using only the final layer or a single intermediate layer, underscoring the value of the full depth‑wise view.

Practical Implications

• Safer LLM APIs – Service providers can expose a “confidence set” alongside each answer, letting downstream applications decide whether to accept, ask for clarification, or fall back to a human.
• Dynamic routing – In multi‑model ensembles, the LI score can act as a gating signal to route uncertain queries to a more specialized model or a retrieval system.
• Monitoring & alerting – Because LI scores are derived from internal activations, they can be logged continuously to detect distribution drift in production without re‑training.
• Regulatory compliance – Finite‑sample guarantees satisfy emerging AI‑risk standards (e.g., EU AI Act) that demand quantifiable error bounds for high‑impact deployments.
• Low overhead – The method only needs a forward pass to collect hidden states; no extra fine‑tuning or external calibration data beyond a modest validation split.

Limitations & Future Work

• Exchangeability assumption – Conformal guarantees hold only when calibration and test data are exchangeable; severe covariate shift may still violate this premise.
• Scalability to very large models – Extracting all layer activations for massive LLMs (e.g., >100 B parameters) could increase latency and memory usage; pruning or low‑rank approximations of LI may be needed.
• Task generality – The study focuses on QA; extending LI‑based conformal prediction to generation, summarization, or code synthesis remains an open question.
• Calibration set size – Small calibration sets can lead to noisy quantile estimates; adaptive or online conformal methods could alleviate this.

Overall, the paper opens a promising avenue: leveraging the rich, depth‑wise information inside LLMs to produce statistically sound, practically useful uncertainty estimates. For developers building trustworthy AI services, it offers a concrete tool that bridges the gap between raw model scores and real‑world reliability guarantees.

Authors

• Yanli Wang
• Peng Kuang
• Xiaoyu Han
• Kaidi Xu
• Haohan Wang

Paper Information

• arXiv ID: 2604.16217v1
• Categories: cs.CL, cs.AI
• Published: April 17, 2026
• PDF: Download PDF