World Model Depth Benefits Vary in Autoregressive Rollouts

Adaptive-compute world models, which dynamically adjust the computational depth per step, operate on the assumption that greater depth leads to better predictions and can be adaptively routed. This research investigates whether the precision gained from depth at each step actually persists when predictions are composed in autoregressive rollouts. Using a pre-registered instrument called the "shallow penalty" (the ratio of error from shallowest-exit rollouts to full-depth rollouts), the study tested nine DeepMind Control tasks. The findings revealed three distinct regimes: for six tasks, depth improved rollouts significantly; for two tasks, shallow exits surprisingly outperformed full-depth models (an "inversion"); and one task showed no difference. The robust inversion observed in tasks like "cheetah" was not inherent to the dynamics but was a consequence of the training methodology. Specifically, supervising early exits only at the first rollout step eliminated this inversion, highlighting a "routability catch-22" where the supervision intended to make exits routable can inadvertently train them to surpass the full stack. The study also found that observation/action dimensionality and one-step model error could partly predict these regimes.