Critical Windows of Complexity Control: When Transformers Decide to Reason or Memorize

TL;DR

This paper identifies a specific training window during which regularization determines whether Transformers learn to reason or memorize, revealing the importance of timing and initialization in training outcomes.

Contribution

It demonstrates that the memorization versus reasoning behavior in Transformers is governed by a sharp, task-specific training window influenced by initialization and regularization timing.

Findings

A 25% training window of weight decay matches full training in OOD accuracy.

Positioning regularization in the middle of training yields significantly higher OOD accuracy.

The critical window's onset is highly sensitive, shifting with as little as 100 steps.

Abstract

Recent work has shown that Transformers' compositional generalization is governed by \emph{complexity control}, initialization scale and weight decay, which steers training toward low-complexity reasoning solutions rather than high-complexity memorization. Existing analyses, however, treat complexity control as a single static hyperparameter choice, leaving open \emph{when} during training this control is actually decisive. We show that the memorization-versus-reasoning fate of a Transformer is determined within a sharp, identifiable window of training. On a controlled compositional task we find that (i)~weight decay applied for a single 25\%-of-training window matches full-training weight decay in out-of-distribution (OOD) accuracy ($0.93$ vs $0.91$); (ii)~holding total regularization budget constant, placing it in the middle of training yields $5{-}9\times$ higher OOD accuracy than placing it early; (iii)~the boundary of the critical window is remarkably sharp, window onset shifted by as little as $100$ optimization steps causes mean OOD to jump from chance ($0.15$) to reasoning-regime ($0.61$); (iv)~the window's position depends systematically on initialization scale, but the basin of attraction for reasoning solutions \emph{shrinks} at small initialization, contradicting the prevailing recommendation that smaller initialization is uniformly better. We further show that the critical-window phenomenon is task-specific: it does not appear on grokking with modular arithmetic, where properly tuned constant weight decay matches scheduled weight decay.