🤖 machine learning

Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models

Das Papier stellt Flash PD-SSM vor, ein speicheroptimiertes strukturiert-sparses Zustandsraummodell, das dynamisch aus einer Menge trainierbarer sparsamer Matrizen auswählt, um die hohe Ausdruckskraft unstrukturierter Modelle zu erreichen und gleichzeitig die für das Training in großem Maßstab erforderliche Recheneffizienz sowie state-of-the-art-Leistung bei langen Sequenzen beizubehalten.

Ursprüngliche Autoren: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

Veröffentlicht 2026-05-20

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, einen superintelligenten Roboter zu bauen, der lange Geschichten lesen, Details merken und vorhersagen kann, was als Nächstes passiert. Dafür benötigt der Roboter ein „Gehirn", das Informationen während des Lesens speichern kann.

Lange Zeit waren die besten Gehirne für diese Aufgabe wie Transformer (die Technologie hinter vielen aktuellen KI-Chatbots). Sie sind unglaublich intelligent, aber sehr schwer und langsam, wie eine Luxuslimousine, die viel Benzin verbraucht und eine riesige Garage benötigt.

Dann erfanden Ingenieure State-Space-Modelle (SSMs). Stellen Sie sich diese wie elektrische Tretroller vor. Sie sind viel schneller und verbrauchen deutlich weniger Energie (Speicher), was sie perfekt für lange Strecken macht. Allerdings hatten frühe Roller ein Problem: Sie waren zu einfach. Sie konnten gerade Straßen bewältigen, aber wenn der Pfad kurvig wurde oder komplexe Logik erforderte (wie das Merken einer spezifischen Regel für einen bestimmten Charakter), gerieten sie in die Irre.

Die Arbeit stellt eine neue Erfindung namens FLASH PD-SSM vor. Es ist so, als würde man diesen elektrischen Roller mit einem intelligenten, modularen Navigationssystem aufrüsten, das ihn genauso gut im Umgang mit komplexen Kurven macht wie die schwere Limousine, ohne dabei seine Geschwindigkeit oder Kraftstoffeffizienz zu verlieren.

Hier ist die Erklärung, wie sie es geschafft haben, aufgeteilt in einfache Konzepte:

1. Das Problem: Die „zu schwere" Karte

Vorherige Versionen dieser intelligenten Roller (wie das Modell namens PD-SSM) versuchten, superintelligent zu sein, indem sie für jeden einzelnen Schritt der Reise eine riesige, detaillierte Karte mitführten.

Das Problem: Das Mitführen dieser riesigen Karte für eine lange Reise nahm so viel Speicherplatz in Anspruch, dass der Roller die Batterie verlor, bevor er weit kam. Es war zu schwer, um für den praktischen Einsatz geeignet zu sein.
Das Ergebnis: Andere Modelle (wie Mamba) entschieden sich dafür, eine winzige, einfache Karte mitzuführen. Sie waren schnell und leicht, konnten aber keine komplexen Rätsel lösen oder komplizierte Regeln merken.

2. Die Lösung: Der „magische Schalter"

Die Autoren dieser Arbeit, FLASH PD-SSM, kamen auf einen klugen Trick. Anstatt für jeden Schritt eine riesige, detaillierte Karte mitzuführen, bauten sie ein Werkzeugkasten mit vorgefertigten, spezialisierten Karten.

Funktionsweise: Bei jedem Schritt der Reise betrachtet der Roboter die aktuelle Situation und betätigt einen Schalter, um die eine perfekte Karte aus dem Werkzeugkasten auszuwählen, die genau für diesen Moment passt.
Die Analogie: Stellen Sie sich vor, Sie fahren Auto. Anstatt jedes Mal, wenn Sie eine Ecke nehmen, eine neue, detaillierte Karte der gesamten Stadt zu zeichnen (was ewig dauert und viel Papier verbraucht), haben Sie einen Satz von 100 vorgezeichneten „lokalen Karten". Sie wählen einfach die aus, die Sie für den nächsten Block benötigen.
Der Vorteil: Dieser Schalter ist unglaublich schnell zu betätigen und nimmt fast keinen Platz ein. Dies ermöglicht dem Roboter die Komplexität der schweren Limousine (er kann schwierige Logikrätsel lösen), behält aber die Geschwindigkeit und Leichtigkeit des Rollers.

3. Was sie bewiesen

Das Team testete diesen neuen „intelligenten Roller" auf drei Hauptwegen:

Der Logiktest (FSA-Emulation): Sie gaben dem Roboter eine Reihe von Logikrätseln (wie das Zählen von Parität oder das Navigieren durch ein Labyrinth). FLASH PD-SSM löste fast alle davon (96 % Genauigkeit), während einfachere Modelle (wie Mamba2) erhebliche Schwierigkeiten hatten. Es bewies, dass der Roboter tatsächlich durch komplexe Regeln „denken" konnte.
Der Langzeitgedächtnistest (Zeitreihen): Sie baten den Roboter, Datenströme zu analysieren, die unglaublich lang waren (über 17.000 Schritte). FLASH PD-SSM war am genauesten darin, vorherzusagen, was als Nächstes passieren würde, und schlug alle anderen Konkurrenten.
Der Sprachtest: Sie setzten dieses neue Gehirn in ein Sprachmodell ein (einen Roboter, der Text schreibt).
- Zustandsverfolgung: Wenn gebeten wurde, Objekte zu verfolgen, die in einer Geschichte zwischen Boxen bewegt wurden (z. B. „Der rote Ball bewegte sich von Box A nach Box B"), war das neue Modell viel besser darin, den endgültigen Standort zu merken als frühere Modelle.
- Schreiben: Wenn es trainiert wurde, wie ein Mensch zu schreiben, schrieb eine hybride Version dieses Modells (eine Mischung aus dem neuen Gehirn und bestehender Technologie) besser und schneller als Modelle, die nur die älteren, einfacheren Gehirne verwendeten.

4. Das Fazit

Die Arbeit behauptet, dass FLASH PD-SSM den größten Kompromiss in der KI derzeit löst: Geschwindigkeit vs. Intelligenz.

Früher: Man musste sich entscheiden zwischen schnell, aber „dumm" (einfache Modelle) oder „intelligent", aber langsam und teuer (komplexe Modelle).
Jetzt: FLASH PD-SSM ist wie ein Hochgeschwindigkeitszug, der auch Berge erklimmen kann. Er läuft so schnell wie die aktuellen Top-Modelle (Mamba2), kann aber viel komplexere Aufgaben bewältigen, und das alles bei geringerem Speicherverbrauch.

Kurz gesagt, sie fanden einen Weg, KI-Modelle zu entwickeln, die leichter, schneller und intelligenter sind, und das alles gleichzeitig, ohne einen größeren, teureren Computer zum Ausführen benötigen.

Technische Zusammenfassung: FLASH PD-SSM

Problemstellung

Zustandsraummodelle (SSMs) bieten eine skalierbare Alternative zu Transformern und ermöglichen eine lineare Skalierung der Rechenleistung sowie einen konstanten Speicherbedarf während der Inferenz. Sie stehen jedoch vor einem fundamentalen Zielkonflikt zwischen Effizienz und Ausdrucksstärke, der durch die Struktur ihrer Übergangsmatrizen bestimmt wird.

Unstrukturierte Matrizen bieten maximale Ausdrucksstärke (fähig zur Emulation beliebiger endlicher Automaten, FSA), verursachen jedoch prohibitiv hohe Rechen- und Speicherkosten.
Strukturierte Matrizen (z. B. diagonal, skalierte Identität) sind hoch effizient, leiden jedoch unter eingeschränkter Ausdrucksstärke und sind theoretisch nicht in der Lage, FSAs mit nicht auflösbaren Transformationshalbgruppen zu emulieren.
Strukturierte Sparse-Matrizen (PD-SSM): Ein neuer Fortschritt (PD-SSM) nutzt strukturierte Sparse-Matrizen, um eine Ausdrucksstärke auf FSA-Niveau zu erreichen und gleichzeitig eine effiziente parallele Berechnung zu ermöglichen. PD-SSM leidet jedoch an einem kritischen Engpass: Die Generierung der inputabhängigen Sparse-Matrizen verursacht Speicherkosten von $O(LN^2)$ (wobei $L$ die Sequenzlänge und $N$ die Größe des versteckten Zustands ist). Dieser Overhead macht PD-SSM deutlich langsamer (40 $\times$ ) und speicherintensiver als effiziente Konkurrenten wie Mamba2, was es für den großskaligen Einsatz ungeeignet macht.

Methodik: FLASH PD-SSM

Die Autoren schlagen FLASH PD-SSM vor, einen speicheroptimierten strukturierten Sparse-SSM, der den Generierungsengpass von PD-SSM löst, während die theoretischen Garantien für die Ausdrucksstärke erhalten bleiben.

1. Neuer Selektionsmechanismus (Harte Selektion)

Die Kerninnovation ersetzt den Soft-Selektionsmechanismus von PD-SSM durch einen harten Selektionsmechanismus:

Original PD-SSM: Generiert eine konvexe Kombination dichter Wörterbuchmatrizen und wendet eine spaltenweise Hardmax-Funktion an, um das Ergebnis zu verdünnen. Dies erfordert die Materialisierung dichter $N \times N$ -Matrizen für jeden Zeitschritt, was zu einem Speicherbedarf von $O(LN^2)$ führt.
FLASH PD-SSM: Berechnet ein Wörterbuch aus $K$ Sparse-Matrizen vor. Zu jedem Zeitschritt generiert der Input $u_t$ über eine lineare Transformation $K$ Scores. Das Modell wählt die Matrix mit dem höchsten Score (argmax) aus und verwendet sie direkt.
Komplexitätsreduktion: Da die ausgewählte Matrix bereits sparse ist (ein One-Hot-Spaltenselektor), sinkt der Speicherbedarf von $O(LN^2)$ auf $O(LN)$. Die Rechenkosten für die Matrixgenerierung reduzieren sich von $O(LN^2K)$ auf $O(LNK)$.

2. Gradientenpropagation

Da die Selektion diskrete Operationen (argmax/hardmax) beinhaltet, verwendet das Modell Straight-Through Estimatoren (STE) mit Slope-Annealing für die Rückwärtspropagation:

Im Vorwärtspass wird eine harte Selektion (argmax) verwendet.
Im Rückwärtspass wird die harte Selektion durch einen temperierten Softmax angenähert, um den Gradientenfluss durch die Selektionsgewichte und die Wörterbuchmatrizen zu ermöglichen.
Proposition 2 in der Arbeit liefert die Ausdrücke für die Surrogatgradienten, die exakt werden, wenn der Temperaturparameter $\tau \to 0$ .

3. Hardwarebewusste Implementierung

Die Autoren implementieren einen benutzerdefinierten Drei-Phasen-CUDA-Kernel (in Triton gekapselt), um die GPU-Parallelität auszunutzen:

Phase A: Berechnet unabhängige sequenzielle Rekursionen innerhalb von Chunks der Sequenz.
Phase B: Aggregiert inter-Chunk-Abhängigkeiten (Prefix-Scan), um Carry-Zustände zu berechnen.
Phase C: Wendet die Carry-Korrekturen auf die Chunk-Zustände an.
Dieses Design minimiert den Speichertraffic und maximiert die Bandbreitennutzung, wodurch eine nahezu maximale GPU-Speicherbandbreite erreicht wird.

4. Architekturintegration

FLASH PD-SSM ist in eine Standard-Pre-Norm-Architektur mit Multi-Head-SSM-Schichten eingebettet. Das Multi-Head-Design ermöglicht die Amortisierung der Parameterkosten für die dichten Wörterbuchmatrizen. Die Autoren zeigen, dass eine Skalierung der Wörterbuchgröße $K$ als $O(\sqrt{D})$ (wobei $D$ die Embedding-Dimension ist) die Anzahl der Parameter handhabbar hält, ohne die Modellgröße „explodieren" zu lassen.

Hauptbeiträge

Algorithmische Verbesserung: Reduziert die Komplexität der Generierung strukturierter Sparse-SSMs von $O(LN^2)$ auf $O(LN)$ durch einen diskreten Selektionsmechanismus und ermöglicht so das Training im großen Maßstab.
Theoretische Ausdrucksstärke: Beweist (Proposition 1), dass FLASH PD-SSM die theoretische Garantie unstrukturierter SSMs beibehält: Es kann jeden deterministischen FSA mit $N$ Zuständen exakt durch eine einzelne Schicht der Größe $N$ darstellen.
Hardware-Optimierung: Liefert eine benutzerdefinierte CUDA-Implementierung, die einen State-of-the-Art-Durchsatz und eine hohe Speichereffizienz erreicht.
Hybridmodell-Integration: Zeigt, dass FLASH PD-SSM als effektiver Drop-in-Ersatz in hybriden LLMs dient und die Fähigkeiten zur Zustandsverfolgung verbessert, ohne den Durchsatz zu beeinträchtigen.

Experimentelle Ergebnisse

Synthetische Zustandsverfolgung und Ausdrucksstärke

FSA-Emulation: Bei synthetischen Aufgaben (Parität, modulares Zählen, Zyklusnavigation) erreicht FLASH PD-SSM eine durchschnittliche Genauigkeit von 96,3 % und übertrifft damit deutlich diagonale SSMs (Mamba2: 58,9 %, Mamba3: 90,7 %) sowie DPLR-Modelle (Gated DeltaProduct: 80,7 %). Dies bestätigt die praktische Realisierung seiner theoretischen Ausdrucksstärke.
Mechanistische Fähigkeiten (MAD-Benchmark): FLASH PD-SSM erzielt 69,2 %, vergleichbar mit Mamba (69,3 %), was auf starke Fähigkeiten im kontextuellen Abruf und der Memorierung hinweist.

Laufzeit- und Speichereffizienz

Durchsatz: Bei einer Sequenzlänge von 5120 ist FLASH PD-SSM 18 % schneller als DeltaNet und 4 % schneller als Mamba2.
Speicher: Es reduziert den Spitzen-Speicherbedarf um >12 % im Vergleich zu DeltaNet und Mamba2 bei Sequenzlängen über 4000.
Skalierbarkeit: Das Modell behält über verschiedene Sequenzlängen hinweg einen hohen Durchsatz bei und adressiert damit den Engpass, der PD-SSM zuvor behinderte.

Natürliche Sprach-Zustandsverfolgung

Evaluierung auf dem Box-Datensatz (erweitert auf 26 Zustandsübergänge) unter Verwendung hybrider Modelle (Granite4, Nemotron-H).
In diese Modelle eingefügte FLASH PD-SSM-Schichten verbesserten die Genauigkeit außerhalb der Verteilung (OOD) um bis zu 10 % gegenüber Mamba2 und um 8 % gegenüber DeltaNet, bei gleichzeitig vergleichbarem Durchsatz.

Sprachmodellierung

DELTA-PD: Ein Hybridmodell, das FLASH PD-SSM und DeltaNet kombiniert (Verhältnis 1:3), wurde auf FineWeb100B vortrainiert.
Leistung: DELTA-PD übertraf das reine DeltaNet sowohl in der Genauigkeit (42,3 % vs. 41,3 % im Durchschnitt) als auch im Durchsatz (253k tok/s vs. 168k tok/s, eine 1,5 $\times$ -Verbesserung).
Es übertraf auch das reine FLASH PD-SSM, was darauf hindeutet, dass die Hybridisierung die Stärken verschiedener SSM-Typen ausbalanciert.

Zeitreihenmodellierung

Erzielte State-of-the-Art (SoTA)-Genauigkeit bei der multivariaten Zeitreihenklassifizierungsaufgabe UEA (Sequenzen >17.000).
Erreichte einen nahezu SoTA-Mittleren Quadratischen Fehler (MSE) bei der Prognoseaufgabe PPG-DaLiA.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass FLASH PD-SSM erfolgreich den Zielkonflikt zwischen Effizienz und Ausdrucksstärke durchbricht, der SSMs bisher limitiert hat. Durch die Reduzierung des Speicher-Overheads strukturierter Sparse-Matrizen ermöglicht es den Einsatz hoch ausdrucksstarker Modelle (fähig zur FSA-Emulation) in dem Maßstab, der für moderne Sprachmodelle erforderlich ist.

Die Autoren positionieren FLASH PD-SSM nicht nur als theoretische Verbesserung, sondern als praktischen, Drop-in-Ersatz für bestehende SSM-Schichten in fortschrittlichen Hybridmodellen. Es nimmt eine einzigartige Position im Leistungslandschaft ein: Es ist das einzige SSM, das im „oberen rechten Quadranten" von Abbildung 1 identifiziert wird, was bedeutet, dass es gleichzeitig hoch ausdrucksstark und schnell/speicher-effizient ist.

Die Arbeit erkennt Einschränkungen bezüglich der Verwendung von Straight-Through-Estimatoren an und stellt fest, dass die theoretische Erklärung für deren Wirksamkeit bei hochdiskreten sequenziellen Methoden ein offenes Forschungsfeld bleibt. Die empirischen Ergebnisse über synthetische, mechanistische und natürliche Sprachaufgaben validieren den Ansatz jedoch als robuste Lösung für skalierbare, ausdrucksstarke Sequenzmodellierung.