Understanding Dynamic Compute Allocation in Recurrent Transformers

Token-level adaptive computation seeks to reduce inference cost by allocating more computation to harder tokens and less to easier ones. However, prior work is primarily evaluated on natural-language benchmarks using task-level metrics, where token-level difficulty is unobservable and confounded with architectural factors, making it unclear whether compute allocation truly aligns with underlying complexity. We address this gap through three contributions. First, we introduce a complexity-controlled evaluation paradigm using existing algorithmic and synthetic language tasks with parameterized difficulty, enabling direct testing of token-level compute allocation. Second, we propose ANIRA, a unified recurrent Transformer framework that supports per-token variable-depth computation while isolating compute allocation decisions from other model factors. Third, we use this framework to conduct a systematic analysis of token-level adaptive computation across alignment with complexity, generalization, and decision timing. Our results show that compute allocation aligned with task complexity can emerge without explicit difficulty supervision, but such alignment does not imply algorithmic generalization: models fail to extrapolate to unseen input sizes despite allocating additional computation. We further find that early compute decisions rely on static structural cues, whereas online halting more closely tracks algorithmic execution state. Code: https://github.com/merlresearch/ANIRA

Large language models usually spend the same amount of computation on every part of a piece of text they generate when solving a problem, even though some parts are harder than others. This can waste computation and make models more expensive to run. We introduce ANIRA, a model architecture that is comprised of a learned module that is applied repeatedly in order to generate the outputs, as well as “adaptive computation”, which is the capability of deciding how many times to apply the learned module for each part of the generated output. Using tasks where difficulty can be controlled and measured, we test whether ANIRA learns to spend more effort on harder problems. We find that it can learn this behavior even without being directly taught which problems are difficult. However, reasonable-looking effort does not always mean the model has learned the exact underlying rule: it can still fail on problems with difficulties that are not seen during training. Our results show that adaptive computation can help make AI systems more efficient, but it should be tested carefully before being treated as evidence of reliable reasoning.