Looped Transformers as Optimizers

2026-09-29

AI Summary

This study addresses the unclear design principles of recurrent transition mechanisms in recurrent Transformers by proposing the OperLoop model. This work interprets the recurrent hidden state as fast weights and establishes a theoretical mapping between recurrent transitions and gradient descent optimization. Through a local gradient update framework, it derives closed-form expressions while introducing explicit weight decay and adaptive step-size mechanisms. Experimental results demonstrate that this approach significantly reduces training loss and improves accuracy on commonsense reasoning tasks. Furthermore, under equivalent computational budgets, OperLoop surpasses existing baselines in generative performance. These findings provide effective theoretical and practical support for test-time compute scaling.

Abstract

Looped Transformers provide a parameter-efficient approach to depth scaling by repeatedly applying shared Transformer blocks. Recent reasoning models have likewise highlighted the value of scaling test-time computation through longer computation trajectories. However, the principles for designing effective loop transitions remain poorly understood. We view the looped hidden state as a fast weight that is updated throughout the depth. We formulate loop transitions as local gradient-based updates, with recurrent blocks predicting implicit targets at each depth. Our framework derives loop transitions in closed form from a projection, a local objective and an optimizer update rule. Mapping representative loop transitions into this framework reveals mismatches between their transitions and projections. We first align the input maps of existing transitions. We then derive OperLoop, which combines explicit weight decay, adaptive step size and a delta objective. The aligned variants reduce training loss and improve average commonsense accuracy. OperLoop improves average generative performance over the compared looped and non-looped baselines under matched training FLOPs. These results support the framework's usefulness for loop design. We extend the analysis to additional loop models and outline a roadmap for future loop transition design.

Problem

• Looped Transformers
• loop transitions
• depth scaling
• test-time computation

Innovation

• Looped Transformers
• Fast Weights
• Gradient-based Updates
• OperLoop
• Parameter-efficient