Saved in:
Bibliographic Details
Main Authors: Li, He, Song, Feichen, Zeng, Boyi, Song, Shixiang, Xu, Zhiqin John, He, Ziwei, Lin, Zhouhan
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2603.02023
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914382473068544
author Li, He
Song, Feichen
Zeng, Boyi
Song, Shixiang
Xu, Zhiqin John
He, Ziwei
Lin, Zhouhan
author_facet Li, He
Song, Feichen
Zeng, Boyi
Song, Shixiang
Xu, Zhiqin John
He, Ziwei
Lin, Zhouhan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2603_02023
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking
Li, He
Song, Feichen
Zeng, Boyi
Song, Shixiang
Xu, Zhiqin John
He, Ziwei
Lin, Zhouhan
Computation and Language
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.
title PonderLM-3: Adaptive Token-Wise Pondering with Differentiable Masking
topic Computation and Language
url https://arxiv.org/abs/2603.02023