- Ausgehend von Softmax-Attention wird schrittweise hergeleitet: lineare Attention mit Zustand fester Größe, DeltaNet, das nur Fehler speichert, Gated DeltaNet, das den gesamten Zustand dämpft, bis hin zu Kimi Delta Attention (KDA), das kanalweise dämpft
- Grundlegende lineare Attention speichert die Summe vergangener Key-Value-Außenprodukte im Zustand (S_t) und arbeitet dadurch linear in der Sequenzlänge, erzeugt aber additive Schreibinterferenz, weil neue Werte nicht zugewiesen, sondern zu bestehenden Zuordnungen addiert werden
- DeltaNet schreibt die mit (\beta_t) gewichtete Differenz zwischen dem aus dem aktuellen Key vorhergesagten Wert und dem Ziel-Value; drei Deutungen — Bedingung der sofortigen Rekonstruktion, Online-Gradientenabstieg und Rang-1-Zustandsupdate — führen auf dieselbe Formel
- Gated DeltaNet dämpft zuerst den gesamten Zustand mit einem Skalar (\alpha_t), und KDA erweitert dies zu einer Diagonalmatrix (D_t=\operatorname{Diag}(\alpha_t)), sodass Informationen pro Key-Kanal mit unterschiedlichen Raten behalten oder gelöscht werden
- Dieselbe KDA-Rekurrenz wird sowohl mit einem fusionierten rekurrenten Triton-Kernel fürs Decoding als auch mit einem Chunk-Verfahren für Training und langes Prefill ausgeführt; das Chunk-Verfahren rekonstruiert Abhängigkeiten innerhalb eines Chunks per Dreieckslösung und formt sie in Matrixmultiplikationen um
Notation und Reihenfolge der Herleitung
- In der Bra-Ket-Notation ist (\lvert q\rangle) ein Spaltenvektor, (\langle k\rvert) ein Zeilenvektor, (\langle k\vert q\rangle) ein Skalar und (\lvert v\rangle\langle k\rvert) eine Matrix
- Es wird ein einzelner kausaler Attention-Head mit reellen Vektoren verwendet; zudem wird angenommen, dass DeltaNet-Keys normalisiert sind und der Zustand vom Key-Raum in den Value-Raum abbildet
- Die Herleitungsreihenfolge ist Softmax-Attention → lineare Attention → DeltaNet → Gated DeltaNet → KDA; zum Schluss folgt die Verbindung zu rekurrenten und chunkbasierten Triton-Implementierungen
- Zwei Varianten aus der DeltaNet-Familie werden in aktuellen Qwen- und Kimi-Modellfamilien eingesetzt
Von quadratischer Komplexität bei Attention zu linearem Zustand
- Gewöhnliche kausale Softmax-Attention berechnet die Ähnlichkeit von Key und Query, normalisiert die Scores über alle vergangenen Keys zu einer Verteilung und gibt die gewichtete Summe der Value-Vektoren aus
- Für eine Sequenz der Länge (T) gibt es (T^2) Key-Query-Paare
- Bei autoregressiver Inferenz lassen sich Keys und Values cachen, aber die Cache-Größe wächst mit der Sequenz
- Auch eine neue Query muss die gesamte Vergangenheit prüfen
- Der Softmax-Nenner hängt gemeinsam von der aktuellen Query und allen vorherigen Keys ab, daher lässt sich die Reihenfolge der Berechnung nicht einfach umstellen
- Entfernt man Softmax, lässt sich die Ausgabe als Summe vergangener Key-Value-Außenprodukte zusammenfassen
- (S_t=\sum_{i\le t}\lvert v_i\rangle\langle k_i\rvert)
- (S_t=S_{t-1}+\lvert v_t\rangle\langle k_t\rvert)
- (\lvert o_t\rangle=S_t\lvert q_t\rangle)
- Die zentrale Identität ist ((\lvert v\rangle\langle k\rvert)\lvert q\rangle=\langle k\vert q\rangle\lvert v\rangle); statt aller vergangenen Keys und Values speichert man die aufsummierten Außenprodukte in einem Zustand fester Größe (d_v\times d_k)
- Da die Tokens nur einmal durchlaufen werden, arbeitet das Verfahren linear in der Sequenzlänge, verliert dafür aber die Normalisierung und Selektivität von Softmax
- Ausgereiftere lineare Attention nutzt Feature-Maps und Normalisierungsterme
Das Problem des additiven Schreibens bei linearer Attention
- Liest man direkt nach dem Schreiben von (\lvert v_t\rangle\langle k_t\rvert) mit dem normalisierten aktuellen Key wieder, ergibt sich (S_t\lvert k_t\rangle=S_{t-1}\lvert k_t\rangle+\lvert v_t\rangle)
- Ein neuer Schreibvorgang weist also nicht zu, dass der Speicher (v_t) zurückgeben soll, sondern addiert ihn im Stil von
+=auf den bestehenden Rückgabewert - Wenn der vorherige Zustand bereits den korrekten Wert liefert, verdoppelt sich derselbe Value; da Keys nicht orthogonal zueinander sind, kann jeder Schreibvorgang bestehende Schreibvorgänge stören
- Lineare Attention bietet komprimierten assoziativen Speicher, führt aber additive Updates statt eines für die Aufgabe näherliegenden
=-artigen Updates aus
DeltaNet: Vorhersagefehler statt Werte schreiben
- DeltaNet liest zunächst die bestehende Vorhersage für den neuen Key (\widehat v_t=S_{t-1}k_t) und schreibt nur die Differenz statt des vollständigen Value
- (e_t=\beta_t(v_t-S_{t-1}k_t))
- (S_t=S_{t-1}+e_tk_t^\mathsf T)
- Die gelernte Schreibstärke (\beta_t) liegt im Bereich ([0,1])
- Liest man sofort erneut mit demselben Key, erhält man ((1-\beta_t)S_{t-1}k_t+\beta_tv_t)
- Bei (\beta_t=1) wird exakt (v_t) zurückgegeben
- Kleinere Werte verschieben die bestehende Vorhersage nur teilweise in Richtung des Ziels
- Das Update ist lokal im Key-Raum
- In Query-Richtungen, die zum aktuellen Key orthogonal sind, ist das Außenprodukt-Update 0, daher ändert sich die Antwort nicht
- Nur die Zuordnung in Richtung des aktuellen Keys wird selektiv ersetzt
-
Herleitung über Rekonstruktionsverlust
- Betrachtet man den Zustand (S) als lineare Abbildung und setzt den Verlust für das aktuelle Key-Value-Paar auf (\frac12\lVert Sk_t-v_t\rVert_2^2), dann ist der Gradient ((Sk_t-v_t)k_t^\mathsf T)
- Führt man ausgehend von (S_{t-1}) einen Gradientenabstiegsschritt der Größe (\beta_t) aus, erhält man genau die Update-Gleichung von DeltaNet
- Dasselbe Update lässt sich auf drei Arten deuten
- In Speicheroperationen ist (\beta_t) die Stärke des Ersetzens bestehender Zuordnungen
- Im Online-Lernen ist (\beta_t) die Lernrate
- In der linearen Algebra ist es das Rang-1-Außenprodukt aus Vorhersagefehler und Key
-
Strukturierter Zustandsübergang
- Entwickelt man das Update aus, ergibt sich (S_t=S_{t-1}(I-\beta_tk_tk_t^\mathsf T)+\beta_tv_tk_t^\mathsf T)
- Für normierte Keys hat (I-\beta_tk_tk_t^\mathsf T) den Eigenwert (1-\beta_t) in Richtung des aktuellen Keys und den Eigenwert 1 in allen orthogonalen Richtungen
- Die Zuordnung in Richtung des bestehenden Keys wird zuerst entfernt und dann die neue hinzugefügt, aber die Verwaltung der Lebensdauer des gesamten Zustands bleibt noch ungelöst
Gated DeltaNet: Zuerst den gesamten Zustand vergessen
- Wird die gesamte Vergangenheit in einer einzigen Matrix komprimiert, lassen sich einzelne bereits im Zustand verschmolzene Tokens nicht selektiv überspringen
- DeltaNet korrigiert in der Umgebung des aktuellen Keys, aber alte Informationen in anderen Richtungen bleiben erhalten und können auch künftig zu Lesevorgängen beitragen
- Gated DeltaNet führt ein gelerntes skalares Retention-Gate (\alpha_t\in[0,1]) ein
- Vergessen mit (\widetilde S_t=\alpha_tS_{t-1})
- Vorhersage mit (\widehat v_t=\widetilde S_tk_t)
- Korrektur mit (e_t=\beta_t(v_t-\widehat v_t))
- Schreiben mit (S_t=\widetilde S_t+e_tk_t^\mathsf T)
- Die Reihenfolge Vergessen → Vorhersage → Korrektur → Schreiben ist wichtig
- Würde man vor der Dämpfung vorhersagen, wären der Speicher für die Fehlerberechnung und der tatsächlich aktualisierte Speicher nicht derselbe
- Die Delta-Regel übernimmt das Ersetzen für den Ziel-Key, das skalare Gate das globale Löschen; beide lösen also unterschiedliche Probleme
- Allerdings wirkt ein einzelnes (\alpha_t) auf die gesamte Matrix, daher müssen alle Key-Kanäle im gleichen Verhältnis behalten oder vergessen werden
Kimi Delta Attention: Kanalweise Dämpfung
- Kimi Delta Attention ersetzt den Skalar (\alpha_t) durch einen Vektor der Dimension (d_k) und bildet (D_t=\operatorname{Diag}(\alpha_t))
- Da der Zustand vom Key-Raum in den Value-Raum abbildet, entsprechen Key-Kanäle den Spalten von (S); die Rechtsmultiplikation (S_{t-1}D_t) wendet daher auf jede Spalte eine andere Retentionsrate an
- KDA arbeitet in der folgenden Reihenfolge
- Kanalweise Dämpfung der Key-Kanäle mit (\widetilde S_t=S_{t-1}D_t)
- Vorhersage mit (\widehat v_t=\widetilde S_tk_t)
- Korrektur mit (e_t=\beta_t(v_t-\widehat v_t))
- Schreiben mit (S_t=\widetilde S_t+e_tk_t^\mathsf T)
- Lesen mit (o_t=S_t(d_k^{-1/2}q_t))
- Der konzeptionelle Übergang von Gated DeltaNet zu KDA besteht nur darin, (\alpha_t) zu (D_t) aufzuwerten, doch damit lässt sich ein Kanal löschen, während andere erhalten bleiben
-
Diagonal-Niedrigrang-Übergang
- Entwickelt man KDA aus, ergibt sich (S_t=S_{t-1}A_t+\beta_tv_tk_t^\mathsf T) mit (A_t=D_t(I-\beta_tk_tk_t^\mathsf T))
- Man kann (A_t=D_t-b_ta_t^\mathsf T) schreiben, mit (b_t=D_tk_t) und (a_t^\mathsf T=\beta_tk_t^\mathsf T); damit entsteht ein Diagonal-Plus-Low-Rank-(DPLR)-Übergang
- DPLR bezeichnet einen (d_k\times d_k)-Übergang im Key-Raum, während der Speicherzustand selbst weiterhin eine (d_v\times d_k)-Matrix ist
- Mit jeder Familienvariante kommt folgende Funktion hinzu
- Lineare Attention: rekurrenter Speicher fester Größe
- DeltaNet: selektives Ersetzen in Zielrichtung
- Gated DeltaNet: Dämpfung des gesamten Zustands
- KDA: Dämpfung pro Key-Kanal
- In Implementierungen speichert man meist (g_t=\log\alpha_t\le0) und berechnet daraus per (\exp(g_t)) die Retentionsrate
- Eine 5-stufige Referenzimplementierung mit transponiertem (d_k\times d_v)-Layout findet sich in
naive_recurrent_kda
Fusionierter rekurrenter Triton-Kernel fürs Decoding
- Für KDA gibt es zwei zentrale Ausführungsarten
- Fusionierte rekurrente Ausführung: geeignet für Decoding, kurze Sequenzen und zustandsbehaftetes Serving
- Chunk-Verfahren: geeignet für Training und langes Prefill
fused_recurrent_kda_fwdstartet pro Sequenz, Value-Head und 32 breitem Value-Tile ein Triton-ProgrammBKdeckt in den üblichen unterstützten Konfigurationen die Key-Dimension ab- Jedes Programm besitzt ein Tile des transponierten Zustands der Form
[BK, BV]und durchläuft die Tokens der Reihe nach - Unterschiedliche Value-Tiles, Heads und Sequenzen laufen unabhängig voneinander
- Der Kernel führt Zustandsdämpfung, Prädiktionsreduktion über den Key, Residualberechnung, Schreiben des Außenprodukts und Lese-Reduktion über die Query genau gemäß der Rekurrenz aus
- Das eignet sich für Decoding, bei dem jeweils nur ein neues Token eingeht, ist aber für Training und langes Prefill ungünstig, weil sich die Vektoroperationen nicht effizient in große, Tensor-Core-freundliche Matrixmultiplikationen umformen lassen
Chunkwise KDA: Die Rekurrenz in Matrixmultiplikationen umstellen
- Chunkwise KDA verarbeitet (C) Tokens gemeinsam und muss dabei exakt dieselben Zustände und Ausgaben erzeugen wie die rekurrente tokenweise Ausführung
- Jeder Chunk berechnet zwei Ergebnisse
- Den Zustand (S_{c+1}) nach Verarbeitung des gesamten Chunks ausgehend vom eingehenden Zustand (S_c)
- Die kausalen Ausgaben aller Tokens innerhalb des Chunks
- Die entscheidende Schwierigkeit ist, dass der Delta-Fehler jedes Tokens von früheren Schreibvorgängen im selben Chunk abhängt
-
Kumulative Dämpfung und temporäre Fehler
- Bezeichne die diagonale Dämpfung von Token (i) mit (D_i) und die kumulative Dämpfung von der Chunk-Grenze bis Token (i) mit (D_{0:i}=D_0D_1\cdots D_i)
- Wenn ein Schreibvorgang von Token (j) bis Token (i) weitergegeben wird, wirkt (D_{j+1:i}); da die Matrizen diagonal sind, kommutieren die Dämpfungsmatrizen untereinander
- Zunächst werden temporäre Fehler parallel berechnet, wobei andere Schreibvorgänge im Chunk ignoriert werden
- (\bar e_i=\beta_i(v_i-S_cD_{0:i}k_i))
- Für alle Tokens außer dem ersten lassen sich diese temporären Fehler nicht direkt verwenden, weil der Einfluss früherer Schreibvorgänge im selben Chunk fehlt
-
Wiederherstellung kausaler Abhängigkeiten
- Der Koeffizient des Einflusses von früherem Token (j) auf den Fehler des aktuellen Tokens (i) sei (\rho_{ij}=\beta_i k_j^\mathsf TD_{j+1:i}k_i)
- Der tatsächliche Fehler hängt dann sequentiell ab als (e_i=\bar e_i-\sum_{j<i}\rho_{ij}e_j)
- Packt man (\rho_{ij}) in eine strikt untere Dreiecksmatrix (R_c), lässt sich die gestapelte Fehlermatrix berechnen als (E_c=\bar E_c(A_c^{kk})^\mathsf T) mit (A_c^{kk}=(I+R_c)^{-1})
- Eine allgemeine dichte Matrixinversion ist nicht nötig
- (I+R_c) ist eine Dreiecksmatrix mit Einsen auf der Diagonale
- Für jeden Value-Kanal genügt eine kausale Dreieckslösung
-
Berechnung des Zustands am Chunk-Ende
- Der eingehende Zustand durchläuft alle Dämpfungen des Chunks, und jeder Schreibvorgang im Chunk durchläuft nur die Dämpfungen nach seiner eigenen Position
- Stapelt man die bis zum Chunk-Ende gedämpften Keys zeilenweise zu (K_c^{\mathrm{end}}), lässt sich der Zustand als folgende Matrixmultiplikation schreiben
- (S_{c+1}=S_cD_{0:C-1}+E_cK_c^{\mathrm{end}})
- Damit werden viele Rang-1-Schreibvorgänge zu einer einzigen Matrixmultiplikation zusammengefasst und der gesamte Chunk-Zustand in einem Schritt fortgeschrieben
-
Berechnung aller Ausgaben innerhalb des Chunks
- KDA liest erst nach dem Schreiben des aktuellen Tokens, daher enthält die Ausgabe von Token (i) auch seinen eigenen Schreibvorgang
- Der Koeffizient des Einflusses eines früheren Schreibvorgangs (j) auf Query (i) sei (\chi_{ij}=s,k_j^\mathsf TD_{j+1:i}q_i), (j\le i)
- Diese Koeffizienten werden in einer unteren Dreiecksmatrix (A_c^{qk}) angeordnet
- Nullen im oberen Dreieck blockieren Beiträge zukünftiger Tokens
- Diagonalelemente spiegeln wider, dass das aktuelle Token nach seinem eigenen Schreibvorgang liest
- Stapelt man die von der Chunk-Grenze bis zu jeder Query gedämpften Vektoren zu (Q_c^{\mathrm{boundary}}), ergibt sich die Gesamtausgabe zu
- (O_c=sS_cQ_c^{\mathrm{boundary}}+E_c(A_c^{qk})^\mathsf T)
- Die erste Matrixmultiplikation liest den gedämpften Eintrittszustand des Chunks, die zweite addiert die kausalen Beiträge der Schreibvorgänge innerhalb des Chunks
Chunkweiser Triton-Pipeline
- Die Chunk-Implementierung ist kein einzelner riesiger Kernel, sondern eine Pipeline aus mehreren Kernel-Aufrufen
- Zuerst werden die kumulativen Log-Dämpfungen innerhalb eines Chunks berechnet
- Über die Differenz zweier Prefix-Summen lässt sich (D_{j+1:i}) darstellen, ohne Retentionsvektoren lang auszumultiplizieren
- Anschließend werden die kausalen Interaktionsmatrizen (A^{qk}) und (A^{kk}) erzeugt und mit (A^{kk}) eine WY-Form für die korrigierten Schreibvorgänge des Chunks aufgebaut
- Der Zustands-Kernel ist der einzige Durchlauf über die Chunks hinweg
- Er erzeugt den Zustand, der in jeden Chunk eingeht
- Er löst die Delta-Fehler des Chunks auf
- Sobald die Eintrittszustände berechnet sind, kann der Output-Kernel Tokens aus verschiedenen Chunks und Tiles parallel verarbeiten
- Die reale Implementierung berechnet zunächst diagonale Interaktionsblöcke mit 16 Tokens und führt dann fusionierte Kerne für Nichtdiagonales und Dreieckslösungen aus
chunk_kda_fwdkoordiniert die Schritte; die zentralen Einstiegspunkte sindchunk_kda_fwd_intra,chunk_gated_delta_rule_fwd_h,chunk_gla_fwd_o_gkv_newim Code ist der aufgelöste Fehlerhist der Eintrittszustand des Chunkskgist der bis zum Chunk-Ende gedämpfte Key
- Rekurrente und chunkweise Ausführung sind keine unterschiedlichen Attention-Arten, sondern zwei Ausführungspläne derselben KDA-Rekurrenz
- Die rekurrente Ausführung sind serielle Vektoroperationen für latenzarmes Decoding
- Die chunkweise Ausführung sind Matrixoperationen für Tensor-Core-zentriertes Training und Prefill
1 Kommentare
Hacker-News-Kommentare
In den vergangenen 15 Jahren brauchte Machine Learning eine einheitliche mathematische Notation, und vermutlich wird es sie weiterhin brauchen. Früher war es noch schlimmer: In Papers von Forschern aus aller Welt tauchten die abenteuerlichsten Notationen auf.
Wenn sich die Notation von Paper zu Paper ändert, entsteht Reibung beim Verstehen. Immerhin erklärt dieser Beitrag seine Notation von Anfang an ausdrücklich; Papers, die das tun, sind eher selten. Anfangs habe ich nicht einmal die Funktion zum Umschalten der Notation bemerkt, aber sie ist sehr nützlich.
∣q⟩gegenüber Ein-Buchstaben-Symbolen oder expliziten Datentypen bevorzugt. Sie mag den Vorteil der Kürze haben, aber Formeln als Pseudocode oder in einer echten Programmiersprache wie Python zu schreiben, wäre wohl viel leichter verständlich.k,qundSsind; ohne entsprechendes Hintergrundwissen bleibt ein Großteil des Textes undurchsichtig.Es heißt zwar „hätte man auch selbst darauf kommen können …“, aber etwas zu schaffen oder zu kombinieren, das es zuvor nicht gab, ist enorm schwierig.
Sobald jemand eine schwierige Aufgabe schließlich veröffentlicht, kommen sofort Reaktionen wie „ist doch gar nicht so schwer“ oder „das hätte ich auch gekonnt“, und alles beginnt simpel zu wirken. Beim Entwickeln denkt man auch häufig, man habe etwas Neues erfunden, nur um später festzustellen, dass es schon in den 1970ern entwickelt wurde und weit verbreitet ist. Es ist einem nur nie begegnet, deshalb wusste man nicht, dass es existiert.
Für mich macht die Bra-Ket-Notation alles einfach und intuitiv. Bei der Vektornotation war ich oft verwirrt, welche Seite Zeile und welche Spalte ist, folgte nur noch den Blöcken und verlor die Konzentration; mit Bra-Ket war das Ganze sehr intuitiv.
Ich habe das Gefühl, dass mir viele gute Texte entgangen sind, und werde versuchen, auch andere Texte in diese Notation zu übertragen. Zur Einordnung: Ich habe einen Doktortitel in Physik und eine leichte Dyslexie.
Bei einem Stil wie „Das äußere Produkt ist eine Matrix und das innere Produkt ist eine Zahl. Statt alle früheren Keys und Values zu speichern, speichert man die Summe der äußeren Produkte in einem Zustand fester Größe
S_t“ bin ich überzeugt, dass der Text von einem LLM geschrieben wurde.–) zu verwenden, kommt so etwas heraus.Es gibt auch ein visualisiertes Tutorial: https://snowchord.com/blog/linear-attention-visualized/
Jedes Mal, wenn ich solche Texte und Titel sehe, empfinde ich tiefe Dankbarkeit und Demut gegenüber den vielen Menschen, die viel klüger sind als ich. In der Schule und im Bachelor galt ich als sehr klug und bin sicher schlauer als der Durchschnitt, aber es gibt bestimmt Millionen Menschen, die mich wie einen Anfänger aussehen lassen würden.
Mit klug meine ich hier die Fähigkeit, riesige und komplexe Konzepte und Systeme im Kopf zu behalten und darüber nachzudenken; eine Begabung, die besonders für Mathematiker wichtig zu sein scheint.
Ein Gedankenexperiment, das ich bei ein paar Drinks mit einem Freund durchgespielt habe, war, Kinder von Bildschirmen und algorithmisch ausgespielten Masseninhalten abzuschirmen und sie in einer lernfreundlichen Umgebung aufzuziehen, in der die Qualität von Medien und Materialien so streng kontrolliert wird wie beim Training modernster Modelle. Eine Art Kloster für Kinder, in dem ihnen das aktuellste Wissen über die Realität durch Mathematik, Ingenieurwesen, Informatik, Deep Learning usw. vermittelt wird.
Letztlich braucht man, um mit fortgeschrittenen KI-Tools die Grenzen des Wissens zu erweitern, weiterhin sehr kluge Menschen, deren Denken nicht stark verunreinigt ist. Die Vorstellung, dass KI Menschen vollständig ersetzen wird, geht in die falsche Richtung.
Zur Info: Der Name Bra-Ket-Notation leitet sich tatsächlich von „bracket“ ab.
https://en.wikipedia.org/wiki/Bra-ket_notation
Anfangs war ich skeptisch, aber dank der Ket-Notation wurden die Operationen deutlich klarer, und sie gefiel mir. Allerdings wäre auch eine kurze Auffrischung zu einigen Variablen wie
d_kbei quadratischer Attention hilfreich gewesen.Zuerst war ich niedergeschlagen, weil mir diese Lösung nicht eingefallen war, aber als mir klar wurde, dass ich sogar Schwierigkeiten habe, eine binäre Suche in JavaScript selbst zu schreiben, war ich sofort beruhigt. Es besteht keinerlei Chance, dass ich Kimi Delta Attention selbst erdacht hätte.
Auch Schleifen werden selten tiefer als zwei oder drei Ebenen; wenn es komplizierter wird, übergibt man es ohnehin besser an eine Bibliothek.