🤖 machine learning

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

This work demonstrates that a Transformer's ability to generalize through reasoning rather than memorization is determined within a sharp, critical training window where the timing of weight decay and the scaling of initialization are decisive, thereby challenging the view of complexity control as a static hyperparameter choice.

Original authors: Sarwan Ali

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sarwan Ali

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart but slightly confused student (a Transformer AI) how to solve a puzzle. The puzzle consists of taking two simple rules and combining them into a new, complex rule.

The student has two ways to learn:

Reasoning: He actually learns the logic of how to combine the rules. This is the "good" path. It enables him to solve new puzzles he has never seen before.
Rote memorization: He simply memorizes the specific answers to the exact puzzles he has practiced on. This is the "bad" path. If you give him a new puzzle, he fails completely.

For a long time, researchers believed the key to getting the student to "reason" rather than "memorize" was simply to use a specific setting called Weight Decay (think of it as a gentle "nudge" or "penalty" that prevents the student from becoming too confident or rigid) and to start with small initial weights (starting the student with an empty, modest mind). The advice was: "Set these controls once at the beginning, and the student will learn to reason."

This work says: "Not so fast. Timing is everything."

The researchers discovered that it does not matter how strongly you nudge the student or when you start with a modest mind. What matters is exactly when you apply this nudge during the training process.

Here is a breakdown of their findings with simple analogies:

1. The "Critical Window" (The Golden Hour)

Imagine training the student as a 20-hour film. The researchers found there is a very specific, short time window—roughly between the 25% and 75% marks of the film—where the student decides his fate.

If you apply the "nudge" (Weight Decay) during this window: The student wakes up, realizes he must learn the logic, and begins to reason. He becomes smart and flexible.
If you apply the nudge before this window (the first 25%): It is like trying to teach a baby to drive a car. The student is too early in his development. The nudge has absolutely no effect. He simply continues to memorize things.
If you apply the nudge after this window (the last 25%): The student has already decided to memorize. It is too late to change his mind. The nudge is ignored.

The shocking discovery: You can apply the nudge only for this middle 25% of the training time, and the student performs just as well as if you had nudged him the whole time. But if you only nudge him at the beginning, he fails.

2. The "Cliff" Edge

The researchers found that the beginning of this "golden hour" is incredibly sharp. It is like a cliff.

If you start the nudge at step 0, the student fails.
If you wait just 100 steps (a tiny fraction of a second in computer time) and then start the nudge, the student suddenly jumps from failure to success.
It is not a gentle slope; it is a switch. In one moment he memorizes things, in the next he draws logical conclusions.

3. The Myth of the "Modest Mind"

The old advice was: "Start with a very small, modest mind (small initialization) to force reasoning."
The work says: This is actually risky.

If the student starts too modestly (very small initial weights), the "golden hour" window shifts, and the student becomes very fragile. It is like a tightrope walker; a small error in the seed (random starting point), and he falls.
The researchers found that a moderately large starting mind is actually safer. It provides the student with a wider "safety net" (basin of attraction) to land in the reasoning zone.

4. It Is Not a Universal Rule

The researchers tested this on other types of learning tasks to see if this "critical window" rule applies everywhere.

Modular Arithmetic (Grokking): In this task, the student must be nudged constantly from start to finish to eventually figure it out. The "golden hour" window does not exist here.
SCAN Task: In this task, the student is simply too dumb (or the architecture is wrong) to ever learn reasoning, no matter when you nudge him.

The Big Picture

Imagine training an AI like baking a cake.

Old View: "Add a pinch of salt (Weight Decay) at the beginning, and the cake will be perfect."
New View: "The cake only rises if you add the salt exactly when the yeast is active. If you add it before the yeast has woken up, or after it has died, the cake will not rise. And you do not need to keep adding salt; you only need to add it in that one critical moment."

In short: The work proves that AI models for certain types of logical puzzles do not just need the right settings; they need the right timing. There is a specific, narrow moment in training where the model decides whether to think or to memorize, and if you miss that moment by even a tiny bit, it will never learn to draw logical conclusions.

Technical Summary: Critical Windows of Complexity Control in Transformers

Problem Statement

Recent literature has demonstrated that Transformers can achieve compositional generalization (inference) rather than memorization of specific tasks by controlling complexity through small initialization scales ( $\gamma$ ) and weight decay parameters ( $\lambda$ ). However, existing analyses treat these hyperparameters as static decisions, assuming that solely the magnitude of regularization determines the model's trajectory. This article investigates a critical gap: whether the timing of complexity control matters. Specifically, it asks whether the decision between memorization and inference solutions is made throughout the entire training process or if it is fixed within a specific, identifiable temporal window.

Methodology

The authors employ a controlled experimental framework based on the Anchor Function Task (introduced by Zhang et al., 2025), where models must compose permutations of anchor symbols.

Architecture: Small 2-layer and 4-layer Pre-Norm Transformers ( $d_{model}=64$ , 2 heads).
Training: Full instrumentation enables step-by-step tracking of Out-of-Distribution (OOD) accuracy, condensation indices, and subspace alignments.
Core Variable: The authors systematically vary the temporal placement of weight decay ( $\lambda$ ), while keeping the cumulative regularization budget ( $\int \lambda(t) dt$ ) and initialization scale ( $\gamma$ ) constant.
Experimental Design: They test "window-based" schedules where weight decay is applied only during specific intervals (e.g., early, mid, late) of training, comparing these against full-training baselines and no-regularization baselines.
Theoretical Analysis: A linearized two-layer attention model is analyzed to derive rate equations for "memorization mass" and "inference mass," predicting a separation of their timescales.

Main Contributions and Results

1. The Existence of a Critical Window

The primary empirical finding is that the fate of a Transformer (inference vs. memorization) is decided within a sharp, narrow training window.

Window Efficacy: Applying weight decay for a single 25% window of training (specifically steps 2,500–17,500 in a 20k-step run) yields OOD accuracy (0.93) comparable to full training with weight decay (0.91).
Inefficacy Outside the Window: Applying the same weight decay budget entirely before this window (steps 0–5,000) or entirely after it results in OOD accuracy at random levels (~0.15), indistinguishable from no regularization.
Budget vs. Timing: When the total regularization budget is held constant, placing it in the middle of training results in 5–9 times higher OOD accuracy than placing it at the beginning. This demonstrates that when regularization is applied is more critical than how much is applied.

2. Sharpness of the Temporal Boundary

The transition into the critical window is remarkably abrupt.

Shifting the window start by only 100 optimization steps causes the mean OOD accuracy to jump from random levels (0.15) to the inference regime (0.61).
After 400 steps, the system reaches the inference plateau (0.93). This suggests a behavior akin to a phase transition phenomenon, where the model's ability to be steered toward inference becomes active only after a specific initialization phase.

3. Initialization Scale and Basin Shrinking

Contrary to the prevailing recommendation that smaller initialization is consistently better for inference, the article finds a nuanced relationship:

Window Position: The critical window shifts forward as $\gamma$ increases.
Basin of Attraction: At very small initialization scales (e.g., $\gamma=0.5$ ), the basin of attraction for the inference solution shrinks significantly. In 12-run experiments, only 8/12 runs at $\gamma=0.5$ achieved OOD > 0.5, whereas 12/12 were successful at $\gamma=1.1$ .
Implication: A moderate initialization ( $\gamma \in [0.7, 1.1]$ ) provides a wider basin of attraction for inference solutions under finite training time, challenging the heuristic "the smaller, the better."

4. Task Specificity and Scope

The phenomenon is not universal for all delayed generalization tasks:

Grokking (Modular Arithmetic): The critical window does not appear. In modular arithmetic, a correctly tuned constant weight decay schedule achieves grokking 5 times faster than temporally localized schedules.
SCAN (add_prim_jump): Vanilla 2-layer Transformers fail to fully reach the basin of attraction of the compositional solution under any schedule, suggesting the phenomenon requires a loss landscape where both memorization and inference basins are accessible.

5. Theoretical Mechanism: Separation of Two Timescales

The authors provide a theoretical explanation based on a linearized model showing that memorization and inference circuits evolve under qualitatively different rate equations:

Memorization Rate: $\mu_m = \Theta(1)$ (independent of $\gamma$ ).
Inference Rate: $\mu_r = \Theta(\gamma^2)$ .
Mechanism: Since inference evolves much slower than memorization at small $\gamma$ , a time interval exists where memorization is saturated while inference is still developing. Weight decay applied during this interval selectively suppresses the saturated memorization path while allowing the inference path to continue growing. Weight decay applied before this interval contracts both paths proportionally (no effect on the ratio), and applied after this interval, the model has already committed to memorization.

6. Robustness

Depth: The phenomenon persists in 4-layer models, although the inference plateau is lower and variance between runs is higher, consistent with the theoretical prediction of basin shrinking with increasing parameter volume.
Optimizer: The phenomenon reproduces under Vanilla SGD with momentum. Notably, constant weight decay schedules perform significantly worse than window-based schedules under SGD, whereas AdamW partially masks this effect through adaptive preconditioning.

Significance and Claims

The article claims to characterize a previously unreported temporally localized phase transition in compositional generalization. Its significance lies in:

Redefining Complexity Control: It shifts the paradigm from a static hyperparameter view to a dynamic, temporal one. The "recipe" for inducing inference is not just "small $\gamma$ + weight decay," but "small $\gamma$ + weight decay, applied during the critical window."
Correcting Previous Recommendations: It provides a caveat to the recommendation of minimal initialization and shows that moderate initialization can be superior for stability (basin width) in finite-time training regimes.
Theoretical Consistency: It offers a theoretical explanation with two timescales that predicts the existence, position, and width of the critical window, as well as the observed basin shrinking.
Defining Scope: It delineates that this mechanism is applicable specifically to settings where the model can access both memorization and inference basins, distinguishing it from other generalization phenomena like grokking.

The authors conclude that complexity control is fundamentally a temporal phenomenon and that existing static frameworks are incomplete in explaining how Transformers choose between inference and memorization solutions.