Model Architectures · Looped Transformers

DeepLoop

Depth Scaling for Looped Transformers

A loop-aware residual-scaling rule for weight-tied looped Transformers: when depth comes from revisiting the same blocks, the DeepNorm exponent $p=1/4$ is no longer the right one, and DeepLoop replaces it with the aligned-visit threshold $p=1/2$.

Shuzhen Li, Yifan Zhang, Jiacheng Guo, Quanquan Gu, Mengdi Wang
Princeton University  ·  Princeton AI Lab  ·  UCLA  ·  arXiv:2607.13491
🤗
Looped TransformersDepth ScalingPost-LN DeepNormHRM · ARC-AGI

Abstract

Overview of the DeepLoop framework.
Figure 1: Overview of the DeepLoop framework. A looped Transformer stores $K$ physical blocks once and revisits them for $R$ rounds, giving unrolled depth $N=KR$; each residual sublayer applies $\mathbf{x}_{i+1}=\mathrm{Norm}(\alpha\,\mathbf{x}_i+f_j(\mathbf{x}_i;\phi_j))$, with $\alpha=(2N)^{1/2}$ scaling the skip connection and $\beta=(8N)^{-1/2}$ the per-matrix initialization gain.

DeepLoop is a residual-scaling rule for weight-tied looped Transformers. Looped models increase effective depth by revisiting the same physical blocks, and the DeepNorm exponent that stabilizes untied depth ($p=1/4$) is no longer the right one. DeepLoop makes the tied-depth effect explicit through a visit-alignment analysis and keeps the Post-LN DeepNorm architecture with a single loop-aware scaling rule:

$$\mathbf{x}_{i+1} = \mathrm{Norm}\big(\alpha\,\mathbf{x}_i + f(\mathbf{x}_i)\big), \qquad \alpha = (2N)^{1/2}, \quad \beta = (8N)^{-1/2},$$

where $N$ is the unrolled (effective) depth and $\beta$ is the initialization gain of each residual-branch output projection. The exponent $p=1/2$ is the aligned-visit threshold; it recovers DeepNorm's $1/4$ when visits decorrelate. The same rule applies to hierarchical recurrent reasoners (HRM), where it improves voted accuracy on ARC-AGI.

Read the Paper Code

One Loop-Aware Scaling Rule

DeepLoop changes nothing about the block itself: the model stays a Post-LN DeepNorm Transformer. What changes is how the residual stream and the residual branches are scaled once the same physical weights are visited $R$ times.

01

Tied depth is not untied depth

A looped model of $L$ physical blocks unrolled $R$ times behaves like a network of effective depth $N$, but its residual-branch updates are produced by the same weights at every visit, so their contributions do not average out the way independent layers do.

02

Aligned-visit threshold

A visit-alignment analysis makes the tied-depth effect explicit: when visits align, update magnitudes accumulate coherently and the stabilizing exponent moves from DeepNorm's $p=1/4$ to $p=1/2$.

03

Two constants, one rule

The rule sets $\alpha=(2N)^{1/2}$ on the residual stream and $\beta=(8N)^{-1/2}$ as the initialization gain of each residual-branch output projection, with $N$ the unrolled effective depth.

Recovering DeepNorm

The aligned-visit exponent $p=1/2$ is a threshold, not a replacement: when visit directions decorrelate, the analysis recovers the standard DeepNorm exponent $p=1/4$. Untied residual networks are the decorrelated special case of the same picture.

Where the Rule Applies

Weight-tied looped language models

GPT-style models (MHA with RoPE) whose blocks are revisited $R$ times per forward pass. The baseline keeps $\alpha=\beta=1$; the DeepLoop variant enables the $\alpha=(2N)^{1/2}$, $\beta=(8N)^{-1/2}$ scaling on every block, with the loop count set in the training config.

Hierarchical recurrent reasoners

The same rule drops into the Hierarchical Reasoning Model (HRM) as a residual-scaling modification, where recurrent high-level and low-level modules revisit their weights many times per reasoning step.

Experiments

The repository provides end-to-end training and evaluation for both settings.

Language modeling (FineWeb-Edu)

Multi-GPU training via torchrun, one config file per experiment, comparing looped baselines ($\alpha=\beta=1$) against DeepLoop scaling at matched loop counts, with lm-evaluation-harness runners for downstream evaluation.

Reasoning (ARC-AGI via HRM)

A DeepLoop residual-scaling drop-in for HRM together with p-sweep launchers, plus a hash-verified vanilla-HRM baseline package for reproducible comparison. DeepLoop improves voted accuracy on ARC-AGI.

FineWeb-Edu validation loss at step 100K versus loop count.
Figure: FineWeb-Edu validation loss at step 100K versus loop count $R$. DeepLoop improves over the $\alpha=\beta=1$ looped baseline on both the 124M and 350M backbones, with the gap widening as $R$ grows (2.7514 vs 2.7700 at $R=7$ on 124M, 2.5280 vs 2.5558 on 350M).

Repository layout

model/                 model definitions (GPT-MHA-RoPE, looped variants, LoopDeepNorm scaling)
config/                training configurations (FineWeb-Edu LM and ARC-AGI)
utils/                 tokenizers, optimizer, checkpoint, distributed, reproducibility utilities
scripts/               training / evaluation launchers
integrations/hrm/      DeepLoop applied to the Hierarchical Reasoning Model (ARC-AGI)
integrations/vanilla_hrm/  reproducible vanilla-HRM baseline package (build + eval)
train_adam_finewebedu.py   FineWeb-Edu language-modeling entry point
train_adam_arc_agi.py      ARC-AGI entry point

Citation

If you find this work useful, please cite:

@article{deeploop2026,
  title  = {DeepLoop: Depth Scaling for Looped Transformers},
  author = {Li, Shuzhen and Zhang, Yifan and Guo, Jiacheng and Gu, Quanquan and Wang, Mengdi},
  journal= {arXiv preprint arXiv:2607.13491},
  year   = {2026}
}