Semantic Token Clustering for Efficient Uncertainty Quantification in LLMs

Paper: 2603.20161 Authors: Qi Cao, Andrew Gambardella, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa Institution: The University of Tokyo Categories: cs.CL

Abstract

Large language models often generate plausible-sounding but incorrect responses with high confidence, making it critical to quantify their uncertainty. Existing methods rely on repeated sampling or auxiliary models, introducing substantial computational overhead. This paper proposes Semantic Token Clustering (STC), an efficient uncertainty quantification method that requires only a single generation and no external models, achieving performance comparable to state-of-the-art while dramatically reducing computational cost.

The Problem: Overconfident LLMs

Why Uncertainty Matters

LLMs demonstrate impressive capabilities but have critical limitations:

1. No guarantee of truthfulness: Outputs can be factually incorrect
2. Overconfidence problem: High confidence even when wrong
3. Plausible-sounding errors: Mistakes are hard to detect
4. High-stakes domains: Healthcare, law, science require reliability

Current Approaches and Their Limitations

Supervised Methods:

• Train probes to predict correctness
• Require labeled data and additional training
• Poor generalization to out-of-distribution data

Unsupervised Methods:

• Logit-based (Perplexity): Ignore semantic consistency
• Sampling-based (Semantic Entropy): Require multiple generations (expensive)
• External models (CCP): Depend on NLI models (computational overhead)

Key Gap: Existing methods either overlook semantic information or require expensive repeated sampling/external models.

The Solution: Semantic Token Clustering

Core Insight

When an LLM is confident about an answer, probability mass is distributed across semantically equivalent tokens:

• “television” vs. “TV” vs. “tv”
• “United States” vs. “USA” vs. “America”

Looking at individual token probabilities underestimates the model’s true confidence. We should aggregate probability mass over semantic clusters.

Method Overview

Pre-computation Stage (offline):

1. Embedding clustering: Group vocabulary tokens by semantic similarity
2. Store cluster assignments for fast lookup

Inference Stage (online):

1. Generate response (single pass)
2. At each decoding step, identify semantic cluster for predicted token
3. Prefix matching: Include tokens with matching prefixes (“television”, “Television”, “televis…”)
4. Aggregate probabilities: Sum probability mass over cluster
5. Compute uncertainty score from aggregated probabilities

Mathematical Formulation

Given input x and generated response y, estimate uncertainty:

U(x, y) = g(p̂(C = 0 | x, y))

where:

• C is binary correctness indicator
• g is monotonically increasing link function
• p̂ is estimated from cluster-aggregated probabilities

At each token position, cluster probability:

P_cluster = Σ P(token) for all tokens in semantic cluster

Uncertainty is computed from the entropy or confidence of cluster probabilities.

Key Advantages

1. Leveraging Internal Representations

• Uses token embedding clustering to capture semantic structure
• Directly exploits LLM’s internal semantic knowledge
• No need for external semantic similarity models

2. Self-Contained Implementation

• No fine-tuning required
• No supervised data needed
• No external models (NLI, etc.)
• Works with any white-box LLM out-of-the-box

3. Computational Efficiency

• Single generation: No repeated sampling
• Offline clustering: Expensive operations done once
• Fast inference: Minimal runtime overhead
• Ideal for resource-constrained and low-latency scenarios

Experimental Results

Datasets

Evaluated on question-answering tasks:

• NQ (Natural Questions)
• TQA (TriviaQA)

Performance (AUROC)

The method achieves comparable or better performance than state-of-the-art baselines:

Method	Semantic Aware	Single Sample	External Model Free	Overhead	Performance
Perplexity	✗	✓	✓	Low	0.70
Semantic Entropy	✓	✗	✗	High	0.88
CCP	✓	✓	✗	High	0.87
STC (Ours)	✓	✓	✓	Low	0.88

Key Findings

1. Matches multi-sample methods with single generation
2. Outperforms single-sample baselines that ignore semantics
3. Dramatically lower computational cost than semantic-aware methods
4. Robust across different model sizes and architectures

Technical Deep Dive

Embedding Clustering

Approach:

• Extract token embeddings from LLM’s embedding layer
• Apply clustering algorithm (K-means, hierarchical, etc.)
• Group semantically similar tokens together

Example clusters:

• Cluster 1: [“television”, “TV”, “tv”, “telly”]
• Cluster 2: [“USA”, “United States”, “America”, “US”]
• Cluster 3: [“yes”, “yeah”, “yep”, “affirmative”]

Prefix Matching

Motivation: Capture morphological variations and capitalization

• “television” → “Television”, “televisions”, “televis…”
• “United States” → “united states”, “United states”

Implementation:

• For predicted token, find all vocabulary tokens with matching prefix
• Add their probabilities to cluster probability

Uncertainty Scoring

Multiple options for computing final uncertainty:

1. Entropy: H = -Σ P_cluster * log(P_cluster)
2. Confidence: U = 1 - max(P_cluster)
3. Margin: U = P_cluster[1] - P_cluster[2]

Paper uses entropy-based scoring for best results.

Comparison with Baselines

vs. Perplexity

• Perplexity: Token-level probability, ignores semantics
• STC: Cluster-level probability, captures semantic equivalence
• Result: STC significantly outperforms

vs. Semantic Entropy

• Semantic Entropy: Multiple generations + NLI model for semantic clustering
• STC: Single generation + embedding clustering
• Result: Similar performance, 10x+ faster

vs. CCP (Claim Conditioned Probability)

• CCP: Single generation but requires NLI model
• STC: Single generation, no external model
• Result: Similar performance, lower overhead

Practical Implications

When to Use STC

Ideal for scenarios requiring:

• Low latency: Real-time applications
• Resource constraints: Limited compute budget
• High throughput: Processing many queries
• Reliability: Detecting potentially incorrect outputs

Applications

1. Conversational AI: Flag uncertain responses for human review
2. Question Answering: Provide confidence scores with answers
3. Content Generation: Identify sections needing fact-checking
4. Decision Support: Highlight high-uncertainty predictions in critical domains

Integration

Easy to integrate into existing LLM pipelines:

# Pseudo-code
def generate_with_uncertainty(prompt):
    # Pre-computed: token_clusters

    response, logits = model.generate(prompt, return_logits=True)

    uncertainties = []
    for token, logit in zip(response, logits):
        cluster = token_clusters[token]
        cluster_prob = sum(logit[t] for t in cluster)
        uncertainty = -log(cluster_prob)
        uncertainties.append(uncertainty)

    avg_uncertainty = mean(uncertainties)
    return response, avg_uncertainty

Limitations and Future Work

Current Limitations

1. White-box requirement: Needs access to token probabilities (not available in API-only models)
2. Clustering quality: Depends on embedding space quality
3. Fixed clusters: Doesn’t adapt clusters to specific contexts
4. Token-level only: Doesn’t capture sentence-level semantic consistency

Future Directions

1. Dynamic clustering: Adapt clusters based on context
2. Hierarchical uncertainty: Combine token, sentence, and document-level uncertainty
3. Black-box adaptation: Extend to API-only models
4. Multi-modal: Apply to vision-language models
5. Calibration: Improve uncertainty calibration for better reliability

Takeaways

1. Semantic clustering is key: Aggregating probability over semantically similar tokens dramatically improves uncertainty estimation
2. Single generation suffices: No need for expensive repeated sampling
3. Internal representations are powerful: LLM embeddings encode rich semantic structure
4. Efficiency matters: Practical deployment requires low-overhead methods
5. Self-contained is better: Avoiding external models simplifies deployment and reduces dependencies

This work demonstrates that efficient and accurate uncertainty quantification is possible by cleverly leveraging the semantic structure already encoded in LLMs. The method’s simplicity, efficiency, and strong performance make it highly practical for real-world deployment.

For applications requiring reliable LLM outputs, STC provides a principled way to identify potentially incorrect responses with minimal computational overhead – a critical capability for deploying LLMs in high-stakes domains.

U(x, y) = g(p̂(C = 0 | x, y))

P_cluster = Σ P(token) 对于语义簇中的所有令牌

# 伪代码
def generate_with_uncertainty(prompt):
    # 预计算：token_clusters

    response, logits = model.generate(prompt, return_logits=True)

    uncertainties = []
    for token, logit in zip(response, logits):
        cluster = token_clusters[token]
        cluster_prob = sum(logit[t] for t in cluster)
        uncertainty = -log(cluster_prob)
        uncertainties.append(uncertainty)

    avg_uncertainty = mean(uncertainties)
    return response, avg_uncertainty