AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth
Authors: Shixiang Song, He Li, Zitong Wang, Boyi Zeng, Feichen Song, Yixuan Wang, Zhiqin John Xu, Ziwei He, +1 more
Organizations: LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai Innovation Institute · Sun Yat-sen University · Shanghai Jiao Tong University · Shanghai AI Laboratory
Test-time scaling via recurrent/iterative Transformers enables large language models to spend more computation at inference, but most pretrained recurrent LMs run a fixed number of iterations, wasting compute on easy tokens and lacking token-wise adaptivity. Following the core idea of Adaptive Computation Time(ACT) and Early Exit(EE), we propose AdaPonderLM, a self-supervised recurrent language model that learns token-wise early exiting during pretraining without manually tuned per-token/per-layer pruning ratios. AdaPonderLM uses iteration-specific MLP gates with a monotonic halting mask to decide when each token stops recurring, and introduces a KV reuse mechanism that reuses cached key/value states for halted tokens, ensuring train--test consistency and practical acceleration. Across Pythia backbones from 70M to 410M (pretraining) and up to 2.8B (continued pretraining), AdaPonderLM reduces inference compute at about 10% while maintaining comparable language modeling perplexity and competitive downstream accuracy. Our analysis shows the learned gates allocate more computation to high-NLL (hard) tokens, exhibiting adaptive computation time behavior in a fully self-supervised setting. Meanwhile, under iso-FLOPs, the learned halting policy consistently outperforms fixed pruning, showing AdaPonderLM allocates compute to the right tokens rather than just reducing average depth.
Figures & tables
Figure 1: Illustration of the gating behavior at inference time. Across recurrent iterations, tokens that are pruned reuse their KV cache from the previous iteration, avoiding redundant computation. In this example, token B is halted by the gating MLP after the first iteration; its KV cache is then reused in the second and third iterations. Tokens D and F are halted in the second iteration and similarly reuse their KV caches thereafter.
Figure 2: Mechanism overview. The same Transformer is executed recurrently across iterations. At iteration i , an iteration-specific MLP produces gate probabilities, which update a persistent mask mi . KV states are aligned token-wise via where : pruned tokens reuse cached KV states, ensuring training–inference consistency.
Figure 3Figure 4
Backbone
Inference FLOPs
#Params
Val Loss ( ↓ )
Val PPL ( ↓ )
AdaPonderLM
3.8×
70M
2.66
14.32
PonderLM
4×
70M
2.67
14.40
Pause Token
4×
70M
2.79
16.27
Loop Transformer
4×
70M
2.71
14.98
AdaPonderLM
3.8×
410M
2.29
9.87
PonderLM
4×
410M
2.27
9.72
Table 1: Validation perplexity (PPL) and loss of baselines.
Backbone
Inference FLOPs
Training#Tokens
#Params
Val Loss ( ↓ )
Val PPL ( ↓ )
AdaPonderLM
3.7×
310B
1.4B
1.92
6.82
PonderLM
4×
300B
1.4B
1.92
6.82
PonderLM
4×
310B
1.4B
1.92
6.82
AdaPonderLM
3.8×
305B
2.8B
1.83
6.21
PonderLM
4×
300B
2.8B
1.83
6.21
Table 2: Validation perplexity (PPL) and loss of AdaPonderLM and PonderLM under different Pythia backbones.
Model (#training tokens)
Lambada OpenAI
ARC -E
Lambada Standard
ARC -C
Wino Grande
PIQA
Hella Swag
SciQ
RACE
Avg acc / Δ acc ↑
0-shot
Pythia-410M (300B)
51.4
52.2
36.4
21.4
53.8
66.9
33.7
81.5
30.9
47.6
PonderLM-410M (25B)
45.8
48.7
33.5
20.3
51.5
65.9
32.4
79.7
29.9
45.3
Loop Teansformer-410M (25B)
44.3
48.1
31.9
20.0
51.3
64.1
32.1
79.3
31.2
44.7
Pause Token-410M (25B)
46.3
50.5
35.7
21.6
52.1
65.0
31.9
81.6
30.3
46.1
AdaPonderLM-410M (25B)
45.4
47.9
32.6
20.1
49.7
64.7
31.8
76.5
30.3
44.3
Table 3: Zero-shot and five-shot performance (accuracy, %) on downstream tasks. For all baselines, we evaluate the released pretrained checkpoints from their official repositories. Δ acc reports the mean accuracy improvement over the matched Pythia baseline.
Figure 3: Hyperparameter Sensitivity. Loss increases monotonically with λ (2.7425 → 2.7455) and k (2.7425 → 2.7454), indicating that lower regularization strengths are preferred for minimizing loss.
Figure 4: Gating mechanism reduces effective compute steps by 10% without compromising representation quality.
Figure 5: Dynamics of pruning across recurrent iterations. Top : Mean token NLL for tokens halted at different forward steps. Bottom : Density map of the token NLL distribution (clipped for clarity).
Figure 6: Adaptive Policy vs. Fixed Distributions. Top : Eval loss on the Pile under Iso-FLOPs. Bottom : Effective keep probability per step for each strategy.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Hyperparameter
Value
Hyperparameter
Value
Batch Size
256
Max Positional Embedding
2048
ZerO Stage
0
Dropout
0
Weight Decay
0.1
BF16
True
Warmup Ratio
0.02
MLP Hidden States
4h
lr scheduler
cosine
Adam Beta1
0.9
Adam Beta 2
0.95
Adam Epsilon
1e-8
Appendix
Table 4: The hyperparameter for the pretrained model.
Test-time scaling has shown that allocating more additional computation at inference can improve generation quality, motivating a natural follow-up question: where should this computation be spent? Building on this insight, we introduce PonderLM-3, a pretraining framework for token-wise adaptive pondering that learns to selectively allocate additional computation under purely self-supervised objectives, built on top of the PonderLM-2 backbone. This makes additional inference computation an allocatable per-token resource, so tokens receive more computation only when it is beneficial, rather than paying a uniform extra cost. To make this allocation learnable while maintaining train-inference consistency, PonderLM-3 injects a differentiable attention mask during pretraining and pairs it with a matching hard pruning rule at inference. PonderLM-3 defines a stronger Pareto frontier: compared with existing recursive or adaptive baselines, it achieves lower pretraining perplexity at equal inference FLOPs. On downstream benchmarks, PonderLM-3 attains comparable performance to fixed-step PonderLM-2 under the same maximum number of additional computation steps, while using fewer inference FLOPs in practice. Overall, PonderLM-3 provides an end-to-end differentiable and train-inference consistent framework for token-wise adaptive computation, enabling additional inference compute to be allocated where it is most useful rather than paid uniformly by every token.
He Li, Feichen Song, Boyi Zeng +4
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai AI Laboratory · Shanghai Jiao Tong University +1
Recurrent-depth language models improve inference-time reasoning by repeatedly applying a shared network to refine internal representations, but the additional recurrent computation also increases inference cost. This overhead is particularly pronounced during long-context prefill, where the full input sequence is repeatedly processed across recurrent depth. Weight-and-activation quantization offers a natural way to accelerate this computation through lower-precision arithmetic, yet repeatedly applying quantized computation introduces numerical perturbations that accumulate along the recurrent trajectory. We systematically analyze how such errors propagate across recurrent depth and show that quantization errors introduced earlier are attenuated by more subsequent recurrent steps and therefore have a smaller effect on the final representation. Matched-budget experiments strongly support this prediction, showing that placing quantized computation earlier in recurrent depth consistently yields the best fidelity. Based on this insight, we introduce RAMP (Recurrent Adaptive Mixed Precision), which performs early recurrent refinement using low-precision weight-and-activation computation and restores high precision for terminal refinement. Across context lengths up to 4K, RAMP preserves near-full-precision prefill quality while achieving up to 1.16x end-to-end first-prefill speedup on an NVIDIA H100, corresponding to a 13.9% latency reduction.
Ke Wan, Chen Chen
Department of Computer Science, University of Virginia
Efficient LLM inference research has largely focused on reducing the cost of each decoding step (e.g., using quantization, pruning, or sparse attention), typically applying a uniform computation budget to every generated token. In practice, token difficulty varies widely, so static compression can over-compute on easy steps and under-compute on hard ones. We study dynamic budget allocation for autoregressive decoding: learning how much computation to spend per token from within a single model. Self-Optimizing Language Models (SOL) pair a frozen LLM with a lightweight policy network that reads the LLM hidden state and selects a discrete efficiency action at each decode step. Actions can jointly control (i) token-level attention sparsity, (ii) structured activation pruning in the MLP, and (iii) activation quantization bit-width, while leaving the base model weights unchanged. We train the policy with group-relative policy optimization on teacher-forced episodes: the token sequence is fixed, while we sample multiple compute schedules (i.e., "counterfactual" schedules that vary only the efficiency actions for the same token path) and compare their likelihoods under the same supervision. Our reward trades off language-model quality against soft penalties that encourage episode-average budget usage to match a requested target. Across model variants and compute regimes, SOL improves quality at matched budget over static allocation and strong random schedule search, offering a complementary axis for inference-efficiency optimization. SOL discovers a better quality-efficiency pareto-front across all our experiments and improves MMLU accuracy by up to 7.3% over uniform budget allocation strategies.