Technical Summary: Adaptive Nucleus Truncation for Long-Form Reasoning

Problem Statement

In long-form language model reasoning, the generation budget is often large enough to support multi-step reasoning, self-refinement, and tool use. In these regimes, the decoding sampler is not merely a post-processing detail but a critical control mechanism that determines reachable reasoning paths and the stability of training trajectories. Existing truncation methods—such as top-$p$, min-$p$, and fixed top-$n_\sigma$ sampling—improve upon unrestricted sampling but suffer from significant limitations:

• Fixed Thresholds: Methods like top-$p$ and min-$p$ rely on fixed probability thresholds that cannot adapt to changes in entropy, task difficulty, training stage, or generation budget.
• Probability-Space Distortions: Probability-based methods operate after softmax and temperature scaling, inheriting distortions from probability normalization. Consequently, the same underlying logit geometry can yield different candidate sets as temperature changes.
• Lack of Adaptability: A single fixed truncation parameter (e.g., a fixed $n$ in top-$n_\sigma$) cannot simultaneously serve decisive contexts requiring sharp pruning and ambiguous contexts requiring wider support. Furthermore, in unstable Reinforcement Learning (RL) phases, aggressive truncation can lead to training instability (overshooting in entropy, KL divergence, and gradient norms).

Methodology: Adaptive Nucleus Truncation Sampling (ANTS)

The authors propose Adaptive Nucleus Truncation Sampling (ANTS), which extends top-$n_\sigma$ sampling from a fixed decoding heuristic into an adaptive, budget-aware rollout-control mechanism. ANTS operates in three stages:

1. Logit-Space Support Estimation

Unlike probability-based methods, ANTS selects the candidate set in the pre-temperature logit space. It defines a neighborhood $N_t(n)$ around the maximum logit $\ell_{t, \text{max}}$ based on the vocabulary-wise standard deviation $\sigma(\ell_t)$:
$$N_t(n) = \{ v \in V \mid \ell_{t,v} > \max_{u \in V} \ell_{t,u} - n \cdot \sigma(\ell_t) \}$$
This selection is invariant to positive logit rescaling (temperature scaling), ensuring the candidate set remains consistent regardless of temperature changes.

2. Entropy-Conditioned Thresholding

The width of the truncation neighborhood, $n_t$, is not fixed but adapts based on the local uncertainty of the model. It is calculated as:
$$n_t = n_0 + \gamma \cdot H(p^{(0)}_t)$$
where $H(p^{(0)}_t)$ is the entropy of the unit-temperature distribution. High-entropy contexts (ambiguous) retain more alternatives (larger $n_t$), while low-entropy contexts (decisive) are pruned more sharply.

3. Online Thompson-Sampling Controller

To determine the optimal truncation strength $\gamma$ dynamically, ANTS employs a Thompson-sampling bandit controller over a finite set of arms:

• Finite Arms: A set of $K$ arms with log-spaced $\gamma$ values ($\gamma_k = 10^{\eta_k}$).
• Fallback Arm: A critical $K+1$-th arm with $\gamma_{K+1} = +\infty$, which effectively disables truncation ($N_t = V$), recovering the baseline sampler.
• Reward Signal: The controller uses an intrinsic reward based on the entropy of the distribution induced by each arm. Arms that preserve sufficient support (higher-than-average entropy) receive larger posterior updates.
• Stability Mechanism: The fallback arm acts as a "safe escape hatch." If truncation becomes unsafe during RL rollouts (causing metrics like KL divergence or gradient norms to drift), the controller can learn to select the fallback arm to stabilize training.

Key Contributions

1. Adaptive Logit-Space Sampler: Formulation of ANTS as a temperature-invariant sampler that separates candidate-set selection from within-set stochasticity.
2. Entropy-Driven Control: Introduction of an entropy-conditioned Thompson controller that adapts truncation strength online and includes a no-truncation fallback arm to prevent training instability.
3. Generation-Budget Scaling Evaluation: Comprehensive evaluation across 8K, 16K, and 32K generation budgets, demonstrating that performance gains generally scale with the budget.
4. Budget-Aware Truncation: Characterization of a specific failure mode in short-budget code generation, motivating the need for truncation policies that are jointly learned with generation budget and task type.

Experimental Results

The method was evaluated on a 33B-total / 4B-active sparse Mixture-of-Experts (MoE) reasoning model.

• Scaling with Budget: Average performance over percentage-based benchmarks improved by +1.9 points at 8K, +3.8 at 16K, and +5.2 at 32K.
• Instruction Following (IFBench): Showed the strongest gains, improving by +10.5 (Loose) and +10.8 (Strict) points at 32K. ANTS prevented drift into verbose or tangential continuations that violate constraints.
• Mathematical Reasoning (AIME 2025): Improved by +7.0 points at 32K, with gains nearly doubling from 8K to 32K.
• Code Generation (Codeforces): Revealed a critical budget interaction. At 8K, ANTS trailed the baseline by -59 ELO despite generating more solution tokens. However, at 16K and 32K, it reversed this gap, achieving +230 and +212 ELO gains, respectively. This suggests truncation is harmful in tightly constrained code regimes but beneficial when the feasible program space is larger.
• Scientific QA (GPQA) & Knowledge (MMLU Pro): Showed near-parity or modest gains, indicating that logit-space truncation offers less leverage for tasks requiring precise single-answer selection compared to long-form generation.
• Token Allocation: ANTS altered token usage patterns. At moderate budgets, it increased "thought" tokens (deliberation). At 32K, it reduced thought tokens while maintaining high accuracy, suggesting it suppresses unnecessary continuation rather than simply generating longer chains.

Significance and Claims

The paper argues that sampler design should be treated not merely as a fixed decoding hyperparameter, but as an integral component of inference-time scaling and rollout control.

• Stability: The inclusion of the fallback arm is central to the method's success in RL, allowing the system to recover from unstable truncation states that fixed methods cannot escape.
• Context Sensitivity: The results demonstrate that the optimal sampling strategy is context-dependent, varying by generation budget, task type (e.g., code vs. math), and training state.
• Efficiency: ANTS often reaches performance bands earlier than the baseline, improving effective compute efficiency in RL settings by reducing the frequency of low-quality or unstable rollouts.

The authors conclude that while ANTS does not universally dominate all existing samplers across all tasks, it exhibits a robust generation-budget scaling pattern, particularly for long-form reasoning and instruction following, where suppressing noisy tails without removing useful reasoning paths is critical.