Marginals Before Conditionals

Marginals Before Conditionals: Why Transformers "Stall" Before Getting the Point

Marginals Before Conditionals

Summary

Problem

Method

Results

Takeaways

Abstract

The paper introduces a "wind tunnel" task to study staged learning in Transformers, where models must disambiguate a mapping $B o A$ using a selector token $z$ . It reveals that models consistently learn marginal distributions $P (A ∣ B)$ before conditionals $P (A ∣ B, z)$ , resulting in a loss plateau at exactly $lo g K$ (where $K$ is the ambiguity level) followed by a sharp, collective transition to zero loss.

TL;DR

Imagine teaching a model that "Bank" can mean a river edge or a financial institution, but it only figures out which one is which (using context) after a long, frustrating plateau where it just guesses both. This paper proves that Transformers learn the "average" (marginal) answer long before they learn to use specific context (conditional) to disambiguate. Surprisingly, more data and more optimization noise actually make this wait longer.

The "Wind Tunnel" of Conditional Learning

Why do models fail to see the obvious? To answer this, the authors constructed a minimal task:

Input: A base string $B$ and a selector $z$ .
Mapping: $B$ maps to $K$ possible targets. Without $z$ , you are guessing (Loss = $lo g K$ ). With $z$ , the answer is unique (Loss = 0).

The study finds a universal training trajectory: a rapid drop to the "marginal" solution ( $lo g K$ ), a long flat plateau, and then a sudden "snap" where the model suddenly "gets it" and loss hits zero.

Insight 1: Duration Depends on Data Size ( $D$ ), Not Complexity ( $K$ )

One might assume that highly ambiguous tasks (large $K$ ) would take longer to solve. The authors found this is false. If you keep the total dataset size $D$ constant, the model takes the same amount of time to escape the plateau regardless of whether $K = 5$ or $K = 36$ .

The bottleneck isn't the complexity of the confusion, but the volume of examples the optimizer must process to align the internal "routing" circuits.

Staged Disambiguation Progress Figure 1: The model hits the $lo g K$ plateau almost immediately and stays there until a sharp "snap."

Insight 2: The Noise Trap (Entropic Stabilization)

In many optimization problems, noise (via high learning rates or small batches) helps "kick" the model out of local minima. Here, the opposite happens.

Fact: Higher learning rates prolong the plateau.
Why?: The marginal solution is a "low-gradient" state where competing directions cancel out. This creates a flat saddle point. Noise acts as an "entropic force" that keeps the model floating in this flat region, preventing it from finding the narrow, shallow "escape hatch" that leads to the conditional solution.

Insight 3: The Collective Snap & Internal Cascade

Is the model learning one group at a time? No. The transition is collective. At half the waiting time ( $a u /2$ ), zero percent of the groups are solved. Then, suddenly, everything solves at once.

Mechanistically, the authors found that a specific "selector-routing" head (L0H3) starts specializing long before the loss actually drops. This "hidden progress" accounts for about 50% of the plateau duration.

Internal Transition Consistency Figure 2: The direction consistency snaps from 0 (random walk) to 0.8 (targeted escape) at the moment of transition.

Directional Asymmetry: The "Arrow of Time"

The paper links this to the Reversal Curse. Learning the mapping $(B, z) o A$ is much faster than the reverse $A o B$ . Because the "backward" direction lacks the group structure (shared $B$ strings) that the Transformer can use to scaffold its circuits, it takes 1.7x to 4.4x longer to learn.

Critical Analysis & Conclusion

This paper provides a rigorous "physics of training" view of why models stall. It falsifies several common myths:

Myth: Plateaus are caused by gradient cancellation across groups. (Result: Falsified; duration doesn't scale with the number of groups).
Myth: Noise helps escape. (Result: Falsified; noise stabilizes the "dumb" solution).
Myth: Models learn one thing at a time. (Result: Falsified; learning is a collective phase transition).

Limitations: The study uses a 4-layer Transformer on synthetic data. While it provides a "wind tunnel" for theory, natural language has many overlapping hierarchies of ambiguity, which might smooth out these sharp transitions or create a "staircase" of plateaus.

The Takeaway for Researchers: If your model is stuck on a plateau and ignoring context, adding more data or increasing the learning rate might actually be counter-productive. You may be inadvertently "stabilizing" the very ignorance you're trying to fix.

Find Similar Papers

Try Our Examples

Search for recent papers investigating the "entropic stabilization" of saddle points in SGD and how it relates to delayed generalization in Transformers.
Which foundational paper first discussed the "Reversal Curse" in LLMs, and how does the concept of Type 2 directional asymmetry qualitatively differ from simple memorization failure?
Explore studies that apply mechanistic interpretability to identify "selector-routing" or "circuit-switching" heads in large-scale language models during the training process.

Contents

Marginals Before Conditionals: Why Transformers "Stall" Before Getting the Point

1. TL;DR

2. The "Wind Tunnel" of Conditional Learning

3. Insight 1: Duration Depends on Data Size ($D$), Not Complexity ($K$)

4. Insight 2: The Noise Trap (Entropic Stabilization)

5. Insight 3: The Collective Snap & Internal Cascade

6. Directional Asymmetry: The "Arrow of Time"

7. Critical Analysis & Conclusion

Marginals Before Conditionals: Why Transformers "Stall" Before Getting the Point