[ARCH-PROP] Looped OLMo

[JT-Ushio]

Architecture Name

Looped OLMo

Parent issue

#6

Motivations

Standard decoder Transformers increase depth by assigning different parameters to every layer. This provides high capacity, but it does not explicitly encourage the model to refine a representation through repeated application of the same update rule. It also ties effective depth directly to parameter count and optimizer-state memory.

Looped Transformers or Looped LM instead reuse a Transformer block across multiple depth iterations. On synthetic in-context learning problems, this recurrent-depth structure learns iterative refinement behavior and matches standard Transformers with fewer than 10% of the parameters. Re-injecting the original input at every iteration and supervising a tail window of loop outputs are important for stable convergence beyond the trained loop count.

The existing results are mainly from small synthetic regression and function-learning tasks. The proposed experiment asks whether the same inductive bias transfers to autoregressive language-model pretraining: can an OLMo 3 model trade unique layer parameters for recurrent depth while retaining language-model quality and improving parameter efficiency?

Proposed Architecture

The model retains the OLMo 3 tokenizer, embeddings, normalization, attention, RoPE, MLP, causal masking, and language-model head. A small group of consecutive Transformer blocks forms a shared loop body $F_\theta$. The body should align with OLMo 3's attention cycle so that one loop preserves the original [SWA, SWA, SWA, Full] pattern.

Let $P$ denote the embedded input sequence and initialize the recurrent depth state as $H_0=0$. The same loop body is applied $b$ times with additive input injection:

$$ H_{t+1}=F_\theta(H_t+P), \qquad t=0,\ldots,b-1. $$

Input injection prevents the original token representation from vanishing as the loop count increases. Each iteration uses the same parameters but represents a different virtual depth. The final normalization and language-model head are shared across loop outputs.

Training applies the causal language-model loss to the final $T$ loop outputs:

$$ L=\frac{1}{T}\sum_{t=b-T+1}^{b}\mathrm{CE}(\mathrm{LMHead}(\mathrm{Norm}(H_t)),y). $$

The first $b-T$ iterations run without gradient tracking, and the recurrent state is detached before the final $T$ differentiable iterations. This bounds activation memory while preserving all $b$ forward iterations. The maximum loop count $b$ is progressively increased during training to avoid optimizing a long recurrence from initialization.

At inference, the model runs a fixed loop count and may optionally be evaluated with fewer or additional loops. Parameter count and optimizer-state memory depend on the loop-body depth rather than effective depth. Compute remains proportional to the total number of loop applications. Autoregressive decoding also requires a KV cache for each virtual layer/iteration, so this proposal does not claim KV-cache savings.

Experiments Plan

1B validation

• Compare standard OLMo 3 with Looped OLMo 3 at matched effective depth and training compute.
• Report both quality-versus-unique-parameter and quality-versus-FLOP comparisons; parameter savings must not be presented as compute savings.
• Compare input injection against plain weight tying and compare tail-window supervision against final-loop-only supervision.
• Sweep the training loop count $b$ and differentiable tail length $T$, then evaluate below, at, and beyond the trained loop horizon.

Key metrics

• Record training loss, evaluation loss/perplexity, downstream-task accuracy, and performance at each loop iteration.
• Monitor gradient norms, activation norms, loss spikes, and post-horizon divergence.
• Report unique parameter count, optimizer-state memory, peak training memory, FLOPs per token, tokens per second, and total training time.
• Report decode latency and KV-cache size for the selected inference loop count.

Scaling decision

• Proceed to 3B and 7B configurations if the 1B model reaches a better quality/parameter trade-off than standard OLMo 3 without unacceptable compute, latency, or stability regressions.