Kronecker-Produkt-Attention (KPA): Schlüssel und Werte als Summen von Kronecker-Produkten kompakter Faktoren, bis zu 10-fache KV-Cache-Kompression, ohne eigenen Rotationszweig

Bei der Inferenz wird für große Modelle ein immer deutlicherer Engpass der KV-Cache. Um Doppelrechnungen zu vermeiden, speichert das Modell die Schlüssel und Werte jedes Tokens in jeder Schicht, und dieser Cache wächst linear mit Sequenzlänge, Batchgröße und Tiefe. Erreichen Kontexte Hunderttausende Token und wachsen Batches, um den Durchsatz zu erhöhen, kann der KV-Cache den Großteil des GPU-Speichers belegen, wenig Platz für Gewichte und Aktivierungen lassen und die Kontextlänge, Batchgröße und Einsatzgröße begrenzen, die man sich leisten kann. Die vorhandenen Abhilfen gehen jeweils einen Kompromiss ein. Multi-Query-Attention (MQA) und Grouped-Query-Attention (GQA) teilen Schlüssel und Werte über Köpfe hinweg auf Kosten der Darstellungsvielfalt. Multi-Head Latent Attention (MLA) speichert komprimierte latente Zustände, braucht aber zusätzliche abgekoppelte Mechanismen zur Positionskodierung. Die Tensor-Produkt-Attention (TPA) faktorisiert mit niedrigrangigen äußeren Vektorprodukten und kann durch eine zu strenge Niedrigrangstruktur eingeschränkt sein. Die Arbeit 2610.12657 stellt vor diesem Hintergrund die Kronecker-Produkt-Attention (KPA) vor.

Die Kernidee von KPA ist, die vollständigen Schlüssel- und Wertematrizen nicht zu speichern, sondern die Schlüssel und Werte jedes Tokens als Summen von Kronecker-Produkten kompakter, tokenabhängiger Matrixfaktoren zu schreiben. Ein Kronecker-Produkt baut aus zwei kleinen Matrizen eine große, kann also Struktur auf Kopfebene und auf der Ebene der Merkmale innerhalb eines Kopfes setzen, während das äußere Vektorprodukt bei TPA nur eine der beiden abdeckt. Die Arbeit merkt an, empirische Studien zeigten, dass sich Attention-Köpfe auf verschiedene Rollen wie syntaktische, semantische und positionelle Muster spezialisieren, sodass eine gute KV-Faktorisierung auf beiden Ebenen zugleich arbeiten sollte. Sie zeigt, dass die Zerlegung in ein äußeres Vektorprodukt von Rang eins bei TPA als „entartete Kronecker-Aufteilung“ von KPA auftritt, KPA TPA also strikt verallgemeinert, und dass nicht entartete Matrixfaktoren reichere Blockmuster darstellen können. Beim Decodieren werden nur die Faktoren gespeichert, und das Kompressionsverhältnis hängt davon ab, wie die Kopfdimension aufgeteilt wird: Je asymmetrischer die Faktoren, desto stärker die Kompression. Diese „Kronecker-Kopfaufteilung“ ist ein Hyperparameter, den man beim Entwurf einstellt, um sich entlang des Kompromisses zwischen Kompression und Genauigkeit zu bewegen, ohne sonstige Änderung der Architektur.

Die Positionskodierung ist der andere Kernpunkt. RoPE wirkt auf den vollständigen Schlüssel, KPA speichert aber nur Faktoren, daher muss die Arbeit klären, wo die Rotation hingehört. Die Antwort von MLA ist ein eigener abgekoppelter RoPE-Zweig. Die Arbeit gibt einen zur Kronecker-Darstellung passenden strukturierten Rotationsoperator an. Er bewahrt die Relativpositions-Eigenschaft von RoPE und dreht nur die kompakten Faktoren, sodass der abgekoppelte Zweig nicht mehr nötig ist. Die Arbeit sagt, KPA übertreffe MLA, obwohl es keinen solchen Zweig benötige.

Bei den Experimenten trainieren die Autoren Modelle im Llama-Stil in vier Größen, etwa 125 Mio., 353 Mio., 773 Mio. und 1,5 Mrd. Parameter, mit SwiGLU und RoPE auf nanoGPT-Basis, auf etwa 50 Mrd. Token FineWeb-Edu, mit Rang RK=RV=2 für Schlüssel und Werte und RQ=6 für Anfragen. Sie werten neun Zero-Shot-Benchmarks mit lm-evaluation-harness aus: ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenbookQA, PIQA, WinoGrande, MMLU und SciQ. Die Arbeit sagt, bei identischen Trainingsbedingungen erreiche oder übertreffe KPA TPA, MHA, MQA, GQA und MLA und erreiche dabei bis zu 10-fache KV-Cache-Kompression. Bei einer Größe erzielt KPA den besten Durchschnitt bei 6,4-facher Kompression, demselben Verhältnis wie TPA, mit 1,49 Punkten mehr Genauigkeit, und der größte Gewinn liegt bei ARC-Easy. Eine andere Konfiguration erreicht 8-fache Kompression. Bei etwa 6,7-facher Kompression übertrifft es die anderen Basislinien, mit klaren Gewinnen bei ARC-Easy, BoolQ und SciQ. Bei 1,5 Mrd. führt es in acht von neun Benchmarks, und bei 125 Mio. übertreffen alle KPA-Konfigurationen die 46,21 % von TPA. Asymmetrische Aufteilungen mit einem größeren Faktor auf Kopfebene schlagen sowohl symmetrische Aufteilungen als auch TPA.

Wie ist die Arbeit zu beurteilen? Ihre Stärken sind eine klare Idee, der Beweis einer strikten Verallgemeinerung, eine Antwort auf das Zusammenspiel von Kompression und Rotations-Positionskodierung und ein einstellbares Kompressionsverhältnis. Die Vorbehalte sind ebenso klar. Erstens sind alle Modelle klein und von den Autoren von Grund auf trainiert, bis zu 1,5 Mrd. Parametern und etwa 50 Mrd. Token, und nicht auf ein bestehendes großes Modell angewendet. Zweitens liegt die Auswertung vor allem bei der Zero-Shot-Genauigkeit. Die Arbeit sagt am Ende selbst, der Bau eines effizienten faktorisierten Decodierkerns und das Testen von KPA bei längerem Kontext seien die nächsten natürlichen Schritte, was heißt, dass es zu Decodiergeschwindigkeit, Latenz und Verhalten bei langem Kontext noch keine Daten gibt. Drittens sind die Ergebnisse selbst berichtet, und wir sahen keine unabhängige Nachprüfung. Wenn Sie die Kosten des KV-Caches im Blick haben, sehen Sie darin ein beobachtenswertes Design und beurteilen, ob es GQA oder MLA ersetzen kann, wenn Kern und Ergebnisse bei langem Kontext vorliegen.