LayerNorm as Implicit Gain Control in Looped Transformers

This paper investigates the stability and memory mechanisms of looped transformers—architectures that reuse a single transformer block across multiple depth iterations. While these models offer significant inference-time savings, they are prone to training instabilities. The author analyzes the interaction between the recurrent 'carry' (a diagonal scaling matrix) and the internal transformer block, specifically focusing on how LayerNorm influences the system's dynamics. The study uses analytical derivations and empirical verification on small-scale models to characterize the stability boundary of these recurrent architectures.

The central contribution is the discovery that LayerNorm functions as an implicit gain controller. By coupling the block's local Lipschitz constant inversely to the activation scale, LayerNorm ensures that the recurrence Jacobian remains asymptotically contractive at fixed points, even when the operator norm exceeds 1. This means the system's stability is governed by the spectral margin (the distance of the spectral radius from 1) rather than the operator-norm bound. The author demonstrates that the 'carry' mechanism, often assumed to be the primary memory component, actually functions as a stabilization tool. In most cross-channel tasks, gradient descent routes memory through the more expressive nonlinear recurrence of the transformer block, leaving the carry at rest.

Understanding these stability properties is critical for deploying looped transformers at scale, where training divergences are prohibitively expensive. By identifying that the stability of these models is structurally tied to the inverse-norm coupling of LayerNorm, the paper provides a theoretical foundation for why certain constraints (like those proposed in the Parcae architecture) work. It also clarifies the boundary of this behavior: in tasks with axis-aligned per-channel structure, the model does recruit the carry, suggesting that the memory-stability tradeoff is highly dependent on the task structure.