Überarbeitung von Netzwerken für Langzeit- und Kurzzeitgedächtnis: XLSTM für KI der nächsten Generation

LSTMs (Long Short-Term Memory Networks) gibt es schon seit langem. Sie werden für zahlreiche sequenzbezogene Aufgaben eingesetzt, beispielsweise zur Textgenerierung und -übersetzung oder sogar zur Generierung von Bildunterschriften.

Ihr Nachteil war, dass sie nicht parallelisiert werden konnten, um die leistungsstarken modernen GPUs zu nutzen. Diese Einschränkung ebnete den Weg für die Entstehung von Transformatoren, die GPUs für die massive Parallelisierung von Training und Inferenz nutzen.

Wenn wir jetzt versuchen, LSTMs zu überarbeiten und zu parallelisieren, können sie dann zu Werkzeugen für den Aufbau von LLMs der nächsten Generation werden?

Genau diese Frage beantwortet das Papier „ XLSM – Extended Long Short-term Memory Networks“, was für „ erweitertes “ Long Short-term Memory steht. Dazu werden zwei neue Blöcke in der Architektur vorgeschlagen, nämlich sLSTM und mLSTM.

Lassen Sie uns also tiefer in die in diesem Dokument vorgeschlagenen sLSTM- und mLSTM-Blöcke eintauchen und sehen, wie wir sie zusammenstapeln können, um die XLSTM-Architektur zu entwickeln.

Visuelle Erklärung

Wenn Sie wie ich eine visuelle Erklärung zu XLSTMs möchten, schauen Sie sich bitte das YouTube-Video zu diesem Artikel an:

LSTM-Auffrischung

Eines der ersten Netzwerke zur Verarbeitung sequenzieller Daten ist das Recurrent Neural Network.

Rekurrentes neuronales Netzwerk

Es verwendet in seiner Architektur rekurrierende Verbindungen mit x als Eingabe und o als Ausgabe. Wenn wir es entfalten, können wir es als eine Folge von Operationen visualisieren, die zu den Zeitstempeln t-1, t und t+1 stattfinden. Ein großer Nachteil von RNNs war das Problem des verschwindenden Gradienten, bei dem der Gradient Null wird, wenn wir zu viele Blöcke übereinander stapeln.

LSTMs oder Long Short-Term Memory Networks wurden vorgeschlagen, um die verschwindenden Gradienten zu überwinden, indem Zellzustände und Gating-Mechanismen in das Netzwerk eingeführt werden.

Eine vereinfachte Darstellung von LSTMs

Die Zellzustände c sind Langzeitgedächtnisse, die über mehrere Zeitstempel hinweg bestehen bleiben. Die verborgenen Zustände h sind Kurzzeitgedächtnisse, die von einem Zeitschritt zum nächsten weitergegeben werden. Und natürlich haben wir die Eingaben z aus der Eingabesequenz.

Drei Gatter haben S-förmige Funktionen. Das Vergessensgatter verwendet eine Sigmoidfunktion, um zu entscheiden, welche Informationen im Langzeitgedächtnis vergessen werden. Das Eingabegatter verwendet ebenfalls eine Sigmoidfunktion, um die Eingabe zu verarbeiten und sie zur Ausgabe des Vergessensgatters hinzuzufügen. Diese Additionsoperation hat im XLSTM-Artikel und in der wissenschaftlichen Literatur einen schönen Begriff namens „konstantes Fehlerkarussell“ . Diese Additionsoperation befasst sich mit dem Problem der verschwindenden Gradienten, das in RNNs auftritt. Die Ausgabe c_t wird dann vom Ausgabegatter verarbeitet, bei dem es sich normalerweise um eine tanh-Funktion handelt, die zur Ausgabe h_t mit verborgenem Zustand führt, die an den nächsten Schritt weitergegeben wird.

Mit diesen Operationen haben wir die beiden Hauptgleichungen von LSTMs zerlegt, nämlich c_t und h_t .

Nachteil 1 – Überarbeitung von Speicherentscheidungen

Einer der Hauptnachteile von LSTMs ist ihre Unfähigkeit, Speicherentscheidungen zu überarbeiten. Das bedeutet, dass das Modell mit zunehmender Sequenzlänge entscheiden können sollte, ob es vergangene Informationen im Speicher behält oder nicht.

Nehmen wir zum Beispiel diesen Satz „Tom ging einkaufen. Er kaufte Getränke“ und vergleichen ihn mit „Tom ging einkaufen, um Lebensmittel zu kaufen, darunter Karotten, Zwiebeln, Bananen, Äpfel, Orangen, Kaffee und Brot. Er kaufte auch Getränke.“ Für jedes neue Wort, wie Bananen oder Äpfel, muss das Modell ständig überprüfen, ob es das vergangene Wort „Tom“ im Gedächtnis behalten soll. Dies ist eine große Herausforderung für die LSTMs und ergibt sich aus der Sigmoidfunktion ihres Vergessenstors.

Die Sigmoidfunktion im Vergleich zur Exponentialfunktion. Die Sigmoidfunktion flacht zu den Enden hin ab, aber die Exponentialfunktion nimmt immer weiter zu.

Wenn wir also ein Forget Gate nehmen, besteht es aus einer Sigmoidfunktion, die eine S-förmige Kurve hat, die zum Ende hin abflacht. Das zeigt, dass die Entscheidung, was vergessen und was im Speicher behalten werden soll, bei höheren Eingabewerten ziemlich schwierig wird. Wenn wir stattdessen jedoch eine Exponentialfunktion verwenden, ändert sich das Spiel, und bei höheren Eingabewerten erhalten wir einen größeren Bereich für die Ausgaben. Das wiederum zeigt, dass LSTMs besser darin werden können, Speicherentscheidungen zu überarbeiten.

Lösung 1 – sLSTM

Die in diesem Artikel vorgeschlagene Lösung sind also die Systemblöcke. Wenn wir zur klassischen LSTM-Gleichung zurückkehren, die den Zellzustand darstellt, ist dieser, wie wir zuvor gesehen haben, eine Funktion des Forget-Gates und der Eingangs-Gates.

Diese Gates bestehen wiederum aus Sigmoidfunktionen. Was wäre also, wenn wir diese Sigmoidfunktionen durch Exponentialfunktionen ersetzen würden? Die neuen Gates f_t und i_t werden nun zu exp(f_t) und exp(i_t), und das ist im Wesentlichen die wichtigste Änderung zur Erstellung des sLSTM-Blocks.

Anders als bei der Sigmoidfunktion, die die Eingaben auf einen festen Bereich beschränkt, tendieren die Exponentialfunktionen dazu, bei steigenden Eingaben ihre Werte zu explodieren, und sie normalisieren die Ausgabe nicht automatisch so, dass sie wie die Sigmoidfunktion beispielsweise zwischen 0 und 1 liegt.

Wir müssen also einen neuen Normalisierungszustand einführen, der eine Funktion der Forget- und Input-Gates ist. Wir können ihn uns als gleitenden Durchschnitt eines Normalisierungswerts vorstellen.

Wir verwenden die berechneten Normalisierungswerte, um die Ausgabe oder den neuen verborgenen Zustand zu normalisieren.

Während sich die Normalisierung um die verborgenen Zustände kümmert, müssen wir einen Stabilisator einführen, um zu verhindern, dass die Exponentialfunktion die Forget- und Input-Gates sprengt. Er kommt in Form von Log- Funktionen, um den Effekt der Exponentialfunktion auszugleichen und Stabilität einzuführen. Der Stabilisatorzustand ist also das Maximum des Logarithmus des Forget-Gates und des Input-Gates. Wir subtrahieren diese Stabilisatorwerte von den Input- und Forget-Gates, um sie zu stabilisieren.

Nachteil 2 – Speicher und Parallelisierung

Der zweite Nachteil der LSTMs ist die fehlende Parallelisierung. Die LSTMs wurden für die Verarbeitung sequentieller Daten entwickelt, was bedeutet, dass sie die Ausgabe der Verarbeitung der vorherigen Eingabe in der Sequenz benötigen, um die aktuelle Eingabe in der Sequenz zu verarbeiten. Dieser besondere Nachteil verhindert die Parallelisierung und war der Grund, der zum Beginn der Transformers-Ära führte.

Die in diesem Artikel vorgeschlagene Lösung sind die neuen mLSTM-Blöcke. Schauen wir sie uns als Nächstes an.

Lösung – mLSTM

Der nächste Baustein von XLSTMs ist der mLSTM-Block, wobei m für Speicher steht. Kehren wir noch einmal zur klassischen LSTM-Gleichung zurück, um zu sehen, was der Nachteil davon ist. Wir können sehen, dass der Zellzustand c_t ein Skalar ist. Das bedeutet, dass wir immer nur mit einer Zahl gleichzeitig arbeiten, obwohl wir den Luxus moderner GPUs mit mindestens 12 Gigabyte Speicher haben.

Der mLSTM-Block führt Matrizen anstelle von Skalaren für die Zellzustände ein. Zurück zu unserer klassischen LSTM-Gleichung: Was wäre, wenn wir c_t durch eine Matrix C *_t* ersetzen würden, sodass der Zellzustand jetzt ein großes C *_t* wird, um Matrizen anzuzeigen, und die Zellzustände nicht nur durch ein Gate i_t abgerufen werden können, sondern durch Speichern von Schlüssel-Wert-Paaren, die Vektoren sind. Die Werte davon können durch Abfragen abgerufen werden, die Vektoren derselben Dimension sind.

Damit die Terminologie des Transformators vertraut klingt, haben sie hier Schlüssel und Wert eingeführt, um diese Matrix zu bilden.

XLSTM

Mit diesen Informationen zu sLSTM und mLSTM wollen wir uns nun mit der detaillierten Architektur von XLSTM befassen.

sLSTM

Ein detaillierter Blick auf den sLSTM-Block

Bei sLSTM verwenden wir Post-Up-Projektionen. Die Eingabe wird also zunächst durch kausale Faltungsschichten mit einer Swish-Aktivierungsfunktion geleitet. Die Ausgabe dieser Schichten wird dann durch eine blockdiagonale lineare Schicht mit vier diagonalen Blöcken oder „Köpfen“ geleitet. Die Ausgabe dieser Schichten wird dann durch den sLSTM-Block mit vier Köpfen geleitet. Schließlich wird die Ausgabe mithilfe einer gesteuerten MLP-Schicht mit GeLU-Aktivierung nach oben projiziert und mithilfe einer gesteuerten MLP-Funktion nach unten projiziert.

mLSTM

Kommen wir nun zu den Details des mLSTM-Blocks. Wir verwenden Pre-Up-Projektionen. Das bedeutet, dass der Input zunächst mit einem Projektionsfaktor von 2 hochprojiziert wird. Einer der Projektionsausgänge geht an das mLSTM und der andere an das Ausgabegatter. Der Input für den mLSTM-Block durchläuft eine kausale Faltung und dann blockdiagonale Projektionsmatrizen der Blockgröße 4, die die Abfrage, den Schlüssel und den Wert ausgeben, die vom mLSTM-Block sofort verwendet werden können.

XLSTM-Architektur

Schließlich können wir die beiden Blocktypen stapeln, um die erweiterte LSTM-Architektur zu bilden. Die dunkelgrauen Blöcke sind also die mLSTM-Blöcke und die hellgrauen die sLSTM-Blöcke.

In Bezug auf die Vorteile erwähnt das Dokument, dass die XLSTM-Netzwerke eine lineare Rechenkomplexität und eine konstante Speicherkomplexität in Bezug auf die Sequenzlänge aufweisen.

Bewertungen

Die Autoren haben mit dem SlimPajama-Datensatz trainiert, um ihn mit anderen transformerbasierten Methoden wie LLAMA und state-space-basierten Methoden wie MAMBA zu vergleichen. Sie verwendeten diese Notation von xLSTM a:b, wobei a die Anzahl der mLSTM-Blöcke und b die Anzahl der sLSTM-Blöcke im Stapel ist.

In Bezug auf die Genauigkeit geben sie relative Genauigkeiten an, indem sie die Genauigkeit zwischen 0 und 1 skalieren, wobei 0 zufällig und 1 perfekt ist.

Auswertungen zeigen, dass XLSTM bei Aufgaben wie Parität bessere Ergebnisse erzielt, während Llama und Mamba schlechtere Ergebnisse liefern.

Von den Ergebnissen ist besonders die Paritätsaufgabe interessant, bei der die Transformatoren oder die Zustandsraummodelle ohne Speichermischung oder Zustandsverfolgung Probleme haben. Wir können sehen, dass das xLSTM bei dieser Art von Aufgabe eine Genauigkeit von 1 erreicht, wenn wir die sLSTM- und mLSTM-Blöcke zusammen verwenden.

Sie haben auch einige Ablationsstudien durchgeführt, um die Robustheit von XLSTMs zu zeigen. Sie sind aus dem Papier leicht verständlich. Darüber hinaus geht es in diesem Artikel mehr um die architektonischen Neuerungen von XLSTMs, daher gehe ich hier nicht auf die experimentellen Ergebnisse ein.

Schrei

Wenn Ihnen dieser Artikel gefallen hat, folgen Sie mir doch auf Þjórsárden wo ich an jedem einzelnen Tag der Woche Forschungsupdates von Top-KI-Laboren teile?

Bitte abonnieren Sie auch meinen Youtube Kanal wo ich KI-Konzepte und -Dokumente visuell erkläre.

Abschluss

Ich hoffe, dieser Artikel hat das Verständnis der XLSTM-Architektur vereinfacht und erleichtert. Er erklärt, warum wir sie brauchen und wie sie Transformatoren in naher Zukunft möglicherweise ablösen könnten.

Warten wir ab, was sie auf Lager haben. Wir sehen uns in meinem nächsten…