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

TL;DR

AdaPonderLM introduces token-wise adaptive depth in language models, enabling efficient inference by dynamically halting computation per token, which improves resource allocation without sacrificing performance.

Contribution

It presents a self-supervised recurrent language model with learned token-wise early exiting and KV reuse, enhancing inference efficiency and adaptivity.

Findings

Reduces inference compute by about 10% with maintained perplexity.

Learns gates that allocate more computation to harder tokens.

Outperforms fixed pruning under equal FLOPs conditions.

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.