Semantic Self-Distillation for Language Model Uncertainty

1. Problem Statement

Large Language Models (LLMs) are increasingly deployed in high-stakes domains (e.g., healthcare, law) and agentic frameworks where reliability is critical. However, quantifying uncertainty in LLMs remains challenging due to:

• Complexity of Outputs: LLMs generate diverse, high-dimensional semantic outputs, making simple scalar uncertainty metrics insufficient.
• Latency Constraints: Current state-of-the-art uncertainty quantification (UQ) methods, such as Semantic Entropy or Semantic Dispersion, rely on generating multiple stochastic samples ($S$) and comparing them (often using Natural Language Inference models). This incurs a prohibitive computational cost ($O(S)$ or $O(S^2)$), rendering them unsuitable for latency-critical applications like real-time agents.
• Information Loss in Probing: Existing lightweight "probing" methods train models to predict a single scalar uncertainty score (e.g., probability of correctness). While fast, these methods collapse the complex topology of the output distribution into a single number, losing valuable structural information needed for tasks like answer verification or selection.

Goal: Develop a method that provides distribution-level uncertainty (capturing the geometry of possible answers) with the inference latency of a single-pass probe, avoiding the need for repeated sampling at deployment.

2. Methodology: Semantic Self-Distillation (SSD)

The authors propose Semantic Self-Distillation (SSD), a framework that distills the sampled semantic distribution of a "teacher" LLM into a lightweight "student" model.

Core Architecture

1. Teacher (Base LLM):

   • Given a prompt $x$, the teacher generates $S$ stochastic completions $y^{(1)}, \dots, y^{(S)}$.
   • These completions are mapped to a semantic embedding space using a specialist embedding model $\Phi$: $z^{(s)} = \Phi(y^{(s)})$.
   • The prompt is encoded into a hidden representation $h$ (extracted from a specific layer $\ell^*$ of the LLM).

2. Student (Mixture Density Network - MDN):

   • The student is a lightweight neural network trained to predict a conditional probability density $q_\phi(z | h)$ over the semantic embeddings $z$, given only the prompt representation $h$.
   • The student models the distribution as a Gaussian Mixture Model (GMM) with $K$ components:
     $$q_\phi(z | h) = \sum_{k=1}^K \pi_k(h) \mathcal{N}(z; \mu_k(h), \text{diag}(\sigma_k^2(h)))$$
   • Training Objective: The student is trained via conditional maximum likelihood to match the empirical distribution of the teacher's sampled embeddings:
     $$\max_\phi \sum_x \frac{1}{S} \sum_{s=1}^S \log q_\phi(z^{(s)} | h(x))$$

Inference Capabilities (Single Forward Pass)

Once trained, the student model provides an explicit density $q_\phi(z | h)$ without generating new tokens from the base LLM. This enables three distinct operations:

1. Pre-generation Risk Estimation (Prompt-Level):
   • Compute the Rényi-2 Entropy of the predicted density analytically. High entropy indicates high semantic dispersion (high uncertainty/hallucination risk).
   • Advantage: Closed-form calculation avoids Monte Carlo sampling.
2. Post-generation Verification (Answer-Level):
   • Given a candidate answer $y'$, embed it to $z'$.
   • Score its compatibility with the prompt using the log-likelihood $\log q_\phi(z' | h)$. Low likelihood suggests the answer is out-of-domain or hallucinated.
3. Latent Sampling & Consensus:
   • Sample from the mixture distribution or compute the mixture mean $\mu_{SSD} = \sum \pi_k \mu_k$ to approximate the "semantic consensus" of the model's belief.

3. Key Contributions

1. Novel Framework (SSD): Introduces a method to distill a language model's sampled semantic distribution into a lightweight, prompt-conditioned density estimator, bridging the gap between expensive sampling and information-losing scalar probes.
2. Competitive Performance with Low Latency: Demonstrates that SSD achieves hallucination prediction performance comparable to (and sometimes better than) sampling-based baselines (like Semantic Entropy) while maintaining the inference speed of single-pass methods.
3. Beyond Scalar Uncertainty: Shows that modeling the full distribution unlocks new reliability primitives unavailable to scalar probes, specifically:
   • Out-of-Domain (OOD) Detection: Identifying mismatched prompt-answer pairs.
   • Multiple-Choice Selection: Selecting the best answer from a set of options based on likelihood.
4. Generalizability: The framework is agnostic to modality and can be applied to any sequence model where outputs can be mapped to meaningful latent representations.

4. Experimental Results

The authors evaluated SSD across 5 model families (Mistral, Llama-3, Qwen, Gemma, SmolLM) on TriviaQA (knowledge-seeking) and MMLU (reasoning-heavy).

• Hallucination Prediction (Pre-generation):

  • SSD outperformed Teacher Dispersion (TD) (the standard deviation of $S=32$ samples) on 4/7 models in AUROC and all 7 in AUPRC on TriviaQA.
  • SSD achieved competitive results against Semantic Entropy (SE) and Semantic Entropy Probes (SEP) but with significantly lower latency (Single-pass vs. $S$ samples).
  • Key Insight: SSD acts as a "smoothing" mechanism, assigning non-zero uncertainty even when the teacher consistently generates the same (potentially wrong) answer, whereas TD assigns zero uncertainty in such degenerate cases.

• Post-generation Capabilities:

  • OOD Detection: SSD log-likelihood achieved AUROC > 0.95 for most models in distinguishing correct prompt-answer pairs from random mismatches.
  • Multiple-Choice Selection: When trained only on open-ended data, the SSD student could select the correct option in MMLU multiple-choice questions with accuracy approaching the base model's open-ended generation performance (e.g., ~36% vs 40% for Ministral 8B), significantly beating the random baseline (25%).

• Ablation Studies:

  • Distillation Fidelity: Performance is strongly correlated with how well the student learns the teacher's uncertainty ordering.
  • Dimensionality: Reducing embedding dimensions via PCA (to $d=16$) optimized the trade-off between semantic resolution and the difficulty of the regression task for the student.
  • Capacity: Increasing student network capacity yielded diminishing returns, suggesting the bottleneck is the information available in the prompt representation, not the student's expressivity.

5. Significance and Future Impact

• Practical Deployment: SSD solves the "latency vs. accuracy" trade-off in UQ. It allows systems to perform rigorous uncertainty checks (both before and after generation) without the computational overhead of generating multiple completions.
• Agentic Systems: By providing pre-generation risk scores and post-generation verification, SSD enables safer agentic workflows where unreliable steps can be detected and mitigated (e.g., via abstention or retrieval) before errors compound.
• Beyond LLMs: The authors highlight the potential for SSD in Electronic Health Records (EHR) and Time-Series Foundation Models, where predicting the distribution of future clinical states or trajectories is more valuable than a single point estimate.
• Theoretical Shift: The work moves the field from treating uncertainty as a scalar "confidence score" to treating it as an explicit geometric structure in semantic space, enabling richer downstream decision-making.

In summary, Semantic Self-Distillation offers a principled, efficient, and versatile approach to uncertainty quantification, transforming the expensive process of sampling-based dispersion analysis into a fast, single-pass density estimation task.