Semantic Self-Distillation for Language Model Uncertainty

1 INTRODUCTION

• In chained agentic tasks, unreliable outputs can compound over long horizons, making low-latency uncertainty estimates important.The introduction identifies healthcare and law as high-stakes deployment domains for LLMs.Source: §1 · p. 1
• Sampling-based semantic-dispersion measures can detect hallucinations, but require multiple stochastic completions and are prohibitively expensive for latency-sensitive applications.Semantic diversity among answers is described as a useful uncertainty proxy.Source: §1 · p. 1
• Lightweight probing methods predict dispersion or hallucination risk [Kadavath et al., 2022, Obeso et al., 2025, Han et al., 2025, Kossen et al., 2024], but scalar scores limit richer reliability signals.The introduction says these methods typically reduce uncertainty to a scalar and simple risk classification.Source: §1 · p. 1
• SSD distills sampled semantic answer embeddings into a lightweight prompt-conditioned density, using entropy for pre-generation risk and likelihood for answer verification.The density is predicted from the prompt representation rather than by generating multiple answers at inference.Source: §1 · pp. 1–2
• SSD matches single-pass probe latency while delivering competitive pre-generation hallucination prediction across multiple model families.The introduction contrasts its inference latency with single-pass probing methods.Source: §1 · p. 2
• The predicted density supports post-generation answer verification and distribution-based multiple-choice selection beyond scalar uncertainty scoring.The introduction presents these as capabilities enabled by retaining the full semantic distribution.Source: §1 · p. 2

2 METHODS

• SSD distils a model’s sampled semantic output distribution into a lightweight density conditioned on the prompt representation, rather than a scalar uncertainty proxy.The method uses sampled completions to train a student density estimator and derives both pre-generation and post-generation signals from it.Source: §2 · p. 2, §2.3 · p. 4

2.1 PROBLEM SETUP

• The setup predicts a distribution over semantic answer embeddings z from the prompt representation h.The prompt representation comes from the language model’s internal states; z represents the completed sequence through a mapping Φ.Source: §2.1 · pp. 2–3

2.1.1 Prompt and Answer Representations

• SSD represents the prompt with the final-token hidden state from a selected layer and embeds answers with Φ, optionally reducing embedding dimensions with PCA.The layer is a hyperparameter selected for its predictiveness of semantic uncertainty.Source: §2.1.1 · p. 3

2.1.2 Student Model

• The student is a mixture density network that predicts a K-component Gaussian mixture over semantic embeddings conditioned on the prompt representation.Training fits the student by conditional maximum likelihood using embeddings of stochastic completions from the base model.Source: §2.1.2 · p. 3

2.2 EXTRACTING UNCERTAINTY SIGNALS

• The predicted density provides prompt-level dispersion through entropy and answer-level compatibility through the likelihood of an embedded candidate response.The density also supports latent sampling and summary statistics such as its mixture mean; this operation is discussed in the appendix.Source: §2.2 · p. 3

2.2.1 Quantifying Dispersion

• Order-2 Rényi entropy provides an analytic, sampling-free dispersion score because it has a closed form for Gaussian mixtures.Higher Rényi-2 entropy corresponds to greater semantic dispersion; Shannon differential entropy generally lacks a closed form for Gaussian mixtures.Source: §2.2.1 · p. 3

2.3 EXTENSION TO GENERIC SEQUENCE DOMAINS

• SSD extends beyond language by distilling a sequence model’s sampled distribution over task-relevant latent representations into a lightweight conditional density.In language tasks, the latent representation captures answer semantics rather than individual tokens or events.Source: §2.3 · p. 4

Uncertainty estimation in language models

• LLM uncertainty methods include single-sample, probing-based, and multi-sample approaches, which often estimate task-oriented proxies; decomposition into aleatoric and epistemic uncertainty has limited consensus.The cited review also reports that uncertainty signals are often only modestly correlated.Source: §3 · p. 4

Distillation in language models

• Language-model distillation commonly transfers logits or token outputs, whereas SSD targets sequence-level output distributions and broader semantic targets.The related work also distinguishes a narrower study of predicted distributions over numerical targets from SSD’s broader semantic targets.Source: §3 · p. 4

Knowledge distillation for uncertainty estimation

• Ensemble-uncertainty distillation approximates ensemble predictive distributions or uncertainty signals in one forward pass, while SSD needs no external ensemble and models semantic output distributions.The cited approaches commonly match logits, predictive variances, or other distributional summaries; probing methods typically compress uncertainty to a scalar.Source: §3 · p. 4

4 EXPERIMENTS

• The experiments ask whether SSD distills sampling-based dispersion into a single-pass model and what answer-level capabilities arise from modeling a full density rather than a scalar score.They evaluate prompt-level hallucination prediction and answer-level verification, including out-of-domain detection and multiple-choice selection.Source: §4 · p. 4

4.1 MODELS AND DATASET GENERATION

• Experiments use TriviaQA and an open-ended MMLU subset, with 4,000 training prompts and 1,000 held-out test prompts across five model families spanning 3B–8B parameters.The datasets represent knowledge-seeking and reasoning-heavy tasks; MMLU is included as a more challenging semantic-dispersion setting.Source: §4.1 · pp. 4–5

4.2 PROMPT-LEVEL HALLUCINATION PREDICTION

• On TriviaQA, SSD beats teacher dispersion (TD) on AUROC for four of seven models and on AUPRC for all seven; on MMLU, it is weaker overall but beats TD on AUPRC for three of seven.Unlike TD, SSD requires no inference-time sampling.Source: §4.2 · pp. 5–6, Table 2 · p. 6
• Compared with specialized probes, SSD has probe-level inference latency and avoids entailment-based target construction, while preserving a density for post-generation likelihood scoring.SE has the highest AUROC on TriviaQA, while PCP is strongest on MMLU; SSD’s pre-generation performance is slightly weaker than specialized probes.Source: §4.2 · p. 6, Table 1 · p. 6
• SSD performance depends on the base model: TriviaQA AUROC is 74.9 for Ministral 8B and 67.7 for Qwen3 8B, with variation linked primarily to distillation fidelity.For some models, predicting the semantic distribution from prompt representations is particularly difficult.Source: §4.2 · p. 6

4.3.1 Out-of-domain Answer Detection

• Student log-likelihood distinguishes prompt-matched from random answers, achieving AUROCs above 0.95 for most base models on both datasets.The same likelihood signal is weaker for post-generation hallucination prediction, with AUROC scores of approximately 0.6, reflecting the limits of the single-pass approximation.Source: §4.3.1 · p. 7, Table 3 · p. 7

4.3.2 Multiple-choice Answer Selection

• For multiple-choice MMLU, SSD ranks the four options by likelihood and exceeds the 0.25 random baseline, approaching base-model open-ended accuracy for Llama 3.2 and Ministral.The student is trained only on open-ended completions, without multiple-choice supervision.Source: §4.3.2 · p. 7, Table 4 · p. 8

4.3 ANSWER-LEVEL RESPONSE EVALUATION

• The explicit density supports pre-generation entropy and post-generation likelihood signals, combining risk estimation and answer verification in one low-latency framework.Figure 2 illustrates entropy differences for correct and hallucinated answers and precise versus underspecified prompts, plus context-aware candidate scoring.Source: Figure 2 · p. 5, §4.3.2 · p. 7, §4.3 · p. 7

5 DISCUSSION & FUTURE WORK

• SSD’s hallucination-prediction effectiveness depends on reconstructing the teacher’s prompt-induced uncertainty ordering from prompt representations alone.The authors suggest scaling distillation data, enriching the student input, or using stronger teachers; they note that representation limits or insufficient supervision may explain performance gaps.Source: §5 · p. 8
• A unified decision policy could combine pre-generation dispersion forecasts with post-generation likelihood checks, such as abstaining under high dispersion or thresholding candidate actions.The authors present this as a natural next step, not an evaluated result.Source: §5 · p. 8
• SSD may suit diffusion-based language models because their prompt-to-output mapping may be more direct without token-wise autoregressive decoding.The authors suggest this could reduce the gap between internal belief states and sampling-induced semantic distributions.Source: §5 · p. 8
• The authors propose applying SSD to EHR and time-series models by representing predicted trajectories in latent spaces and estimating uncertainty over future states.For EHR models, the proposed density could yield low-latency forecasts of plausible patient states, with entropy quantifying uncertainty over long-term outcomes.Source: §5 · p. 8

6 CONCLUSION

• SSD distills sampled semantic answer distributions into prompt-conditioned predictive densities that preserve semantic output-space geometry.The density supports analytic dispersion estimates and likelihood-based evaluation in a single forward pass.Source: §6 · p. 8
• SSD delivers competitive hallucination prediction relative to sampling-based semantic dispersion baselines while avoiding multiple-completion inference overhead.The full density also supports answer-level reliability scoring and distribution-based selection unavailable to scalar probes.Source: §6 · p. 8
• SSD shifts computational cost from inference to training to provide distribution-level uncertainty for latency-critical language systems and other sequence domains.The conclusion frames this as a practical mechanism for integrating uncertainty over latent outcome representations.Source: §6 · p. 8

A DERIVATIONS

• For a Gaussian mixture, the quadratic functional ∫q(z)^2 dz decomposes into pairwise component overlaps, using a Gaussian product identity and a collision matrix.The collision matrix represents the pairwise overlap terms.Source: A DERIVATIONS · p. 13
• The collision-matrix computation is O(K^2) in the number of mixture components and is negligible relative to the transformer backbone.Source: A DERIVATIONS · p. 13

B FURTHER ANALYSES

• Further analyses are reported on TriviaQA unless otherwise specified, because SSD showed its strongest performance variation across base models there.Source: B FURTHER ANALYSES · p. 13

B.1 DISTILLATION FIDELITY VS. DETECTION PERFORMANCE

• Distillation fidelity is the Spearman rank correlation over 1,000 test prompts between student-predicted entropy and teacher dispersion.Higher fidelity indicates more faithful recovery of the teacher’s prompt-induced uncertainty ordering.Source: §B.1 · p. 13
• ρmeta = 0.65 indicates a strong positive relationship between distillation fidelity and hallucination-detection performance across models and student settings.The figure plots detection AUROC against fidelity; failure cases such as Qwen3-4B have near-zero fidelity, suggesting the uncertainty signal may not be accessible from the chosen prompt representation.Source: §B.1 · p. 14, Figure A1 · p. 14

B.2 SEMANTIC CONSENSUS FROM THE PREDICTED MIXTURE

• The centroid comparison treats the empirical mean of 32 stochastic answer embeddings as reference and uses Euclidean mean squared distance, where lower values indicate better recovery.Source: §B.2 · p. 14
• Over 76% of incorrect-answer test cases show improved centroid recovery with the SSD mixture mean versus a single default decoding pass.The evaluation uses the empirical mean of 32 stochastic samples as the reference semantic centroid and mean squared distance as the measure.Source: §B.2 · p. 14

C.1 PCA DIMENSION

• SSD performance peaks at dpca = 16, while teacher-dispersion detection improves with dimensionality; OOD performance is similar at 16 and 32 before dropping at higher dimensions.The main experiments use dpca = 16, except OOD, which uses dpca = 32; all PCA ablations use Ntrain = 4000.Source: §C.1 · p. 15, Table A2 · p. 15, Figure A2 · p. 15

C.2 STUDENT CAPACITY

• Increasing student width or depth yields negligible mean AUROC gains, suggesting the prompt representation’s available information, rather than student expressivity, limits performance.The authors use K = 10 mixture components but do not increase it further, given S = 32 teacher samples and the risk to generalization.Source: §C.2 · p. 16, Table A3 · p. 16, Figure A3 · p. 16

D MODEL CONFIGURATION AND SELECTION

• The model-configuration table reports TriviaQA accuracy, total layers, selected PCP and SEP probe layers, and each model’s best SSD architecture.The best SSD configuration specifies hidden size, depth, and mixture-component count.Source: D MODEL CONFIGURATION AND SELECTION · p. 16, Table A4 · p. 17

D.1 LAYER SELECTION PROTOCOL

• Following Kossen et al. [2024], the layer-selection heuristic chooses the layer with the highest validation accuracy for predicting binarized Semantic Entropy.The heuristic assumes that layers predictive of scalar dispersion also contain rich signal for learning the full semantic density; selected layers tend to be deeper in the network.Source: §D.1 · p. 17

• Sampling-based Semantic Entropy and teacher-dispersion methods require S stochastic sequences per prompt, with autoregressive decoding and pairwise entailment comparisons adding inference cost.The passage gives typical S as 5–20 and notes that scaling with S can be unacceptable in latency-sensitive applications.Source: E COMPUTATIONAL COMPLEXITY ANALYSIS · p. 17
• Scalar probes use a cached prompt representation from one base-model forward pass, but compress uncertainty to a scalar and discard semantic-distribution geometry.The probe itself is a lightweight classifier or regressor and needs no additional sampling or decoding.Source: E COMPUTATIONAL COMPLEXITY ANALYSIS · p. 17
• SSD matches scalar probes’ inference complexity while retaining distributional structure, with analytical entropy enabling O(1) uncertainty estimation without sampling at deployment.Its MDN head adds a slightly larger parameter set than a linear probe, described as negligible beside the transformer backbone; the computational burden shifts to training.Source: E COMPUTATIONAL COMPLEXITY ANALYSIS · p. 17

• The full results tables cover all methods across datasets and provide metrics for individual models.Source: F FULL HALLUCINATION PREDICTION RESULTS · p. 18
• Table A5 reports hallucination-prediction AUROC means over 1,000 test-set bootstrap resamples, with standard deviations shown as subscripts.The table marks the best method per dataset and model in bold.Source: Table A5 · p. 18, F FULL HALLUCINATION PREDICTION RESULTS · p. 18
• Table A6 reports hallucination-prediction AUPRC means over 1,000 test-set bootstrap resamples, with standard deviations shown as subscripts.The table marks the best method per dataset and model in bold.Source: Table A6 · p. 18, F FULL HALLUCINATION PREDICTION RESULTS · p. 18