Nahom

N300

[Larry Dial @classiclarryd]

Always enjoy the work from this group. This sort of "“role ossification” may be particularly relevant when later changing the data mix or context length. (where the optimal role under the new regime may not match the old one)

[Nahom @Nahomes_15]

There is some incredible dynamic ui here... trust: https://t.co/Bv1YSCRb4b

[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.

[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.

[Nahom @Nahomes_15]

If you force every head in one layer share an attention map, you will notice that median token loss barely moves. The probability of “Andrews” collapsed even though “Amy Andrews” was already in context.

[Nahom @Nahomes_15]

I think we mistake our ability to scale pretraining for an understanding of it. Attention settles into seed-specific patterns early. The loss can keep falling while the model grows more dependent on choices it made near init

[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.

[Nahom @Nahomes_15]

Attention allocates its own credit (a token that gets more attention contributes more to the value gradient). That can make it more useful, which gives the head more reason to look there again.

[Hyperstition @hyperstition_cc]

Attention patterns vary more across innit seeds than they do across training!

[Shivi @Shivi_Raisurana]

Huge alpha in gay spies.

[Nahom @Nahomes_15]

@agentsettles @hyperstition_cc Yes.

[Nahom @Nahomes_15]

@_thomasip @hyperstition_cc The serving stack is for ultra low latency inference. Will release the megakernels for higher concurrencies for Deepseek, Kimi and GLM soon :)

[Nahom @Nahomes_15]

@jimmyyangrz @hyperstition_cc Just you wait.

[Nahom @Nahomes_15]

@utsav_khairnar @hyperstition_cc Roofline😔

[Nahom @Nahomes_15]

@098uyiubgvgcss @hyperstition_cc What's the point here exactly? This is the same vLLM default config. They also store 203GB of Engram tables in host ram.

[Nahom @Nahomes_15]

@SoumilBhandari @hyperstition_cc You can mostly predict per expert load during training from a running average. Every n steps you redistribute the experts so each GPU gets roughly the same load. For extreme spikes you need to redo that stage completely with a bigger buffer (but this only happens near init).

[Nahom @Nahomes_15]

@jus_eng_njainam @hyperstition_cc Yes.

[elie @eliebakouch]

very cool work, megakernels are still a relatively new field and i'm glad to see more effort in the open!!

big fan of this team and excited for what they will ship next

https://t.co/H8nHxbxeWC

[Nahom @Nahomes_15]

This is actually a generational result. Win @ElliotGlazer

[Elliot Glazer @ElliotGlazer]

hello there ℤ is Diophantine in ℚ, thanx to my close friends Astra and Claude

The polynomial characterizing ℤ isn’t quite Tweet-sized; it has m=292,890,231,905 variables, is of degree d &lt; 2^^7, and all its coefficients have absolute value &lt; 2^^8.

[Kyle Vedder @KyleVedder]

@hyperstition_cc @arcinstitute @physical_int this is very exciting!

one of the reasons for using a denoising head is simply because FAST decodes are too slow on typical research runtimes

[Nahom @Nahomes_15]

@hyperstition_cc Fewer projects were more fun to work on and as important. more to come

[Jerry Tworek @MillionInt]

Megakernels are the future and incredibly cracked teams seem to be pushing the boundaries of possibility each day getting further and further.

Looking forward to seeing more from this team 😀