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

Mar 2, 2026

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

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.

Shixiang Song, He Li, Zitong Wang +6

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

Ask PDF

May 11, 2026

Compute Where it Counts: Self Optimizing Language Models

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.

Yash Akhauri, Mohamed S. Abdelfattah

Cornell University

LLM Inference Efficiency RL for Language Models LLM Pruning LLM Inference Scheduling Inference-Time Compute Allocation

Ask PDF

Oct 3, 2026

ALoDLM: Adaptively Looped Diffusion Language Models

Diffusion language models (DLMs) enable fast generation by predicting multiple tokens in parallel, but their practical adoption remains limited by a persistent quality gap relative to comparably sized autoregressive (AR) models. We attribute this gap to a computation-difficulty mismatch: within a partially observed sequence, some unknown tokens are easy to predict, while others require substantially more computation. Existing DLMs nevertheless apply uniform computational depth to all unknown positions at each denoising step. We introduce ALoDLM, which replaces uniform computation with token-adaptive latent recurrence. At each denoising step, ALoDLM iteratively refines latent representations and allocates computation according to token difficulty. Tokens ready to commit are fed back as discrete context, while unresolved tokens retain and further refine their latent states through additional recurrent passes. To learn token prediction and computation allocation jointly, we formulate token-wise computation schedules as latent variables and derive a conditional negative evidence lower bound (NELBO). We train ALoDLM at 1.7B and 8B parameter scales. Across eleven benchmarks, ALoDLM outperforms all evaluated DLMs and the corresponding AR baselines in average benchmark score at both scales. ALoDLM also retains fast parallel decoding, yielding a strong quality-efficiency trade-off among evaluated autoregressive and diffusion models under optimized inference engines.

Liancheng Fang, Zhuowei Li, Youngeun Kim +10

University of Illinois Chicago · Amazon AGI · Korea University

Diffusion Language Models Diffusion Language Model Inference Diffusion Language Model Decoding

Ask PDF