TT

AI research on efficient deep learning, build things I need then share them.

[TT @timothibo]

How come K3 is so powerful ?

🍋 Kimi Delta Attention (KDA)
🍋 Attention residual (AttnRes)
🍋 Stable latent MoE

🍋🍋🍋 that’s a jackpot…

All three topics covered and explained in details:

KDA: https://t.co/4wQiPdO7qV

AttnRes : https://t.co/QwgJQap29z

Stable Latent MoE: https://t.co/sNbsz6uUB3

[TT @timothibo]

How did Moonshet make K3's MoE architecture THAT efficient ?

- Stable Latent MoE
- Quantile balancing

How do you actually train and serve a 2.8T MoE where every token only touches 1.8% of the experts without half the model going permanently unused ?
without activation explosions ?
and without the all-to-all communication killing throughput ?

Here is a clear breakdown:

[TT @timothibo]

Extreme sparsity only works when routing, communication, activation stability and load balance are treated as first-class citizens instead of afterthoughts. Stable Latent MoE is one of the cleanest recent answers to that engineering stack.
If you're serious about learning => https://t.co/ZXZKpBCibT

[TT @timothibo]

How did Moonshet make K3's MoE architecture THAT efficient ?

- Stable Latent MoE
- Quantile balancing

How do you actually train and serve a 2.8T MoE where every token only touches 1.8% of the experts without half the model going permanently unused ?
without activation explosions ?
and without the all-to-all communication killing throughput ?

Here is a clear breakdown:

[TT @timothibo]

Well, that’s the theory…
Cause getting all 16 experts in distinct GPUs to compute ri,j, gather them to order them is terribly expensive.

Instead we define an interval [Rmin, Rmax] that we divide into H equal pieces and just ask each GPU to tell you how many ri,j it has in each piece.

You then attribute bj to the upper bound of the piece that contains your 1.79% quantile.

To achieve 1.78% precision, you need 56 pieces.

Of course it doesn’t ensure a perfect repartition since we set weight for token i based on the k+1 highest scores of the token i-1 step.
But that is a good iterative and low cost approach that ensures a balanced usage of all experts of the model.

[TT @timothibo]

"An LSTM is a ResNet rotated 90 degrees" => Rotate a LSTM gives you residuals => Attention rotated 90° gives you Residual Attention.

Zhilin Yang, Moonshot's CEO.

Deep dive in Kimi's Residual Attention: why, how, results.

In basic architectures, the input of layer l is output of layer l-1, but that causes too much drift from the original input after too few steps due to interference and dilution of the signal. Hence the abstract understanding of the inputs is built on shaky ground, or put otherwise, there is not enough signal passing through each layer for the next ones to capture more abstract features.

Hence we use residuals: we feed the layer l with the output of layer l-1 + the original input of layer l-1. That way, the very first input x0 is always passed to a layer. It can rewrite as input l = x0 + sum output j for j = 1 to l-1.

That highlight an issue: simply adding as many outputs as we have dilute information and signal.

Such a sum with implicit weights = 1 for each term increase input's variance linearly with L (number of layers), hence increase output's variance.

One could say that it's not a problem since proportions and ratios are safe: if feature1 is present 50% of the time and feature2 10% of the time and their vector norms gets +1 each time they appears, after 100 steps, feature 1 vector has norm 50, feature 2 has norm 10.
That way, we always have the same ratio so adding more and more outputs of layers just increase the values (cardinality) but not the ratios (ordinality).
Problem is that with non linear activations and non linear operations like softmax, a 5:1 ratio in logits can end up in a 1000:1 ratio in probabilities for prediction. Softmax temperature/scaling gets completely skewed if the input variance grows with L, so cardinality matters.

So our problem is to find a way to pass enough signal through layers by providing layer l with a less abstract input (residual) in addition to the output of the previous layer. It appears clearly that we should not overload that residual otherwise it would be too noisy hence fail to his mission to strengthen signal.

Moonshot's solution: residual attention. Instead of passing the sum of all previous layer's output, we pass a combination of transformation of all of them. Like attention but for previous states instead of tokens. Meaning that we store all previous outputs (that we call hidden states as they are intermediary (hidden) states of comprehension of the word by the model between input and output).

Then each layer has its own Wq, Wk, Wv matrices that allow to compute Q,K,V.
Here a difference with token attention matters hidden states don't attend to themselves: if we are at layer l, Q = Ql or queries vectors layer l while K and V contains keys and values K1,...,Kl-1 and V1,...,Vl-1.

Why ?

Well because hidden state l is gonna ask with Ql how strongly should it attend to each layer's hidden state based on their key Kj and that strength will actually multiply the value Vj of that hidden state.
Why don't a hidden state attend to himself: since Ql and Kl would be generated from the same state hl, the model could learn to be lazy by setting Ql = Kl and hl just mainly attend to himself.

Here is the gold mine: this setup allow each layer to cherry pick from previous hidden state useful signal for him. That signal differ from a layer to another as they tend to specialise into different features detection.
They all get personalised memory based on their needs, no redundant information, no signal lost in the middle (lost under over accumulation of redundant information). So we don't actually use the hidden states hj themselves, just their value vectors Vj. Careful, Vj of layer l differs from Vj of layer l-1 since they have different Wv, Wk,Wq matrices.
This combination of attended hidden states's values are called residuals. It is the new form of memory we feed a layer with to provide with stronger signal.

Attention on tokens is heavy O(N^2), that's fatal when you have 15T training tokens but that's manageable when you have 100 layers and compute attention on previous layers.

Still, Moonshot reported that such architecture would cause a memory bottleneck on the GPU SRAM.

So they used the same trick than with KDA (Kimi Delta Attention): they divided the task in small blocks to run intra block operation in parallel (in KDA, not possible here) and only implement the full sequential workflow on the block level.

Concretely: they divide their network of L layers in N blocks of B layers (N*B = L) and allow each layer to only attend to previous layers of its block and final output of previous blocks (and not all the output of all the previous blocks).
First bloc: input0 goes through layer 1, get output1, layer 2 gets fed with output1 + Value vector of input0 get output2. Layer3 gets fed with output2 + residuals from output1 and input0 etc until getting outputB of the last layer of the block.
Then for the next block, we repeat the same operation but instead of getting residuals on the whole set of all outputs, each layer can attend to outputs of layers of his block + final output of each preceding block. That makes it O(L^1.5) over all layers instead of O(L^2).

That makes echo "an LSTM is a ResNet rotated 90 degrees". Here we have 2 propagations going on:

- horizontal: classic attention for tokens, at layer 1, token 5 attend to tokens 1,2,3,4 and 5. at layer l, the input is xl then xl,5 attend to xl,1,..., xl,5.

- vertical: attention for hidden states, hidden state l can attend to hidden states 0 to l-1. Since there is no matmul that mixes hidden states values, there is no cross tokens contamination. Meaning that for a given token i, the layer l can attend to its abstract representation made by layers 1 to l-1 but can't attend token's j previous representations.

The dense cross layer activity creates short gradients path from layer l to layers 1,...,l-1. Hence, a query dependant residual lets the backward pass route gradients to any earlier layer that actually produced a useful feature. That's golden.

All the engineering lies into the how to equip each input with the signal needed only (source and strength) at a sustainable compute cost.

I was tired of having to launch dozens of llm chats to learn something on the long run. Context reset at 0 every time, no continuity, no plan, no repetition, no direct coding exercise. Now I have everything I need all in one place, optimised. If you're serious about learning => https://t.co/ZXZKpBCibT

[TT @timothibo]

All the hard work in modern architectures is ultimately about the same question:
How do you equip every computation with exactly the signal it needs — right source, right strength — at sustainable cost?
Attention over tokens solved the sequence axis.
Attention Residuals is the same move applied to depth.
Beautiful, minimal, and it actually works at scale.

I was tired of having to launch dozens of llm chats to learn something on the long run. Context reset at 0 every time, no continuity, no plan, no repetition, no direct coding exercise. Now I have everything I need all in one place, optimised. If you're serious about learning => https://t.co/ZXZKpBCibT

[TT @timothibo]

Residual connections let us train absurdly deep nets.
They are the reason Transformers didn’t collapse under their own depth.
But they have a quiet failure mode that gets worse the deeper we go.
=> Kimi Residual Attention: why, how, results

[TT @timothibo]

Results (Kimi Linear 48B total / 3B active, 1.4T tokens): consistent gains across the board. Biggest jumps on multi-step reasoning and code. Scaling laws show Block AttnRes matching a baseline trained with ~1.25× more compute. Output magnitudes stay bounded; gradients distribute more evenly across depth.
The engineering insight is simple and deep: the residual stream is a memory. Stop writing to it with permanent unit markers. Let each consumer decide what (and how strongly) to read.