Technical Summary: Embedding Perturbation for Intermediate-Step Uncertainty in LLM Reasoning

Problem Statement

Large Language Models (LLMs) have achieved significant breakthroughs in reasoning tasks but remain prone to producing unreliable or misleading outputs. While Uncertainty Quantification (UQ) techniques exist to estimate a model's confidence in its final answers, they are often inadequate for intermediate-step uncertainty in reasoning chains. Existing methods face several limitations:

• Sampling-based approaches (e.g., self-consistency) require generating multiple responses to check for agreement. While effective for final answers, they are computationally expensive and struggle to pinpoint the specific intermediate step where an error originates.
• Probability-based metrics (e.g., token probability, entropy, margin) often fail to distinguish between low-probability tokens that are linguistically rare but correct, and those that represent genuine reasoning errors. They also struggle to capture long-term contextual dependencies necessary for auditing reasoning errors.
• Bayesian-based methods (e.g., weight perturbation) can be computationally burdensome, particularly for large models.

The core problem addressed is the lack of efficient, fine-grained metrics capable of identifying the specific tokens or steps within a reasoning chain where the model's uncertainty is highest, thereby indicating a potential error.

Methodology

The authors propose Embedding Perturbation as a novel approach to quantify intermediate uncertainty. The core hypothesis is that incorrect reasoning steps contain tokens that are highly sensitive to small perturbations in the embeddings of preceding tokens. This sensitivity reflects the model's internal hesitation among multiple competing continuations (i.e., the model is "unsure" whether to proceed with the current token or a rival token).

The method operates on the premise that LLMs are autoregressive; the generation of a token $x_t$ depends solely on preceding tokens $x_{1:t-1}$. The authors define the token generation probability as a function of the preceding token embeddings $H_{1:t-1}$. They measure the model's sensitivity to perturbations on these embeddings using two specific metrics:

1. Random Perturbation (Rand. Pert.):

   • Random Gaussian noise $\Delta \sim \mathcal{N}(0, \sigma^2 I)$ is added to the embeddings of all tokens in the input and response.
   • The metric is defined as the variance of the log-probability of the target token $x_t$ over multiple noise samples:
     $$\text{Pert.}(x_t) = \text{Var}_{\Delta} [\log \text{Pr}(x_t | \tilde{H}_{1:t-1})]$$
   • This captures the instability of the token selection under slight context variations.

2. Adversarial Perturbation (Adv. Pert.):

   • Instead of random sampling, embeddings are modified along the direction that maximally decreases the total log-probability of the generated response (using a signed gradient).
   • The metric is the difference in log-probability before and after perturbation:
     $$\text{Pert.}(x_t) = \log \text{Pr}(x_t | H_{1:t-1}) - \log \text{Pr}(x_t | \hat{H}_{1:t-1})$$
   • This identifies tokens that are most disrupted by an adversarial shift, indicating a proximity to the decision boundary between the selected token and a "rival token."

Theoretical Insight:
The paper provides a theoretical analysis showing that these perturbation sensitivities inversely approximate the geometric distance from the current embedding to the decision boundary between the selected token and a rival token. A high sensitivity score implies the embedding is close to a boundary where a different token (potentially leading to a different reasoning path) would be selected, signaling high uncertainty. Crucially, the authors argue this metric is less biased by natural language token frequencies compared to raw probability-based metrics.

Key Contributions

1. Novel UQ Metrics: Introduction of perturbation-based metrics (Random and Adversarial) specifically designed to capture intermediate-step uncertainty in LLM reasoning, distinct from final-answer confidence.
2. Theoretical Justification: A formal demonstration that perturbation sensitivity approximates the gradient norm, which in turn inversely approximates the distance to the decision boundary between competing reasoning paths.
3. Efficiency: The proposed methods are shown to be computationally efficient. Adversarial perturbation requires only a single backward and forward pass, while random perturbation can reuse a single set of perturbations for the entire sequence, avoiding the massive overhead of multiple sampling trials.
4. Empirical Validation: Extensive experiments demonstrating that these metrics outperform existing baselines in identifying the first erroneous step in reasoning chains.

Experimental Results

The authors evaluated their method on Math Reasoning (MATH dataset) and Logical Reasoning (BIG-Bench Hard) tasks across four models (Llama-3.1, OpenMath, Qwen, Mistral).

• Performance: Both Random and Adversarial perturbation metrics consistently outperformed baselines, including:
  • Token-level: Negative Log-Likelihood (NLL), Entropy, Probability Margin, and the Bayesian-based TokUR.
  • Sentence-level: CCP, TokenSAR, and Graph Uncertainty (multiple-sampling).
  • In the MATH dataset, Adversarial Perturbation achieved a top-3 detection rate of 0.65 for Llama (vs. 0.55 for TokUR) and 0.61 for Qwen (vs. 0.56 for TokUR).
• Efficiency: On a single NVIDIA GH200 GPU, Adversarial Perturbation added only 0.28 seconds per generation, and Random Perturbation (with 20 iterations) added 0.90 seconds. This is significantly faster than multiple-sampling approaches.
• Hallucination Detection: When tested on the FactScore dataset for factual hallucination detection, perturbation-based methods underperformed compared to multiple-sampling methods. The authors attribute this to the fundamental difference in error patterns: hallucinations often arise from spurious correlations or weak grounding that remain stable under perturbation, whereas reasoning errors involve genuine instability in pathway selection.

Significance and Claims

The paper claims that embedding perturbation provides a more effective and efficient way to audit the reasoning process of LLMs compared to existing methods. By identifying tokens with high sensitivity to embedding perturbations, practitioners can:

• Pinpoint the exact location where a model begins to deviate from a correct reasoning path.
• Enable fine-grained interventions, such as reflection-based mechanisms or self-correction, triggered specifically at uncertain steps.
• Facilitate the development of verification models by providing interpretable uncertainty signals derived from the generation process itself.

The authors modestly acknowledge limitations: the method may not detect hallucinations driven by spurious correlations, and it does not always perfectly disentangle uncertainty regarding correctness from uncertainty arising from the selection among multiple valid reasoning pathways. However, the study establishes that perturbation-based metrics are a superior tool for detecting procedural reasoning errors compared to probability-based or sampling-based baselines.