Understanding Dynamic Compute Allocation in Recurrent Transformers

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Moosa, Ibraheem Muhammad, Lohit, Suhas, Wang, Ye, Chatterjee, Moitreya, Yin, Wenpeng
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911434971021312
author Moosa, Ibraheem Muhammad
Lohit, Suhas
Wang, Ye
Chatterjee, Moitreya
Yin, Wenpeng
author_facet Moosa, Ibraheem Muhammad
Lohit, Suhas
Wang, Ye
Chatterjee, Moitreya
Yin, Wenpeng
contents 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 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.
format Preprint
id arxiv_https___arxiv_org_abs_2602_08864
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Understanding Dynamic Compute Allocation in Recurrent Transformers
Moosa, Ibraheem Muhammad
Lohit, Suhas
Wang, Ye
Chatterjee, Moitreya
Yin, Wenpeng
Computation and Language
Artificial Intelligence
Machine Learning
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 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.
title Understanding Dynamic Compute Allocation in Recurrent Transformers
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2602.08864