Was ist ein Mamba-Modell?

Mamba-Modelle sind vielleicht die erste Deep-Learning-Architektur, die mit der Effizienz von Transformer-Modellen bei der Aufgabe mithalten kann, für die Transformer ursprünglich bekannt wurden: Sprachmodellierung. Vor allem hat die Mamba-Architektur bei prominenten LLM-Benchmark-Bewertungen die Fähigkeit unter Beweis gestellt, Transformer gleicher Größe abzugleichen, während sie in Bezug auf Latenz und Speicheranforderungen oft deutlich effizienter ist.

Die Mamba-Architektur wurde erstmals von Tri Dao und Albert Gu in dem Artikel „Mamba: Linear-Time Sequence Modeling with Selective State Spaces“ aus dem Jahr 2023 vorgestellt. Ein Jahr später ergänzten sie das ursprünglichen Mamba-Paper mit einem weiteren Paper, das sowohl die Verbindungen zwischen SSMs und Transformern erkundete als auch eine verfeinerte, deutlich schnellere Version der Mamba-Architektur vorstellte, die sie Mamba-2 nannten.

Zwar sind Transformer in den zwei Jahren nach der Veröffentlichung des ursprünglichen Mamba-Papers die vorherrschende Art von LLM geblieben, jedoch wurde die Architektur in eine wachsende Anzahl von Open-Source-Modellen integriert. Einige, wie z. B. Codestral Mamba von Mistral KI, sind reine Mamba-Modelle. Viele weitere, einschließlich der Jamba-Serie von AI2I und IBM Granite 4.0, sind hybride Modelle, die sowohl Aufmerksamkeitsschichten (Transformer) als auch SSM-Schichten (Mamba) enthalten. Zusätzlich zu ihren Vorteilen hinsichtlich der Leistung verspricht die Verbreitung von Mamba-basierten Modellen eine Demokratisierung des Zugangs zur KI, da sie auf vergleichsweise preiswerter Hardware reibungslos laufen.

Zustandsraummodelle (State Space Models, SSMs) wurden ursprünglich entwickelt, um den nächsten Zustand einer kontinuierlichen Sequenz, wie z. B. ein elektrisches Signal, ein Wettermuster oder die Flugbahn eines sich bewegenden Objekts, auf der Grundlage einer Eingabe vorherzusagen. Konzeptionell und mathematisch sind sie mit den rekurrenten neuronalen Netzen (RNNs) verwandt, die vor der Einführung der Transformer im Jahr 2017 die Verarbeitung natürlicher Sprache (NLP) beherrschten, sowie mit anderen Algorithmen des maschinellen Lernens, wie z. B. faltende neuronale Netze (convolutional neural networks, CNNs) und Hidden Markov-Modellen (HMMs).

Wie der Name schon sagt, treffen SSMs Vorhersagen über den nächsten Zustand in einem dynamischen System, indem sie den Zustandsraum modellieren: eine mathematische Darstellung aller Zustandsvariablen, die den Zustand eines Systems und die Bandbreite der Möglichkeiten für jede dieser Variablen beschreiben erzielen.

Ein SSM nimmt eine Eingangssequenz x(t) und ordnet sie einer latenten Zustandsdarstellung h(t) zu – analog zum verborgenen Zustand eines RNN – um eine Ausgangssequenz y(t) vorherzusagen. Den Kern eines jeden SSM bilden zwei Gleichungen:

• Die Zustandsgleichung,  $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• Die Ausgangsgleichung $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

Die wichtigsten Parameter des Modells sind A, B, C und D, die typischerweise die Form einer Gewichtungsmatrix haben. In den Bereichen, in denen SSMs konventionell verwendet werden, wie in der Kontrolltheorie, wird oft davon ausgegangen, dass diese Matrizes fest sind: Sie repräsentieren die Dynamik eines etablierten Systems, und das SSM wird verwendet, um die Eingaben x zu finden, die zu wünschenswerten Ausgaben y führen.In moderneren Konzepten von SSMs sind diese Matrizes selbst Parameter, die durch maschinelles Lernen optimiert werden müssen. In Deep-Learning-Modellen werden diese Matrizes durch die erlernbaren Gewichte eines neuronalen Netzwerks dargestellt.

Die Zustandsgleichung beschreibt, wie sich der Zustand ändert. Die Werte in Matrix A bestimmen, wie sich jede Zustandsvariable im Laufe der Zeit entwickelt, wenn man sie sich selbst überlässt. Die Werte in Matrix B bestimmen, wie die Eingabe – z. B. das nächste Token in einer Textsequenz – jede Zustandsvariable beeinflusst.

Bei der Sprachmodellierung stellt der aktuelle Status den Kontext einer Textsequenz dar, der nach jedem Token aktualisiert wird. Seine Rolle entspricht der des KV-Cache in einem Transformer-Modell.

Die Ausgangsgleichung beschreibt, wie der aktuelle Zustand die Ausgabe beeinflusst (vermittelt durch Matrix C) und wie die Eingabe die Ausgabe direkt beeinflusst (vermittelt durch Matrix D). Da die Matrix D im Wesentlichen außerhalb der Modellierung von h(t) selbst liegt, wird sie in Diagrammen und Diskussionen über SSMs oft weggelassen, um sich auf die Kernmatrizes A, B und C zu konzentrieren.

In einem Mamba LLM wird die Ausgangsgleichung verwendet, um das nächste Token zu generieren.

Herkömmliche SSMs sind für die Modellierung kontinuierlicher Eingaben konzipiert, aber Textsequenzen (und die meisten anderen Datenmodalitäten, die mit modernen Deep-Learning-Modellen verarbeitet werden) sind diskrete Eingaben. Die Verwendung von SSMs zur Modellierung einer diskreten Sequenz erfordert ein Mittel zur Darstellung ihrer einzelnen, spezifischen Zeitschritte als Teil eines kontinuierlichen Signals.

Konzeptionell läuft die Diskretisierung darauf hinaus, den Wert einer kontinuierlichen Funktion zu bestimmten Zeitpunkten zu erfassen. Dies beinhaltet die Einführung eines neuen Parameters – der Schrittgröße, geschrieben als ∆–, der bestimmt, wie lange dieser Wert bei jedem diskreten Zeitschritt t abgetastet oder „gehalten“ wird. Anpassungen an ∆ sind vergleichbar mit Änderungen an Eigenschaften wie der Datenauflösung (bei Zeitreihendaten ) oder der Bildrate (bei Videodaten). Es gibt mehrere „Diskretisierungs“-Methoden, aber die meisten modernen SSM-Varianten (einschließlich Mamba) verwenden die einfache Zero Order Hold (ZOH)-Methode.

Durch die Diskretisierung eines SSM kann es wie ein RNN für sequenzielle Aufgaben verwendet werden. Die Parameter und Gleichungen eines diskreten SSM werden in der Regel umgeschrieben, um sie von ihren zeitkontinuierlichen Äquivalenten zu unterscheiden, wobei die für RNNs üblicherweise verwendete Indexnotation verwendet wird. In dieser Notation stellt h_t den aktualisierten Zustandsraum dar, den das Modell generiert, und h_t-1 stellt den Zustand davor dar, d. h. den aktuellen Zustandsraum.

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

Die Modellierung von Textdaten mit diskreten Standard-SSMs ist aufgrund einer Reihe von Mängeln, die sie mit RNNs teilen, unpraktisch. Zwei der Mängel wurden durch die Einführung von strukturierten Zustandsraumsequenzmodellen (oder „S4-Modellen“) durch Albert Gu et al. im Jahr 2021 behoben: die Ineffizienz ihres Trainings und ihre Unfähigkeit, lange Sequenzen zu modellieren.

Der Erfolg der S4-Modelle – und ihrer vielen Derivate wie diagonale SSMs (DSS), diagonale S4-Modelle (S4D) und H3-Modelle – ebnete jedoch direkt den Weg für das, was zu Mamba wurde.

Der Nutzen von diskretisierten SSMs als Äquivalent zu einer spezifischen Instanz eines RNN ist, dass RNNs bei der Inferenz extrem schnell sind. Der Nachteil ist jedoch, dass RNNs extrem langsam zu schulen sind.

Glücklicherweise haben diskretisierte SSMs eine wichtige Eigenschaft, die sie von anderen RNNs unterscheidet: Sie modellieren ausschließlich lineare Abhängigkeiten. Mit anderen Worten, sie verwenden nur einfache, unkomplizierte Multiplikations- und Additionsoperationen. Wie das S4-Paper zeigt, können diese einfachen, wiederholten und voneinander abhängigen linearen Wiederholungen in eine eindimensionale Faltungsmatrix umgewandelt werden, die die Eingabe x in einem einzigen Schritt direkt der Ausgabe y zuordnet: . Dies kann mit Hilfe der schnellen Fourier-Transformation sehr effizient berechnet werden.

Der einzige Haken ist, dass dies nur möglich ist, wenn jeder Schritt der gesamten Eingabesequenz bekannt ist. Während der Inferenz ist dies nicht möglich, aber während des Trainings schon. Ein strukturiertes SSM bietet daher das Beste aus beiden Welten: Während des Trainings kann es als CNN sehr effizient betrieben werden, während es während des Trainings sehr effizient als CNN betrieben werden kann; während der Inferenz kann es als RNN sehr effizient betrieben werden.

Wie die meisten RNNs sind Standard-SSMs von Natur aus schwach bei der Modellierung von Abhängigkeiten über große Entfernungen. Mit anderen Worten: Sie sind nicht gut darin, die Beziehung zwischen weit voneinander entfernten Schritten in einer Sequenz zu verstehen, z. B. Wörter am Anfang und am Ende eines Absatzes – was sie bei der Modellierung langer Sequenzen insgesamt schwach macht.

Um das zu lösen, verwendeten Gu und seine Co-Autoren einer von ihnen war Tri Dao) eine Technik namens HiPPO – kurz für High-Order Polynomial Projection Operators –, um das Verhalten der A- und B-Matrizes zu definieren, indem sie ihre Anfangswerte mithilfe einer von orthogonalen Polynomen abgeleiteten Formel strukturierten . Dies steht im Gegensatz zur Standardpraxis des maschinellen Lernens, bei der Modellgewichte zu Beginn des Modelltrainings zufällig initialisiert werden. Für S4 schlugen Dao und Gu Initialisierungsschemata vor, die aus Legendre-Polynomen abgeleitet wurden. Sie erkundeten zusätzliche Formeln in einem Follow-up-Papier mit dem Titel „How to Train Your HiPPO.“¹

Im S4-Paper heißt es, dass die einfache Änderung eines SSM von einer Zufallsmatrix A auf [die HiPPO-Matrix] seine Leistung beim sequenziellen MNIST-Benchmark von 60 % auf 98 % verbesserte, wodurch das Problem des Langzeitspeichers von SSMs effektiv gelöst wurde. Spätere Variationen von strukturierten SSMs, wie DSS, S5 und Mamba, verwenden unterschiedliche (oft einfachere) Initialisierungsschemata für A und B, die jedoch die Kernprinzipien von HiPPO beibehalten: die Implementierung einer diagonalen Struktur, die stabile Aktualisierungen und ein gewisses Maß an Unabhängigkeit zwischen den einzelnen Werten in der Matrix erzwingt.

Den Kern der Mamba-Architektur bilden zwei Innovationen. Das erste ist das selektive Zustandsraummodell, das Mamba eine entscheidende Funktion verleiht, die bisher nur Transformer-Modelle besaßen: die Funktion, sich selektiv auf bestimmte Teile der vergangenen Eingabe zu konzentrieren oder diese zu ignorieren, basierend auf ihrer gegenwärtigen Relevanz. Der andere ist der hardwarebewusste parallele Scan, ein Algorithmus, der die Art und Weise optimiert, wie eine Grafikverarbeitungseinheit (GPU) die Berechnungen des Modells in ihrer Speicherhierarchie verarbeitet, um Geschwindigkeit und Recheneffizienz zu maximieren.

Bei Transformern wird diese Fähigkeit durch den Aufmerksamkeitsmechanismus bereitgestellt, der die Aufmerksamkeitsgewichtungen anpasst, die den Einfluss jedes vorherigen Tokens betonen oder abschwächen, basierend auf seiner Relevanz für das aktuelle Eingabe-Token. Gewöhnliche SSMs sind explizit so konzipiert, dass sie die Eingabe der Ausgabe unter Verwendung des gesamten Eingabeverlaufs zuordnen. Dies ist für einige Aufgaben der Sequenzmodellierung akzeptabel oder sogar wünschenswert, aber ein erhebliches Handicap für die meisten fortgeschrittenen Aufgaben der Sprachmodellierung.

Um dem Unvermögen entgegenzuwirken, bestimmte Teile ihrer Eingabe dynamisch wegzulassen oder hervorzuheben, schlugen Dao und Gu eine neue Klasse von Zustandsraummodellen mit einem „selektiven Scan“ vor. In dem Mamba-Artikel bemerken die Autoren, dass sie „selektive SSMs manchmal als S6-Modelle abkürzen, weil es sich um S4-Modelle mit einem Auswahlmechanismus handelt, die mit einem Scan berechnet werden“. Sie gaben ihrer auf S6 basierenden Architektur den Spitznamen „Mamba“, unter anderem, weil die vielen „S“ wie das Zischen einer Schlange klingen.

Mamba kann am besten als Neural Networks-Architektur verstanden werden, deren Kern das selektive Zustandsraummodell ist. Als Analogie: Mamba verhält sich zu selektiven SSMs wie das Transformer-Modell zum Aufmerksamkeitsmechanismus.

Ein herkömmliches SSM hat eine feste Dynamik: Die Regeln, die bestimmen, wie sich der verborgene Zustand von einem Schritt zum nächsten entwickelt – die Modellparameter– sind für jede Eingabe und für jeden Schritt in der Sequenz dieselben. Diese Eigenschaft wird als lineare Zeitinvarianz (linear time invariance, LTI) bezeichnet. Um SSMs die Möglichkeit zu geben, bestimmte Informationen aus der Vergangenheit auf der Grundlage des aktuellen Kontextes selektiv zu priorisieren oder zu depriorisieren, haben Dao und Gu ihr SSM so umkonfiguriert, dass die Werte der wichtigsten Modellparameter für verschiedene Eingaben unterschiedlich sind.

Genauer gesagt machen selektive SSMs die Schrittgröße ∆_t und die Matrizes B_t und C_t zu direkten Funktionen des aktuellen Tokens xt. Dies wird erreicht, indem die Vektoreinbettung von x_t durch drei parallele lineare Projektionsschichten durchläuft – mit anderen Worten: Standard-Forward-Schichten eines neuronalen Netzwerks (oder MLP-Schichten). Dies entspricht der Art und Weise, wie die parallelen Abfrage-, Schlüssel- und Wertköpfe die entsprechenden Q-, K- und V-Vektoren einer Eingabe in einem Transformer-Modell erzeugen.

Die Multiplikation der Vektoreinbettung von x_tmit den Gewichtungs-und Bias-Termen in diesem linearen Projektionsnetz liefert die resultierenden Werte von ∆_t, B_t und C_t. Die Gewichtungs- und Verzerrungsbegriffe der linearen Projektionsschichten selbst werden während des Vortrainings des Modells anhand umfangreicher Datensätze von Textproben gelernt und dann (optional) durch anschließende Feinabstimmung verfeinert.

• Der Wert von ∆_t bestimmt das Ausmaß des Einflusses von x_t auf die Erinnerung des Modells an den Kontext, den es bisher gegeben hat: mit anderen Worten, wie groß die Aktualisierung vom verborgenen Zustand h_t-1 auf h_t sein wird.Eine größere Schrittgröße ∆_t ergibt größere Änderungen und beschleunigt den Zerfall – mit anderen Worten, das „Vergessen“ – älterer Informationen, die im Zustand enthalten sind. Umgekehrt führt eine kleinere Schrittgröße zu einer kleineren Aktualisierung. Bei einer ausreichend kleinen Schrittgröße hat die aktuelle Eingabe überhaupt keinen Einfluss auf den verborgenen Zustand.
  
  
• Änderungen an der Matrix B_k bestimmen,wie das aktuelle Eingabetoken den ausgeblendeten Zustand aktualisiert. Wenn x_t beispielsweise ein Token für das Wort „gestern“ ist, könnte B_t so angepasst werden, dass der Zustand so aktualisiert wird, dass der folgende Kontext wahrscheinlich zur Vergangenheit gehört.
  
  
• Änderungen an der Matrix C_t bestimmen, wie diese Kontextinformationen umgesetzt werden, um den Output des Modells y_t zu beeinflussen. Um das Beispiel fortzufahren, in dem x_k ein Token für „gestern“ ist, könnte C_k so beeinflusst werden, dass alle Verben, die später als Ausgabe vom Modell ausgegeben werden, in der Vergangenheitsform konjugiert werden.

Bemerkenswert ist, dass keine derartigen input-basierten Anpassungen an der A-Matrix vorgenommen werden. Seine Rolle bleibt die gleiche wie bei den S4-Modellen: den gesamten Verlauf früherer Eingaben effizient zu speichern. Die Rolle der Bestimmung,welche Teile dieser Historie zu einem bestimmten Zeitpunkt genutzt werden sollen, wird von den Matrizes B und C übernommen.

Sobald das Modell jedoch nicht mehr zeitinvariant ist, kann es die Faltungsverknüpfung während des Trainings nicht mehr verwenden, da der Übergangskern nicht mehr konstant ist: Die Krux des Selektivitätsmechanismus besteht darin, dass der Übergang von h_t-1 zu h_t nun kontextabhängig ist.

Mamba nutzt nun stattdessen eine clevere Umgehung, um ähnliche Vorteile zu erzielen. Da das SSM nur Multiplikation und Addition verwendet, unterliegen seine Berechnungen der bekannten assoziativen Eigenschaft der Mathematik: Sie können auf unterschiedliche Weise gruppiert werden, ohne das Endergebnis zu verändern. Dadurch können die vielen sequentiellen Berechnungen in kleine, unabhängige Abschnitte unterteilt werden, die von einer GPU durch einen Scan der parallelen Präfixsummen verarbeitet werden können.

Darüber hinaus werden die Ergebnisse auf eine spezielle hierarchische Weise kombiniert, die verschiedenen Arten von Hardwarespeichern auf einer GPU optimal effizient nutzt. Dabei werden Prinzipien angewendet, die den FlashAttention-Techniken ähneln, die ebenfalls von Tri Dao entwickelt wurden und heute in modernen LLMs allgegenwärtig sind.  

Innerhalb der Mamba-Architektur dient das S6-Modell als ein Modul des größeren „Mamba-Blocks“, ähnlich wie der Aufmerksamkeitsmechanismus als ein Modul innerhalb des größeren „Aufmerksamkeits-Blocks“ dient. Es kombiniert das S6-Modul mit einer Gated Neural Network-Architektur. Mamba-Modelle umfassen in der Regel mehrere Mamba-Blöcke – d. h. eine Reihe aufeinanderfolgender Mamba-Schichten in einem neuronalen Netzwerk – vor der Ausgabeschicht, die die endgültige Ausgabe des Modells vornimmt.

Vor dem Eintritt in den Mamba-Block wird eine Kopie der Eingabe als Residualverbindung direkt an das Ende gesendet.Das Innenleben des Mamba-Blocks dient nicht nur dazu, festzustellen, welche Teile des größeren Kontexts für diese Eingabe relevant sind, sondern auch dazu, zu bestimmen , wie stark diese Kontextinformationen die ursprüngliche Bedeutung der Eingabe verändern sollen.

Innerhalb des Mamba-Blocks wird der ursprüngliche Eingabevektor wie folgt verarbeitet:

• Zunächst wird die Eingabe durch eine lineare Schicht geleitet, die doppelt so breit ist wie der Eingabevektor selbst, und in einen höherdimensionalen Raum projiziert. Wenn das Modell beispielsweise ursprünglich jedes Eingabe-Token x als 512-dimensionale Vektoreinbettung darstellt, wird es durch Multiplikation von x mit den Gewichten der linearen Projektionsschicht zu einem 1024-dimensionalen Vektor erweitert.
  
  
• Als nächstes wird der erweiterte Vektor in zwei Teile geteilt. Eine Hälfte (die wir x_proj nennen) wird in den Pfad eingespeist, der durch das SSM führt, und die andere Hälfte (die wir z_proj nennen) wird in einen separaten Pfad eingespeist, der durch einen Gating-Mechanismus führt. Der Übersichtlichkeit halber wird der vorherige Erweiterungsschritt in der Regel so dargestellt, dass er von 2 parallelen linearen Schichten ausgeführt wird.
  
  
• Bevor x_proj das SSM erreicht, wird es in eine 1-dimensionale Faltungsschicht eingespeist. Diese Faltungsschicht extrahiert lokale Muster (wie Abhängigkeiten zwischen benachbarten Token, z. B. einfache Verb-Subjekt-Paare). Dadurch kann sich das SSM auf das kontextuelle Verständnis weitreichender, globaler Abhängigkeiten „fokussieren“.
  
  
• Der Ausgang der Faltungsschicht dient als Eingabe für eine nichtlineare Aktivierungsfunktion. Die Einführung von Nichtlinearität ist ein Markenzeichen aller neuronaler Netze, die es ihnen ermöglicht, komplexere Muster zu erfassen. Das Mamba-Papier verwendet die Sigmoid Linear Unit (SiLU). Wir nennen den daraus resultierenden Vektor x_act.
  
  
• Unterdessen wird z_proj im Pfad des separaten Gating-Mechanismus auch als Eingabe in eine nichtlineare Aktivierung eingegeben, was z_act ergibt.
  
  
• Im SSM-Pfad wird x_act in drei parallele lineare Projektionsschichten eingegeben, die die jeweiligen Werte für ∆ _x, B_x und C_x ergeben.
  
  
• Das SSM verwendet diese eingabeabhängigen Parameter (sowie die A- und D-Matrizes), um die Zustandsraumaktualisierung und den Output y des SSM zu berechnen.
  
  
• Die Ausgabe y des SSM wird nun elementweise mit der Ausgabe z_act des Gating-Pfades multipliziert. Im Wesentlichen verhält sich jedes Element in z_act wie ein Lautstärkeregler an einem Audiomischpult: Wenn ein bestimmtes Element von z_act nahe Null liegt, ergibt die Multiplikation mit dem entsprechenden Teil von y einen Wert, der näher an Null liegt, und sein Einfluss wird verringert. Umgekehrt wird, wenn ein gegebenes Element von zact groß ist, durch die Multiplikation mit y der Einfluss seiner Kontextinformationen verstärkt.
  
  
• Der daraus resultierende Vektor wird wieder auf seine ursprüngliche Größe projiziert. Er kann als Vektor gewichteter kontextueller Aktualisierungen (oder Nicht-Aktualisierungen) für jedes der Elemente des ursprünglichen Eingabevektors verstanden werden.
  
  
• Schließlich wird dieser Aktualisierungsvektor zu der Kopie des ursprünglichen Eingabevektors hinzugefügt, der als Residualverbindung direkt an das Ende des Blockes gesendet wurde.
  
  
• Der ursprüngliche Eingabevektor wurde nun aktualisiert, um das kontextuelle Verständnis des selektiven SSM widerzuspiegeln. Es kann nun an die nächste Mamba-Schicht gesendet werden oder in den letzten Schichten des Modells als Eingabe für eine Softmax-Funktion dienen, die die entsprechende Wahrscheinlichkeit als Ausgabe ausgibt, dass der vollständig aktualisierte Vektor jedem Wort im Vokabular des Modells entspricht.

Ein Jahr nach dem ursprünglichen Mamba-Bericht veröffentlichten Dao und Gu „Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality“. Dieses Paper enthielt drei wesentliche Beiträge:

• Eine Erkundung der theoretischen Verbindungen zwischen Mamba und Transformer und ein gemeinsames Vokabular für die beiden Architekturen
• Eine Reihe von Klarstellungen und Erkundungen verschiedener Designentscheidungen für Mamba-Modelle
• Eine modifizierte Architektur, Mamba-2, die auf diesen Designuntersuchungen basiert und verbessert wurde

Der Mamba-2-Algorithmus ist deutlich schneller und einfacher zu implementieren als das ursprüngliche Mamba: Die Autoren stellten eine „minimale SSD“-Codebasis zur Verfügung, die das selektive SSM in etwa 25 Codezeilen implementiert.² Diese Effizienz ermöglicht es Mamba-2, viel größere Dimensionen versteckter Zustände zu verwenden, ohne das Modell zu verlangsamen, was die Entwicklung größerer, leistungsfähigerer und ausdrucksstärkerer Modelle mit der Architektur ermöglicht. Bei den Tests entsprachen die Mamba-2-Modelle bei einer Reihe von nachgelagerten Aufgaben definitiv den Mamba- und Transformer-Modellen entsprechender Größe oder übertrafen diese.

Wie in der Einleitung des Artikels dargelegt, bestand das Hauptziel von Dao und Gu darin, „eine umfassende Sammlung theoretischer Zusammenhänge zwischen strukturierten SSMs und Varianten der Aufmerksamkeit zu entwickeln“. Dadurch entstand ein neuer konzeptioneller Framework, der die beiden vereint, den sie „State Space Duality“ (SSD) nannten.³ Auf diese Weise ermöglichten sie Mamba, einen Vorteil von der jahrelangen Erforschung und Optimierung der Transformer-Architektur zu haben.

Ein bemerkenswerter Nutzen war die Entwicklung eines Mamba-Äquivalents für Mehrkopf-Aufmerksamkeit (MHA), bei dem ein Mamba-Block in mehrere „Mamba-Köpfe“ aufgeteilt werden kann, ähnlich den multiplen „Aufmerksamkeitsköpfen“ in Transformer. Eine Variante dieses Ansatzes, die sie als analog zur gruppierten Abfrageaufmerksamkeit betrachteten, ermöglicht noch mehr Effizienz durch Tensor-Parallelität in GPUs.

Im Mamba-2-Block – den sie als parallelen Mamba-Block bezeichnet (im Gegensatz zum ursprünglichen „sequentiellen Mamba-Block“) – werden die eingabeabhängigen Parameter ∆, B und C in der anfänglichen Projektionsschicht parallel erzeugt. Insbesondere B und C werden durch einfaches Kopieren von Teilen von x_proj abgeleitet, anstatt durch Multiplikation von x_proj durch dedizierte lineare Schichten. Neben der Vereinfachung und Reduzierung der Gesamtmodellparameter ermöglicht diese Parallelität ein deutlich effizienteres Training im großen Maßstab.⁴

Sowohl Mamba als auch Transformers haben ihre jeweiligen Stärken, aber Mamba-basierte Modelle sind in allen Belangen, die mit der Speichernutzung und der Geschwindigkeit zu tun haben, im Allgemeinen überlegen: Laut dem Mamba-Papier bietet Mamba einen 5-mal höheren Durchsatz als entsprechende Transformers.

Transformer sind unglaublich präzise und vielseitig, stellen aber auch enorme Anforderungen an die Rechenressourcen. Während des Vortrainings (und der Feinabstimmung) skalieren die Gedächtnisanforderungen der Selbstaufmerksamkeit quadratisch mit der Sequenzlänge: Wenn Sie die Kontextlänge einer Sequenz verdoppeln, verbraucht der Aufmerksamkeitsmechanismus das Vierfache der Ressourcen. Dieser „quadratische Engpass“ drosselt die Geschwindigkeit und Speicherverfügbarkeit zunehmend, wenn das Kontextfenster größer wird. Während der Inferenz skaliert ihr Speicherbedarf linear.

Während des Trainings skaliert der Speicherverbrauch eines Mamba-Modells nur während des Trainings linear. Noch wichtiger ist, dass der Speicherverbrauch während der Inferenz konstant ist: Unabhängig davon, wie viele Token das Modell gesehen hat, behält der SSM eine Darstellung seiner Eingabehistorie mit fester Größe bei. Dies ermöglicht theoretisch eine unbegrenzte Kontextlänge, die nur durch Hardwarebeschränkungen eingeschränkt ist.

Allerdings hat die speicherintensivere und rechenredundantere Methode von Transformern ihre eigenen Vorteile. So hat beispielsweise die Forschung gezeigt, dass Transformer sowohl Mamba als auch Mamba-2 bei Aufgaben, die kontextbezogenes Lernen erfordern (z. B. Few-Shot-Prompting), Kopieren oder Long-Context-Reasoning immer noch übertreffen.

Glücklicherweise schließen sich die jeweiligen Stärken von Transformern und Mamba nicht gegenseitig aus. Das Mamba-2-Paper legt nahe, dass ein Hybridmodell sowohl reine Transformer als auch SSMs übertreffen könnte – eine Annahme, die später im Jahr 2024 durch NVIDIA-Forschung formal bestätigt wurde.⁵ Im Großen und Ganzen scheinen Hybridmodelle die Effizienz-Nutzen von Mamba mit der Nuancen- und Kontextlernleistung zu kombinieren, die durch den ressourcenintensiveren Aufmerksamkeitsmechanismus der Transformer bereitgestellt wird.

Um dies weiter zu untersuchen, arbeitete IBM Research mit Dao und Gu sowie mit Minjia Zhang von der University of Illinois in Urbana-Champaign (UIUC) an Bamba und Bamba V2 zusammen. Bamba wiederum hat viele der architektonischen Elemente von IBM Granite 4.0 beeinflusst.

Die Erforschung hybrider Modelle ist weiterhin ein Bereich aktiver Forschung, insbesondere innerhalb der Open-Source-Community.