Abstract
Looped LMs are parameter efficient and promise dynamic computation (saving memory and FLOPs on easy tokens). However, state-of-the-art open Looped LMs trained with this dynamic computation capability (Ouro models) do not realize it in practice as each loop iteration (depth) requires its own level of KV-cache, necessitating all loop computations. Moreover, Ouro's early-exit prior is enforced on each token equally, which results in static lower-depth like processing of all tokens regardless of difficulty. In this work, we propose a simple "best-available" KV caching strategy that works out-of-the-box, creating a new frontier in the performance vs depth space. Our approach enables up to 30% reduction in FLOPs and KV memory while retaining full-depth performance, showing the true flexibility of Looped LMs. Furthermore, training looped LMs with awareness about this KV caching strategy improves performance and efficiency. Finally, we apply a small but effective fix to the early-exit prior enforcement objective that makes tokens exit at truly heterogeneous depths based on effort. Our findings are validated on Ouro models as well as smaller looped LMs pre-trained from scratch.
Keywords
Subject
Publication details
- Journal
- Not available
- Open access
- Green open access
Cite this article
APA 7
Mishra, A., Masdemont, A. P., Vendrell, V. C., Ros-Giralt, J., Behboodi, A., & Massoli, F. V. (2026). Enabling Dynamic Computation in Looped LMs. https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms
MLA 9
Mishra, Aayush, et al. "Enabling Dynamic Computation in Looped LMs." https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms.
Chicago (author–date)
Mishra, Aayush, Arnau Padrés Masdemont, Victor Conchello Vendrell, Jordi Ros-Giralt, Arash Behboodi, and Fabio Valerio Massoli. 2026. "Enabling Dynamic Computation in Looped LMs." https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms.
Harvard
Mishra, A., Masdemont, A. P., Vendrell, V. C., Ros-Giralt, J., Behboodi, A. and Massoli, F. V. (2026) 'Enabling Dynamic Computation in Looped LMs', Available at: https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms.
Vancouver
Mishra A, Masdemont AP, Vendrell VC, Ros-Giralt J, Behboodi A, Massoli FV. Enabling Dynamic Computation in Looped LMs. https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms
IEEE
A. Mishra, A. P. Masdemont, V. C. Vendrell, J. Ros-Giralt, A. Behboodi, and F. V. Massoli, "Enabling Dynamic Computation in Looped LMs," https://omanscience.com/en/articles/enabling-dynamic-computation-in-looped-lms.