Hochmoderne Techniken, die KI ohne Mehrkosten beschleunigen

Linktabelle

Zusammenfassung und 1 Einleitung

2 Zustandsraummodelle

3 Selektive Zustandsraummodelle und 3.1 Motivation: Selektion als Mittel zur Kompression

3.2 Verbesserung der SSMs durch Selektion

3.3 Effiziente Umsetzung selektiver SSM

3.4 Eine vereinfachte SSM-Architektur

3.5 Eigenschaften von Selektionsmechanismen

3.6 Weitere Modelldetails

4 Empirische Auswertung und 4.1 Syntheseaufgaben

4.2 Sprachmodellierung

4.3 DNA-Modellierung

4.4 Audiomodellierung und -generierung

4.5 Geschwindigkeits- und Speicherbenchmarks

4.6 Modellablationen

5 Diskussion

6 Schlussfolgerung und Referenzen

Eine Diskussion: Auswahlmechanismus

B Verwandte Arbeiten

C Mechanik selektiver SSMs

D Hardwarebewusster Algorithmus für selektive SSMs

E. Experimentelle Details und zusätzliche Ergebnisse

3.3 Effiziente Umsetzung selektiver SSM

Hardwarefreundliche Architekturen wie Faltungen (Krizhevsky, Sutskever und Hinton 2012) und Transformer (Vaswani et al. 2017) finden weit verbreitete Anwendung. Hier zielen wir darauf ab, selektive SSMs auch auf moderner Hardware (GPU) effizient zu machen. Der Auswahlmechanismus ist recht natürlich, und frühere Arbeiten versuchten, spezielle Fälle der Auswahl einzubeziehen, wie z. B. die zeitliche Variation von ∆ in wiederkehrenden SSMs (Gu, Dao et al. 2020). Wie bereits erwähnt, ist jedoch ihre Rechenleistung eine wesentliche Einschränkung bei der Verwendung von SSMs, weshalb S4 und alle Derivate LTI-Modelle (nicht selektive Modelle) verwendeten, am häufigsten in Form globaler Faltungen.

3.3.1 Motivation früherer Modelle

Wir gehen zunächst noch einmal auf diese Motivation ein und geben einen Überblick über unseren Ansatz zur Überwindung der Einschränkungen früherer Methoden.

• Auf hohem Niveau balancieren rekurrierende Modelle wie SSMs immer einen Kompromiss zwischen Ausdrucksstärke und Geschwindigkeit: Wie in Abschnitt 3.1 erläutert, sollten Modelle mit größerer verborgener Zustandsdimension effektiver, aber langsamer sein. Daher möchten wir die verborgene Zustandsdimension maximieren, ohne Geschwindigkeits- und Speicherkosten in Kauf zu nehmen.

• Beachten Sie, dass der rekurrierende Modus flexibler ist als der Faltungsmodus, da letzterer (3) aus der Erweiterung des ersteren (2) abgeleitet wird (Gu, Goel und Ré 2022; Gu, Johnson, Goel et al. 2021). Dies würde jedoch die Berechnung und Materialisierung des latenten Zustands ℎ mit der Form (B, L, D, N) erfordern, der viel größer ist (um einen Faktor N, die SSM-Zustandsdimension) als der Eingang x und der Ausgang y der Form (B, L, D). Daher wurde der effizientere Faltungsmodus eingeführt, der die Zustandsberechnung umgehen und einen Faltungskern (3a) von nur (B, L, D) materialisieren kann.

• Frühere LTI-SSMs nutzen die dualen rekurrierenden-konvolutionalen Formen, um die effektive Zustandsdimension um den Faktor Nx (≈ 10 − 100) zu erhöhen, was deutlich mehr ist als bei herkömmlichen RNNs, ohne dass es zu Effizienzeinbußen kommt.

3.3.2 Übersicht über den selektiven Scan: Hardwarebewusste Statuserweiterung

Der Auswahlmechanismus soll die Einschränkungen von LTI-Modellen überwinden. Gleichzeitig müssen wir uns daher erneut mit dem Berechnungsproblem von SSMs befassen. Wir gehen dies mit drei klassischen Techniken an: Kernelfusion, paralleler Scan und Neuberechnung. Wir machen zwei Hauptbeobachtungen:

• Die naive rekurrente Berechnung verwendet O(BLDN) FLOPs, während die Faltungsberechnung O(BLD log(L)) FLOPs verwendet, und erstere hat einen niedrigeren konstanten Faktor. Daher kann der rekurrente Modus für lange Sequenzen und eine nicht zu große Zustandsdimension N tatsächlich weniger FLOPs verwenden.

• Die beiden Herausforderungen sind die sequentielle Natur der Rekurrenz und der hohe Speicherbedarf. Um Letzteres anzugehen, können wir, genau wie beim Faltungsmodus, versuchen, den vollständigen Zustand ℎ nicht wirklich zu materialisieren.

Die Grundidee besteht darin, die Eigenschaften moderner Beschleuniger (GPUs) zu nutzen, um den Zustand ℎ nur in effizienteren Ebenen der Speicherhierarchie zu materialisieren. Insbesondere sind die meisten Operationen (außer der Matrixmultiplikation) durch die Speicherbandbreite begrenzt (Dao, Fu, Ermon et al. 2022; Ivanov et al. 2021; Williams, Waterman und Patterson 2009). Dies schließt unsere Scan-Operation ein, und wir verwenden Kernelfusion, um die Anzahl der Speicher-IOs zu reduzieren, was zu einer deutlichen Beschleunigung im Vergleich zu einer Standardimplementierung führt.

Um die sequentielle Rekurrenz zu vermeiden, stellen wir fest, dass es trotz der fehlenden Linearität immer noch mit einem arbeitseffizienten Parallel-Scan-Algorithmus parallelisiert werden kann (Blelloch 1990; Martin und Cundy 2018; Smith, Warrington und Linderman 2023).

Schließlich müssen wir auch das Speichern der Zwischenzustände vermeiden, die für die Backpropagation erforderlich sind. Wir wenden sorgfältig die klassische Technik der Neuberechnung an, um den Speicherbedarf zu reduzieren: Die Zwischenzustände werden nicht gespeichert, sondern im Rückwärtsdurchlauf neu berechnet, wenn die Eingänge vom HBM in den SRAM geladen werden. Infolgedessen hat die fusionierte selektive Scan-Schicht den gleichen Speicherbedarf wie eine optimierte Transformer-Implementierung mit FlashAttention.

Einzelheiten zum fusionierten Kernel und zur Neuberechnung finden Sie in Anhang D. Die vollständige selektive SSM-Schicht und der Algorithmus sind in Abbildung 1 dargestellt.