TRE-K/TRE-P: Restricting Entropy in LLMs

• The paper introduces TRE-K and TRE-P methods that restrict entropy regularization to a trusted subset of actions, effectively reducing tail-noise in large language models.
• It proposes a re-normalized entropy formulation over plausible tokens, ensuring controlled exploration while maintaining stability in long generation tasks.
• Empirical evaluations demonstrate that TRE methods outperform global entropy approaches on reasoning and alignment tasks, highlighting improved performance and robustness.

Restricting entropy maximization to plausible actions in reinforcement learning with LLMs addresses the failure modes observed when applying naïve entropy regularization in vast action spaces. Standard entropy bonuses—formulated as global Shannon entropy over the entire vocabulary—are detrimental in LLMs because the action space (often 10^5 tokens) and long generation horizons lead to the accumulation of probability mass on semantically invalid tokens. Over many decoding steps, this “tail noise” severely degrades coherent reasoning by injecting unstructured randomness at every step. Trust Region Entropy (TRE), and specifically its instantiations TRE-K (“top-K”) and TRE-P (“top-p”), restrict the entropy maximization process to a dynamically selected set of plausible actions, thereby directing exploration to high-confidence regions of the token distribution while preserving stability and reasoning fidelity (Huang et al., 3 Feb 2026).

1. Motivation: Limitations of Global Entropy Regularization

Naïve entropy regularization in RL is conventionally used to encourage policy exploration via a global entropy term added to the standard surrogate (e.g., vanilla PPO loss). For LLMs, this is written as:

$$L_{\text{total}} = L_{\text{surr}} + \beta(-H(\pi(\cdot|s))),$$

where $\pi(a|s)$ outputs probabilities for actions $a$ given state $s$, and $H(\cdot)$ denotes Shannon entropy over the full token set $A$. Because $|A|$ can be $\gtrsim 10^5$, even a minuscule allocation of mass to the tail can result in substantial aggregate leakage, $$\epsilon = \sum_{a\notin \text{valid}} \Delta \pi(a)$$. Over a sequence of length $T$, the probability of not sampling an invalid token decays as $\pi(a|s)$0, rapidly diminishing to zero for long horizons. This compounds incoherence and degrades performance on long reasoning tasks (Huang et al., 3 Feb 2026).

2. Trust Region Entropy: Core Principle and Formalism

Trust Region Entropy (TRE) is designed to remedy the global tail risk by maximizing entropy only within a carefully chosen subset of plausible tokens—the “trust region”—at each decision point. Let $\pi(a|s)$1 denote this subset for state $\pi(a|s)$2. The policy is re-normalized over $\pi(a|s)$3:

$\pi(a|s)$4

and $\pi(a|s)$5 otherwise. The local entropy within this trust region is

$\pi(a|s)$6

To ensure comparability in scale with global entropy, the local entropy is rescaled by $\pi(a|s)$7:

$\pi(a|s)$8

If $\pi(a|s)$9, the entropy bonus is omitted ($a$0).

The overall RL objective at step $a$1 becomes

$a$2

preserving the PPO surrogate and implicit KL-regularization properties but constraining entropy-driven exploration to trusted actions. This formulation ensures exploration without mass drifting into implausible or invalid regions (Huang et al., 3 Feb 2026).

3. TRE-K and TRE-P: Instantiations of the Trust Region

Two practical variants are introduced:

• TRE-K (“top-K”): Here, $a$3 comprises the $a$4 tokens with highest logits at step $a$5. The entropy bonus is computed as

$a$6

where $a$7 is the entropy of the top-$a$8 re-normalized policy. The algorithm selects the $a$9 largest-logit tokens, re-normalizes, computes the entropy, and applies the scaled bonus.

• TRE-P (“top-p” or nucleus): Here, $s$0 is the smallest subset of $s$1 whose cumulative softmax probability exceeds a threshold $s$2. The corresponding bonus is

$s$3

with $s$4 the selected prefix tokens. This variant adapts the region size to the model’s confidence—expanding when uncertain (large $s$5), shrinking when confident (small $s$6)—yielding smoother policy confidence dynamics and more effective regulation of exploration (Huang et al., 3 Feb 2026).

4. Empirical Performance and Comparative Analysis

Empirical evaluation demonstrates the superiority of both TRE-K and TRE-P over global entropy regularization, vanilla PPO, Forking-Tokens, and covariance-based KL penalties (KL-Cov) across mathematical reasoning (MATH), combinatorial (Countdown), and preference alignment (HH) tasks. Key results from Table 1 in (Huang et al., 3 Feb 2026):

Method	MATH Pass@1 ↑	Countdown Pass@1 ↑	HH Reward ↑
PPO (vanilla)	57.04%	64.12%	3.24
Entropy (global)	56.64%	63.20%	3.19
Forking-Tokens	57.16%	62.82%	3.32
KL-Cov	58.23%	66.50%	3.39
TRE-K ($s$7)	58.26%	66.28%	3.82
TRE-P ($s$8)	58.28%	66.96%	3.88

TRE-P specifically led to more pronounced improvements on larger models (Qwen2.5-7B), with smoother entropy regularization and better maintenance of exploratory capacity over training, as indicated by non-saturated top-token probabilities and resilience against premature convergence (Huang et al., 3 Feb 2026).

5. Algorithmic Implementation and Complexity

Integration of TRE-K or TRE-P into standard RL fine-tuning loops for LLMs is direct:

• Compute actor logits for each step in the rollout.
• Apply TRE-K (select top-$s$9 logits) or TRE-P (greedily accumulate tokens until cumulative softmax mass $H(\cdot)$0).
• Re-normalize the selected logits and compute the corresponding entropy.
• Scale and add the TRE loss to the PPO surrogate.
• Backpropagate through the combined objective.

The additional computational cost is $H(\cdot)$1 for top-$H(\cdot)$2 and $H(\cdot)$3 for top-$H(\cdot)$4, but can be mitigated via partial-sorting and prefix-sum optimizations. The wall-clock overhead remains modest relative to model forward-pass costs (Huang et al., 3 Feb 2026).

6. Connections to Related Entropy Regularization Approaches

Restricting entropy maximization to plausible actions has convergent motivation with methods such as AEnt (“clamped entropy”) (Shen, 3 Sep 2025) and SIREN (“selective entropy regularization”) (Jiang et al., 29 Sep 2025). All adopt the principle of restricting the entropy bonus to a reduced, high-confidence subset of the action space:

• AEnt evaluates clamped entropy on a subset of top $H(\cdot)$5-percent tokens, adaptively adjusting its entropy coefficient to maintain bounded entropy within this plausible region (Shen, 3 Sep 2025).
• SIREN applies a two-step masking strategy: top-p masking for output token selection and a peak-entropy mask (selecting only positions with top entropies), in combination with a self-anchored entropy penalty that stabilizes entropy drift while targeting exploration (Jiang et al., 29 Sep 2025).

These convergent approaches highlight the consensus that global entropy regularization is fundamentally unsuited to large-scale, sparse-reward reasoning tasks, and that targeted, trust region-style entropy is required for stable and effective LLM RL fine-tuning.

7. Practical Considerations and Hyperparameterization

Best empirical results using TRE observed with $H(\cdot)$6. Increasing $H(\cdot)$7 or $H(\cdot)$8 reintroduces tail noise, while $H(\cdot)$9 (i.e., minimum entropy) collapses exploration and underperforms vanilla PPO. Proper hyperparameter tuning over a small validation set is critical.

Both TRE-K and TRE-P exhibit stability compatible with PPO’s clipped surrogate, and their locality ensures bounded per-step KL divergence. This design maintains trust-region goals of controlling policy divergence while enabling high-quality exploration strictly within high-confidence neighborhoods of the prior network’s generative distribution (Huang et al., 3 Feb 2026).

A plausible implication is that as LLMs scale further, explicit restriction of entropy maximization to trusted action subsets is increasingly essential for effective and stable policy improvement in RL finetuning—an insight reflected by several distinct research programs converging on similar methodology.