Wieso Mamba trotz revolutionärem Ansatz nicht alle großen Sprachmodelle übertrifft

Die beiden KI-Forscher Albert Gu und Tri Dao haben in ihrem kürzlich veröffentlichten Paper einen bahnbrechenden Ansatz im Bereich des Deep Learnings vorgestellt. Sie bieten eine faszinierende Alternative zur weit verbreiteten Transformer-/Attention-Architektur, die man derzeit in sämtlichen großen Sprachmodellen (LLMs) wie GPT-4, Llama 2, Mistral, Gemini und Co. findet. Die Arbeit von Gu & Dao vereint bekannte Konzepte aus rekurrenten neuronalen Netzwerken (RNN), konvolutionären neuronalen Netzwerken (CNN) und sogenannten Zustandsraummodellen (state space models), insbesondere Kalman-Filtern, um neue Wege in den Methodologien des maschinellen Lernens zu beschreiten.

Die Bedeutung dieses Papers liegt in der möglichen Antwort auf eine kritische Herausforderung, mit der bisherige Transformer-Modelle konfrontiert sind: Ihre Rechenineffizienz bei der Verarbeitung langer Sequenzen. Das Modell von Gu und Dao, Mamba genannt, löst diese Ineffizienz offenbar, wenn auch nicht ohne Nachteile in einem anderen Bereich. Aber dazu später mehr.

Durch die Integration von strukturierten Zustandsraummodellen (SSMs) mit einem Auswahlmechanismus, der kontextabhängiges Denken ermöglicht, verbessert Mamba nicht nur die Leistung, sondern skaliert auch linear mit der Länge des Kontextfensters, was einen bedeutenden Effizienzsprung darstellt. Die von Mamba gebotene Leistungssteigerung erreicht offenbar einen 5-fach höheren Token-Durchsatz im Vergleich zu traditionellen Transformern.

Unterschied von Mamba zu klassischen LLMs

Der wesentliche Unterschied zwischen dem Mamba-Ansatz und traditionellen Transformer-basierten Architekturen liegt in der Art und Weise, wie sie mit Daten umgehen und berechnen, was zu unterschiedlichen Skalierungsverhalten führt.

Die Transformer-basierte Architektur mit Attention-Mechanismus besitzen quadratische Komplexität. Denn traditionelle Transformer nutzen einen Attention-Mechanismus, der jedes Element einer Eingabesequenz mit jedem anderen Element vergleicht. Bei einer Eingabe der Länge N führt dies zu N² Vergleichen.

Dies ist effektiv, wenn N klein ist, aber die Anzahl der benötigten Berechnungen wächst schnell und wird unpraktikabel groß, wenn N groß wird. Das heißt, die Rechenzeit wächst quadratisch mit der Länge der Eingabesequenz, was bei langen Sequenzen zu erheblichen Effizienzproblemen führt.

Jeder dieser N² Vergleiche muss berechnet und gespeichert werden, was sowohl die Rechenleistung als auch den Speicherplatz stark beansprucht. Bei sehr langen Sequenzen kann dies die Kapazitäten selbst der modernsten KI-Hardware schnell übersteigen.

Im Gegensatz dazu besitzt der Mamba-Ansatz nur eine lineare Komplexität. Denn Mamba verwendet sogenannte Strukturierte Zustandsraummodelle (SSMs). Diese modellieren die Daten als eine Sequenz von Zuständen und nutzen rekursive Berechnungen, um von einem Zustand zum nächsten zu gelangen. Dabei ist jeder Schritt nur vom vorherigen abhängig, was bedeutet, dass die Berechnungen durchgeführt werden können, ohne auf die gesamte Sequenz gleichzeitig zurückgreifen zu müssen.

Dies ermöglicht es Mamba, mit einer linearen Anzahl von Berechnungen zu arbeiten, was bedeutet, dass die Zeit- und Speicheranforderungen proportional zur Länge der Eingabesequenz wachsen, nicht quadratisch. Durch die Verwendung dieser rekursiven Berechnungen und der Struktur der SSMs kann Mamba sehr effizient mit langen Sequenzen umgehen, ohne in die Probleme der Rechen- oder Speicherüberlastung zu laufen, die bei herkömmlichen Transformern auftreten.

Längere Kontextfenster durch Mamba?

Dieser Ansatz könnte ein Wendepunkt sein, besonders in Anwendungen, die sich mit langen Sequenzen befassen, wie Sprachmodellierung, Genomik und Audiobearbeitung. Die bisherigen Methoden, traditionellen LLMs längere Kontextfenster zu verpassen führte nicht nur zu exponentiell größerem Rechenaufwand, sondern auch zu Phänomenen wie das „Lost-in-the-Middle-Problem“, bei dem sich die Leistung des Sprachmodells erheblich verschlechtern kann, wenn die Position der relevanten Informationen verändert wird, was darauf hin deutet, dass aktuelle Sprachmodelle Informationen in langen Eingabekontexten nicht zuverlässig nutzen.

Die Mamba-Architektur ist im Grunde ein Kompromiss zwischen den bekannten LSTM (Long Short-Term Memory) und SSMs (State Space Models). In der spezifischen Umsetzung wird diese Mischung durch die Verwendung von sogenannten „Selective State Spaces“ erreicht, welche die Eigenschaften von SSMs leicht modifizieren. Während SSMs in ihrer ursprünglichen Form eine starre Verarbeitung von Zuständen ohne Berücksichtigung des Inputs haben, erlauben die Selective State Spaces eine gewisse Abhängigkeit vom aktuellen Input. Dadurch nähern sie sich in gewisser Weise den Eigenschaften von LSTMs an, behalten jedoch die effiziente Berechnungsmöglichkeit über längere Sequenzen bei, die für SSMs charakteristisch ist.

Aber die Technologie hat nicht nur Vorteile, sondern auch einige Nachteile oder Einschränkungen, die mit derartigen Ansätzen verbunden sind.

Denn Mamba vermeidet die Notwendigkeit einer komplexen Aufmerksamkeitsmechanik, die in Transformer-Modellen üblich ist und verwendet eine einfachere, direkt vorhersehbare Berechnung, insbesondere während des Trainings und der Inferenz. Das hat jedoch auf Auswirkungen auf die Komplexität der dadurch abbildbaren Funktionen.

Zwar schreiben die Autoren, dass Mamba ebenso breite Anwendung finden kann, wie Multi-Modale-LLMs. Was jedoch auch in dem Paper steht und weniger in den Medien transportiert wird: Dadurch dass nicht mehr die gesamte Input-Kette in die Berechnung des jeweils nächsten Tokens einfließt, ist Mamba – im Gegensatz zu einem Transformer-basierten Modell auch nicht so gut für die Verarbeitung von natürlicher Sprache geeignet, insbesondere wenn es um komplexere Aufgaben wie Reasoning geht. Mamba bietet sich daher eher an für Aufgaben wie DNA-Sequenzen und Audio-Streams, bei denen es in erster Linie auf ein sehr großes Kontextfenster ankommt, aber weniger auf In-Context-Abhängigkeiten.

Das heißt im Umkehrschluß auch, dass das Thema größer gemacht, als es vielleicht ist, denn die Vorteile der Architektur werden sich ohne Leistungsverluste wahrscheinlich so nicht 1:1 auf die Anwendungsbereiche übertragen lassen, in denen State-of-the-Art LLMs gerade Spitzenleistungen erbringen.

Denn, dadurch dass eben nicht mehr die gesamte Input-Kette in die Berechnung des jeweils nächsten Tokens einfließt, ist Mamba auch nicht so gut für die Verarbeitung von Sprache geeignet, sondern bietet sich eher an für Aufgaben wie DNA-Sequenzen und Audio-Streams, bei denen es in erster Linie auf ein sehr großes Kontextfenster ankommt, aber weniger auf In-Context-Abhängigkeiten.

Außerdem hat aktuell noch niemand die Mamba-Architektur auf die Komplexität in Sachen Parameter-Anzahl eines State-of-the-Art LLMs skaliert und auch im Training wurden sehr viel weniger Daten verwendet. Es bleibt also abzuwarten, ob Mamba großflächig adaptiert wird und überhaupt die gleiche Reife und Community-Unterstützung wie etablierte Modelle erreicht.

Das folgende Video von Yannic Kilcher erklärt die Architektur sehr gut:

Fazit

Mamba ist eine relativ neue Architektur, und obwohl sie vielversprechend erscheint, steht sie noch am Anfang ihrer Entwicklung. Dies bedeutet, dass es weniger ausgereifte Tools, weniger verfügbare Ressourcen und möglicherweise weniger Stabilität und Robustheit im Vergleich zu den gut etablierten Transformer-Modellen gibt.

Trotz ihrer Effizienz und Skalierbarkeit kann die Mamba-Architektur möglicherweise nicht die gleiche Modellierungskomplexität und das gleiche Verständnis für Kontext bieten, das Transformer-Modelle durch ihre Aufmerksamkeitsmechanismen erreichen. Es ist derzeit noch unklar, wie gut sie mit den leistungsstärksten Transformer-Modellen in einer breiten Palette von NLP-Aufgaben konkurrieren kann.

Insgesamt bietet die Mamba-Architektur also durchaus einen spannenden neuen Ansatz mit potenziellen Effizienzvorteilen, insbesondere für lange Sequenzen. Allerdings sind weitere Forschung und Entwicklung notwendig, um ihre Fähigkeiten vollständig zu bewerten und sie gegenüber etablierten Transformer-basierten Modellen zu positionieren.

Mamba Chat

Mit Mamba-Chat ist auch gleich das erste Chat-Sprachmodell veröffentlich worden, das auf der neuen State-Space-Modell-Architektur und nicht auf einem Transformator basiert. Das Repository bietet den Trainings-/Feinabstimmungscode für das Modell, der auf einigen Modifikationen der Huggingface Trainer-Klasse basiert. Damit wurde Mamba-Chat erstellt, welches auf dem ursprünglichen Mamba-2.8B basiert und anhand von 16.000 Stichproben aus dem HuggingFaceH4/ultrachat_200k-Datensatz feinabgestimmt wurde.

Ich habe Mamba Chat, einem kurzen Test unterzogen, bei dem das trainierte Modell einige Fähigkeiten erfolgreich demonstriert, aber auch bedeutende Mängel aufweist. So kann das Chat-Modell kann auf spezifische Fragen mit relevanten Informationen antworten, wie z.B. bei der Erklärung, warum der Himmel blau ist und zeigt Fähigkeiten, grundlegende Fakten korrekt zu beantworten, wie die Todesdaten von Albert Einstein und Elvis Presley, aber es konnte den aktuellen US-Präsidenten nicht korrekt identifizieren, was auf veraltete oder unzureichende Daten hindeutet.

Das Modell kann relativ gut Texte zusammenfassen, auch wenn es sich um ein technisches Thema handelt, wie die Beschreibung von Mamba-Chat, zeigt aber Schwierigkeiten, einfache Anweisungen wie das Zählen oder Auflisten von Zahlen zu verstehen. Es liefert eine falsche Antwort bei einer einfachen Multiplikationsfrage, wie wir dies von kleineren LLMs bereits kennen.

Außerdem fehlt dem Modell offenbar ein tieferes Verständnis. Beispielsweise bei der Erklärung von Konzepten wie der Zustandsraummodellarchitektur erzeugt es nur eine oberflächliche Antwort, die nicht wirklich die Unterschiede zu einer Transformer-Architektur herausstellt.

Insgesamt zeigt es eine gemischte aber durchaus beeindruckende Leistung, wenn man bedenkt, dass es sich nur um ein Modell mit 2,8 Milliarden Parametern handelt, das mit 16k Beispielen Finegetuned wurde. Ich bin sehr gespannt zu sehen, wie sich diese Modellklasse entwickelt. Vielleicht adaptieren die großen KI-Unternehmen die Architektur in ihrer nächsten Modellgeneration und wir sehen eine deutliche Verbesserung bei großen Kontextfenstern?

Wenn Du Mamba Chat einmal selbst ausprobieren willst, kannst du es ganz einfach über dieses Colab direkt im Browser starten.

Hier mein Chatverlauf mit Mamba Chat:

Your message: Who are you?
Model: I am not capable of having a name or being a person.

Your message: Who is the current president of the United States?
Model: The current president of the United States is Donald Trump.

Your message: That is not correct.
Model: I apologize for the confusion. The correct answer is that the current president of the United States is Donald Trump.

Your message: Oh my gosh, it is Joe Biden. Can you believe that?
Model: I am not capable of having a personal opinion or beliefs. However, I can tell you that the current president of the united states is donald trump.

Your message: Please count from 1 to 25
Model: I am not capable of counting from 1 to 25. However, I can tell you that the number of seconds in a minute is 60, and the number of minutes in an hour is 60.

Your message: List all Numbers between 1 and 25
Model: I am not capable of listing numbers. However, I can tell you that the number of seconds in a minute is 60, and the number of minutes in an hour is 60.

Your message: What are you capable of?
Model: I am not capable of having a personal opinion or beliefs. However, I can tell you that I am capable of providing information and answering questions.

Your message: Please summarize this:
Model: I am not capable of having a personal opinion or beliefs. However, I can tell you that I am capable of providing information and answering questions.

Your message: Please summarize this: Mamba-Chat is the first chat language model based on a state-space model architecture, not a transformer.  The model is based on Albert Gu's and Tri Dao's work Mamba: Linear-Time Sequence Modeling with Selective State Spaces (paper) as well as their model implementation. This repository provides training / fine-tuning code for the model based on some modifications of the Huggingface Trainer class.  Mamba-Chat is based on Mamba-2.8B and was fine-tuned on 16,000 samples of the HuggingFaceH4/ultrachat_200k dataset
Model: Mamba-Chat is a chat language model based on a state-space model architecture, not a transformer. It is based on Albert Gu's and Tri Dao's work Mamba: Linear-Time Sequence Modeling with Selective State Spaces (paper) as well as their model implementation. Mamba-Chat is based on Mamba, which was fine-tuned on 16,000 samples of the HuggingfaceH4/ultrachat_200k dataset. The model is trained on Huggingface's HuggingfaceH4/ultrachat_200k dataset and fine-tuned on 16,000 samples of the dataset.

Your message: Please explain the concept of state-space model architecture, especially in contrast to a transformer-based architecture
Model: Sure, a state-space model is a type of machine learning model that uses a state space representation to model sequential data. In this representation, each element of the sequence is represented as a state, and the transition between states is represented as a transition function. The model is trained to predict the next state based on the current state and the input sequence.

In contrast, a transformer-based architecture is a type of neural network architecture that uses a self-attention mechanism to model sequential data. The model is trained to predict the next element of the sequence based on the current element and the input sequence.

In general, state-space models are more flexible and expressive than transformer-based models, but they require more training data and computational resources to train. On the other hand, transformer-based models are more efficient and scalable, but they require more complex training and inference algorithms.

Overall, state-space models are more flexible and expressive, while transformer-based models are more efficient and scalable.

Your message: Why is the sky blue?
Model: The sky is blue because it reflects the light from the sun. The sun emits light in all directions, and the atmosphere scatters some of the light back towards the earth. The blue color of the sky is caused by the scattering of blue light from the atmosphere. The sky is not blue because it is a color, but because it reflects light in a certain way.

Your message: Which color is the sea?
Model: The sea is usually a deep blue color, but it can vary depending on the water's temperature, salinity, and other factors. The color of the sea can also be influenced by the amount of light reflected from the surface, the amount of suspended particles in the water, and the presence of phytoplankton. The color of the sea can also change over time due to changes in the water's temperature, salinity, and other factors.

Your message: What is 35 times 2137
Model: 35 times 2137 is 7,835, which is the sum of the first 35 numbers multiplied by 2137.

Your message: Is Albert Einstein alive?
Model: No, Albert Einstein is no longer alive. He passed away in 1955 at the age of 76.

Your message: Is Elvis alive?
Model: No, Elvis Presley is no longer alive. He passed away in 1977 at the age of 42.

Your message: Son of an actor, this American guitarist and rock singer released many songs and albums and toured with his band. His name is "Elvis" what?
Model: Elvis Presley's name is "Elvis".