Table of Contents
Fetching ...

Semantic Token Clustering for Efficient Uncertainty Quantification in Large Language Models

Qi Cao, Andrew Gambardella, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Abstract

Large language models (LLMs) have demonstrated remarkable capabilities across diverse tasks. However, the truthfulness of their outputs is not guaranteed, and their tendency toward overconfidence further limits reliability. Uncertainty quantification offers a promising way to identify potentially unreliable outputs, but most existing methods rely on repeated sampling or auxiliary models, introducing substantial computational overhead. To address these limitations, we propose Semantic Token Clustering (STC), an efficient uncertainty quantification method that leverages the semantic information inherently encoded in LLMs. Specifically, we group tokens into semantically consistent clusters using embedding clustering and prefix matching, and quantify uncertainty based on the probability mass aggregated over the corresponding semantic cluster. Our approach requires only a single generation and does not depend on auxiliary models. Experimental results show that STC achieves performance comparable to state-of-the-art baselines while substantially reducing computational overhead.

Semantic Token Clustering for Efficient Uncertainty Quantification in Large Language Models

Abstract

Large language models (LLMs) have demonstrated remarkable capabilities across diverse tasks. However, the truthfulness of their outputs is not guaranteed, and their tendency toward overconfidence further limits reliability. Uncertainty quantification offers a promising way to identify potentially unreliable outputs, but most existing methods rely on repeated sampling or auxiliary models, introducing substantial computational overhead. To address these limitations, we propose Semantic Token Clustering (STC), an efficient uncertainty quantification method that leverages the semantic information inherently encoded in LLMs. Specifically, we group tokens into semantically consistent clusters using embedding clustering and prefix matching, and quantify uncertainty based on the probability mass aggregated over the corresponding semantic cluster. Our approach requires only a single generation and does not depend on auxiliary models. Experimental results show that STC achieves performance comparable to state-of-the-art baselines while substantially reducing computational overhead.
Paper Structure (23 sections, 5 equations, 3 figures, 13 tables)

This paper contains 23 sections, 5 equations, 3 figures, 13 tables.

Figures (3)

  • Figure 1: Illustration of the proposed method. Token embedding clustering is performed in the pre-computation stage. During inference, we aggregate next-token probability mass over embedding-clustered and prefix-matched tokens to quantify uncertainty.
  • Figure 2: Performance comparison between our method and baseline approaches across different models and datasets. For clarity, only the top seven methods ranked by performance are shown in each subfigure, with methods ordered by descending performance. Metric: AUROC.
  • Figure 3: Efficiency comparison across methods. The performance (AUROC) and relative execution time overhead (%) are plotted on the y-axis and x-axis, respectively, illustrating the efficiency of the proposed method. The relative execution time overhead represents the additional execution time required for uncertainty quantification relative to basic inference. For clarity, only the top ten methods ranked by performance are shown in each subfigure.