GPT-2 から Kimi K3 まで、AI モデルの進化を解説

✕

GPT-2

GPT-2 is 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 receives token and positional embeddings:

✕

Each transformer block, zoomed in, looks like this:

1class Block(nn.Module):
2    def __init__(self, config):
3        super().__init__()
4        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
5        self.attn = CausalSelfAttention(config)
6        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
7        self.mlp = MLP(config)
8
9    def forward(self, x):
10        x = x + self.attn(self.ln_1(x))
11        x = x + self.mlp(self.ln_2(x))
12        return x

✕

The attention process looks like this:

1class CausalSelfAttention(nn.Module):
2
3    def __init__(self, config):
4        super().__init__()
5        assert config.n_embd % config.n_head == 0
6        # key, query, value projections for all heads, but in a batch
7        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
8        # output projection
9        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
10        # regularization
11        self.attn_dropout = nn.Dropout(config.dropout)
12        self.resid_dropout = nn.Dropout(config.dropout)
13        self.n_head = config.n_head
14        self.n_embd = config.n_embd
15        self.dropout = config.dropout
16
17    def forward(self, x):
18        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
19
20        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
21        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
22        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
23        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
24        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
25
26        # manual implementation of attention
27        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
28        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
29        att = F.softmax(att, dim=-1)
30        att = self.attn_dropout(att)
31        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
32        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
33
34        # output projection
35        y = self.resid_dropout(self.c_proj(y))
36        return y

Once the final hidden-state matrix is produced, the language-model head maps it into vocabulary logits. During autoregressive decoding, only the logits at the final position are needed to select the next token.

This is an inefficiency of decoder-only generation: the model computes representations for every input position, but each decode step consumes only the final position’s logits. Without caching, much of that work would be repeated for the next token.

✕

Without a cache, we then repeat the process for token n+1.

KV Cache Invention

The KV cache comes from a straightforward observation: after appending the generated token to the input, the model would otherwise recompute projections for all previous tokens. Storing their key and value vectors avoids that redundant work.

That storage is the KV cache. It retains vectors for the previous N-1 tokens and can become large enough to create a memory-bandwidth bottleneck.

✕

Size and scale

Overall, with about 50k possible tokens, 12 blocks, 12 heads, and an embedding dimension of 768, our baseline model is about 124M parameters.

vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
n_layer: int = 12
n_head: int = 12
n_embd: int = 768

At 2.8 trillion parameters, one Kimi K3 model contains roughly as many parameters as 22,580 GPT-2 models.

Linear attention

Softmax attention applies its nonlinearity after the q·k product, coupling every query to every key. Linear attention instead applies a feature map, such as ELU+1, to q and k separately. This makes the product reassociable, so the growing set of K and V vectors can be folded into a fixed D×D state.

The paper’s O(N²) framing is easy to misread without its 2020 context. At the time, training commonly materialized the full N×N attention matrix, FlashAttention did not exist, and reference autoregressive implementations often recomputed the token history without a KV cache.

✕

A KV cache avoids recomputation, but each decode step still reads all previous keys and values. The per-step cost grows as O(ND), the cache grows as O(ND), and generating N tokens costs O(N²D) in total. In practice, repeatedly streaming that cache from HBM is the bottleneck.

Linear attention replaces the growing cache with a fixed D×D state. Across N tokens, the work becomes O(ND²). Decode remains sequential, but when N is much larger than D, the fixed state can reduce memory traffic and deliver substantial speedups. Claims of thousand-fold improvements generally compare against the older cacheless baseline; gains over a modern KV-cached implementation are more modest.

The implementation difference is easiest to see by starting with the conventional KV cache:

1def forward(self, x, mask=None, past_kv=None):
2  # x is b,t,d
3  b,t,d=x.shape
4  d_head=d//self.num_heads
5  h=self.num_heads
6  qkv=self.qkv_proj(x)
7
8  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
9  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
10  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
11
12  # at prefill, q,k,v have shapes b,h,t,d
13  # at decode, shape is b, h, 1, d
14  # so i cat at the t dimension, dim(2)
15
16  if past_kv is not None:
17    k_past=past_kv[0]
18    v_past=past_kv[1]
19    k=torch.cat((k_past, k), dim=2)
20    v=torch.cat((v_past, v), dim=2)
21
22  scores=(q@k.transpose(-1,-2))/math.sqrt(d_head)
23  if past_kv is None: #we're in prefill and need to mask
24    causal_mask=torch.ones(t,t,dtype=bool, device=q.device)
25    causal_mask=torch.triu(causal_mask, diagonal=1)
26    scores=scores.masked_fill(causal_mask, float('-inf'))
27
28  if mask is not None:
29    scores=scores.masked_fill(~mask, float('-inf'))
30
31  #get attn (bhtt x bhtd)
32  attn=scores.softmax(-1)#bhtt
33  o=attn@v #bhtd
34  o=o.transpose(1,2).contiguous().view(b,t,d)  #b,t,d
35
36  # use x to get qkv
37  o_proj=self.o_proj(o)
38  past_kv=(k, v)
39  return o_proj, past_kv

The same process is easier to see visually. Each decode step performs two ND reads and two 1D writes to HBM, while the KV cache grows linearly, in O(N), with the sequence length.

✕

Notice the excessive reads and writes, which this paper replaces with:

1def forward(self, x, mask=None, cache=None):
2  # x is b,t,d
3  b,t,d=x.shape
4  d_head=d//self.num_heads
5  h=self.num_heads
6  qkv=self.qkv_proj(x)
7
8  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
9  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
10  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
11
12  k=F.elu(k)+1
13  k=k.transpose(-1,-2)
14  q=F.elu(q)+1
15
16  S,z=cache if cache is not None else (0.0, 0.0)
17  S=S+k@v
18  z=z+k
19
20 o=q@S #bhtd
21 denom=q@z
22 o_scaled=o/denom
23 o_scaled=o_scaled.transpose(1,2).contiguous().view(b,t,d)
24 o_proj=self.o_proj(o_scaled)
25 cache=(S,z)
26
27 return o_proj, cache

There is a trade-off. Here, we replace the exponential used by softmax with ELU+1 applied separately to q and k before they interact. Both approaches normalize the resulting scores, but the feature map used by linear attention is a less expressive approximation of the softmax kernel. That approximation can reduce fidelity, although the practical accuracy loss depends on the architecture and workload.

✕

Notice that we still divide by the sum of qk, which is omitted from the diagram for simplicity. At a high level, attention consists of three steps:

• Make the qk scores non-negative. Linear attention uses ELU+1, while softmax uses exponentiation.
• Divide by the sum.
• Compute the weighted average of the values.

This preserves the basic attention contract, but uses a less expressive feature map to make the QK scores non-negative.

DeltaNet (Fast weight programmers)

A finite cache must overwrite or combine with information already stored. The state from token i-1 does not receive its own slot; it is added to the same D by D matrix. New queries can therefore no longer retrieve a perfectly isolated representation of each earlier token.

That addition is also the source of the efficiency gain. Updating the cache additively rather than by concatenation prevents it from growing in O(N), but the same operation causes information to interfere. DeltaNet addresses this loss of recoverability.

✕

Eloquently put by Schlag’s paper (Fast Weight Programmers): “when the sequence length exceeds storage capacity, the model may end up in an overcapacity regime. To properly operate under such a regime, the model should learn to dynamically interact with the memory contents and selectively decide which key-value associations to keep and which ones to delete. The purely additive instruction may be inappropriate for this purpose…. endlessly adding new associations to a memory of finite size, as in Eq. 17, inevitably will reach a limit.“

The regime that makes linear attention attractive, where N is much larger than D, also exposes its main limitation. Once the state exceeds its effective capacity, associations begin to interfere because the update is additive and nothing leaves the cache.

1def forward(self, x, mask=None, cache=None):
2  # x is b,t,d
3  b,t,d=x.shape
4  d_head=d//self.num_heads
5  h=self.num_heads
6  qkv=self.qkv_proj(x)
7
8  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
9  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
10  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
11
12  q = F.normalize(F.silu(q), dim=-1)
13  k = F.normalize(F.silu(k), dim=-1)
14  beta = torch.sigmoid(self.w_beta(x)).view(b, 1, t, 1)
15  # new: per-token write strength
16
17  S = cache if cache is not None else 0.0
18
19  v_old = k @ S # read the board at this key
20  u = beta * (v - v_old) # the delta: only what's actually new
21  S = S + k.transpose(-1, -2) @ u # same outer-product write as before
22
23  o = q @ S # read, no denominator
24  o = o.transpose(1, 2).contiguous().view(b, t, d)
25  return self.o_proj(o), S

A visual example makes this easier to follow.

✕

Take a single association written as S = k.T @ v. Read it back with the same key and you get k @ (k.T @ v), which is (k @ k.T) v, which is the squared norm of k times v. So the read comes back scaled by the key's squared norm, and if you normalize k to unit length, or just divide the result by that norm, you get v back exactly. This is why F.normalize sits in the code: it's what makes the read exact and what makes the erase term below an actual projection. Q is also a learned pointer. Wq and Wk read the same residual stream, and the query for a fact points at the key direction that fact was written into. The update first asks what information the current key retrieves from the cache. It subtracts that existing information from the value we want to store, multiplies the key by the difference, and adds the result back. Old information is removed and new information is written in its place.

DeltaNet (Parallelizing linear transformers with delta rule)

This is the most difficult section of the post. It took me about seven hours to develop a working understanding of it, so I will build the explanation from the implementation. In short, DeltaNet implements a first-order linear recurrence with generalized Householder transition matrices, enabling chunk-wise parallel forward passes for hardware-efficient linear-time training. It splits the inputs and outputs into several chunks of size C, and computes outputs for each chunk based on the final state of the previous chunk and the query key value blocks of the current chunk.

The practical problem is prefill. A direct implementation of the Delta rule over a sequence of T tokens would look like this:

1S = torch.zeros(b, h, dh, dh) if cache is None else cache
2outs = []
3for i in range(t):
4    k_i = k[:, :, i:i+1]
5    v_i = v[:, :, i:i+1]
6    b_i = beta[:, :, i:i+1]
7    v_old = k_i @ S
8    u_i  = b_i * (v_i - v_old)
9    S = S + k_i.transpose(-1, -2) @ u_i # write
10    outs.append(q[:, :, i:i+1] @ S)
11o = torch.cat(outs, dim=2)

Unlike standard attention, this formulation requires a correction at every key vector, so the path to a parallel matrix multiplication is not immediately obvious. Even without the Delta rule, a direct linear-attention prefill remains sequential:

1S = torch.zeros(b, h, dh, dh) if cache is None else cache
2outs = []
3for i in range(t):
4    q = q[:, :, i:i+1]
5    k = k[:, :, i:i+1]
6    v = v[:, :, i:i+1]
7
8    S=S_old+k@v
9    o=q@S #bhtd
10    o=self.norm(o)
11    o=o.transpose(1, 2).contiguous().view(b, t, d)
12
13    out=self.o_proj(o)
14    cache=S
15    outs.append(out)
16
17o = torch.cat(outs, dim=2)

A chunked formulation provides a more efficient approach. The mechanics are easier to understand through an example:

✕

Setting C=N recovers standard O(N^2) attention, while C=1 gives regular linear attention. Intermediate values trade additional within-chunk work for better hardware utilization. In practice, C is often 64 or 128 because tensor-core instructions operate efficiently at that granularity; UMMA is one example.

The intermediate tiles are folded into S as part of the state update:

✕

1S = torch.zeros(b, h, dh, dh) if cache is None else cache
2outs = []
3for i in range(t//C):
4    q_c = q[:, :, i*C:(i+1)*C]
5    k_c = k[:, :, i*C:(i+1)*C]
6    v_c = v[:, :, i*C:(i+1)*C]
7
8    o_prev=q_c@S #this is everything up to this block
9
10    attn=(q_c@k_c.transpose(-1,-2)).tril() #masked attention
11    o_curr=attn@v_c
12
13    o=o_prev+o_curr
14
15    S_new=k_c.transpose(-1,-2)@v_c #recurrent attention
16    S=S+S_new
17    outs.append(o)
18
19o = torch.cat(outs, dim=2)

Within a block, we do q(kᵀv). This is score first, the normal attention order with masking. Across blocks, we follow (kᵀv)q, so we’re doing recurrent order, state first.

Attention grows in O(N²) and this does not. Inside a block I do real attention (the masked QKᵀ times V), and across blocks I fold everything into the state and read it back with one matmul. So the cost splits in two. There's a fixed piece, 2Ld², which is the state work and doesn't care about C at all. And there's a growing piece, 2LCd, which is the score matrices sitting on the diagonal. Full attention is just the case where C equals L, and then that second term becomes 2L²d, quadratic. So the smaller I make C, the fewer FLOPs I do.

C=1 is the cheapest option in pure FLOP terms, but not necessarily in wall-clock time. A GPU can complete more arithmetic faster when the work maps efficiently onto its matrix-multiply hardware.

The next step is to extend the same approach to DeltaNet.

✕

The mathematical detail is substantial, but the underlying issue is simple: the chunking method used for purely additive attention does not directly apply to the delta updates:

v_old = k_i @ S
u_i  = b_i * (v_i - v_old)

This means we need every single state in order to compute the information that needs to be subtracted out. We can't parallelize it the same way without some mathematical reparameterization.

The authors therefore rewrite the delta updates from:

u=v_new-v_old
S_t= S_(t-1)+K.T@u
o=q@S_T

Here, a sequential loop computes one delta per iteration. The reparameterized form is: