AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth

Shixiang Song; He Li; Zitong Wang; Boyi Zeng; Feichen Song; Yixuan Wang; Zhiqin John Xu; Ziwei He

cs.CLMar 2, 2026

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

Abstract

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 $m^{i}$ . KV states are aligned token-wise via where : pruned tokens reuse cached KV states, ensuring training–inference consistency.

Figure 3

Figure 4

Backbone	Inference FLOPs	#Params	Val Loss ( $\downarrow$ )	Val PPL ( $\downarrow$ )
AdaPonderLM	$3.8\times$	70M	2.66	14.32
PonderLM	$4\times$	70M	2.67	14.40
Pause Token	$4\times$	70M	2.79	16.27
Loop Transformer	$4\times$	70M	2.71	14.98
AdaPonderLM	$3.8\times$	410M	2.29	9.87
PonderLM	$4\times$	410M	2.27	9.72

Table 1: Validation perplexity (PPL) and loss of baselines.

Backbone	Inference FLOPs	Training#Tokens	#Params	Val Loss ( $\downarrow$ )	Val PPL ( $\downarrow$ )
AdaPonderLM	$3.7\times$	310B	1.4B	1.92	6.82
PonderLM	$4\times$	300B	1.4B	1.92	6.82
PonderLM	$4\times$	310B	1.4B	1.92	6.82
AdaPonderLM	$3.8\times$	305B	2.8B	1.83	6.21
PonderLM	$4\times$	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 / $\Delta$ acc $\uparrow$
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.

\Delta

acc reports the mean accuracy improvement over the matched Pythia baseline.

Figure 3: Hyperparameter Sensitivity. Loss increases monotonically with $\lambda$ (2.7425 $\rightarrow$ 2.7455) and $k$ (2.7425 $\rightarrow$ 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.

Explore similar work

PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking
Mar 2, 2026He Li, Feichen Song, Boyi Zeng +4Adaptive Computation Language Model Pretraining
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai AI Laboratory · Shanghai Jiao Tong University +1
RAMP: Recurrent Adaptive Mixed Precision for Efficient Prefill in Recurrent-Depth Language Models
Oct 9, 2026Ke Wan, Chen ChenLLM Inference Acceleration Mixed-Precision Quantization
Department of Computer Science, University of Virginia
Compute Where it Counts: Self Optimizing Language Models
May 11, 2026Yash Akhauri, Mohamed S. AbdelfattahLLM Inference Efficiency RL for Language Models
Cornell University

Mar 2, 2026J/K move · Enter open · S save

PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking

He Li, Feichen Song, Boyi Zeng, Shixiang Song+3

LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai AI Laboratory · Shanghai Jiao Tong University+1

Adaptive Computation Language Model Pretraining Test-Time Scaling LLM Inference Acceleration

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking

RAMP: Recurrent Adaptive Mixed Precision for Efficient Prefill in Recurrent-Depth Language Models

Compute Where it Counts: Self Optimizing Language Models