Semantic Token Clustering for Efficient Uncertainty Quantification in Large Language Models

1. Problem Statement

Large Language Models (LLMs) exhibit a critical limitation: they often generate plausible-sounding but factually incorrect responses (hallucinations) while displaying high confidence. This "overconfidence" makes it difficult to identify unreliable outputs.

• Current Limitations: Existing uncertainty quantification (UQ) methods generally fall into two categories, both of which have significant drawbacks:
  1. Sampling-based methods (e.g., Semantic Entropy, EigenScore): These require generating multiple responses to measure semantic dispersion. This introduces substantial computational overhead and latency.
  2. Auxiliary model-based methods (e.g., CCP, P(True)): These rely on external Natural Language Inference (NLI) models or probes to verify correctness, adding inference costs and dependency on external resources.
• Goal: The authors aim to develop a UQ method that is computationally efficient (single generation), self-contained (no external models), and semantically aware (leveraging internal LLM representations).

2. Methodology: Semantic Token Clustering (STC)

The proposed method, Semantic Token Clustering (STC), quantifies uncertainty by aggregating probability mass over semantically consistent token groups rather than individual tokens. It operates in two stages:

A. Pre-computation Stage (Offline)

Before inference, the model's vocabulary is processed to create semantic clusters.

• Embedding Extraction: Input and output embeddings for each token are concatenated to form a unified semantic representation.
• Clustering: An unsupervised clustering algorithm (specifically Agglomerative Clustering with cosine distance) groups tokens into semantically consistent clusters (e.g., grouping "TV", "television", "TV", and "televis" into one cluster).
• Optimization: Function words and Arabic numerals are excluded to reduce noise. This step is performed once per model, incurring no runtime cost during inference.

B. Inference Stage (Online)

During a single generation pass, the method calculates uncertainty for each generated token $y_i$:

1. Semantic Aggregation: Instead of looking at the probability of the single generated token $P(y_i)$, the method identifies a set of candidate tokens $T_i$ that are semantically consistent with $y_i$.
2. Dual Filtering: The set $T_i$ is formed by the union of:
   • Embedding-Clustered Tokens ($T^e_i$): Tokens belonging to the same pre-computed semantic cluster as $y_i$.
   • Prefix-Matched Tokens ($T^p_i$): Tokens that serve as a prefix for the subsequent generation (e.g., if the model generates "tele", tokens like "television" or "televis" are included to capture surface-form consistency).
3. Probability Mass Aggregation: The probability mass is summed over this set:
   $$\hat{p}_c(y_i | x, y_{<i}) = \sum_{t \in T_i} p(t | x, y_{<i})$$
4. Final Score: The overall uncertainty score $S(x, y)$ for the sequence is calculated as one minus the product of these aggregated probabilities across all steps:
   $$S(x, y) = 1 - \prod_{i=1}^{n} \hat{p}_c(y_i | x, y_{<i})$$
   A higher score indicates higher uncertainty (lower confidence).

3. Key Contributions

1. Efficiency: STC requires only a single generation and no external models. It eliminates the need for repeated sampling or auxiliary NLI models, reducing inference-time overhead by approximately 98% compared to methods like CCP.
2. Semantic Awareness: Unlike logit-based metrics (e.g., Perplexity) that treat tokens as discrete symbols, STC leverages the internal semantic structure of the LLM. It recognizes that different tokens (e.g., "TV" vs. "television") can represent the same meaning, preventing the model from appearing uncertain due to surface-level variations.
3. Self-Contained Implementation: The method requires no fine-tuning, supervised data collection, or labeled data. It works directly on off-the-shelf white-box LLMs.
4. Offline Pre-computation: The computationally intensive clustering is done offline, ensuring that the inference process remains lightweight and suitable for resource-constrained or low-latency environments.

4. Experimental Results

The authors evaluated STC against 14 baselines (including Perplexity, Semantic Entropy, CCP, and EigenScore) across multiple models (Llama-2/3, Mistral, Qwen2.5) and datasets (NQ, TQA, WQ).

• Performance (AUROC): STC achieves performance comparable to State-of-the-Art (SOTA) baselines. In many cases, it outperforms sampling-based methods and significantly outperforms logit-based metrics like Perplexity.
  • Example: On the TQA dataset with Llama-3-8B, STC achieved an AUROC of 0.882, comparable to Semantic Entropy (0.882) and EigenScore (0.882), but with drastically lower cost.
• Efficiency:
  • Execution Time: STC adds negligible overhead (e.g., ~1.7% relative overhead on Llama-3-8B) compared to sampling-based methods which often add 200%–400% overhead.
  • Hardware: STC can run efficiently on CPUs without requiring GPUs, whereas many baselines require GPU acceleration for sampling.
• Ablation Study: Removing either embedding clustering or prefix matching leads to performance degradation, confirming that both components are necessary for capturing semantic and surface-level consistency.

5. Significance and Limitations

Significance:
STC represents a paradigm shift in uncertainty quantification by moving away from expensive sampling or external verification toward leveraging the inherent semantic knowledge already encoded in the LLM's weights. It makes high-quality uncertainty estimation feasible for real-time applications and resource-constrained environments.

Limitations:

1. Access Requirements: Like most internal-state methods, it requires access to token logits and embeddings, making it inapplicable to closed-source models (black-box APIs) without modification.
2. Static Embeddings: The method uses static token embeddings. It may struggle with polysemy (words with multiple meanings) where context-independent clustering might group divergent meanings. The authors suggest future work could integrate context-aware embeddings.
3. Calibration: The raw uncertainty scores are not explicitly calibrated (though they can be post-calibrated if needed).

In conclusion, Semantic Token Clustering offers a highly efficient, accurate, and self-contained solution for detecting unreliable LLM outputs, bridging the gap between performance and computational cost.