Search papers, labs, and topics across Lattice.
The paper introduces AdaPonderLM, a recurrent language model that learns token-wise early exiting during pretraining using iteration-specific MLP gates and a monotonic halting mask, eliminating the need for manual tuning of pruning ratios. AdaPonderLM incorporates a KV reuse mechanism to maintain train-test consistency and accelerate inference by reusing cached key/value states for halted tokens. Experiments on Pythia models (70M-2.8B) demonstrate that AdaPonderLM achieves a 10% reduction in inference compute with comparable perplexity and downstream accuracy, allocating more computation to high-NLL tokens.
Key contribution not extracted.
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.
