22580: Von GPT-2 bis KimiK3, erklärt

Zweiundzwanzigtausendfünfhundertachtzig. So viele GPT-2-Modelle (2019) passen in ein KimiK3 (2026). Wir haben in sieben Jahren um den Faktor 22.580 skaliert. Aber ist es nur … Skalierung?

In diesem Worklog zeige ich, wie wir hierhergekommen sind und wie viel – oder wie wenig – sich seitdem tatsächlich geändert hat. Wir verfolgen die wichtigsten architektonischen Entwicklungen, die zu KimiK3 geführt haben.

GPT-2

GPT-2 ist eine Decoder-only-Architektur:

1tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
2pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
3x = self.transformer.drop(tok_emb + pos_emb)
4for block in self.transformer.h:
5    x = block(x)
6x = self.transformer.ln_f(x)
7logits = self.lm_head(x)
8return logits

Die Eingabe erhält Token- und Positions-Embeddings:

Jeder Transformer-Block sieht vergrößert so aus:

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

Der Aufmerksamkeitsprozess:

1        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
2
3        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
4        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
5        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
6        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
7        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
8
9        # manual implementation of attention
10        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
11        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
12        att = F.softmax(att, dim=-1)
13        att = self.attn_dropout(att)
14        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
15        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
16
17        # output projection
18        y = self.resid_dropout(self.c_proj(y))
19        return y

Sobald die finale Hidden-State-Matrix erzeugt ist, bildet der Language-Model-Head sie auf Vokabular-Logits ab. Während der autoregressiven Dekodierung werden nur die Logits an der letzten Position benötigt, um das nächste Token auszuwählen.

Dies ist eine Ineffizienz der Decoder-only-Generierung: Das Modell berechnet Repräsentationen für jede Eingabeposition, aber jeder Dekodierschritt verbraucht nur die Logits der letzten Position. Ohne Caching würde ein Großteil dieser Arbeit für das nächste Token wiederholt werden.

Der KV-Cache entsteht aus einer einfachen Beobachtung: Nachdem das generierte Token an die Eingabe angehängt wurde, müsste das Modell sonst die Projektionen für alle vorherigen Token neu berechnen. Das Speichern ihrer Key- und Value-Vektoren vermeidet diese redundante Arbeit.

Dieser Speicher ist der KV-Cache. Er behält die Vektoren der vorherigen N-1 Token und kann groß genug werden, um einen Memory-Bandwidth-Engpass zu verursachen.

Insgesamt hat unser Basismodell bei etwa 50k möglichen Token, 12 Blöcken, 12 Heads und einer Embedding-Dimension von 768 etwa 124 Millionen Parameter.

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

Mit 2,8 Billionen Parametern enthält ein KimiK3-Modell ungefähr so viele Parameter wie 22.580 GPT-2-Modelle.

Lineare Aufmerksamkeit

Softmax-Aufmerksamkeit wendet ihre Nichtlinearität nach dem q·k-Produkt an und koppelt so jede Query an jeden Key. Lineare Aufmerksamkeit wendet stattdessen eine Feature-Map, wie ELU+1, getrennt auf q und k an. Dadurch wird das Produkt reassoziierbar, sodass die wachsende Menge von K- und V-Vektoren in einen festen D×D-Zustand gefaltet werden kann.

Die O(N²)-Darstellung des Papiers hat mich verwirrt. Es stimmt nicht, dass „die Kosten pro Zeitschritt für Transformer mit dem Quadrat der aktuellen Sequenzlänge skalieren". Das behebt Flash Attention … dann sah ich, dass es 2020 veröffentlicht wurde.

Zu der Zeit materialisierte das Training oft die gesamte N×N-Aufmerksamkeitsmatrix, FlashAttention existierte nicht, und Referenzimplementierungen für autoregressive Modelle berechneten die Token-Historie oft ohne KV-Cache neu.

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

Derselbe Prozess ist visuell leichter zu erkennen. Jeder Dekodierschritt führt zwei ND-Lese- und zwei 1D-Schreibvorgänge im HBM durch, während der KV-Cache linear mit der Sequenzlänge wächst, in O(N).

Beachten Sie die übermäßigen Lese- und Schreibvorgänge, die dieses Papier durch Folgendes ersetzt:

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

Es gibt einen Kompromiss.

Hier ersetzen wir die von Softmax verwendete Exponentialfunktion durch ELU+1, das getrennt auf q und k angewendet wird, bevor sie interagieren. Beide Ansätze normalisieren die resultierenden Scores, aber die von der linearen Aufmerksamkeit verwendete Feature-Map ist eine weniger ausdrucksstarke Approximation des Softmax-Kerns. Diese Approximation kann die Genauigkeit verringern, obwohl der praktische Genauigkeitsverlust von der Architektur und der Arbeitslast abhängt.

Beachten Sie, dass wir immer noch durch die Summe von qk dividieren, die der Einfachheit halber im Diagramm weggelassen wurde. Auf hoher Ebene besteht Aufmerksamkeit aus drei Schritten:

1. Die qk-Scores nichtnegativ machen. Lineare Aufmerksamkeit verwendet ELU+1, während Softmax Exponentiation verwendet.
2. Durch die Summe dividieren.
3. Den gewichteten Durchschnitt der Werte berechnen.

Dies bewahrt den grundlegenden Aufmerksamkeitsvertrag, verwendet aber eine weniger ausdrucksstarke Feature-Map, um die QK-Scores nichtnegativ zu machen.

DeltaNet (Fast Weight Programmers)

Ein endlicher Cache muss bereits gespeicherte Informationen überschreiben oder damit kombinieren. Der Zustand von Token i-1 erhält keinen eigenen Slot; er wird zur selben D×D-Matrix hinzugefügt. Neue Abfragen können daher keine perfekt isolierte Repräsentation jedes früheren Tokens mehr abrufen.

Diese Addition ist auch die Quelle des Effizienzgewinns. Die additive Aktualisierung des Caches anstelle einer Verkettung verhindert, dass er in O(N) wächst, aber dieselbe Operation führt dazu, dass Informationen interferieren. DeltaNet adressiert diesen Verlust der Wiederherstellbarkeit.

Wie Schlag in seinem Paper (Fast Weight Programmers) treffend formuliert: „Wenn die Sequenzlänge die Speicherkapazität überschreitet, kann das Modell in ein Überkapazitätsregime geraten. Um in einem solchen Regime ordnungsgemäß zu funktionieren, sollte das Modell lernen, dynamisch mit den Speicherinhalten zu interagieren und selektiv zu entscheiden, welche Key-Value-Assoziationen behalten und welche gelöscht werden sollen. Die rein additive Anweisung ist für diesen Zweck möglicherweise ungeeignet … Das endlose Hinzufügen neuer Assoziationen zu einem Speicher endlicher Größe, wie in Gleichung 17, wird unweigerlich eine Grenze erreichen."

Das Regime, das lineare Aufmerksamkeit attraktiv macht – bei dem N viel größer als D ist – legt auch ihre Hauptlimitation offen. Sobald der Zustand seine effektive Kapazität überschreitet, beginnen Assoziationen zu interferieren, weil die Aktualisierung additiv ist und nichts den Cache verlässt.

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

Ein visuelles Beispiel macht dies leichter verständlich.

Nehmen Sie eine einzelne Assoziation, geschrieben als S = k.T @ v. Wenn Sie mit demselben Key zurücklesen, erhalten Sie k @ (k.T @ v), was (k @ k.T) v ist, also die quadrierte Norm von k mal v. Das Lesen gibt also skaliert mit der quadrierten Norm des Keys zurück, und wenn Sie k auf Einheitslänge normalisieren oder das Ergebnis einfach durch die Norm dividieren, erhalten Sie v exakt zurück.

Q ist ebenfalls ein gelernter Pointer. Wq und Wk lesen denselben Residualstrom, und die Query für einen Fakt zeigt auf die Key-Richtung, in die dieser Fakt geschrieben wurde. Die Aktualisierung fragt zunächst, welche Informationen der aktuelle Key aus dem Cache abruft. Sie subtrahiert diese vorhandenen Informationen von dem Wert, den wir speichern möchten, multipliziert den Key mit der Differenz und addiert das Ergebnis zurück. Alte Informationen werden entfernt und neue Informationen an ihrer Stelle geschrieben.

DeltaNet (Parallelizing Linear Transformers with Delta Rule)

Dies ist der schwierigste Abschnitt des Beitrags. Ich habe etwa sieben Stunden gebraucht, um ein funktionierendes Verständnis davon zu entwickeln, daher werde ich die Erklärung aus der Implementierung aufbauen. Kurz gesagt, implementiert DeltaNet eine lineare Rekurrenz erster Ordnung mit generalisierten Householder-Übergangsmatrizen, die parallele chunkweise Forward-Pässe für hardwareeffizientes lineares Training ermöglicht. Es teilt die Eingaben und Ausgaben in mehrere Chunks der Größe C auf und berechnet die Ausgaben für jeden Chunk basierend auf dem Endzustand des vorherigen Chunks und den Query-Key-Value-Blöcken des aktuellen Chunks.

Das praktische Problem ist der Prefill. Eine direkte Implementierung der Delta-Regel über eine Sequenz von T Token würde so aussehen:

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)

Im Gegensatz zur Standardaufmerksamkeit erfordert diese Formulierung eine Korrektur bei jedem Key-Vektor, daher ist der Weg zu einer parallelen Matrixmultiplikation nicht sofort ersichtlich. Selbst ohne die Delta-Regel bleibt ein direkter linearer Aufmerksamkeits-Prefill sequentiell:

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)

Eine chunkweise Formulierung bietet einen effizienteren Ansatz. Die Mechanik ist anhand eines Beispiels leichter zu verstehen:

C=N gesetzt ergibt die Standard-O(N²)-Aufmerksamkeit, während C=1 die reguläre lineare Aufmerksamkeit ergibt. Zwischenwerte, die wir interpolieren, tauschen zusätzliche Arbeit innerhalb eines Chunks gegen eine bessere Hardwareauslastung. In der Praxis ist C oft 64 oder 128, weil Tensor-Core-Instruktionen auf dieser Granularität effizient arbeiten; UMMA ist ein Beispiel.

Die Zwischenkacheln werden als Teil der Zustandsaktualisierung in S gefaltet:

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)

Innerhalb eines Blocks machen wir q(kᵀv). Das ist zuerst Score, die normale Aufmerksamkeitsreihenfolge mit Maskierung. Über Blöcke hinweg folgen wir (kᵀv)q, also die rekurrente Reihenfolge, zuerst Zustand. Aufmerksamkeit wächst in O(N²) und das hier nicht. Innerhalb eines Blocks mache ich echte Aufmerksamkeit (die maskierte QKᵀ mal V), und über Blöcke hinweg falte ich alles in den Zustand und lese es mit einem Matmul zurück. Die Kosten teilen sich also in zwei Teile. Es gibt ein festes Stück, 2Ld², das die Zustandsarbeit ist und sich nicht um C kümmert. Und es gibt ein wachsendes Stück, 2LCd, das sind die Score-Matrizen auf der Diagonalen. Volle Aufmerksamkeit ist nur der Fall, wenn C gleich L ist, und dann wird dieser zweite Term zu 2L²d, quadratisch. Je kleiner ich C mache, desto weniger FLOPs habe ich.

C=1 ist die günstigste Option in reinen FLOPs, aber nicht unbedingt in der Wanduhrzeit. Eine GPU kann mehr Arithmetik schneller erledigen, wenn die Arbeit effizient auf ihre Matrix-Multiplikations-Hardware abgebildet wird.

Der nächste Schritt ist, denselben Ansatz auf DeltaNet zu erweitern.

Das zugrundeliegende Problem ist einfach: Die für rein additive Aufmerksamkeit verwendete Chunking-Methode ist nicht direkt auf die Delta-Updates anwendbar:

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

Wir brauchen jeden einzelnen Zustand, um die Informationen zu berechnen, die subtrahiert werden müssen. Wir können es nicht auf die gleiche Weise parallelisieren, ohne eine mathematische Umparametrisierung. Die Autoren schreiben daher die Delta-Updates um von:

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

Hier berechnet eine sequentielle Schleife ein Delta pro Iteration. Die umparametrisierte Form ist:

1S_t = S_{t-1}(I − β_t k_t k_tᵀ)  +  β_t v_t k_tᵀ
2o_t = S_t q_t

Diese Formulierung ermöglicht es dem chunkweisen Code, alle C Deltas auf einmal zu berechnen:

1def chunk_delta_rule_forward(Q, K, V, beta, C):
2        # L: sequence length, d: head dimension
3        L, d = Q.shape
4        # chunking
5        Q, K, V = map(lambda x: x.reshape(-1,C,d), [Q, K, V])
6        beta = beta.reshape(-1, C)
7        K_beta = K * beta.unsqueeze(-1)
8        V_beta = V * beta.unsqueeze(-1)
9
10        # compute eq. 10 with vectorized forward substitution for fast inverse
11        T = -(K_beta @ K.t()).tril(-1)
12        for i in range(1, C):
13                T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2)
14
15        T += torch.eye(C)
16        W = T @ K_beta
17        U = T @ V_beta
18
19        # chunkwise parallel. Eq. 8-9
20        S = torch.zeros(d, d)
21        O = torch.empty_like(V)
22
23        for i in range(L//C):
24                q_i, k_i, w_i = Q[i], K[i], W[i]
25                u_i = U[i] - w_i @ S # the corrections, all of one chunk
26                o_inter = q_i @ S
27                A_i = (q_i @ k_i.t()).tril() #qk.t
28                o_intra = A_i @ u_i # attention @ v (with corrections, so u)
29                S += k_i.t() @ u_i # update state with addition
30                O[i] = o_intra + o_inter #update output with flash + recurrent
31        return O.reshape(L, d)

Damit erreichen wir unseren ersten Vergleichspunkt: MHA vs. DeltaNet Transformers:

Gated Delta Net

Wir haben nun eine Methode, um präzise Änderungen am Cache vorzunehmen. Mit jedem neuen Fakt (jedem neuen Key-Vektor) können wir genau die alten Informationen betrachten, die an dieser Stelle gespeichert sind, und sie durch die neuen Informationen ersetzen, auf die wir uns konzentrieren möchten.

Dieser Mechanismus kann jedoch nur eine Assoziation vergessen, für die er einen spezifischen Ersatz hat. Er kann nicht effizient mehrere Assoziationen während eines Kontextwechsels löschen oder das Gedächtnis allgemein abbauen, um Kapazität freizugeben.

Wenn wir rein additive lineare Aufmerksamkeit betreiben würden:

Das Hinzufügen der Fähigkeit zu vergessen wäre einfach. Wir bräuchten nur einen Parameter, der den vergesslichen Zustand steuert:

1S_old=cache
2S_new=k@v
3# cache=S_old+S_new
4cache=alpha * S_old + S_new

Dies ist der Beitrag von Mamba-2. Wir dekayieren den vorherigen Cache und fügen dann den neuen Cache in voller Stärke hinzu, wodurch verhindert wird, dass der Zustand unbegrenzt wächst.

Das gleichmäßige Dekayieren aller Key-Value-Assoziationen zu jedem Zeitschritt um ein dynamisches Verhältnis ist ein funktionierender Ansatz, und genau das macht Mamba. Es berücksichtigt jedoch nicht die unterschiedliche Wichtigkeit verschiedener Key-Value-Assoziationen.

Das heißt, wenn das Modell eine bestimmte Assoziation vergessen muss, werden alle Assoziationen gleichermaßen vergessen. Die Delta-Regel hingegen kann einen einzelnen Fakt aktualisieren, hat aber keine Möglichkeit, die restlichen Fakten abklingen zu lassen.

Die Gated-Delta-Regel kombiniert daher Mambas Gated-Update-Regel mit der Delta-Regel. Sie fügt einen Parameter alpha hinzu, der bei Einstellung auf eins zur reinen Delta-Regel wechselt und bei Null den Speicher löscht. Die Herausforderung besteht darin, dies mit derselben Parallel-Chunks-Methode zu implementieren.

Die Implementierung verwendet dieselbe in der letzten Sektion beschriebene DeltaNet-Umparametrisierung. Die Mathematik ist nahezu identisch, mit einer Ergänzung: einem datenabhängigen Skalar zwischen null und eins, der das Dekay des vorherigen Zustands steuert. Dies kombiniert effektives Lernen von Key-Value-Assoziationen mit adaptivem Speichermanagement.

Die entsprechenden Codeänderungen sind unten dargestellt:

Der γʳ/γⁱ-Term berücksichtigt das kumulative Dekay. Ein Token, das zum Zeitpunkt x geschrieben und bei x+t gelesen wird, wurde mit αₓαₓ₊₁αₓ₊₂…αₓ₊ₜ multipliziert. Dies ist das multiplikative Analogon zu einer Prefix-Sum-Berechnung.

Die resultierende Architektur sieht so aus:

KDA/Kimi Linear

Zu diesem Zeitpunkt begannen Forscher, mit hybriden Modellen zu experimentieren, die mehrere Formen von Aufmerksamkeit innerhalb einer Architektur kombinieren, wie Gated DeltaNet mit Mamba.

Kimi Linear erregte Aufmerksamkeit durch eine zentrale Behauptung: Unter kontrollierten Vergleichen übertraf es die volle Aufmerksamkeit. Die Autoren präsentierten es als Drop-in-Architekturersatz mit besserer Qualität und bis zu 6-fachem Dekodierungsdurchsatz.

Kimi Linear verbessert Gated DeltaNet durch die Einführung von feinkörnigem Gating. Anstelle eines einzelnen skalaren Dekays lernt es einen separaten Dekay-Wert für jeden Kanal.

Die KDA-Update-Regel bleibt ähnlich, aber der Code sieht jetzt eher so aus:

Hier erfasst alpha.reshape(nb, C, d) den bedeutendsten Beitrag des Papiers: feinkörnige Kontrolle über den Speicherverfall.

Neben dem DeltaNet Transformer führt die Kimi Linear-Architektur drei größere Änderungen ein:

1. Sie verwendet ein Hybridsystem, das Multi-head Latent Attention (MLA)-Schichten einfügt.
2. Sie ersetzt das MLP durch eine Mixture-of-Experts (MoE)-Schicht.
3. Sie fügt DeltaNet durch die Alpha-Projektion Kapazität hinzu.

Die späteren Abschnitte behandeln MLA und MoE detaillierter. Vorerst ist der wichtige Punkt, dass dies keine blinde Skalierung ist. Die zusätzliche Kapazität hat einen spezifischen mathematischen Zweck: Die kanalweise Skalierung gibt dem Modell eine feinere Kontrolle über den Speicherverfall.

Skalierungsgesetze bleiben relevant, aber Kapazität muss an der richtigen Stelle und in einer Form hinzugefügt werden, die das System nutzen kann. Jede Architektur in dieser Progression fügt Kapazität hinzu, um eine konkrete Einschränkung des Vorgängersystems zu adressieren.

Kimi K3

Letztendlich ähnelt das KimiK3-Sprachbackbone dem obigen Kimi Linear-Modell. Es enthält 23 vierschichtige Makrozyklen. In jedem Makrozyklus verwenden drei Schichten Kimi Delta Attention und die vierte Multi-head Latent Attention. Die erste Schicht verwendet ein dichtes Feed-Forward-Netzwerk; jede verbleibende Schicht verwendet eine latente Mixture-of-Experts.

Auf den ersten Blick erscheinen die Änderungen gegenüber Kimi Linear bescheiden:

• Eine deutliche Skalierungserhöhung
• Blockweises AttnRes alle 12 Schichten
• MLA-Query-LoRA und Output-Gating
• Latent-Space MoE
• SiTU-Aktivierungen
• Gated MLA

KDA liefert einen konstanten rekurrenten Speicher, während periodische MLA-Schichten den vollständigen Softmax-Abruf über den Kontext beibehalten. Die folgende vereinfachte Visualisierung bietet eine nützliche Referenz für die unten diskutierten Änderungen.

Wir beginnen mit den direkteren Änderungen: Gated MLA, Latent-Space MoE und SiTU-Aktivierungen.

Gated MLA bestimmt, wie viel jedes abgerufene Merkmal von MLA in den Residualstrom gelangt. Dies geschieht durch elementweise Multiplikation mit einem Gate, das aus der Eingabe projiziert wird.

In einem konventionellen MoE verwendet ein gelernter Router die Dot-Produkt-Ähnlichkeit, um jedes Token an eine Teilmenge von Expertennetzwerken zu senden. KimiK3 hat insgesamt 898 Experten. Zwei sind gemeinsam genutzt und verarbeiten jedes Token; von den restlichen 896 wählt der Router 16 für jedes Token aus.

KimiK3 ändert auch die Expertenaktivierung. Anstatt SiLU auf die Up-Projektion anzuwenden, sie elementweise mit dem Gate zu multiplizieren und dann die Down-Projektion anzuwenden, verwendet es SiTU:

1d = x.shape[-1] // 2
2gate = x[..., :d].to(torch.float32)
3up = x[..., d:].to(torch.float32)
4situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
5if self.linear_beta is not None:
6    up = self.linear_beta * torch.tanh(up / self.linear_beta)
7return (situ_a * up).to(x.dtype)

Das Modell projiziert die Eingaben für die gemeinsamen Experten auch herunter und ihre endgültige Summe hoch:

Dies veranschaulicht eine wiederkehrende Herausforderung bei der Modellinferenz. Ohne einen fusionierten Kernel ist die neue Aktivierung fast 3x langsamer als der ursprüngliche Pfad. Eine ausgleichende Optimierung ist, dass die Experten in einem komprimierten latenten Raum arbeiten, was ihren Forward-Pass viel schneller macht und die FLOPs fast halbiert.

Die verbleibenden Änderungen sind MLA-Query-LoRA, Output-Gating und blockweise Attention Residuals alle 12 Schichten. AttnRes fügt etwa 2% Inferenzlatenz hinzu, bietet aber zwei wichtige Vorteile:

• Selektiver Abruf früherer Repräsentationen, der die Residualverdünnung und das Wachstum des verborgenen Zustands mildert
• Ein 1,25-facher Rechenvorteil

AttnRes und MLA adressieren dieselbe grundlegende Einschränkung aus unterschiedlichen Richtungen. KDA-Schichten arbeiten mit einem Zustand konstanter Größe und müssen zwangsläufig Informationen verwerfen. MLA greift auf den Token-Kontext zurück, während AttnRes auf frühere tiefenbezogene Darstellungen zugreift.

AttnRes

Dank @chloey3k für die Hilfe bei diesem Abschnitt. Bei jedem Vorwärtsdurchlauf durchläuft die Eingabe einen Stapel von Schichten. Hier besteht jede Schicht aus einem Aufmerksamkeitsblock (KDA oder MLA) und einem MLP- oder MoE-Block. Normalerweise ist die Eingabe für jede Schicht die Summe der ursprünglichen Einbettung und der Ausgabe aller vorherigen Schichten, alle gleich gewichtet.

hl = h1 + ∑_{i=1}^{l-1} fi(hi)

Hier ist hi die Eingabe für Schicht i, h1 die Einbettung des aktuellen Tokens (des letzten Tokens in der Sequenz bisher) und fi(hi) die Ausgabe von Schicht i (ein Aufmerksamkeits- oder MLP-Block).

Das Problem ist der fehlende selektive Zugriff. Verschiedene Schichttypen erhalten denselben aggregierten Zustand, obwohl sie von unterschiedlichen Gewichtungen profitieren könnten. Da die Rekursion rein additiv ist, müssen spätere Schichten zudem zunehmend größere Ausgaben lernen, um den akkumulierten Residuum zu beeinflussen, was das Training destabilisieren kann. Anstatt alle Schichten gleich zu behandeln, multipliziert AttnRes jeden Term dieser Summe mit einer spezifischen Gewichtung, wodurch das Modell den Schichten, die im Kontext am nützlichsten sind, mehr Bedeutung beimessen kann.

hl = α0 · h1 + ∑_{i=1}^{l-1} αi · fi(hi)

Jede Gewichtung αi wird aus einem Query-Key-Skalarprodukt berechnet. Der Query wird für jede Schicht gelernt, während die Keys und Values aus früheren Residual-Stream-Zuständen stammen. Die Scores werden normalisiert, sodass sie sich zu eins summieren, und dann zur Bildung einer gewichteten Kombination dieser Zustände verwendet.

Das Modell muss sich also nicht nur auf seinen unmittelbaren Vorgänger stützen. AttnRes gibt jeder Schicht selektiven Zugriff auf frühere Schichtausgaben, sodass ihr gelernter Query die für die aktuelle Berechnung nützlichsten Darstellungen abrufen kann.

Der Pseudocode unten wendet dieselbe Idee auf Blockebene an. Ein Block ist die elementweise Summe der Aufmerksamkeits- und MLP-Ausgaben, die über 12 Decoder-Schichten akkumuliert wurden und als einzelne Tiefendarstellung für spätere AttnRes-Mischung gespeichert wird.

Die Anwendung von Residual-Aufmerksamkeit auf jeder Schicht würde zu hohe Trainings- und Inferenzkosten verursachen. Die Anwendung nur an festen Blockgrenzen bringt den Großteil des Nutzens zu geringeren Kosten. In KimiK3 erfolgt jede Grenze nach 12 Decoder-Schichten. Über 23 vierschichtige Makrozyklen hinweg entstehen so acht AttnRes-Blöcke, was unsere Inferenzgeschwindigkeit erhöht.

Dies ist möglicherweise der wichtigste Teil der Funktion block_attn_res

1V = torch.stack(blocks + [partial_block]) # [N+1, B, T, D]
2K = norm(V)
3logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)
4h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
5return h

Damit ist der Fortschritt von GPT-2 zu KimiK3 abgeschlossen.

Die zentrale Veränderung ist nicht allein der Maßstab. Jeder architektonische Schritt verändert, was das Modell speichert, wie es diesen Zustand aktualisiert oder wie es Informationen abruft, die ein Zustand fester Größe nicht bewahren kann.

KimiK3 kombiniert ein rekurrentes Gedächtnis mit konstantem Zustand, periodischen Softmax-Abruf, sparse Expert-Kapazität und selektiven tiefenbezogenen Residualzugriff. Das Ergebnis ist ein System, das zusätzliche Kapazität dort einsetzt, wo es eine spezifische funktionale Rolle gibt.

Im Wesentlichen benötigt ein assoziatives Gedächtnis mit fester Kapazität (feste Dimensionen) eine Räumungsstrategie, da eine rein additive lineare Operation irgendwann Interferenzen verursacht, sobald die Kapazität erreicht ist. Daher ist gelernte Selektion wie Gating, Routing oder Decay notwendig, und Aufmerksamkeit ist der effektivste selektive Lesemechanismus.