Forking Paths in Neural Text Generation

TL;DR

Forking Paths in Neural Text Generation investigates token-level uncertainty in LLM generation, arguing that a single token can drastically alter subsequent text. The authors introduce the Forking Tokens Hypothesis and the Forking Paths Analysis, a three-stage sampling pipeline that builds token-level outcome distributions and applies Bayesian Change Point Detection and survival analysis to identify forking tokens. Across 7 tasks in 4 domains, they demonstrate dynamic uncertainty and forking tokens including punctuation, challenging static final-token uncertainty measures. The work has implications for LLM evaluation, safety, and guidance during inference, and points to future directions for efficiency and control of text generation.

Abstract

Estimating uncertainty in Large Language Models (LLMs) is important for properly evaluating LLMs, and ensuring safety for users. However, prior approaches to uncertainty estimation focus on the final answer in generated text, ignoring intermediate steps that might dramatically impact the outcome. We hypothesize that there exist key forking tokens, such that re-sampling the system at those specific tokens, but not others, leads to very different outcomes. To test this empirically, we develop a novel approach to representing uncertainty dynamics across individual tokens of text generation, and applying statistical models to test our hypothesis. Our approach is highly flexible: it can be applied to any dataset and any LLM, without fine tuning or accessing model weights. We use our method to analyze LLM responses on 7 different tasks across 4 domains, spanning a wide range of typical use cases. We find many examples of forking tokens, including surprising ones such as punctuation marks, suggesting that LLMs are often just a single token away from saying something very different.

Figures (30)

• Figure 1: Forking paths in text generation: can a single token alter the outcome? At each step of next-word text generation, an LLM has some probability of sampling a variety of possible next tokens. This raises the question: are there specific forking tokens, where choosing a certain token over other probable alternatives results in a distinct outcome? To test this hypothesis, we estimate uncertainty in text generation by systematically re-sampling alternate completions at each token in a single base path, such as a greedily decoded sequence, to identify if there are forking tokens where two paths diverge into different outcomes. We find forking tokens where expected content words (2021/2024 in this example) lead to a different final answers (Queen Elizabeth or King Charles), but we also find forking tokens in unexpected places (e.g. that/who).
• Figure 2: (Top) Systematically sampling alternate paths in text generation. Our data collection proceeds in three stages. First, we decode a base path text completion $x^*$ from an LLM given some prompt and record the most probable alternate tokens $w$ at each step $t$. Next, we re-sample $S$completions $x_{>t}^{(s)}$ by appending $x_{<t}^*$ to the original prompt, along with each alternate token $x_t = w$. Finally, we extract Cyan outcome vector representations $R( \cdot )$ for each sample. In our experiments, for $R$ we use a different LLM to extract the final answer from a path. (Bottom) Probability-weighted outcome distributions We aggregate outcome representations $R( \cdot )$ into outcome distributions $o_t$ and $o_{t,w}$ by weighting each outcome with next-token probabilities for the forking token $p(x_t = w \ | \ x^*_{<t})$ and for the sampled path completion $p( x^{(s)}_{>t} \ | \ x_{<t}^* , x_t = w )$. Also see Eq. \ref{['eq:outcome-dist']}.
• Figure 3: Uncertainty dynamics time series (Top) The outcome distribution $o_t$ is equivalent to a multivariate time series, where each possible outcome (e.g. King Charles or Queen Elizabeth) can be plotted as a separate line. To simplify our modeling, we convert $o_t$ into a univariate time series using a semantic drift transformation $y_t$. Bayesian Change Point Detection model. (Bottom, Left) We use a Bayesian Change Point Detection (CPD) model to identify sudden shifts in $y_t$. This model fits $y_t$ by splitting it into $m$ segments at times $\tau_i$, and fitting separate linear models with parameters $\beta_i, \delta_i$ to each segment $i$. (Bottom, Right) Our CPD model uses Monte Carlo sampling to infer posterior distributions that help us interpret analysis results: $p(\tau = t \ | \ y)$, or how likely it is that a change $\tau$ happened at time $t$, and $p(m \ | \ y)$, or how many change points, if any, occurred in the time series.
• Figure 4: The outcome distribution can collapse after a single step in chain-of-thought reasoning. (Left, Top) The outcome distribution $o_t$ and (Left, Bottom) change point probabilities $p(\tau = t | y)$ for a single question from HotpotQA and a greedily decoded base path $x^*$. (Right) Tokens in $x^*$, highlighted according to $p(\tau = t | y)$, where yellow indicates low probability of a change point, red indicates high probability; the initial prompt is shown above $x^*$ with no coloring. We see striking uncertainty dynamics in $o_t$: the outcome distribution remains stable with the top single answer being Robin Tunney (the correct answer) until the forking token $t=34$, when the distribution suddenly collapses to Mia Sara (a hallucinated wrong answer).
• Figure 5: Further examples of forking tokens Two examples of outcome distributions $o_t(x^*)$ with forking tokens: a physics question from MMLU (Top; Correct Answer: B) and a mathematical reasoning question from GSM8k (Bottom; Answer: 21). In MMLU-12, we see a similar pattern to Fig. \ref{['fig:cpd-panel']}, where $o_t$ remains stable for most of the base path, before suddenly collapsing when the answer token is first specified "a = 9.8 m/s 2". In GSM8k-59 we see multiple forking tokens, both occurring at unexpected places. E.g. for the second change point, $o_{>t}$ changes depending on whether the token ' ( ' or 'by' is sampled. We also note that, similar to Fig. \ref{['fig:cpd-panel']}, at $t=0$ the most probable outcome is the correct answer (21) but this answer disappears from $o_t$ part-way through.