FLEET: From Logits Entropy to Enhanced Trajectories in Text Generation

AI summary

Overview

Research area: Machine learning, specifically inference-time decoding and test-time scaling for large language models.

Technical level: Advanced. The paper combines information-theoretic uncertainty metrics (entropy, varentropy), Monte Carlo Tree Search (pUCT), Online/Intermediate hidden-state clustering, and benchmark evaluation with verifiable rewards.

Scope: The paper proposes FLEET, a memory-augmented decoding method that replaces memoryless temperature sampling with a deterministic, search-like procedure that penalizes empirically suboptimal tokens at high-uncertainty decoding states, and evaluates it on LiveCodeBench and GSM8K.

What This Paper Is About

Sampling multiple completions from an LLM with temperature is one of the most common ways to improve accuracy at inference time, but it is memoryless: each draw ignores earlier draws and their evaluations, so as more samples are taken the method produces a growing share of semantically duplicate answers and returns diminish. The authors argue this is a structural flaw at "branching points" — the small subset of tokens that determine whether a whole reasoning trajectory succeeds — and propose FLEET, which keeps a memory of past trajectories and uses it to steer new generations. The goal is to reach higher accuracy under the same sampling budget, with a deterministic and lightly-tuned procedure.

Key Contributions

1. Vector Disjoint Set Union (VectorDSU): An online data structure that maps continuous hidden states into unified discrete search states, preventing trajectory duplication and preserving state utility history without storing every member vector.
2. FLEET: A memory-augmented, deterministic search paradigm that replaces memoryless temperature sampling with targeted exploration of the completion space, intervening as a soft logit penalty rather than replacing the decoder.
3. Empirical analysis: A comparison of FLEET against repeated sampling on mathematical and coding problems, using both ground-truth verifiers and outcome reward models (ORMs), reporting that FLEET's accuracy is always higher under the same budget on these tasks.
4. A principled calibration pipeline: A single-pass hyperparameter selection procedure (layer selection score, Gaussian Mixture Model threshold derivation, resampling-temperature calibration, and VectorDSU similarity threshold calibration) presented as an alternative to ad hoc grid search over temperature, top-k, top-p, and min-p.

Main Findings

• LiveCodeBench improvement: LiveCodeBench Pass@32 increases from 59.9% to 66.2% under the same budget, reported as a substantial gain on complex coding tasks.
• Efficiency parity and speed: FLEET achieves the same accuracy as the repeated sampling baseline with a 3x speedup.
• Determinism and low calibration cost: In the greedy-decoding configuration evaluated, the approach is deterministic and uses a single calibration pass to derive its principal hyperparameters, requiring only minimal modifications to existing LLM pipelines. The temperature-sampling baseline, by contrast, required repeated generations during temperature optimization through Bayesian search.
• Diagnosis of existing sampling methods: The authors report that current logit-processing strategies (top-k, top-p, min-p) can mitigate the lack of selective exploration and hyperparameter brittleness, but not sample inefficiency, because they are unaware of the rewards associated with generated completions.
• Entropy plus varentropy as the trigger: High conditional entropy alone is described as susceptible to false positives from redundant tokens and deterministic reasoning steps, so FLEET requires both entropy and varentropy to exceed their thresholds before treating a state as a branching point.
• Calibration values reported: In the example calibration diagnostics for Llama 3.2-3B, the resampling temperature is identified at approximately T ≈ 2.9 at the intersection where the exploratory tail mass surpasses the top-1 mass, with k = 32 in that illustration; the VectorDSU threshold is defined as the minimal similarity value beyond which normalized divergence is bounded below an error tolerance of 0.05.
• Cube-root exploration schedule: The target fraction of token states processed by FLEET follows ρ(i) = min(i^(1/3), 100)%, with thresholds relaxed via a running entropy–varentropy buffer when the observed hit rate falls below target.
• Penalty and prior settings: The logit penalty λ was dynamically set to the maximum logit in the experiments, and explored actions without an external task signal receive a default multiplicative prior of 0.5, described as an exploration bias rather than a learned domain probability.
• Reported limitations: The method currently requires directly patching the underlying model architecture; it can struggle in induction or deep-reflection cognitive states, where deterministic selection is characterized as actively detrimental; generated text can occasionally contain noise or gibberish; and cross-session/cross-domain transfer and parallel speedup benefits are explicitly not evaluated.
• GSM8K numeric results: The introduction to the experiments states that GSM8K is used for mathematical reasoning with a sparse binary reward, but the numerical GSM8K outcomes are not visible in the provided content, so no figures are reported here.

Methodology in Plain English

The authors start from the observation that greedy decoding gives zero sample variance, while temperature sampling injects noise uniformly across all steps, including steps that are purely syntactic and should stay deterministic. Their alternative is to detect only the genuinely uncertain decision points and remember what happened there.

At each generation step, the method inspects an intermediate hidden layer rather than the final layer, on the reasoning that final-layer logits are smoothed. It projects that intermediate state into vocabulary space using the language model head — the Logit Lens technique — and computes two uncertainty measures from the resulting distribution: entropy (how uncertain the distribution is overall) and varentropy (how spread out the log-probabilities are around the mean). When both exceed thresholds, the step is treated as a branching point.

The hidden state at that step is then matched against stored representative vectors using cosine similarity. If it is close enough to an existing representative, it is assigned to that cluster; otherwise a new cluster is created. This is what VectorDSU does: it keeps only representatives and compact metadata (the token taken, the next state reached, and the reward observed) rather than every vector, which keeps memory small. Because state retrieval relies on internal activations, the method can use a Monte Carlo Tree Search-style objective, specifically pUCT, where the language model's own probabilities serve as the prior and the observed rewards serve as the exploitation term.

FLEET then computes which tokens are empirically suboptimal and applies a logit penalty to them, leaving the configured decoder in place. Unexplored high-probability tokens are grouped into a single "exploration action" so they remain available. Hyperparameters are set once in a calibration pass: a composite score picks the intermediate layer, a Gaussian Mixture Model over the joint entropy–varentropy space sets the thresholds, the resampling temperature is annealed to where the exploratory tail outweighs the top-1 candidate, and the similarity threshold is chosen where measured divergence stays under tolerance.

Why This Matters

Impact on research: The paper reframes test-time scaling from "sample more and pick the best" to "search the completion space with memory." If the reported results generalize, it challenges the practice of tuning temperature, top-k, top-p, and min-p as the primary knobs for inference-time improvement, and it links decoding to the trajectory-level machinery of tree search. It also supplies token-level attributed trajectory data (which actions were tried, what rewards they earned, how they compared) that the authors suggest could feed supervised fine-tuning or reinforcement learning, though they do not evaluate that use.

Real-world applications (as motivated by the paper's own framing):

• Code generation in agentic or autonomous problem-solving pipelines, where the paper reports a LiveCodeBench Pass@32 gain from 59.9% to 66.2%.
• Mathematical reasoning with verifiable answers such as GSM8K.
• Deployment settings without ground truth, where an outcome reward model supplies the signal instead of a verifier.
• Inference pipelines that need deterministic, reproducible outputs, since the greedy-decoding configuration is deterministic.

Industry relevance: The method is presented as a soft constraint that does not replace the configured decoder, can batch structural updates between generations, and can run with production inference backends that compile the computation graph. The 3x speedup at matched accuracy and the single-calibration-pass setup address the practical costs that make repeated sampling expensive at scale. However, the requirement to patch the model architecture and the reported noise in generated text are adoption barriers the authors acknowledge.

Future Directions

• Cross-session and cross-domain transfer: The decoupled VectorDSU memory could let state-action statistics or priors learned on one worker or task be supplied to others, but the authors state this benefit is not evaluated in the present experiments.
• Parallelization and speedup measurement: Workers could reserve or be assigned distinct high-scoring branches in shared memory to reduce overlap, but quantifying the resulting parallel speedup is listed as future work.
• Compatibility with reasoning models: The authors report that the method may not synergize well with "thinking" or reflection-focused models, because deterministic selection is actively detrimental in reflective states — an open problem for the next generation of models.
• Output noise mitigation: A secondary lightweight language model is proposed as a way to rewrite syntactically or formatting-corrupted output, but this mitigation is not evaluated.
• Logit-only variants and training value: Whether FLEET can be applied purely to output logits (rather than patching the architecture) and whether its attributed trajectory data improves SFT or RL are both left open.

Target Audience

Researchers and engineers working on LLM inference, decoding strategies, and test-time scaling; practitioners who deploy repeated sampling and want higher accuracy for the same compute; and readers interested in the intersection of information-theoretic uncertainty signals, state-space search, and hidden-state clustering. The paper assumes familiarity with autoregressive decoding, temperature sampling, and tree-search notation, so beginners may find the methodology section demanding despite the accessible motivation.

Authors’ abstract

Solutions based on large language models (LLMs) often rely on temperature sampling to improve accuracy and stability by aggregating multiple samples from the completion distribution. However, this memoryless approach is inherently suboptimal: because it lacks awareness of prior generations and their evaluations, it produces an increasing proportion of semantically duplicate answers as more samples are drawn, leading to diminishing returns. To address this limitation, we introduce FLEET, a novel method that integrates a memory mechanism into the generation process. FLEET represents each generation as a sparse trajectory through states whose entropy exceeds a predefined threshold and uses these trajectories to infer per-token utility scores that adjust the logits. Benchmark evaluations demonstrate that FLEET achieves the same accuracy as the repeated sampling baseline, with a 3x speedup, and substantially improves accuracy on complex coding tasks (LiveCodeBench Pass@32 increases from 59.9% to 66.2%) under the same budget. Furthermore, in the greedy-decoding configuration evaluated here, the approach is deterministic and uses a single calibration pass to derive its principal hyperparameters, requiring only minimal modifications to existing LLM pipelines.