Pretraining with Token-Level Adaptive Latent Chain-of-Thought

Paper: 2602.08220 Authors: Boyi Zeng, Yiqin Hao, He Li, Shixiang Song, Feichen Song, Zitong Wang, Siyuan Huang, Yi Xu, ZiWei He, Xinbing Wang, Zhouhan Lin Categories: cs.CL, cs.LG

Abstract

This work enhances language models by increasing the computational work performed per token without adding parameters. The model generates variable-length internal reasoning sequences before predicting each token. Crucially, the model naturally learns to allocate more reasoning steps to challenging tokens and fewer to simpler ones — a behavior that emerges during standard pretraining without explicit supervision. An adaptive token-wise stopping mechanism reduces computational demands. Experiments on Llama models show consistent improvements even with fewer training FLOPs than prior recurrent baselines.

Key Contributions

• Token-level adaptive latent reasoning: Variable-length internal computation before each token prediction
• Emergent difficulty-aware allocation: The model learns to spend more compute on harder tokens without explicit signals
• Adaptive stopping: Token-wise early termination reduces unnecessary computation
• FLOP-efficient improvements: Better performance than recurrent baselines with fewer training FLOPs
• Compatible with standard pretraining: No special training procedures or auxiliary losses required

Background: Compute Per Token

Standard transformers allocate exactly the same computation to every token — a single forward pass through all layers. This is clearly suboptimal: predicting “the” after “in” requires far less reasoning than predicting a complex technical term or the next step in a mathematical derivation.

Several approaches have attempted to address this:

• Chain-of-thought prompting: Explicit reasoning tokens, but increases sequence length and requires supervision
• Recurrent depth: Repeating transformer blocks, but typically with fixed iteration counts
• Adaptive computation: Early exit or variable depth, but usually at the sequence level

This work introduces token-level adaptive latent reasoning that operates in the model’s hidden space.

Method

Latent Reasoning Sequences

Before predicting each output token $x_t$, the model generates a sequence of $k_t$ latent reasoning steps in hidden space:

$z_t^{(1)}, z_t^{(2)}, \ldots, z_t^{(k_t)} = f_{\theta}(h_t, z_t^{(1)}, \ldots, z_t^{(k_t-1)})$

where $h_t$ is the hidden state from the main transformer and $f_{\theta}$ is a shared reasoning module. Each step refines the representation before the final token prediction.

Adaptive Token-Wise Stopping

The number of reasoning steps $k_varies per token. A halting mechanism learns when sufficient reasoning has been performed:

$p_{\text{halt}}^{(j)} = \sigma(w_{\text{halt}}^T z_t^{(j)})$

When the cumulative halting probability exceeds a threshold, reasoning stops for that token. This naturally allocates:

• Few steps for predictable tokens (function words, common phrases)
• Many steps for challenging tokens (rare words, logical conclusions, mathematical operations)

Integration with Pretraining

The latent reasoning module is integrated directly into the pretraining objective. The standard next-token prediction loss is computed after the adaptive reasoning:

$\mathcal{L} = -\sum_t \log P(x_t | x_{<t}, z_t^{(1:k_t)})$

A small regularization term encourages computational efficiency by penalizing excessive reasoning steps:

$\mathcal{L}_{\text{reg}} = \lambda \sum_t k_t$

Emergent Difficulty-Aware Allocation

One of the most striking findings is that the model spontaneously learns to allocate more reasoning steps to tokens that are objectively more difficult to predict. Without any explicit difficulty labels or curriculum:

• High-entropy tokens (multiple plausible continuations) receive more reasoning steps
• Low-entropy tokens (highly predictable) receive fewer steps
• The allocation pattern correlates with human intuitions about linguistic complexity

This suggests that adaptive computation is a natural inductive bias that emerges when the architecture permits it.

Results

• Consistent improvements on Llama models across multiple benchmarks
• Fewer training FLOPs than prior recurrent baselines while achieving better performance
• Graceful degradation: Limiting maximum reasoning steps reduces compute with minimal quality loss
• Scaling behavior: Benefits increase with model size, suggesting the approach is complementary to parameter scaling

Significance

1. Efficient scaling: Increases effective model capacity without adding parameters
2. Natural adaptive compute: Demonstrates that difficulty-aware computation can emerge from standard training
3. Practical deployment: Adaptive stopping keeps inference costs manageable
4. Theoretical implications: Suggests that fixed-depth transformers are fundamentally limited in their per-token reasoning capacity
5. Pretraining integration: Unlike post-hoc methods, this approach benefits from being present during the entire pretraining process