One email each morning — yesterday's AI, sortedGet it in your inbox

METAL LAB

Forking Fast: Efficiently Estimating Uncertainty Dynamics in Text Generation

arXiv:2608.196112026-08-21

A statistical smoothing trick lets you find where an LLM's reasoning could branch into different answers without resampling hundreds of times at every step

When an LLM answers a question, resampling its reasoning at every token or sentence reveals moments where the answer could branch off differently, but doing this accurately normally requires hundreds of costly resamples at every position. The authors built a statistical model that segments stable regions from sharp branch points and smooths noisy low-sample data, matching high-sample accuracy at a fraction of the cost. They validated this on Llama-3-8B-Instruct and DeepSeek-R1-Distill-Llama-8B answering tinyMMLU questions, collecting 1.77 billion tokens of data.

What they did

  1. Problem: tracking how an LLM's eventual answer could change at each generation step (Forking Paths Analysis) normally needs tens or hundreds of resampled continuations at every single position, which is extremely expensive in tokens
  2. Key finding: with very heavy resampling (S=1000), the distribution of possible outcomes becomes smooth almost everywhere except at sharp 'forking points' – meaning the noisy jumpiness seen with few samples is mostly just sampling noise, not real model sensitivity
  3. Solution: a three-stage statistical model – change point detection (PELT) to locate stable segments versus forks, then Gaussian-kernel-weighted pooling to smooth within each stable segment
  4. Result: with the model, low sample counts (S=5–30) achieve accuracy equivalent to collecting 3.3x to 5x more raw samples, and total token budget can be cut by up to 8x with only a small accuracy loss
  5. Limitation: tested only on multiple-choice tinyMMLU questions with two 8B-parameter models, and the method is slightly less accurate than raw data at very large answer shifts (forking threshold 0.20 or above)
Figure 1: Uncertainty dynamics ot for one question (tinyMMLU question 39; Llama-3-8B-Instruct). Colors represent the fraction of rollouts at t that end with different final answers. (a) When collecting a massive number of resampled rollouts (S=1000) at every token (N=1), the outcome distribution becomes increasingly smooth, except at forking points. (c) Uncertainty dynamics with a lower-sample analysis (S=20, N=2) using 1% of the reference’s sampling cost. (b) Smoothing the lower-sample data ot closely approximates the high-sample data and recovers the same forking points. More examples in Appendix B and our interactive dashboard (https://github.com/ericb-goodfire/forking-fast).
Figure 1: Uncertainty dynamics ot for one question (tinyMMLU question 39; Llama-3-8B-Instruct). Colors represent the fraction of rollouts at t that end with different final answers. (a) When collecting a massive number of resampled rollouts (S=1000) at every token (N=1), the outcome distribution becomes increasingly smooth, except at forking points. (c) Uncertainty dynamics with a lower-sample analysis (S=20, N=2) using 1% of the reference’s sampling cost. (b) Smoothing the lower-sample data ot closely approximates the high-sample data and recovers the same forking points. More examples in Appendix B and our interactive dashboard (https://github.com/ericb-goodfire/forking-fast).
Figure 2: Accuracy versus number of resampled rollouts. We find that with raw data observations ot (dotted lines), greater resample spacing N leads to slightly lower accuracy with equal S, with considerably more difference at higher S. This effect is significantly more pronounced with the smoothed ot (solid lines), which also shows that increasing log⁡S has a consistent linear effect on log-TVD. Larger N having decreasing accuracy for a given S has a relatively simple explanation, since this means fewer samples collected overall. Accuracy on the y-axis is measured by total variation distance (TVD) relative to the S=200 reference ot.
Figure 2: Accuracy versus number of resampled rollouts. We find that with raw data observations ot (dotted lines), greater resample spacing N leads to slightly lower accuracy with equal S, with considerably more difference at higher S. This effect is significantly more pronounced with the smoothed ot (solid lines), which also shows that increasing log⁡S has a consistent linear effect on log-TVD. Larger N having decreasing accuracy for a given S has a relatively simple explanation, since this means fewer samples collected overall. Accuracy on the y-axis is measured by total variation distance (TVD) relative to the S=200 reference ot.
Figure 3: Accuracy versus token cost. We find that overall, with raw ot data, wider spacing between resampled points N leads to better accuracy with lower overall token cost (dotted lines). Our statistical model significantly increases the accuracy for most S and N, but also negates the improvement of larger N increasing accuracy with equal cost (solid lines). TVD between the S=200 reference and each low-sample estimator is shown on the y-axis, with colors representing different N. The black star represents the sampling parameters used in (1).
Figure 3: Accuracy versus token cost. We find that overall, with raw ot data, wider spacing between resampled points N leads to better accuracy with lower overall token cost (dotted lines). Our statistical model significantly increases the accuracy for most S and N, but also negates the improvement of larger N increasing accuracy with equal cost (solid lines). TVD between the S=200 reference and each low-sample estimator is shown on the y-axis, with colors representing different N. The black star represents the sampling parameters used in (1).
Figure 4: Reconstructions for three questions from the full tinyMMLU evaluation set using Llama-3-8B-Instruct. From top to bottom, the rows show the S=200 every-token reference, the Full Model’s reconstruction from a nested S=15 run at 4-token spacing, and Raw Data from the same reduced run.
Figure 4: Reconstructions for three questions from the full tinyMMLU evaluation set using Llama-3-8B-Instruct. From top to bottom, the rows show the S=200 every-token reference, the Full Model’s reconstruction from a nested S=15 run at 4-token spacing, and Raw Data from the same reduced run.

Why it matters

Interpretability research that tries to pinpoint which reasoning steps decide an LLM's final answer has been limited by the huge cost of resampling, so this method lowers that cost barrier substantially. That opens the door to studying more questions and larger models to understand why a model answers the way it does.

Figure 5: DeepSeek counterpart to Figure 2. Pooled TVD to the full S=200 reference is plotted against the number of samples per position, S, with one line per observation spacing in sentences. Dotted lines are results for raw data, solid lines are smoothed data.
Figure 5: DeepSeek counterpart to Figure 2. Pooled TVD to the full S=200 reference is plotted against the number of samples per position, S, with one line per observation spacing in sentences. Dotted lines are results for raw data, solid lines are smoothed data.
Figure 6: DeepSeek counterpart to Figure 3. Pooled TVD is plotted against the number of tokens generated per question; colors match Figure 5. Dotted lines are results for raw data, solid lines are smoothed data. The star marks the baseline setting of S=30 at 1-sentence spacing.
Figure 6: DeepSeek counterpart to Figure 3. Pooled TVD is plotted against the number of tokens generated per question; colors match Figure 5. Dotted lines are results for raw data, solid lines are smoothed data. The star marks the baseline setting of S=30 at 1-sentence spacing.
Figure 7: Forking region counterpart to Figure 2. TVD to the S=200 reference is evaluated only at forking regions, using forking thresholds of 0.10 (left) and 0.15 (right). Solid lines show Smoothed Data, and dotted lines show Raw Data; colors match those in the main text.
Figure 7: Forking region counterpart to Figure 2. TVD to the S=200 reference is evaluated only at forking regions, using forking thresholds of 0.10 (left) and 0.15 (right). Solid lines show Smoothed Data, and dotted lines show Raw Data; colors match those in the main text.
Figure 8: Fork-region counterpart to Figure 3. TVD in forking regions is plotted against the number of tokens generated per question at forking thresholds of 0.10 and 0.15. Solid lines show Smoothed Data, and dotted lines show Raw Data.
Figure 8: Fork-region counterpart to Figure 3. TVD in forking regions is plotted against the number of tokens generated per question at forking thresholds of 0.10 and 0.15. Solid lines show Smoothed Data, and dotted lines show Raw Data.

Terms in this paper

  • Forking Paths Analysis · a method that resamples continuations at each generation step to see how the model's eventual answer could branch
  • forking point · a position where the distribution of possible final answers suddenly and sharply changes
  • PELT (Pruned Exact Linear Time) · an algorithm that detects points in a sequence where the underlying statistical pattern changes
  • kernel pooling · smoothing neighboring data points using Gaussian-weighted averaging to reduce noise
  • Total Variation Distance (TVD) · a measure of how different two probability distributions are; lower means more similar
Figure 9: Paired per-question TVD differences between each ablation variant and the Full Model for DeepSeek-R1-Distill-Llama-8B at 1-sentence spacing. Error bars show 95% t-intervals over tinyMMLU, and the dashed zero line denotes the Full Model. Results are shown for pooled accuracy (left) and forking region accuracy at threshold 0.10 (right). Removing segmentation harms accuracy near forks but not pooled accuracy; fixing the detection cost to L2 is indistinguishable from the Full Model in both panels.
Figure 9: Paired per-question TVD differences between each ablation variant and the Full Model for DeepSeek-R1-Distill-Llama-8B at 1-sentence spacing. Error bars show 95% t-intervals over tinyMMLU, and the dashed zero line denotes the Full Model. Results are shown for pooled accuracy (left) and forking region accuracy at threshold 0.10 (right). Removing segmentation harms accuracy near forks but not pooled accuracy; fixing the detection cost to L2 is indistinguishable from the Full Model in both panels.
Figure 10: Ablation overlays for two example questions. Reconstructions use S=30 on the densest spacing N and are evaluated against leave-replicate-out S=200 references; dashed lines mark reference forks. The left column shows Llama-3-8B-Instruct question 46 using the ablation variants, cross-validation, and operating point from the recorded S=200 store. The right column shows the recorded DeepSeek example, where the fork occurs when continuations stop resolving to an answer. Removing segmentation (Kernel Pooling) blurs the forks, whereas removing kernel pooling (Segment+Pool) reduces gradual drift to discrete steps. The L2 detection-cost variant is visually indistinguishable from the Full Model.
Figure 10: Ablation overlays for two example questions. Reconstructions use S=30 on the densest spacing N and are evaluated against leave-replicate-out S=200 references; dashed lines mark reference forks. The left column shows Llama-3-8B-Instruct question 46 using the ablation variants, cross-validation, and operating point from the recorded S=200 store. The right column shows the recorded DeepSeek example, where the fork occurs when continuations stop resolving to an answer. Removing segmentation (Kernel Pooling) blurs the forks, whereas removing kernel pooling (Segment+Pool) reduces gradual drift to discrete steps. The L2 detection-cost variant is visually indistinguishable from the Full Model.
Figure 11: Replicate TVD versus the number of samples per position, S, on log–log axes. The measured curve remains within 2% of the exact i.i.d. multinomial null through S=200. Triangles show the full dataset means for both models.
Figure 11: Replicate TVD versus the number of samples per position, S, on log–log axes. The measured curve remains within 2% of the exact i.i.d. multinomial null through S=200. Triangles show the full dataset means for both models.
Figure 12: Number of detected segments as a function of the PELT penalty for S=30 dense-grid counts on the two development questions. Curves show the exact multinomial cost, a trend-aware cost, and the silently substituted L2 cost. The dashed line marks the ideal segment count.
Figure 12: Number of detected segments as a function of the PELT penalty for S=30 dense-grid counts on the two development questions. Curves show the exact multinomial cost, a trend-aware cost, and the silently substituted L2 cost. The dashed line marks the ideal segment count.

Original abstract (English)

LLM reasoning is stochastic, and so understanding a model requires grappling with the distribution of reasoning chains that it might produce for a given question, i.e., its uncertainty. Resampling-based analyses characterize this distribution, revealing which steps of a rollout determine how the model arrives at its answer. However, a major limitation of these approaches is that resampling text sequences at every token or sentence in a reasoning chain is very costly. Our work strives to make resampling analysis more computationally efficient, while also shedding light on an important scientific question: what is the right statistical model for explaining uncertainty dynamics in text generation? We show that when resampling many reasoning chains, uncertainty dynamics converge to stable patterns, and noise is largely an artifact of sampling rather than an LLM's sensitivity to each individual token or reasoning step. We develop a statistical model for smoothing noisy low-sample rollout data to better approximate high-sample data, allowing us to significantly cut sampling costs.

Authors · Eric Bigelow, Amir Zur, Satchel Grant, Tal Haklay, Can Rager, Owen Lewis, Thomas McGrath, Jack Merullo, Ekdeep Singh Lub

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: Eric Bigelow et al., arXiv:2608.19611, CC BY 4.0