Parcae: A Dynamical Systems Perspective to Stable Looped LLMs

Traditional fixed-depth architectures scale performance by increasing training FLOPs, typically through expanded parameterization at the expense of higher memory footprints. Recent work has suggested a promising alternative — looped architectures — which increases FLOPs while keeping parameters constant by looping layers in the model. While initial results are promising, existing recipes for training looped architectures can be unstable, suffering from residual explosion and loss spikes. To address these challenges, we recast layer looping as a nonlinear time-variant dynamical system. Linearizing this system, we observe strict theoretical conditions for stability on the parameters of the system. With these conditions in mind, we propose Parcae, a novel looped Transformer that enforces stability on the looped layers. When trained on the same data, Parcae converges up to 1.25x faster than prior looped models, while reducing downstream perplexity by up to 27% and 64% in comparison to similarly-sized prior looped models and transformers, respectively. Finally, we investigate recurrence as a medium to increase FLOPs, observing a new scaling axis in training and providing a means of scaling test-time compute and improving quality with constant memory overhead.