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

2026-03-02

AI Summary

This work addresses the challenge of dynamically allocating computation during inference to avoid uniformly increasing overhead across all tokens. The authors propose an end-to-end differentiable, self-supervised pretraining framework built upon the PonderLM-2 architecture, which introduces a learnable, differentiated attention mask coupled with a hard-pruning inference rule. This approach achieves, for the first time, token-level adaptive computation with consistent training and inference behavior. Under identical inference FLOPs, the method substantially reduces pretraining perplexity compared to uniform baselines, while matching or exceeding the downstream task performance of fixed-step methods—all with lower actual computational cost.

Technology Category

Machine Learning: Hardware-aware MLNatural Language Processing: (Large) Language ModelsComputer Vision: Diffusion Models for Vision

Application Category

Economics, Online Markets and Human Computation: Cost models of using LLMs in production systemsWeb Mining and Content Analysis: Large pretrained models with web dataSearch and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• adaptive computation
• token-wise pondering
• test-time scaling
• inference efficiency
• differentiable masking

Innovation

• adaptive computation
• differentiable masking
• token-wise pondering
• train-inference consistency
• test-time scaling