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

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.

• The paper investigates how Transformers' solutions to compositional tasks are influenced by temporally-targeted regularization during training.
• Empirical data reveals a crucial 25% training window where weight decay maximizes generalization, challenging the notion that constant low initialization is always beneficial.
• Theoretical analysis outlines the dynamics of memorization and reasoning circuits, showing weight decay's selective impact during specific training intervals.

Critical Windows of Complexity Control: Temporal Structure of Reasoning and Memorization in Transformers

Introduction

The study "Critical Windows of Complexity Control: When Transformers Decide to Reason or Memorize" (2605.04396) provides a rigorously controlled empirical and theoretical analysis of how regularization—specifically, weight decay and initialization scale ($\gamma$)—determines whether Transformers converge to reasoning or memorization solutions in compositional tasks. Conventional wisdom suggests that compositional generalization is linked to hyperparameter selection, especially favoring constant low initialization and static weight decay. This work shifts focus to the temporal aspect of regularization, positing and validating that the decisive influence of weight decay is localized to a critical, time-bounded window during training. The paper presents key empirical findings, robust theoretical modeling, and important caveats that challenge earlier dogma, especially regarding the universal benefit of smaller $\gamma$. It further delineates the scope of these dynamics, showing task specificity and optimizer/depth robustness.

The Anchor-Function Task and Experimental Methodology

The empirical backbone of the paper is the anchor-function task [zhang2025complexity], a controlled compositional benchmark where a model predicts the result of composing two anchor-defined permutations applied to a key. The Transformer models used have two layers and $d=64$, with precisely tracked per-step OOD accuracy and order parameters, enabling high-resolution diagnostics.

Three core metrics are captured during training:

1. Condensation Index (participation ratio of singular values for value matrices), measuring low-rank convergence.
2. Bridge Alignment (subspace overlap between layer circuits), inspired by induction-head circuit analysis [song2025composition].
3. Weight Norm (for comparison with grokking phenomena [power2022grokking]).

Weight decay schedules are manipulated in windowed, constant, or absent forms across matched regularization budgets ($\int \lambda(t) dt$). Thousands of runs, each with controlled random seeds, yield statistically robust characterizations of solution regimes.

Static Complexity Control: Revisiting the Phase Diagram

The initial phase diagram (Figure 1) demonstrates that compositional generalization on anchor-functions is sharply delimited in the $(\gamma, \lambda)$ plane. Notably, high OOD accuracy arises in a horizontal stripe at modest $\lambda$ and moderate $\gamma$. Contrary to frequent prior recommendations [zhang2025complexity; chizat2018global], further decreasing $\gamma$ does not strictly improve generalization; in fact, excessively small $\gamma$ degrades performance by shrinking the reasoning solution basin.

Figure 1: Phase diagram of OOD accuracy across $(\gamma, \lambda)$; the reasoning regime forms a horizontal stripe at moderate $\gamma$0, and is robust across a range of $\gamma$1, refuting the claim that smaller $\gamma$2 is always superior.

The Existence and Sharpness of the Critical Window

The pivotal empirical finding is that when weight decay is applied is more crucial than how much is applied. When a single window of weight decay (spanning just 25% of total training steps) is positioned during a specific interval—after the initial phase of optimization—models achieve OOD accuracy indistinguishable from continuous weight decay or full-budget constant regularization. However, the same regularization applied exclusively early or late in training is ineffective.

Figure 2: OOD accuracy for 5000-step windowed weight decay placed at variable onsets. Only windows during the critical interval (middle placements) reach the reasoning regime.

Budget-matched experiments (Figure 3) confirm this: the relative placement of fixed total regularization budget drastically alters OOD generalization, with middle placements producing up to 9× higher OOD accuracy than early placements.

Figure 3: Same total regularization budget, different placements; only middle windows elicit reasoning, confirming that cumulative regularization is not a sufficient statistic for generalization outcome.

Further resolution reveals a step-function transition in OOD outcome as the onset of the weight-decay window is shifted by as little as 100 steps (batch size 128): the cliff between the memorization and reasoning regimes is exceedingly sharp (Figure 4).

Figure 4: OOD accuracy as window onset is swept with 500- and 100-step granularity, revealing a step-function transition at the early boundary.

Dependence on Initialization Scale and Basin Shrinkage

Systematic variation of $\gamma$3 shows critical window positions do shift as predicted, but more strikingly, the basin of attraction for reasoning solutions shrinks dramatically at smaller $\gamma$4, evidenced by increasing seed-level variance and higher probability of collapsing into memorization solutions (Figure 5, Figure 6).

Figure 5: Critical window location and reasoning-plateau height degrade at small $\gamma$5, accompanied by increased variance.

Figure 6: At $\gamma$6, only $\gamma$7 of seeds reach the reasoning solution, while at $\gamma$8 all do, highlighting the non-monotonicity of the basin width as a function of initialization scale.

This empirical result contradicts the conventional view that smaller initialization is always preferable and calls for a more nuanced recommendation: there is a trade-off between initial speed, basin size, and probability of attaining the desired reasoning solution.

Theoretical Analysis: Two-Timescale Dynamics

The linearized theory introduced in the paper provides mechanistic clarity. The key insight is that the respective evolutionary rates of memorization and reasoning circuits under gradient flow are distinct:

• Memorization mass grows at a rate independent of $\gamma$9.
• Reasoning mass grows at a rate proportional to $d=64$0.

Weight decay exerts selective pressure only in the window between the characteristic timescales of these two trajectories, thereby steering convergence toward the reasoning solution. Outside this window, regularization is ineffective in altering solution fate.

The theory rigorously predicts:

• Critical window onset and width: Scaling as $d=64$1 and $d=64$2 respectively.
• Basin shrinkage at small $d=64$3 due to both slower reasoning-path growth and increased initialization sensitivity under finite training steps.

This theoretical framework precisely matches empirical outcomes, including the sharp window, the window’s $d=64$4-dependence, and the degradation of the reasoning solution’s attainability at small $d=64$5.

Robustness and Limitations

Depth and optimizer robustness: The critical window persists at 4 layers and under both AdamW and SGD optimizers, although the reasoning plateau is somewhat reduced, and seed variance increases with depth, in line with theoretical predictions (Figure 7, Figure 8).

Figure 7: The critical-window effect persists at 4 layers, but reasoning-plateau height is reduced and variance increases, as predicted by basin-shrinkage theory.

Figure 8: Critical-window phenomenon persists under SGD, which recovers the gradient-flow regime more cleanly than AdamW; constant weight decay is less effective than scheduled regularization under SGD.

Task specificity: The phenomenon is observed only in tasks for which the model architecture can reach both reasoning and memorization basins under standard training. For grokking (modular arithmetic) and SCAN add_prim_jump, appropriately tuned constant weight decay suffices or windowing has no effect, confirming that the critical-window effect emerges only in the presence of accessible solution basins of both types.

Figure 9: Online condensation index and bridge-alignment metrics are poor monotonic predictors of OOD outcome; the relationship is categorically predictive but non-monotonic.

Figure 10: On modular-arithmetic grokking, scheduled weight decay slows generalization compared to optimal constant weight decay; the critical-window effect does not generalize to this task.

Implications and Future Directions

These findings necessitate a new practical and theoretical framing for compositional generalization in Transformers:

• Complexity control must be considered a temporal, not merely static, phenomenon. The same hyperparameter configuration can yield drastically different outcomes depending on the timing of regularization application.
• Recommendations for small $d=64$6 must be caveated. While a small initialization scale can facilitate reasoning-bias, it simultaneously narrows the basin of attraction, diminishing the probability of attaining the desired outcome under typical (finite) training durations.
• Task dependency is critical. Generalization about solution selection dynamics cannot be straightforwardly transferred to arbitrary settings; care is required in identifying the structural properties of tasks and architectures for which critical-window control is operable.

Emerging questions highlighted include:

• How does the softmax nonlinearity and full MLP stack influence the two-timescale separation in larger models?
• What are the implications for scaling laws when increasing model depth or parameter count?
• Can practical diagnostics for critical-window identification be developed for more complex or real-world data?

Conclusion

This work provides definitive evidence that solution selection in Transformer training on compositional tasks is governed by a sharply localized critical window in training where complexity control exerts its principal influence. Both the existence and properties of this window are analytically tractable and experimentally robust, but strongly task- and initialization-dependent. Importantly, recommendations in the prior literature must be revised: static hyperparameter selection is insufficient to ensure generalization via reasoning, and constant regularization is, in some regimes, suboptimal or even detrimental. The implications for mechanistic interpretability and the design of training protocols in deep learning are substantial, opening a productive line of inquiry into the temporal dynamics of inductive bias and compositionality in modern architectures.