Rezension von Transformer mit einer ausführlichen Erklärung der Grundsätze der alternativen Mamba (illustriert)

Die Transformer-Architektur ist eine Schlüsselkomponente für den Erfolg von großen Sprachmodellen (LLMs). Fast alle großen Sprachmodelle, die heute verwendet werden, nutzen diese Architektur, von Open-Source-Modellen wie Mistral zu Closed-Source-Modellen wie ChatGPT.

Um das große Sprachmodell weiter zu verbessern, wurden neue Architekturen entwickelt, die sogar über die Transformer-Architektur hinausgehen können. Ein solcher Ansatz ist die Mambaeine Art von Zustandsraummodell.

Mamba in der Zeitung Mamba: Linear-Time Sequence Modelling mit selektiven Zustandsräumen1 in der sie vorgestellt wurde. Mehr darüber erfahren Sie in der Repository Die offizielle Umsetzung und die Musterprüfpunkte finden Sie in.

In diesem Beitrag führe ich in das Gebiet der Zustandsraummodellierung im Kontext der Sprachmodellierung ein und erkunde Schritt für Schritt verschiedene Konzepte, die zum Verständnis dieses Gebiets beitragen. Anschließend werden wir erörtern, wie Mamba die Transformer-Architektur herausfordern könnte.

Als visueller Leitfaden werden in diesem Artikel zahlreiche Visualisierungen vorgestellt, die das Verständnis von Mamba und Zustandsraummodellen erleichtern!

Um zu verdeutlichen, was für eine interessante Architektur Mamba ist, wollen wir zunächst kurz auf Transformers zurückblicken und eine seiner Schwächen untersuchen.

Der Transformer behandelt jede Texteingabe so, als wäre sie eine Texteingabe aus dem Token konstituierend Reihenfolge.

Einer der Hauptvorteile von Transformers ist, dass es unabhängig von der Eingabe, die es erhält, zu allen früheren Token in der Sequenz zurückgehen kann, um seine Darstellung abzuleiten.

Denken Sie daran, dass der Transformer aus zwei Strukturen besteht, einem Satz von Encoder-Blöcken zur Darstellung von Text und einem Satz von Decoder-Blöcken zur Erzeugung von Text. Kombiniert können diese Strukturen für mehrere Aufgaben, einschließlich der Übersetzung, verwendet werden.

Wir können diese Struktur nutzen und generative Modelle erstellen, indem wir nur Decoder verwenden. Dieses Transformator-basierte Modell, dasGenerative Pre-training Transformer(GPT), das Decoderblöcke verwendet, um einen Teil des eingegebenen Textes zu vervollständigen.

Mal sehen, wie das funktioniert!

Ein einzelner Decoder-Block besteht aus zwei Hauptkomponenten, dem maskierten Selbstaufmerksamkeitsmechanismus und dem neuronalen Feedforward-Netzwerk.

Der Mechanismus der Selbstbeobachtung ist ein wichtiger Grund, warum diese Modelle so gut funktionieren. Er ermöglicht einen unkomprimierten Blick auf die gesamte Sequenz und ist schnell zu trainieren.

Wie funktioniert es also?

Es wird eine Matrix erstellt, in der jede Token Vergleiche mit jedem vorherigen Token. Die Gewichtung in der Matrix hängt von der Relevanz der Tokenpaare zueinander ab.

Diese Matrix wird während der Ausbildung auf einmal erstellt. "Meine"und"Name"Die Aufmerksamkeit zwischen muss nicht berechnet werden, bevor sie berechnet werden kann".Name"und"ist"Achtung zwischen.

Sie realisiertParallelisierungwas die Ausbildung erheblich beschleunigt!

Allerdings gibt es einen Nachteil. Bei der Erzeugung des nächsten Tokens müssen wir diegesamte SequenzAufmerksamkeit, auch wenn wir bereits einige Token generiert haben.

Erzeugen Sie die Länge derLDie Sequenz erfordert etwaL²mal, was bei zunehmender Länge der Sequenz sehr rechenintensiv sein kann.

Diese Notwendigkeit, die gesamte Sequenz neu zu berechnen, ist einer der größten Engpässe der Transformer-Architektur.

Sehen wir uns an, wie eine "klassische" Technik, das rekurrente neuronale Netz (RNN), dieses langsame Inferenzproblem löst.

Das rekurrente neuronale Netz (RNN) ist ein sequenzbasiertes Netz. Es erhält in jedem Zeitschritt zwei Eingaben, d. h. ZeitschritttEingabe und dem vorherigen Zeitschrittt-1des verborgenen Zustands wird verwendet, um den nächsten verborgenen Zustand zu erzeugen und die Ausgabe vorherzusagen.

Das RNN verfügt über einen Schleifenmechanismus, der es ihm ermöglicht, Informationen von einem Schritt zum nächsten weiterzugeben. Wir können diesen Visualisierungsprozess "erweitern", um ihn deutlicher zu machen.

Bei der Erzeugung der Ausgabe muss das RNN nur den vorherigen versteckten Zustand und die aktuelle Eingabe berücksichtigen. Es vermeidet das Problem der Neuberechnung aller vorherigen versteckten Zustände, die der Transformer benötigt.

Mit anderen Worten: Ein RNN ist zu schnellen Schlussfolgerungen fähig, weil es linear mit der Sequenzlänge skaliert! Theoretisch könnte es sogarUnbegrenzte Kontextlänge.

Zur Veranschaulichung wenden wir das RNN auf den Eingabetext an, den wir zuvor verwendet haben.

Jeder verborgene Zustand ist eine Zusammenfassung aller vorherigen verborgenen Zustände, normalerweise eine komprimierte Ansicht.

Allerdings gibt es hier ein Problem ......

Beachten Sie, dass bei der Generierung des Namens "Maarten", wenn der letzte verborgene Zustand keine Informationen mehr über das Wort "Hallo"RNNs neigen dazu, Informationen im Laufe der Zeit zu vergessen, weil sie nur den letzten Zustand berücksichtigen.

RNNs sind zwar schneller in Bezug auf Training und Inferenz, aber ihnen fehlt die Genauigkeit, die Transformer-Modelle bieten können.

Daher untersuchen wir Zustandsraummodelle, um RNNs (und manchmal Faltung) effizient zu nutzen.

Zustandsraummodelle (State Space Models, SSMs) wie Transformers und RNNs befassen sich mit Sequenzen von Informationen, wie Text und Signale. In diesem Abschnitt stellen wir die grundlegenden Konzepte von SSMs vor und zeigen, wie sie sich auf Textdaten beziehen.

Ein Zustandsraum enthält die Mindestanzahl von Variablen, die zur vollständigen Beschreibung eines Systems erforderlich sind. Er ist eine Möglichkeit, ein Problem mathematisch darzustellen, indem die möglichen Zustände eines Systems definiert werden.

Vereinfachen wir das Ganze. Stellen wir uns vor, wir reisen durch ein Labyrinth. "Zustandsraum"Es ist eine Karte mit allen möglichen Orten (Zuständen). Jeder Punkt steht für einen bestimmten Ort im Labyrinth mit spezifischen Details, z. B. wie weit man vom Ausgang entfernt ist.

"Zustandsraumdarstellung" ist eine vereinfachte Beschreibung dieser Karte. Sie zeigt Ihren aktuellen Standort (aktueller Zustand), den nächsten Ort, zu dem Sie reisen können (möglicher zukünftiger Zustand), und wie Sie zum nächsten Ort gelangen können (Bewegung nach rechts oder links).

Das Zustandsraummodell verwendet zwar Gleichungen und Matrizen, um dieses Verhalten nachzuvollziehen, aber es ist eigentlich nur eine Methode, um festzustellen, wo man ist, wohin man gehen kann und wie man dorthin kommt.

Die Variablen beschreiben einen Zustand, in unserem Fall die X- und Y-Koordinaten, und die Entfernung zum Ausgang, die wie folgt ausgedrückt werden kann "Zustandsvektor".

Kommt Ihnen das bekannt vor? Das liegt daran, dass Einbettungen oder Vektoren auch in Sprachmodellen häufig verwendet werden, um den "Zustand" einer Eingabesequenz zu beschreiben. Ein Vektor, der Ihre aktuelle Position beschreibt (ein Zustandsvektor), könnte zum Beispiel so aussehen:

In neuronalen Netzen bezieht sich der Begriff "Zustand" in der Regel auf den verborgenen Zustand des Systems, der einer der wichtigsten Aspekte bei der Erzeugung neuer Token in einem großen Sprachmodell ist.

SSM (State Space Models) sind eine Klasse von Modellen, die zur Beschreibung dieser Zustandsdarstellungen und zur Vorhersage ihres nächsten Zustands verwendet werden, wobei die Vorhersagen auf bestimmten Eingaben basieren.

Traditionell, mit der Zeit tSSM:

Anstatt jedoch diskrete Sequenzen zu verwenden (z. B. einmaliges Verschieben nach links), akzeptiert SSM kontinuierliche Sequenzen als Eingabe und sagt die Ausgabesequenz voraus.

SSM geht davon aus, dass dynamische Systeme (z. B. Objekte, die sich im 3D-Raum bewegen) zeitlich t des Zustands und zwei Gleichungen zu dessen Vorhersage.

Wir gehen davon aus, dass wir durch die Lösung dieser Gleichungen statistische Prinzipien für die Vorhersage des Systemzustands auf der Grundlage der beobachteten Daten (Eingangssequenzen und vorherige Zustände) aufdecken können.

Ziel ist es, diese Zustandsdarstellung zu finden h(t)die es uns ermöglicht, von Eingabe- zu Ausgabesequenzen zu gelangen.

Diese beiden Gleichungen sind das Herzstück des Zustandsraummodells.

Auf diese beiden Gleichungen wird im gesamten Leitfaden Bezug genommen. Um sie intuitiver zu machen, werden dieWir verwenden Farbkodierungso dass Sie sie schnell zitieren können.

Zustandsgleichung beschreibt, wie der Zustand in Abhängigkeit von den Eingaben geändert wird (über die Matrix B) Auswirkungen auf den Zustand (über Matrix A) und Veränderung.

Wie bereits erwähnt.h(t) bezieht sich auf einen bestimmten Zeitpunkt t Die potenzielle Zustandsdarstellung desx(t) bezieht sich auf eine bestimmte Eingabe.

Ausgangsgleichung beschreibt, wie der Zustand durch die Matrix C umgewandelt werden, und wie die Eingaben in Ausgaben umgewandelt werden Matrix D Beeinflusst die Ausgabe.

Wenn man diese beiden Gleichungen visualisiert, erhält man die folgende Architektur:

Schauen wir uns nun Schritt für Schritt an, wie diese Matrizen den Lernprozess beeinflussen.

Angenommen, wir haben einige Eingangssignale x(t)Das Signal wird zunächst mit dem Matrix B Multiplizieren.Matrix B Beschreibt, wie sich Eingaben auf das System auswirken.

Der aktualisierte Zustand (ähnlich wie der verborgene Zustand eines neuronalen Netzes) ist ein Potenzialraum, der das zentrale "Wissen" über die Umgebung enthält. Wir beziehen den Zustand auf die Matrix A Miteinander multipliziert beschreibt diese Matrix die Verbindungen zwischen allen internen Zuständen und stellt die zugrunde liegende Dynamik des Systems dar.

Wie Sie vielleicht schon bemerkt haben.Matrix A Anwendung vor der Erstellung der Zustandsdarstellung und erneut nach der Aktualisierung der Zustandsdarstellung.

Dann verwenden wir die Matrix C um zu beschreiben, wie Zustände in Ausgaben umgewandelt werden.

Schließlich können wir die Matrix D Liefert ein direktes Signal vom Eingang zum Ausgang. Dies wird auch allgemein bezeichnet als Sprungverbindung.

aufgrund von Matrix D Ähnlich wie bei Sprungverbindungen wird SSM oft als die folgende Form betrachtet, die keine Sprungverbindungen enthält.

Um zu unserer vereinfachten Perspektive zurückzukehren, können wir uns nun auf die Matrix A, undB im Gesang antworten C als Kernstück der SSM.

Wir können die ursprünglichen Gleichungen aktualisieren (und einige schöne Farben hinzufügen), um zu markieren, wofür jede Matrix verwendet wird, wie wir es zuvor getan haben.

Zusammen sollen diese beiden Gleichungssätze den Zustand des Systems aus den beobachteten Daten vorhersagen. Da die Eingaben kontinuierlich sind, lautet die Hauptdarstellung von SSM kontinuierliche Zeitskala.

Wenn Sie ein kontinuierliches Signal haben, finden Sie die Zustandsdarstellung h(t) schwierig zu parsen. Da wir in der Regel diskrete Eingaben haben (z. B. Textsequenzen), möchten wir das Modell außerdem diskretisieren.

Zu diesem Zweck verwenden wir die Null-Ordnung-Halteverfahren. Es funktioniert folgendermaßen: Wenn wir ein diskretes Signal erhalten, halten wir zunächst seinen Wert, bis wir ein neues diskretes Signal erhalten. Dieser Prozess erzeugt ein kontinuierliches Signal für SSM:

Wir behalten den Wert der Zeit durch einen erlernbaren neuen Parameter Herzschrittmacher ∆ Zeigt an. Sie gibt die Auflösung des Eingangs an.

Nachdem wir nun ein kontinuierliches Signal für den Eingang erzeugt haben, können wir als nächstes einen kontinuierlichen Ausgang erzeugen und diese Werte entsprechend dem Zeitschritt des Eingangs abtasten.

Diese abgetasteten Werte sind das Ergebnis unserer Diskretisierung!

Mathematisch gesehen können wir die Null-Ordnung wie folgt anwenden:

Zusammen ermöglichen sie den Übergang von einer kontinuierlichen SSM zu einer diskreten SSM. an diesem Punkt ist das Modell nicht mehr Funktion zu Funktion x(t) → y(t)Stattdessen ist es Sequenz zu Sequenz xₖ → y_ₖ::

Hier wird die Matrix A im Gesang antworten B Bezeichnen Sie nun die diskretisierten Parameter des Modells.

Wir verwenden k anstelle von tzu unterscheiden, wann es sich um kontinuierliche SSM und wann um diskrete SSM handelt.

Da wir nun eine Formel für die diskretisierte Darstellung haben, wollen wir nun untersuchen, wie man tatsächlich zählen Das Modell.

Unser diskretisiertes SSM ermöglicht es uns, das Problem in einem bestimmten Zeitschritt und nicht in einem kontinuierlichen Signal zu konstruieren. Wie wir bereits bei RNNs gesehen haben, sind rekursive Methoden hier sehr nützlich.

Wenn wir diskrete Zeitschritte anstelle von kontinuierlichen Signalen betrachten, können wir das Problem mit Hilfe von Zeitschritten neu formulieren:

In jedem Zeitschritt berechnen wir die aktuelle Eingabe (Bxₖ), wie es sich auf den vorherigen Zustand auswirkt (Ahₖ₋₁) und berechnen dann die vorhergesagte Leistung (Chₖ).

Diese Darstellung ist Ihnen vielleicht schon bekannt! Wir können sie analysieren, wie wir es bei der vorherigen Behandlung von RNNs getan haben.

Wir können wie folgt expandieren (oder in eine Reihe von Zeitschritten erweitern):

Es sei darauf hingewiesen, dass wir eine diskretisierte Version davon verwenden können, indem wir den grundlegenden Ansatz eines RNN verwenden.

Zur Darstellung von SSM können wir die Faltung verwenden. Erinnern Sie sich daran, dass wir bei klassischen Bilderkennungsaufgaben Filter anwenden (Kerne), um aggregierte Merkmale zu extrahieren:

Da wir es mit Text und nicht mit Bildern zu tun haben, müssen wir eine eindimensionale Ansicht verwenden:

Der Kernel, den wir zur Darstellung dieses "Filters" verwenden, ist von der SSM-Formel abgeleitet:

Schauen wir uns an, was dieser Kernel in der Praxis leistet. Wie bei der Faltung können wir den SSM-Kernel verwenden, um jeden Satz von Token zu durchlaufen und die Ausgabe zu berechnen:

Dies verdeutlicht auch die Auswirkungen, die Auffüllungen auf die Ausgabe haben können. Ich habe die Reihenfolge der Auffüllung geändert, um die Darstellung zu verbessern, aber normalerweise wird die Auffüllung am Ende des Satzes vorgenommen.

Im nächsten Schritt bewegt sich der Kernel einmal, um den nächsten Schritt der Berechnung durchzuführen:

Im letzten Schritt können wir die volle Wirkung des Kernels sehen:

Einer der Hauptvorteile der Darstellung von SSMs als Faltungen ist, dass sie wie Faltungsneuronale Netze (CNNs) parallel trainiert werden können. Aufgrund der festen Kernelgröße ist ihre Inferenz jedoch nicht so schnell und uneingeschränkt wie bei RNNs.

Diese drei Darstellungen -Verlauf, undrekursiv (Berechnung) im Gesang antworten Faltung Beide haben unterschiedliche Vor- und Nachteile:

Interessanterweise können wir nun rekursive SSMs für eine effiziente Inferenz verwenden, während wir für das parallele Training Faltungs-SSMs einsetzen.

Mit diesen Darstellungen können wir einen cleveren Trick anwenden, indem wir je nach Aufgabe eine Darstellung wählen. Beim Training verwenden wir eine Faltungsrepräsentation, die parallelisiert werden kann, während wir bei der Inferenz eine effiziente rekursive Repräsentation verwenden:

Dieses Modell ist bekannt als Lineare Zustandsraumschicht (LSSL). 2

Diese Darstellungen haben eine wichtige Eigenschaft gemeinsam, nämlichlineare Zeitinvariante (math.)（LTI bezeichnet die Parameter von SSM. A, undB im Gesang antworten C ist in allen Zeitschritten fest. Dies bedeutet, dass die Matrix A, undB im Gesang antworten C ist für jedes erzeugte Token gleich.

Mit anderen Worten, egal welche Reihenfolge Sie SSM geben.A, undB im Gesang antworten C Die Werte aller dieser Daten bleiben gleich. Wir haben eine statische Darstellung, die sich nicht um den Inhalt kümmert.

Bevor wir uns ansehen, wie Mamba dieses Problem löst, wollen wir uns das letzte Teil des Puzzles ansehen - dieMatrix A.

Einer der wichtigsten Aspekte der SSM-Formel ist wohl, dass die Matrix A. Wie wir bereits bei der rekursiven Darstellung gesehen haben, erfasst sie Informationen über die vorherige Informationen über den Zustand zu konstruieren, um die erneuert Status.

Im Wesentlichen.Matrix A Erzeugen Sie einen verborgenen Zustand:

Daher ist die Schaffung Matrix A kann darüber entscheiden, ob wir uns nur an einige wenige frühere Token erinnern oder alle Token erfassen können, die wir bisher gesehen haben. erinnert daran, dass Vorheriger Zustand.

Wie man so erstellt, dass viel Speicherplatz erhalten bleibt (Kontextgröße) Matrix A?

Wir verwenden Hungry Hungry Hippo! oder HiPPO3 zu realisierenIhr (Ehrentitel)Treppe wie (in welchem Umfang)Begriff (math.) sich in die Hände vonFilm Glückoperator.HiPPO versucht, alle Eingangssignale, die es bisher gesehen hat, in einen Vektor von Koeffizienten zu komprimieren.

Sie verwendet Matrix A um eine Zustandsdarstellung zu konstruieren, die neuere Token besser erfasst und ältere Token abklingen lässt, kann die Formel wie folgt ausgedrückt werden:

Angenommen, wir haben ein Quadrat Matrix ADas gibt uns:

Bauen mit HiPPO Matrix A erwies sich als wesentlich besser als die Initialisierung mit einer Zufallsmatrix. Sie ist also besser in der Lage, die Daten zu rekonstruieren. Update Der Aspekt der Signalisierung (nächstgelegener Token) ist wichtiger als ältere Das Signal (Initial Token) ist genauer.

Die Kernidee der HiPPO-Matrix besteht darin, einen verborgenen Zustand zu erzeugen, der sich an seine Geschichte erinnert.

Mathematisch geschieht dies durch die Verfolgung Legendre-Polynom Koeffizienten, um dies zu erreichen, was eine Annäherung an alle historischen Aufzeichnungen ermöglicht.4

HiPPO wird dann auf die rekursiven und gefalteten Darstellungen angewandt, die wir bereits gesehen haben, um weitreichende Abhängigkeiten zu behandeln. Das Ergebnis. Strukturierter Zustandsraum für Sequenzen (S4)Die SSM ist eine Klasse von SSMs, die lange Sequenzen effizient verarbeiten kann.5

Sie besteht aus drei Teilen:

Diese Art von SSM hat je nach gewählter Darstellung (rekursiv vs. faltig) mehrere Vorteile. Sie kann auch lange Textsequenzen verarbeiten und Speicher effizient speichern, da sie auf HiPPO-Matrizen basiert.

Wir haben nun alle Grundlagen behandelt, die notwendig sind, um zu verstehen, was Mamba so besonders macht. Zustandsraummodelle können zur Modellierung von Textsequenzen verwendet werden, aber es gibt immer noch eine Reihe von Nachteilen, die wir gerne vermeiden möchten.

In diesem Abschnitt werden wir die beiden Hauptbeiträge von Mamba erörtern:

Diese beiden bilden zusammen die Selektive Zustandsraummodellierung vielleicht S6 Modell, das wie die Selbstaufmerksamkeit genutzt werden kann, um Mamba-Block.

Bevor wir uns mit diesen beiden Hauptbeiträgen befassen, wollen wir zunächst untersuchen, warum sie notwendig sind.

Zustandsraummodelle, selbst S4 (strukturierte Zustandsraummodelle), sind bei bestimmten Schlüsselaufgaben der Sprachmodellierung und -generierung unzureichend, nämlich Fähigkeit, sich auf bestimmte Eingaben zu konzentrieren oder sie zu ignorieren.

Wir können dies mit zwei synthetischen Aufgaben illustrieren, nämlich selektives Kopieren im Gesang antworten Inducer-Kopf.

existieren selektives Kopieren Das Ziel von SSM in einer Aufgabe ist es, die Eingabeteile zu kopieren und sie in der richtigen Reihenfolge auszugeben:

Da SSM jedoch eine Linear zeitinvariant(in Form eines Nominalausdrucks)dass sie bei dieser Aufgabe schlecht abschneidet. Wie wir bereits gesehen haben, ist die Matrix A, undB im Gesang antworten C Jeder generierte Token ist derselbe für SSM.

Infolgedessen ist SSM nicht in der Lage, die Inhaltsbezogene ArgumentationDas ist ein Problem, weil wir wollen, dass SSM über die Eingaben (Hinweise) nachdenkt. weil es jedes Token aufgrund einer festen Matrix von A, B und C gleich behandelt. Das ist ein Problem, weil wir wollen, dass SSM über die Eingaben (Hinweise) nachdenkt.

SSM schneidet bei einer anderen Aufgabe schlecht ab, nämlich bei der Inducer-Kopfdessen Ziel es ist, die in der Eingabe gefundenen Muster zu reproduzieren:

Im obigen Beispiel führen wir im Wesentlichen eine einmalige Aufforderung durch, bei der wir versuchen, dem Modell "beizubringen", bei jedem "[...]" denselben Hinweis zu verwenden.Q." nach der Bereitstellung eines "A." Antwort. Da SSM jedoch zeitinvariant ist, kann es nicht wählen, welche früheren Token aus der Historie abgerufen werden sollen.

Dazu sollten wir uns auf folgende Punkte konzentrieren Matrix B um diesen Punkt zu veranschaulichen. Unabhängig von der Eingabe x Was es ist.Matrix B ist immer gleich, also mit x Irrelevant:

Gleichermaßen.A im Gesang antworten C bleibt auch immer gleich, unabhängig von der Eingabe. Dies legt nahe, dass das, was wir bisher über die SSM gesehen haben Arbeitsniederlegung Merkmale.

Im Gegensatz dazu sind diese Aufgaben für den Transformer relativ einfach, da sie auf der Eingangssequenz beruhen dynamisch (Wissenschaft) Wechselnde Aufmerksamkeit. Sie können selektiv auf verschiedene Teile der Sequenz "schauen" oder sich darauf "konzentrieren".

Die schlechte Leistung von SSM bei diesen Aufgaben verdeutlicht die potenziellen Probleme der zeitinvarianten SSM, Matrix A, undB im Gesang antworten C Der statische Charakter der Inhalt Sinn des Themas.

Die rekursive Darstellung von SSM erzeugt einen kleineren Zustand, der sehr effizient ist, weil er die gesamte Historie komprimiert. Im Vergleich zum Transformer-Modell komprimiert das Transformer-Modell die Historie jedoch nicht (über die Aufmerksamkeitsmatrix), so dass es leistungsfähiger ist.

Mamba will das Beste aus beiden Welten bieten. Ein kleiner Staat, der so mächtig ist wie ein Transformer-Staat:

Wie bereits erwähnt, geschieht dies durch selektive Komprimierung von Daten in Zustände. Bei einem Eingabesatz gibt es in der Regel einige Informationen, wie z. B. ein Stoppwort, die nicht viel Sinn ergeben.

Um Informationen selektiv komprimieren zu können, müssen die Parameter von der Eingabe abhängig sein. Zu diesem Zweck wollen wir zunächst die Dimensionen der Eingaben und Ausgaben von SSM während des Trainings untersuchen:

In dem strukturierten Zustandsraummodell (S4) wird die Matrix A, undB im Gesang antworten C sind unabhängig von den Inputs, da ihre Dimensionen N im Gesang antworten D ist statisch und ändert sich nicht.

Stattdessen kombiniert Mamba die Länge der Eingabesequenz und die Losgröße, um die Matrix B im Gesang antworten Csogar Herzschrittmacher ∆_, abhängig von der Eingabe:

Das bedeutet, dass wir jetzt für jedes Eingabe-Token verschiedene B im Gesang antworten C Matrix, die das Problem der Inhaltswahrnehmung löst!

zusammen Selektiv Wählen Sie aus, was verborgen bleiben soll und was ignoriert werden soll, da es nun von der Eingabe abhängig ist.

kleiner Herzschrittmacher ∆ dazu führt, dass bestimmte Wörter zugunsten des vorherigen Kontextes ignoriert werden, während größere Herzschrittmacher ∆ Stattdessen liegt der Schwerpunkt eher auf dem Eingangsvokabular als auf dem Kontext:

Da diese Matrizen nun dynamisch (Wissenschaft) s, können sie nicht mit der Faltungsdarstellung berechnet werden, die von einer Fixierung des Faltungskerns. Wir können nur rekursive Darstellungen verwenden, wodurch der Parallelisierungsvorteil, den die Faltung bietet, verloren geht.

Um eine Parallelisierung zu erreichen, wollen wir untersuchen, wie wir die Rekursion zur Berechnung der Ausgabe nutzen können:

Jeder Zustand ist der vorherige Zustand (multipliziert mit A) mit dem Stromeingang (multipliziert mit B) von und. Dies ist bekannt als die ScanvorgangDies kann leicht mit einer for-Schleife berechnet werden.

Im Gegensatz dazu scheint eine Parallelisierung unmöglich, da jeder Zustand erst berechnet werden kann, wenn ein vorheriger Zustand verfügbar ist. Mit Mamba ist es jedoch möglich, den Zustand jedes Zustands zu parallelisieren, indem Parallelabtastung Algorithmen machen dies möglich.

Sie geht davon aus, dass die Reihenfolge, in der wir die Operationen durchführen, unwichtig ist, und macht sich eine Eigenschaft des Vereinigungsgesetzes zunutze. So können wir die Sequenz in Teilen berechnen und sie dann iterativ kombinieren:

dynamische Matrix B im Gesang antworten C und parallele Abtastalgorithmen zusammenkommen, um eine Selektiver Abtastalgorithmusum die dynamische und schnelle Natur der Verwendung rekursiver Darstellungen darzustellen.

Einer der Nachteile aktueller Grafikprozessoren ist die begrenzte Übertragungsgeschwindigkeit (IO) zwischen ihrem kleinen, aber effizienten SRAM und ihrem großen, aber etwas weniger effizienten DRAM. Das häufige Kopieren von Informationen zwischen SRAM und DRAM kann zu einem Engpass werden.

Mamba ähnelt Flash Attention insofern, als es versucht, die Anzahl der Wechsel vom DRAM zum SRAM und zurück zu begrenzen. Dies geschieht durch Kernel-Fusion Dadurch kann das Modell das Schreiben von Zwischenergebnissen verhindern und die Berechnung so lange durchführen, bis sie abgeschlossen ist.

Konkrete Beispiele für die Zuweisung von DRAM und SRAM können wir sehen, wenn wir uns die Grundarchitektur von Mamba ansehen:

Hier werden die folgenden Elemente zu einem Kernel verschmolzen:

Der letzte Teil des hardwarenahen Algorithmus lautet Neuberechnung.

Zwischenzustände werden nicht gespeichert, sind aber für die Berechnung des Gradienten im Umkehrdurchgang erforderlich. Stattdessen berechnen die Autoren diese Zwischenzustände während des Umkehrdurchgangs neu.

Dies mag zwar ineffizient erscheinen, ist aber viel billiger als das Lesen all dieser Zwischenzustände aus dem relativ langsamen DRAM.

Wir haben nun alle Komponenten seiner Architektur behandelt, die unten in seinem Artikel abgebildet ist:

Selektive SSM. Abgerufen von: Gu, Albert und Tri Dao, "Mamba: lineare Zeitreihenmodellierung mit selektiven Zustandsräumen". arXiv-Vorabdruck arXiv:2312.00752 (2023).

Diese Architektur wird oft bezeichnet als Selektive SSM vielleicht S6 Modell, da es sich im Wesentlichen um ein S4-Modell handelt, das mit dem Algorithmus der selektiven Abtastung berechnet wurde.

Was wir bisher erforscht haben Selektive SSM kann als Modul implementiert werden, so wie wir die Selbstaufmerksamkeit in einem Decodermodul darstellen können.

Wie ein Decoder können wir mehrere Mamba-Module stapeln und deren Ausgang als Eingang für das nächste Mamba-Modul verwenden:

Es beginnt mit einer linearen Projektion, um die Eingabeeinbettung zu erweitern. Dann wird die Selektive SSM Um zu verhindern, dass die Token unabhängig voneinander berechnet werden, wird zuvor eine Faltung durchgeführt.

Selektive SSM hat die folgenden Merkmale:

Wenn wir uns die Code-Implementierung ansehen, können wir diese Architektur noch weiter ausbauen und untersuchen, wie ein End-to-End-Beispiel aussehen könnte:

Beachten Sie einige Änderungen, wie z. B. die Hinzufügung einer Normalisierungsschicht und einer Softmax für die Auswahl von Ausgabewörtern.

Alles zusammen ergibt eine schnelle Inferenz und ein schnelles Training, selbst bei unendlichem Kontext. Bei der Verwendung dieser Architektur stellten die Autoren fest, dass die Leistung mit der eines Transformer-Modells derselben Größe übereinstimmte und diese manchmal sogar übertraf!

Dies ist der Abschluss unserer Erkundung von Zustandsraummodellen und der unglaublichen Mamba-Architektur mit selektiven Zustandsraummodellen. Ich hoffe, dieser Beitrag hat Ihnen ein besseres Verständnis von Zustandsraummodellen und insbesondere von Mamba vermittelt. Wer weiß, ob dies Transformers ersetzen wird, aber im Moment ist es erstaunlich zu sehen, dass eine so unterschiedliche Architektur die Aufmerksamkeit erhält, die sie verdient!

Weitere Visualisierungen zum Thema Big Language Modelling und zur Unterstützung dieses Newsletters finden Sie in dem Buch, das ich gemeinsam mit Jay Alammar verfasst habe.