From GPT-2 to Kimi K3: Seven Years of Architecture Upgrades That Go Far Beyond Parameter Count
Recently, Moonshot AI's Kimi officially open-sourced the full model weights of Kimi K3. Kimi K3 is a Mixture-of-Experts (MoE) large model with a total parameter count of 2.8 trillion and a context window of 1 million tokens. It is also the world's first nearly 3-trillion-parameter open-source large model to be deployed, sparking widespread discussion in the industry.
One blogger, ali@waterloo_intern, realized that from the 1.24-billion-parameter GPT-2 released by OpenAI in 2019 to the 2.8-trillion-parameter Kimi K3 in 2026, only seven years have passed, yet the two models differ in scale by a factor of 22,580!
A simple conversion shows that roughly 22,580 GPT-2 Smalls can fit inside a single Kimi K3.
This sparked his curiosity: "But is it really just about getting bigger?"
In response, ali said he spent about 48 hours reading Kimi K3's modeling code and 8 papers, ultimately clarifying the complete technical lineage from 2019's GPT-2 to Kimi K3. "I will take you through how we got here step by step, and exactly how much the model has changed from GPT-2 to Kimi K3? And what things have actually never changed? We will follow this technological evolution path and sort out the key architectural upgrades that ultimately led to Kimi K3."
Let's take a look together.
GPT-2
GPT-2 uses a decoder-only architecture:
tok_emb = self.transformer.wte (idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.transformer.wpe (pos) # position embeddings of shape (t, n_embd)
x = self.transformer.drop (tok_emb + pos_emb)
for block in self.transformer.h:
x = block (x)
x = self.transformer.ln_f (x)
logits = self.lm_head (x)
return logits
The input first adds token embeddings and position embeddings:
Zooming into each Transformer module, its structure is as follows:
class Block (nn.Module):
def __init__(self, config):
super ().__init__()
self.ln_1 = LayerNorm (config.n_embd, bias=config.bias)
self.attn = CausalSelfAttention (config)
self.ln_2 = LayerNorm (config.n_embd, bias=config.bias)
self.mlp = MLP (config)
def forward (self, x):
x = x + self.attn (self.ln_1 (x))
x = x + self.mlp (self.ln_2 (x))
return x
The attention calculation process is as follows:
B, T, C = x.size () # batch size, sequence length, embedding dimensionality (n_embd)
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
q, k, v = self.c_attn (x).split (self.n_embd, dim=2)
k = k.view (B, T, self.n_head, C //self.n_head).transpose (1, 2) # (B, nh, T, hs)
q = q.view (B, T, self.n_head, C //self.n_head).transpose (1, 2) # (B, nh, T, hs)
v = v.view (B, T, self.n_head, C //self.n_head).transpose (1, 2) # (B, nh, T, hs)
# manual implementation of attention
att = (q @ k.transpose (-2, -1)) * (1.0 /math.sqrt (k.size (-1)))
att = att.masked_fill (self.bias [:,:,:T,:T] == 0, float ('-inf'))
att = F.softmax (att, dim=-1)
att = self.attn_dropout (att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose (1, 2).contiguous ().view (B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_dropout (self.c_proj (y))
return y
When the final hidden state matrix is generated, the language model head maps it to logits over the vocabulary. During autoregressive decoding, the model only needs the logits of the last position to select the next token.
This is also an inefficiency of the decoder-only generation method: the model computes representations for every position in the input sequence, but at each decoding step, only the logits of the last position are actually used. Without a caching mechanism, a large amount of computation needs to be re-executed when generating the next token.
KV Cache originates from a very direct observation: when a newly generated token is appended to the input sequence, the model originally needs to recompute the projections of all previous tokens. By saving the Key and Value vectors corresponding to these tokens, this part of the repeated computation can be avoided.
This saved data is the KV Cache. It retains the vectors of the previous N-1 tokens, and its scale can become very large, even forming a memory bandwidth bottleneck.
Overall, with a vocabulary size of about 50,000, 12 Transformer modules, 12 attention heads, and an embedding dimension of 768, this baseline model has approximately 124 million parameters.
Linear Attention
Softmax attention applies a non-linear transformation after the q·k product is completed, so each Query is coupled with every Key. Linear attention, on the other hand, applies a feature map, such as ELU+1, to q and k separately. This way, the matrix multiplication can be re-associated, and the continuously growing K and V vectors can be compressed into a fixed-size D×D state.
The author stated that the description of O(N²) in the paper once confused him. Strictly speaking, "the computational cost of a Transformer at each time step grows quadratically with the current sequence length" is not accurate. FlashAttention solves exactly this problem... He later discovered that this paper was published in 2020.
At that time, training typically explicitly constructed the complete N×N attention matrix, FlashAttention had not yet appeared, and many reference-level autoregressive implementations did not use KV Cache, requiring repeated computation of the historical states of all previous tokens.
def forward (self, x, mask=None, past_kv=None): # x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj (x)
q=qkv [:, :, :d].view (b,t,h,d_head).transpose (1,2)
k=qkv [:, :, d:2*d].view (b,t,h,d_head).transpose (1,2)
v=qkv [:, :, 2*d:].view (b,t,h,d_head).transpose (1,2)
# at prefill, q,k,v have shapes b,h,t,d
# at decode, shape is b, h, 1, d
# so i cat at the t dimension, dim (2)
if past_kv is not None:
k_past=past_kv [0]
v_past=past_kv [1]
k=torch.cat ((k_past, k), dim=2)
v=torch.cat ((v_past, v), dim=2)
scores=([email protected] (-1,-2))/math.sqrt (d_head)
if past_kv is None: #we're in prefill and need to mask
causal_mask=torch.ones (t,t,dtype=bool, device=q.device)
causal_mask=torch.triu (causal_mask, diagonal=1)
scores=scores.masked_fill (causal_mask, float ('-inf'))
if mask is not None:
scores=scores.masked_fill (~mask, float ('-inf'))
#get attn (bhtt x bhtd)
attn=scores.softmax (-1)#bhtt
o=attn@v #bhtd
o=o.transpose (1,2).contiguous ().view (b,t,d) #b,t,d
# use x to get qkv
o_proj=self.o_proj (o)
past_kv=(k, v)
return o_proj, past_kv
Through visualization, this process becomes more intuitive. Each decoding step requires two (ND)-scale reads from HBM and two 1D-scale writes; at the same time, the size of the KV Cache grows linearly with the sequence length at O(N).
It can be seen that this process involves a large number of read and write operations, and this paper replaces them with the following method:
def forward (self, x, mask=None, cache=None): # x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj (x)
q=qkv [:, :, :d].view (b,t,h,d_head).transpose (1,2)
k=qkv [:, :, d:2*d].view (b,t,h,d_head).transpose (1,2)
v=qkv [:, :, 2*d:].view (b,t,h,d_head).transpose (1,2)
k=F.elu (k)+1
k=k.transpose (-1,-2)
q=F.elu (q)+1
S,z=cache if cache is not None else (0.0, 0.0)
S=S+k@v
z=z+k
o=q@S #bhtd
denom=q@z
o_scaled=o/denom
o_scaled=o_scaled.transpose (1,2).contiguous ().view (b,t,d)
o_proj=self.o_proj (o_scaled)
cache=(S,z)
return o_proj, cache
This is a trade-off relationship.
Here, instead of using the exponential operation in Softmax, the author applies ELU+1 to q and k separately before they interact. Both methods normalize the final scores, but the feature map used by linear attention has a weaker approximation capability for the Softmax kernel.
This approximation may reduce the fidelity of the results, but the actual precision loss depends on the specific model architecture and workload.
It should be noted that we still divide by the sum of qk scores, but this step is omitted in the diagram for simplicity.
Overall, the attention mechanism can be divided into three steps:
- Convert qk scores to non-negative numbers; linear attention uses ELU+1, while Softmax attention uses exponential operations;
- Divide by the sum of all scores to complete normalization;
- Perform a weighted average of the Values based on the normalized weights.
Linear attention retains the basic computational form of the attention mechanism, but to make QK scores non-negative, it uses a feature map with relatively weaker expressive power.
DeltaNet (Fast Weight Programmers)
A cache with limited capacity inevitably needs to overwrite existing information or merge new information with old information. The state from the i-1th token does not get an independent storage slot but is written into the same D×D matrix. Therefore, a new Query can no longer retrieve completely isolated representations of each historical token from it.
This cumulative writing is the source of the efficiency gain. By updating the cache through addition rather than continuously concatenating new content, the cache size does not grow with the sequence length at O(N). But the same operation also causes different information to interfere with each other. DeltaNet attempts to solve exactly this problem of information being difficult to recover.
There is a very concise description of this in Schlag's paper "Fast Weight Programmers":
When the sequence length exceeds the storage capacity, the model may enter a capacity overload state. To operate normally in this state, the model should learn to dynamically interact with the memory content and selectively decide which key-value associations to retain and which to delete. Simple cumulative instructions may not be suitable for this goal... As shown in Formula 17, endlessly adding new associations to a memory with limited capacity will eventually reach a limit.
The most attractive scenario for linear attention is when N is much larger than D, but this also exposes its main limitation. Once the state exceeds its effective capacity, different associations will begin to interfere with each other because the update method is only continuous accumulation, and no information is removed from the cache.
def forward (self, x, mask=None, cache=None): # x is b,t,d
b,t,d=x.shape
d_head=d//self.num_heads
h=self.num_heads
qkv=self.qkv_proj (x)
q=qkv [:, :, :d].view (b,t,h,d_head).transpose (1,2)
k=qkv [:, :, d:2*d].view (b,t,h,d_head).transpose (1,2)
v=qkv [:, :, 2*d:].view (b,t,h,d_head).transpose (1,2)
q = F.normalize (F.silu (q), dim=-1)
k = F.normalize (F.silu (k), dim=-1)
beta = torch.sigmoid (self.w_beta (x)).view (b, 1, t, 1) # new: per-token write strength
S = cache if cache is not None else 0.0
v_old = k @ S # read the board at this key
u = beta * (v - v_old) # the delta: only what's actually new
S = S + k.transpose (-1, -2) @ u # same outer-product write as before
o = q @ S # read, no denominator
o = o.transpose (1, 2).contiguous ().view (b, t, d)
return self.o_proj (o), S
Through a visual example, this process can be understood more easily.
Suppose a set of associations is written: S = k.T@v, and then read using the same Key, resulting in k @ (k.T @ v), which is (k @ k.T) v, effectively equal to the squared norm of k multiplied by v. Therefore, the read result is scaled by the squared norm of the Key. Normalizing k to unit length, or simply dividing the result by its norm, can accurately recover v.
Q can also be seen as a learned pointer. Both Wq and Wk read information from the same residual stream, and when the model queries a fact, the corresponding Query points in the direction of the Key that was used when this fact was originally written.
When updating the state, the model first checks what information the current Key can read from the cache. Then, it subtracts the old information currently read from the new Value it wants to store, multiplies this difference by the Key, and adds the result back to the state matrix.
In this way, old information is removed, and new information is written in its original place.
DeltaNet (Parallelizing Linear Transformers via the Delta Rule)
This is the most difficult part of this article to understand. The blogger stated that to truly grasp it, it took about seven hours, so the following will start from the specific implementation and explain step by step.
Simply put, DeltaNet implements a first-order linear recurrence and uses a generalized Householder transformation matrix as the state transition matrix, thereby supporting chunked parallel execution of the forward pass, achieving linear-time training more suitable for hardware.
It divides the input and output into several chunks of size C, and calculates the output of the chunk based on the final state of the previous chunk and the Query, Key, and Value matrices in the current chunk.
The actual problem to be solved is prefill, which is the context pre-filling stage.
If the Delta Rule is implemented directly on a sequence of length T, the calculation process is roughly as follows:
S = torch.zeros (b, h, dh, dh) if cache is None else cache
outs = []
for i in range (t):
k_i = k [:, :, i:i+1]
v_i = v [:, :, i:i+1]
b_i = beta [:, :, i:i+1]
v_old = k_i @ S
u_i = b_i * (v_i - v_old)
S = S + k_i.transpose (-1, -2) @ u_i # write
outs.append (q [:, :, i:i+1] @ S)
o = torch.cat (outs, dim=2)
Unlike standard attention, this form requires a correction for each Key vector, so it is not intuitive how to convert it into parallelizable matrix multiplication.
Even without considering the Delta Rule, directly implementing the prefill process of linear attention is still serial:
S = torch.zeros (b, h, dh, dh) if cache is None else cache
outs = []
for i in range (t):
q = q [:, :, i:i+1]
k = k [:, :, i:i+1]
v = v [:, :, i:i+1]
S=S_old+k@v
o=q@S #bhtd
o=self.norm (o)
o=o.transpose (1, 2).contiguous ().view (b, t, d)
out=self.o_proj (o)
cache=S
outs.append (out)
o = torch.cat (outs, dim=2)
Using a chunked form, a more efficient implementation can be obtained. Through an example, the calculation mechanism is easier to understand:
When C=N, this method degenerates into standard O(N²) attention; when C=1, it corresponds to ordinary linear attention.
Choosing different C between the two is equivalent to a trade-off between computation and hardware utilization: inside the chunk, it adds some computation but can utilize the hardware more fully.
In practice, C is usually set to 64 or 128, because Tensor Core instructions can run efficiently at this granularity, with UMMA being one example.
The matrix blocks generated during intermediate calculations are folded into the state S during the state update process:
S = torch.zeros (b, h, dh, dh) if cache is None else cache
outs = []
for i in range (t//C):
q_c = q [:, :, i*C:(i+1)*C]
k_c = k [:, :, i*C:(i+1)*C]
v_c = v [:, :, i*C:(i+1)*C]
o_prev=q_c@S #this is everything up to this block
attn=(q_c@k_c.transpose (-1,-2)).tril () #masked attention
o_curr=attn@v_c
o=o_prev+o_curr
S_new=k_c.transpose (-1,-2)@v_c #recurrent attention
S=S+S_new
outs.append (o)
o = torch.cat (outs, dim=2)
Inside a chunk, we calculate q(kᵀv), first computing the attention scores using the standard attention calculation order, combined with a causal mask.
Between different chunks, the blogger stated that (kᵀv)q is used, which is the recurrent calculation order: first build the state, then use the Query to read information from the state. The computational cost of standard attention grows at O(N²), while this method does not.
Specifically, inside each chunk, true attention calculation is still performed, which is the masked QKᵀ multiplied by V; between chunks, all historical information is compressed into the state, and then read out through a single matrix multiplication.
Therefore, the overall computational cost can be split into two parts:
- The first part is the fixed cost 2Ld²: this comes from the calculation of the state matrix and is independent of the chunk size C;
- The second part grows with C, 2LCd: it corresponds to the intra-chunk attention score matrices distributed on the diagonal of the matrix.
Full attention is just the special case of C=L, at which point the second term becomes 2L²d, and the computational complexity returns to quadratic.
Therefore, from the perspective of FLOPs, the smaller C is, the fewer computations need to be performed.
When C=1, the theoretical FLOPs are the lowest, but it may not bring the shortest actual runtime. As long as the computation task can be efficiently mapped to the GPU's matrix multiplication hardware, the GPU can often complete more arithmetic operations in a shorter time.
The next step is to extend the same method to DeltaNet.
The core problem here is actually very simple: the chunking method used for purely cumulative attention cannot be directly applied to Delta updates:
v_old = k_i @ S
u_i = b_i * (v_i - v_old)
To calculate the existing information that needs to be subtracted at each step, each intermediate state must be obtained sequentially. Without mathematical reparameterization, these calculations cannot be parallelized in the same way.
Therefore, the paper's authors rewrote the Delta update from this form:
u=v_new-v_old
S_t= S_(t-1)+K.T@u
o=q@S_T
In the original form, the model calculates one Delta per iteration through a serial loop.
The reparameterized form is as follows:
S_t = S_{t-1}(I − β_t k_t k_tᵀ) + β_t v_t k_tᵀ
o_t = S_t q_t
With this expression, the chunked implementation can calculate all C Deltas in the current chunk at once:
def chunk_delta_rule_forward (Q, K, V, beta, C):
# L: sequence length, d: head dimension
L, d = Q.shape
# chunking
Q, K, V = map (lambda x: x.reshape (-1,C,d), [Q, K, V])
beta = beta.reshape (-1, C)
K_beta = K * beta.unsqueeze (-1)
V_beta = V * beta.unsqueeze (-1)
# compute eq. 10 with vectorized forward substitution for fast inverse
T = -(K_beta @ K.t ()).tril (-1)
for i in range (1, C):
T [i, :i] = T [i, :i] + (T [i, :, None] * T [:, :i]).sum (-2)
T += torch.eye (C)
W = T @ K_beta
U = T @ V_beta
# chunkwise parallel. Eq. 8-9
S = torch.zeros (d, d)
O = torch.empty_like (V)
for i in range (L//C):
q_i, k_i, w_i = Q [i], K [i], W [i]
u_i = U [i] - w_i @ S # the corrections, all of one chunk
o_inter = q_i @ S
A_i = (q_i @ k_i.t ()).tril () #qk.t
o_intra = A_i @ u_i # attention @ v (with corrections, so u)
S += k_i.t () @ u_i # update state with addition
O [i] = o_intra + o_inter #update output with flash + recurrent
return O.reshape (L, d)
At this point, we can finally make the first direct comparison: MHA Transformer vs. DeltaNet Transformer.
Gated DeltaNet
Now, there is a method that can precisely modify the cache: whenever a new fact appears, i.e., a new Key vector, the model can accurately check the information originally stored at that location and replace it with the new information it wants to focus on later.
But this mechanism can only forget associations that have explicit new content to replace them. When the context switches, it cannot efficiently clear multiple groups of associations at once, nor can it globally decay memory to free up storage capacity.
Suppose purely cumulative linear attention is used: then adding forgetting capability is not complicated, just introducing a parameter to control the degree of state forgetting:
S_old=cache
S_new=k@v
# cache=S_old+S_new
cache=alpha * S_old + S_new
This is exactly the improvement brought by Mamba-2, which first decays the previous cache and then writes the new cache at full strength, thereby preventing the state from growing indefinitely.
The method adopted by Mamba is to use a dynamic ratio at each time step to uniformly decay all key-value associations. This method is indeed effective, but it does not consider that the importance of different key-value associations is not the same.
In other words, when the model only needs to forget a specific association, all associations are forgotten to the same degree.
In contrast, the Delta Rule can update a single fact individually but cannot allow other facts to decay naturally.
Therefore, Gated Delta Rule combines Mamba's gated update rule with the Delta Rule. It introduces the parameter alpha: when alpha=1, the update degenerates to the pure Delta Rule; when alpha=0, the memory is completely cleared.
The difficulty here is how to continue using the chunked parallel method introduced earlier to implement this mechanism.
The specific implementation still uses the DeltaNet reparameterization method introduced in the previous section. The overall mathematical form is almost the same, with only one addition: a scalar dynamically determined by the data, ranging from 0 to 1, used to control the decay degree of the old state.
In this way, the model simultaneously possesses the ability to effectively learn key-value associations and the ability to adaptively manage memory.
The corresponding code changes are as follows:
Among them, the γʳ/γⁱ term is used to calculate the cumulative decay.
Assuming a token is written at time step x and read at x+t, the cumulative scaling it experiences is: αₓαₓ₊₁αₓ₊₂…αₓ₊ₜ.
This can be seen as the multiplicative counterpart of a prefix sum calculation.
The final architecture obtained is as follows:
KDA / Kimi Linear
At this stage, researchers began experimenting with hybrid architectures: combining multiple attention mechanisms in the same model, such as combining Gated DeltaNet with Mamba.
The reason Kimi Linear has attracted attention lies in its core conclusion: in controlled variable comparison experiments, Kimi Linear's performance surpassed that of the full attention architecture.
The paper's authors describe it as an architectural solution that can directly replace traditional attention, not only with better model performance but also with decoding throughput potentially increased up to 6 times.
The main improvement of Kimi Linear over Gated DeltaNet is the introduction of a more fine-grained gating mechanism.
Previously, the model only used a single scalar to control overall decay; Kimi Linear learns a separate decay value for each channel.
KDA's update rule is still similar, but the corresponding code becomes as follows:
Among them, alpha.reshape (nb, C, d) embodies the most important contribution of this paper: fine-grained control over memory decay.
Compared with the DeltaNet Transformer, the Kimi Linear architecture mainly introduces three changes:
- Adopts a hybrid architecture, alternately inserting Multi-head Latent Attention (MLA) layers in the model;
- Uses Mixture-of-Experts (MoE) layers to replace traditional MLPs;
- Adds extra capacity to DeltaNet through alpha projection.
The key point to understand is that this is not simply, blindly scaling up the model. The newly added capacity has a clear mathematical purpose: the per-channel scaling mechanism allows the model to control memory decay more finely.
Scaling Law still holds, but model capacity must be added in the right places and in a form the system can truly utilize. On this architectural evolution path, each new architecture adds capacity to solve a specific limitation of the previous generation system.
Kimi K3
Finally, the language model backbone of Kimi K3 is quite similar to the Kimi Linear model introduced earlier.
The model contains a total of 23 macro-cycles composed of four layers. In each macro-cycle, the first three layers use Kimi Delta Attention, and the fourth layer uses Multi-head Latent Attention.
The first layer of the model uses a dense feedforward network, and all remaining layers use latent-space Mixture-of-Experts networks.
At first glance, the changes from Kimi Linear to Kimi K3 do not seem to be many:
- Model scale has grown significantly;
- Chunked AttnRes operations are added every 12 layers;
- MLA Query LoRA and output gating;
- Latent-space MoE;
- SiTU activation function;
- Gated MLA;
KDA provides recurrent memory with a constant state size, while the periodically inserted MLA layers retain the Softmax retrieval capability based on the full context.
The simplified architecture diagram below can serve as a reference for understanding the subsequent modifications.
Let's start with a few relatively direct changes: Gated MLA, Latent-space MoE, and the SiTU activation function.
Gated MLA is used to decide how much of each feature retrieved from MLA can enter the residual stream.
The specific approach is: project a gating vector from the input, and then multiply it element-wise with the retrieved features.
In traditional MoE, a learned router assigns each token to a group of expert networks based on dot-product similarity.
Kimi K3 has a total of 898 experts, two of which are shared experts that process all tokens; among the remaining 896 experts, the router selects 16 for each token.
Kimi K3 also changed the activation function in the expert networks. The traditional approach is: first apply SiLU to the up-projection result, then multiply element-wise with the gating branch, and finally perform down-projection.
Kimi K3 instead uses SiTU:
d = x.shape [-1] // 2
gate = x [..., :d].to (torch.float32)
up = x [..., d:].to (torch.float32)
situ_a = self.beta * torch.tanh (gate /self.beta) * torch.sigmoid (gate)
if self.linear_beta is not None:
up = self.linear_beta * torch.tanh (up /self.linear_beta)
return (situ_a * up).to (x.dtype)
The model also first down-projects the input to the shared expert space, and after aggregating the shared experts' outputs, up-projects it back:
This reveals a recurring difficulty in model inference: without fused operators, the new activation function runs almost 3 times slower than the original computation path.
One optimization that can offset some of the overhead is to have the expert networks run in a compressed latent space, which can significantly speed up the forward pass of the expert networks and reduce FLOPs by almost half.
The remaining changes include MLA Query LoRA, output gating, and adding a chunked Attention Residual (AttnRes) every 12 layers. AttnRes increases inference latency by about 2%, but it brings two important benefits:
- Selectively retrieves earlier representations, thereby alleviating information dilution in the residual stream and the problem of growing hidden state magnitudes;
- Achieves approximately a 1.25x computational advantage.
AttnRes and MLA solve the same underlying limitation from different directions.
KDA layers use a fixed-size state, so they inevitably need to discard some information. MLA retrieves information from the token context, while AttnRes retrieves from earlier representations along the network depth direction.
AttnRes (Attention Residual)
In each forward pass, the input goes through a series of stacked network layers. Here, each layer consists of an attention module (KDA or MLA) and an MLP or MoE module.
Normally, the input to a layer is the sum of the original embedding and the outputs of all previous layers, with all parts having the same weight:
Where represents the input of the i-th layer,
is the embedding of the current token, i.e., the last token in the sequence so far; and
is the output of the i-th layer, which can come from an attention module or an MLP module.
The problem with this approach is the lack of selective access capability.
Different types of network layers receive the same aggregated state, even though they might be better suited to use different weight combinations.
Furthermore, because this recursion relies entirely on addition, later layers must learn to produce outputs with increasingly larger magnitudes to have sufficient impact on the accumulating residual, which can lead to unstable training.
AttnRes no longer treats all layers equally but multiplies each term in the summation expression by a specially calculated weight, allowing the model to pay more attention to the most useful network layers based on the current context:
Each weight alpha_i is calculated through the dot product of Query and Key.
Here, each layer has a learned Query, while the Key and Value come from earlier residual stream states. The model normalizes the scores so that their sum is 1, and then uses these weights to perform a weighted combination of the previous states.
Therefore, the model is no longer limited to relying on the immediately preceding layer.
AttnRes allows each layer to selectively access the outputs of earlier layers and, through the learned Query, retrieve the representation most needed for the current computation.
The pseudocode below implements the same idea at the chunk granularity.
Here, a chunk is the element-wise sum of the outputs of the attention modules and MLP modules in 12 consecutive decoder layers. This result is stored as a unified depth representation for subsequent AttnRes mixing. Applying residual attention at every layer would incur excessively high training and inference overhead. Applying it only at fixed chunk boundaries can retain most of the benefits at a lower cost.
In Kimi K3, such a boundary is formed every 12 decoder layers. The model contains a total of 23 four-layer macro-cycles, resulting in 8 AttnRes chunks and improving overall inference efficiency.
This is probably the most important part of the block_attn_res function:
V = torch.stack (blocks + [partial_block]) # [N+1, B, T, D]
K = norm (V)
logits = torch.einsum ('d, n b t d -> n b t', proj.weight.squeeze (), K)
h = torch.einsum ('n b t, n b t d -> b t d', logits.softmax (0), V)
return h
At this point, the architectural evolution process from GPT-2 to Kimi K3 is complete.
The most core change among them is not just the increase in scale.
Each architectural upgrade changed what information the model stores, how it updates its state, or how it re-retrieves information that a fixed-size state cannot fully retain.
Kimi K3 combines fixed-state recurrent memory, periodic Softmax retrieval, sparse expert capacity, and selective access to depth-direction residual representations.
The final result is a system that invests extra capacity into clearly defined functional positions.
Ultimately, an associative memory system with fixed capacity, i.e., a memory system whose dimensions remain unchanged, must have some kind of eviction strategy.
Because when the memory reaches its capacity limit, purely cumulative linear operations will inevitably cause different information to interfere with each other.
Therefore, the system must introduce learned selection mechanisms such as gating, routing, or decay; and the attention mechanism remains the most effective selective reading method currently available.
For more information, check out the full article!
Reference links: