Mason Eyler

[Mason Eyler @eyler_mason]

@Traansistor @hyperstition_cc Surely grokking is for specific circuits rather than a regime change for the whole model though? Like there can be a large spectrum of when different things grok. Also if we use a proper optimiser/batch size, we intentionally stay out of the lazy regime.

[Mason Eyler @eyler_mason]

@Traansistor @hyperstition_cc Not really - feature learning continues the whole time in a power law. Also if you really thought that most training is searching for "near minima," then we wouldn't see the returns to more data that we do. It's only the attention patterns that get stuck, not the entire network!

[Mason Eyler @eyler_mason]

@CameronMThacker @hyperstition_cc The head comparisons are greedy matched for this exact reason

[Mason Eyler @eyler_mason]

@brown_balaklava @hyperstition_cc If pretraining with all its chaos can’t shift attention patterns then RL has no chance. We ablated all gradient flow through dq and dk (so only MLP and dv flow) and were more or less on par with the control. So shifting attn doesn’t contribute to RL.

[Mason Eyler @eyler_mason]

This is init dependence made visceral. Training runs with identical data streams end up with wholly different attention patterns due to seed noise. In recall based continuations, this can be priced: transplanting the correct attention patterns into the worse model closes the gap.

[Hyperstition @hyperstition_cc]

Today we report a startling pretraining finding:

Attention circuits STOP meaningfully evolving much sooner than we expect.

In other words, the model has decided where to look very early on. It is known that different seeds create starkly different results; here we isolate a specific mechanism that causes it.

The setup:

We perform two identical training runs with the same data order and different initializations and find that by the first thousand steps each run’s attention resembles its own step-10k pattern MORE than the individual finished runs resemble each other.

The initialization plays a greater role in the attention patterns than training dynamics, which implies that the final pattern is randomly chosen.

NOTE: Randomly chosen, not from the space of all the possible attention patterns that it could be, but from the different attention pattern basins that it could converge to—which one it ends up in appears to be init dependent.

Some differences in the loss are attributable to these different attention patterns.

This means the final performance is somewhat dependent on a random choice. Depending on how lucky with pattern formation one gets in the first thousand steps, we could end up much better or worse in different contexts.

This further implies training could simply find the optimal patterns, or at least better ones, if the attention could evolve instead of being stuck.

A story for why attention evolves less than we’d hope falls out of the softmax derivative. We can observe a positive feedback loop we call the credit ratcheting.

Let oᵢ be the head output and δᵢ = ∂L/∂oᵢ the loss gradient at that output. Define key k’s advantage for query i as

Aᵢₖ = −(vₖ − oᵢ)ᵀδᵢ

i.e. how much the loss drops if oᵢ moves toward vₖ, relative to the current attention-weighted mix. Thus, the gradient with respect to sᵢₖ = kₖᵀqᵢ / √d is

∂L/∂sᵢₖ = −aᵢₖAᵢₖ

A key can have a huge advantage and have its score still barely move, because the advantage is multiplied by the attention it already gets. It’s the same form as a softmax policy gradient, π(a)·A(a): attention is the policy, and it mostly learns about the actions it already takes.

In aggregate, the query gradient is gᵢ = Jᵢᵀδᵢ where

Jᵢ = (1/√d) ∑ₖ aᵢₖ(vₖ − oᵢ)kₖᵀ.

Across 166,656 Jacobians in Qwen2.5-1.5B, ≥99% of tokens in each of the 3 heads we probed had squared-spectrum effective rank <10, in a 128-dim head. The update to the query is being filtered through an exceedingly narrow channel.

The values have the same issue. Key k’s share of the value gradient is Gᵢₖ = aᵢₖδᵢxₖᵀ, so at similar δᵢxₖᵀ, 10× the attention ≈ 10× the credit. This means the values mostly adapt to positions already looking at them.

This compounds! Attention allocates credit, credit updates values, and values can pull attention further toward the winners.

The most natural approach to fix this is to precondition the query gradient with the Fisher of the attention distribution (NATQ):

Fᵢ = Covₐᵢ(k) / d

uᵢ = (Fᵢ + λᵢI)⁻¹gᵢ

g̃ᵢ = uᵢ · (‖gᵢ‖₂ / ‖uᵢ‖₂)

Even this blunt version does a lot. At 124M, entropy rose from 3.75 to 4.69 nats, top-key attention fell from 0.238 to 0.149, and in layers 4–7, keys the baseline gave under 0.1% attention got 67% of the gain.

Validation loss rose by 0.082 nats. A more targeted intervention need not make this concession.

Here’s the complete blog: https://t.co/bLPjndoxWE

We frame this as an open problem that offers a stark contrast in the way we perceive pretraining. If we can find a way to allow for attention to evolve throughout training, opposed to just the very start, we may find inordinate gains on model capability.

Internally we half joke that the alternative read is that attention might simply not be real past the first few steps, and it’s all the MLPs. If training is able to work as well as it does without really changing the attention patterns much, maybe the MLP is the real workhorse that’s driving the learning.