Adaptive Nucleus Truncation for Long-Form Reasoning

Abstract: Sampling plays an important role in long-form language-model reasoning. Over thousands of decoding steps, small changes in the candidate token set can compound into different reasoning trajectories, stability profiles, and final answers. Existing truncation methods such as top-$p$, min-$p$, and fixed top-$nσ$ sampling improve over unrestricted sampling, but they rely on fixed thresholds that cannot adapt to changes in entropy, task difficulty, training stage, or generation budget. We introduce Adaptive Nucleus Truncation Sampling (ANTS), which extends top-(nσ) sampling from a fixed decoding rule into an adaptive rollout-control mechanism for long-form generation. ANTS selects standardized neighborhoods around the maximum logit before temperature scaling, adapts the truncation width using an entropy-conditioned controller, and retains a no-truncation fallback arm to stabilize training when truncation becomes unsafe. On a 33B-total / 4B-active sparse Mixture-of-Experts reasoning model, ANTS improves average performance over percentage-based benchmarks by +1.9, +3.8, and +5.2 points at 8K, 16K, and 32K generation budgets, respectively. The strongest gains appear on instruction following and mathematical reasoning, with IFBench improving by more than 10 points at 32K and AIME 2025 improving by 7 points. Code generation reveals an important budget interaction. On Codeforces, ANTS trails the baseline at 8K, but reverses this gap and substantially improves ELO at 16K and 32K. These results suggest that sampler design should be treated not just as a decoding hyperparameter, but as part of how we stabilize and scale long-budget reasoning.

• The paper introduces ANTS, an adaptive logit-space sampling method that integrates context entropy with a Thompson bandit controller for improved long-form reasoning.
• It demonstrates significant scaling gains on benchmarks like AIME and IFBench, with performance improvements up to +7.0% at a 32K token generation budget.
• It highlights the importance of budget-aware strategies in sampling, revealing trade-offs between applications such as code generation and instruction following.

Budget-Aware Logit-Space Sampling for Long-Form LLM Reasoning

Motivation and Background

The sampling mechanism is central to how LLMs explore and exploit token-level possibilities during incremental decoding. In long-form reasoning tasks, where generation budgets span thousands of tokens, compounding stochasticity and support truncation dictate both the diversity and trajectory of reasoning paths. Classical sampling approaches (top-$k$, nucleus, min-$p$) restrict the candidate set in probability space, but these depend heavily on temperature scaling and probability normalization, leading to unpredictable variance in the candidate set composition across token positions and tasks.

Recent developments indicate that informative and noisy tokens are more cleanly separated in pre-softmax logit space; for instance, top-$n\sigma$ sampling retains tokens within a fixed standard-deviation window of the maximum logit [tang-etal-2025-top]. However, fixed thresholds fail to adapt to local context entropy and task-specific requirements, especially as the generation budget increases. The paper introduces Adaptive Nucleus Truncation Sampling (ANTS), which transforms logit-space sampling into an online, budget-aware rollout controller, explicitly incorporating context entropy and an RL-style Thompson bandit controller with an explicit fallback to the untruncated baseline. This method decouples candidate-set selection from stochasticity induced by temperature, providing a more stable and adaptive sampling strategy for RL rollouts in LLMs.

Methodology Overview

ANTS extends top-$n\sigma$ into an adaptive, entropy-conditioned controller with stability guarantees. At each generation step $t$, the logit-space nucleus $\mathcal{N}_t(n)$ contains tokens whose logits exceed $\max \ell_{t,u} - n_t \sigma(\ell_t)$. Here, the width $n_t$ is contextually set using the base entropy $\mathcal{H}(p_t^{(0)})$ of the unit-temperature probability vector, such that $n_t = n_0 + \gamma \mathcal{H}(p_t^{(0)})$, where $p$0 is selected online by a Thompson bandit over a finite arm set plus a fallback arm that disables truncation. Beta posteriors track standardized entropy-based rewards for each arm; the controller updates all arms per step, not just the selected arm, enabling efficient online adaptation. The fallback arm is critical for stabilizing RL rollouts, preventing runaway entropy, gradient norm, or importance ratio statistics by reverting to the baseline when truncation becomes unsafe.

Empirical Evaluation Across Generation Budgets

ANTS is evaluated on a sparse MoE reasoning model (33B total, 4B active) across 8K, 16K, and 32K generation budgets on diverse benchmarks: AIME 2024/2025 (mathematical), GPQA (scientific QA), MMLU Pro (knowledge retrieval), Codeforces ELO (code generation), and IFBench/IFEval (instruction following). The primary empirical claim is a scaling gain: ANTS achieves average improvements of +1.9, +3.8, and +5.2 percentage points at 8K, 16K, and 32K, respectively. Strongest gains appear in instruction following and mathematical reasoning, with IFBench improving by +10.8 points and AIME 2025 by +7.0 at 32K. However, ANTS is not uniformly beneficial under all budgets; at 8K on Codeforces, ANTS trails the baseline by 59 ELO, but reverses to +230/+212 ELO at 16K/32K, exposing a critical interaction between truncation strategy and completion length.

Figure 1: ANTS improves AIME and IFBench at 8K, is near-parity on GPQA, and trails Codeforces ELO, indicating limitation under tight budgets.

Figure 2: At 16K, ANTS achieves broad gains, including a large Codeforces ELO reversal and strong IFBench improvements.

Figure 3: At 32K, ANTS sustains a stable IFBench lead, improves AIME 2025 by +7.0 points, and maintains Codeforces ELO gains.

Analysis of Task-Specific Effects

Mathematical Reasoning

AIME 2024 peaks at the lowest budget (+4.2 points at 8K) but does not scale monotonically, while AIME 2025 demonstrates robust scaling, from +3.7 points at 8K to +7.0 at 32K. This reflects different solution lengths and distributional shifts, implying that tasks requiring extensive multi-step reasoning benefit more from budget-aware truncation.

Instruction Following

ANTS demonstrates pronounced gains on IFBench: Loose Prompt improves by +10.5 and Strict by +10.8 at 32K. These gains are accompanied by lower variance and more stable compliance, suggesting that logit-space truncation reduces tail events leading to constraint violation. On IFEval, gains are smaller due to near-saturation at low budgets, but ANTS continues to improve stability and mean accuracy at higher budgets.

Code Generation and Budget Dependency

Codeforces presents a non-monotonic effect: ANTS underperforms at 8K despite increased solution token usage, indicating that the gap is not a simple completion-length artifact, but likely due to less favorable distributional shaping under constrained budgets. At 16K/32K, ANTS reverses this trend, suggesting that truncation should always be interpreted in conjunction with generation budget, task type, and training state.

Figure 4: At 8K, ANTS increases thought-token usage and produces more Codeforces solution tokens despite lower ELO.

Figure 5: At 16K, ANTS uses more thought tokens across all tasks, producing fewer solution tokens, consistent with longer deliberation and more compact answers.

Figure 6: At 32K, ANTS uses fewer thought tokens while maintaining benchmark gains, indicating budget-dependent effects.

Token Allocation and Controller Dynamics

Token statistics reveal non-trivial allocation behavior: at 8K and 16K, ANTS generally increases intermediate reasoning length (thought tokens), with 18–19% increases for AIME 2025 and up to 44% for Codeforces at 16K. At 32K, the effect reverses; ANTS uses fewer thought tokens while sustaining accuracy gains, corroborating the argument that ANTS is not simply producing longer traces but reallocating budget adaptively. The entropy-driven reward structure encourages arms preserving broad support, preventing aggressive candidate collapse, and the fallback arm enables recovery when excessive pruning destabilizes rollouts.

Related Work and Theoretical Implications

ANTS builds on the geometric insights of logit-space separation ([tang-etal-2025-top], [ma2025estimatingllmuncertaintyevidence]) and extends nucleus/min-$p$1 sampling ([Holtzman2020The], [nguyen2025turningheatminpsampling]) by decoupling candidate-set construction from temperature scaling. The bandit-style online controller with fallback embodies RL-style adaptive exploration and exploitation ([Thompson1933ONTL], [pmlr-v23-agrawal12]), enabling sampler-level rollout shaping without policy-gradient augmentation. Empirical scaling laws in RLHF ([10.5555/3618408.3618845], [guo2025deepseek]) echo the importance of sampler design for avoiding reward model overoptimization and stabilizing RL signals.

Practical and Theoretical Implications

The chief practical implication of ANTS is that sampler design cannot be regarded as a fixed decoding hyperparameter for LLM inference and RL fine-tuning. Rather, it must form an integral part of rollout allocation, especially in long-form reasoning regimes with large budgets. Theoretical implications extend to support estimation, entropy-driven diversity control, and RL stability; the explicit fallback arm and adaptive width conditioning align with recent advances in dynamic temperature scheduling and rollout allocation ([liao-etal-2025-enhancing]). In the context of MoE architectures and RL-based scaling, adaptive sampling mechanisms such as ANTS may become central to efficient and robust RLHF and reward model optimization.

Future directions include budget-conditioned fallback policies, stronger calibration and coverage guarantees for rare token reachability ([awad2026lostsamplingassessinglexical]), and integration with reward-model auditing to prevent undesirable support collapse. The efficacy of entropy-driven bandit controllers in other generative modalities (e.g. diffusion, vision-language) and specialized RL rollouts is a prospective research avenue.

Conclusion

ANTS introduces a robust adaptive logit-space sampling method for long-form LLM reasoning, decoupling candidate selection from temperature and enriching sampler design with entropy conditioning, online Thompson-style control, and a fallback safety mechanism. Empirical results confirm scaling gains with generation budget and task specificity, with critical reversals in code generation exposing the necessity for budget-aware policies. These findings reframe sampling as an integral component of inference-time scaling and RL rollout control, rather than a mere decoding detail (2606.13982).