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

He Li; Feichen Song; Boyi Zeng; Shixiang Song; Zhiqin John Xu; Ziwei He; Zhouhan Lin

cs.CLMar 2, 2026

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

Authors: He Li, Feichen Song, Boyi Zeng, Shixiang Song, Zhiqin John Xu, Ziwei He, Zhouhan Lin

Organizations: LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai AI Laboratory · Shanghai Jiao Tong University · Shanghai Innovation Institute

Abstract

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.

Figures & tables

Figure 1: From fixed compute to token-wise allocation. Standard LMs decode with one forward pass per token. PonderLM-2 applies fixed pondering steps, turning extra compute into a uniform per-token tax. PonderLM-3 makes refinement depth token-dependent, allocating additional internal updates only when they provide non-trivial marginal gains.

Method	Training Paradigm	Adaptivity Source	Constraints / Priors
PonderNet	Pretraining	Learned stopping distribution	Geometric prior
MoR / MoD / ITT	Pretraining	Router-based token selection	Fixed continuation ratio
Think-at-Hard	SFT (LoRA)	Supervised step targets	External step labels
Learning to Ponder	RL (GRPO)	Learned controller	Frozen backbone; reward-shaped budget
CALM / DAT	Post-hoc	Confidence-based stopping	Calibration of threshold
LayerSkip	Pretraining	Draft–verify	Static exit policy

Table 1: Adaptive-compute methods vs. ours.

Figure 2: Token-wise adaptive pondering in PonderLM-3. (1) A lightweight router predicts the step distribution $s_{t,k}$ from $h_{t}^{(0)}$ , and obtains a monotone mask score $w_{t,k}$ via the tail-CDF of $s$ ; $w_{t,k}\in[0,1]$ serves as a soft attention mask (color intensity indicates $w$ ). (2) the same LLM performs $K$ pondering steps; at each step, the current hidden state is inserted as a latent token at its corresponding position to form an interleaved sequence, and $w_{t,k}$ is applied in attention as a soft mask. Then we use Jacobi iterations to approximate the sequential inference dynamics, and the output uses the $s_{t,k}$ -weighted fused hidden state. (3) Inference: token-by-token decoding uses the same attention masking and fusion mechanism, and early exits once $w_{t,k}<\tau$ , skipping remaining pondering steps.

Figure 3: Performance–compute Pareto curve. The x-axis is the average number of executed additional computation steps per token (an inference FLOPs proxy). PonderLM-3 yields a better compute–quality trade-off than the baselines.

Model	Inf. FLOPs (G/token)	Lambada OpenAI	ARC -E	Lambada Standard	ARC -C	Wino Grande	PIQA	Hella Swag	SciQ	RACE	Avg.
5-shot
LoopedLM (4 loops)	9.84	34.2	48.2	26.6	20.6	47.0	64.2	30.7	83.4	27.8	42.5
Pause (3 pauses)	9.84	30.7	46.9	26.3	21.8	51.6	63.4	30.1	81.8	28.5	42.3
MoR ( max 3 steps)	8.44	31.5	48.5	25.3	20.0	51.0	63.7	30.6	84.7	28.0	42.6
PonderLM-1 (3 steps)	9.84	34.0	49.1	28.4	21.8	49.7	64.1	30.9	84.4	29.5	43.5
PonderLM-2 (3 steps)	9.84	37.6	50.6	34.0	21.8	53.5	65.0	32.5	87.5	31.2	46.0

Table 2: Downstream accuracy under reported inference compute (estimated inference FLOPs using

6ND

Panel of Figure 4 — Figure 4: Where extra compute helps. (a) Bucket-wise $\Delta CE_{i}$ over executed additional computation steps. (b) Intrinsic difficulty $\ell_{t}$ versus executed additional computation steps.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

	$x_{1}$	$h_{1}^{(0)}$	$h_{1}^{(1)}$	$x_{2}$	$h_{2}^{(0)}$	$h_{2}^{(1)}$	$x_{3}$	$h_{3}^{(0)}$	$h_{3}^{(1)}$
$x_{1}$	$1$
$h_{1}^{(0)}$	$1$	$w_{1,0}$
$h_{1}^{(1)}$	$1$	$w_{1,0}$	$w_{1,1}$
$x_{2}$	$1$	$w_{1,0}$	$w_{1,1}$	$1$
$h_{2}^{(0)}$	$1$	$w_{1,0}$	$w_{1,1}$	$1$	$w_{2,0}$
$h_{2}^{(1)}$	$1$	$w_{1,0}$	$w_{1,1}$	$1$	$w_{2,0}$	$w_{2,1}$

Appendix

Figure 7: Soft attention mask over an interleaved sequence. Rows correspond to queries , and columns correspond to keys . Entries show the multiplicative mask values applied to attention weights. Blank entries are masked out by causality or padding.

Figure 8: RMSE convergence of Jacobi-style parallel iterations. Average RMSE between training-iteration latent states

h^{(n)}

and the inference-time fixed point

h^{*}

over iteration index

n

(log scale). The curve decays rapidly and is well approximated by an exponential fit (effective

L\approx 0.433

R^{2}\approx 0.9959

), indicating fast approach to the inference-time solution.

Table 10

Explore similar work

AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth
Mar 2, 2026Shixiang Song, He Li, Zitong Wang +6LLM Inference Acceleration Language Modeling
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai Innovation Institute · Sun Yat-sen University +2
Compute Where it Counts: Self Optimizing Language Models
May 11, 2026Yash Akhauri, Mohamed S. AbdelfattahLLM Inference Efficiency RL for Language Models
Cornell University
ALoDLM: Adaptively Looped Diffusion Language Models
Oct 3, 2026Liancheng Fang, Zhuowei Li, Youngeun Kim +10Diffusion Language Models Diffusion Language Model Inference
University of Illinois Chicago · Amazon AGI · Korea University

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

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

Shixiang Song, He Li, Zitong Wang, Boyi Zeng+5

LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Shanghai Innovation Institute · Sun Yat-sen University+2

LLM Inference Acceleration Language Modeling Adaptive Computation

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Compute Where it Counts: Self Optimizing Language Models

ALoDLM: Adaptively Looped Diffusion Language Models