Die Entwicklung von Large Language Models (LLMs), die aus reinen Decoder-Transformer-Modellen aufgebaut sind, hat eine entscheidende Rolle bei der Transformation des Bereichs Natural Language Processing (NLP) sowie bei der Weiterentwicklung verschiedener Deep-Learning-Anwendungen gespielt Verstärkung lernen, Zeitreihenanalyse, Bildverarbeitung und vieles mehr. Trotz ihrer Skalierbarkeit und starken Leistung weisen LLMs, die nur aus Decoder-Transformatormodellen bestehen, jedoch immer noch erhebliche Mängel auf. Obwohl ausdrucksstark, erfordert der Aufmerksamkeitsmechanismus in von Transformatoren abgeleiteten LLMs hohe Rechenressourcen sowohl während der Inferenz als auch beim Training, was einen erheblichen Speicherbedarf für die Sequenzlänge und quadratische FLOPs erfordert. Dieser hohe Rechenaufwand schränkt die Kontextlänge von Transformatormodellen ein, was autoregressive Generierungsaufgaben mit der Skalierung proportional teurer macht und das Lernen aus kontinuierlichen Datenströmen und die Möglichkeit einer wirklich unbegrenzten Sequenzverarbeitung behindert.

In letzter Zeit, Zustandsraummodelle (SSMs) haben bemerkenswerte Fähigkeiten und Leistung bewiesen, indem sie in groß angelegten Modellierungs-Benchmarks mit Transformator-Architekturmodellen konkurrierten und gleichzeitig Speicherkomplexität als Funktion der Sequenzlänge und linearen Zeit erreichten. Darüber hinaus hat Mamba, ein kürzlich veröffentlichtes State Space Model, eine herausragende Leistung bei einer Reihe von Sprachmodellierungs- und Langsequenzverarbeitungsaufgaben gezeigt. Gleichzeitig haben auch Mixture of Expert (MoE)-Modelle eine beeindruckende Leistung gezeigt und gleichzeitig die Latenz und die Rechenkosten der Inferenz deutlich reduziert, allerdings auf Kosten eines größeren Speicherbedarfs. Aufbauend auf Mamba- und MoE-Modellen wird in diesem Artikel BlackMamba erörtert, eine neuartige Architektur, die das Mamba State Space Model mit MoE-Modellen kombiniert, um die Vorteile beider Frameworks zu nutzen. Experimente mit BlackMamba haben gezeigt, dass es das vorhandene Mamba-Framework und die Transformer-Baselines sowohl bei Trainings-FLOPs als auch bei Inferenzen übertreffen kann. Die außergewöhnliche Leistung des BlackMamba-Frameworks zeigt, dass es die Fähigkeiten der Mamba- und MoE-Frameworks effektiv kombinieren kann und eine schnelle und kostengünstige Inferenz aus MoE mit der Erzeugung linearer Komplexität aus Mamba bietet.

Ziel dieses Artikels ist es, das BlackMamba-Framework ausführlich zu behandeln. Wir untersuchen den Mechanismus, die Methodik und die Architektur des Frameworks sowie seinen Vergleich mit modernsten Bild- und Videogenerierungs-Frameworks. Lass uns anfangen.

BlackMamba: Eine Einführung in MoE für Zustandsraummodelle

Die Weiterentwicklung von Large Language Models (LLMs), insbesondere solchen, die auf reinen Decoder-Transformer-Architekturen basieren, hat die Entwicklung erheblich beeinflusst Natürliche Sprachverarbeitung (NLP) Bereich erweitert und auf verschiedene Deep-Learning-Anwendungen ausgeweitet, darunter Reinforcement Learning, Zeitreihenanalyse, Bildverarbeitung und mehr. Trotz ihrer Skalierbarkeit und robusten Leistung stehen diese transformatorbasierten LLMs, die nur auf einem Decoder basieren, vor erheblichen Herausforderungen. Der Aufmerksamkeitsmechanismus, ein Schlüsselmerkmal der transformatorbasierten Methode LLMss, erfordert umfangreiche Rechenressourcen sowohl für die Inferenz als auch für das Training. Dies erfordert einen mit der Sequenzlänge wachsenden Speicherbedarf und quadratisch zunehmende Rechenoperationen (FLOPs). Solche intensiven Rechenanforderungen schränken die Kontextlänge der Modelle ein, erhöhen die Kosten autoregressiver Generierungsaufgaben bei der Skalierung des Modells und behindern die Fähigkeit der Modelle, aus kontinuierlichen Datenströmen zu lernen oder Sequenzen unbegrenzter Länge effizient zu verarbeiten. 

In den letzten Jahren wurden erhebliche Anstrengungen unternommen, um diese Einschränkungen zu überwinden, und die Aufmerksamkeit wurde auf die Entwicklung architektonischer Alternativen zu den kanonischen Dense-Attention-Transformer-Modellen verlagert, wobei SSMs und MoE-Modelle die vielversprechendsten Kandidatenarchitekturen sind. Der Hauptvorteil, der durch die Bevorzugung von Zustandsraummodellen gegenüber Modellen mit Transformatorarchitektur erzielt wird, ist die lineare Rechenkomplexität in Bezug auf die Länge der Eingangssequenz, die SSMs bieten, im Gegensatz zur quadratischen Komplexität, die Transformatoren bieten. Theoretisch ermöglicht die lineare Rechenkomplexität in Bezug auf die Länge der Eingabesequenz es State Space Models, größere Sequenzen als Transformer-Architektur-Modelle für ein gegebenes FLOPS- oder Gleitkommaoperations-pro-Sekunden-Budget zu verarbeiten und die autoregressive Generierung in der Berechnung ohne KV-Cache konstant zu machen. Kürzlich entwickelte Zustandsraummodelle, darunter Mamba, RetNet und einige andere, haben eine effiziente Inferenz und Schulung über lange Sequenzen sowie eine wettbewerbsfähige Leistung bei Sprachmodellierungsaufgaben gegenüber Transformatoren mit ähnlichen Skalierungseigenschaften gezeigt. Andererseits werden Mixture-of-Expert-Modellarchitekturen als Alternative zu dichten Transformatoren immer beliebter, da sie eine erhebliche Reduzierung der Inferenz- und Trainings-FLOPs ermöglichen, die für das Erreichen einer mit einem dichten Modell vergleichbaren Qualität unerlässlich sind. MoE-Modelle (Mixture of Experts) funktionieren, indem sie während eines einzelnen Vorwärtsdurchlaufs nur eine spärliche Auswahl der Gesamtparameter aktivieren. Sie nutzen eine Routing-Funktion, um basierend auf dem gegebenen Kontext zu bestimmen, welche „Experten“ zum Einsatz kommen. Dieser Ansatz schafft eine Trennung zwischen dem Rechenaufwand der Inferenz und der Gesamtzahl der Parameter und ermöglicht so eine verbesserte Leistung innerhalb eines festen Inferenzbudgets, allerdings mit einer größeren Anzahl von Parametern und einem größeren Speicherbedarf.

Dieser Fortschritt in der Architektur bietet bemerkenswerte Vorteile gegenüber herkömmlichen Transformatoren und stellt eine spannende Richtung für die weitere Entwicklung dar. Wir gehen davon aus, dass die Integration dieser Verbesserungen in ein kombiniertes Mamba-MoE-Modell die Sprachmodellierungsfunktionen und -effizienz erheblich steigern könnte, die über die von Standardtransformatormodellen hinausgehen. Zu den erwarteten Vorteilen einer Mamba-MoE-Architektur im Vergleich zu einem herkömmlichen Dense-Transformer-Modell gehören:

Mamba: Erreicht eine lineare Rechenkomplexität im Verhältnis zur Länge der Eingabesequenz sowohl für die Trainings- als auch für die Inferenzphase. Es ermöglicht die autoregressive Generierung in einem konstanten Zeitrahmen und mit konstanter Speichernutzung.

MoE: Bietet die Inferenzgeschwindigkeit und Trainingsrecheneffizienz, die mit einem kleineren, dichten Basismodell vergleichbar sind, und behält gleichzeitig ein Modellqualitätsniveau bei, das mit dem eines Modells mit der gleichen Anzahl von Parametern wie die dichtere Version mithalten kann.

Vor diesem Hintergrund muss unbedingt festgestellt werden, dass Transformatorarchitekturmodelle immer noch auf dem neuesten Stand sind und eine konsistente und bemerkenswert starke Leistung bei Sprachmodellierungsaufgaben und Sequenzverarbeitungsaufgaben gezeigt haben. Im Kern nutzt die Transformatorarchitektur Selbstaufmerksamkeit, die einen quadratischen Gesamtvergleich der Skalarproduktähnlichkeiten zwischen den Einbettungen verschiedener Token in einer Sequenz durchführt und eine lineare Abbildung auf einen Ausgabevektor durchführt. Das Transformatormodell besteht aus Selbstaufmerksamkeitsblöcken, die zwischen MLP- oder Multi-Layer-Perceptron-Blöcken gestapelt sind, die außerdem aus einem zweischichtigen MLP mit einer bestimmten Aktivierungsfunktion bestehen. 

BlackMamba: Architektur und Methodik

Zustandsraummodelle

Zustandsraummodelle gehören zur Gruppe der Sequenzmodelle mit linearer Komplexität in Bezug auf die Länge der Eingabesequenz. Die Architektur von State Space Models orientiert sich eher an wiederkehrenden neuronalen Netzen und Faltungs-neuronalen Netzen als an aufmerksamkeitsbasierter Architektur und ist von einem kontinuierlichen dynamischen System inspiriert, das eine eindimensionale Funktion durch einen impliziten latenten Raum abbildet. Ein lineares dynamisches System macht parallele Berechnungen mithilfe eines assoziativen Scans oder eines Faltungsscans effizient. In praktischen Szenarien war die wiederkehrende Natur von Zustandsraummodellen der Grund dafür, dass sie immer noch auf hochparalleler KI-Hardware wie GPUs übernommen werden müssen. Das Aufkommen von SSMs wie RWKV und Mamba haben parallele Scan-Kernel verwendet, um wiederkehrende Vorgänge effizient auf GPUs abzubilden und so das Training neuartiger Architekturen mit einer Effizienz zu ermöglichen, die mit der von Transformatormodellen erreichten Effizienz vergleichbar ist. 

Die inhärente quadratische Komplexität im Verhältnis zur Sequenzlänge innerhalb von Transformatoren ist eine bekannte Einschränkung, die das Denken und Verstehen in sehr langen Kontexten erschwert. Jüngste Innovationen haben die Idee eingeführt, die Kontextlänge zu erweitern, wodurch Transformatoren in einem realisierbaren Umfang trainiert werden können, bevor sie während der Inferenz auf viel längere Kontexte angewendet werden. Trotz dieser Fortschritte erfordert der Inferenzprozess immer noch eine beträchtliche Menge an Rechenressourcen und Speicher, insbesondere für die Verwaltung des Schlüsselwert-Cache (KV), was ihn zu einem ressourcenintensiven Unterfangen macht. Jüngste Forschungsbemühungen konzentrierten sich auf die Verbesserung der Ausdrucksfähigkeiten von Zustandsraummodellen durch die Einbeziehung eingabeabhängiger Gating-Mechanismen, ähnlich den Abfrage-, Schlüssel- und Wertmatrizen (QKV), die in Aufmerksamkeitsmechanismen zu finden sind. 

Diese Bemühungen zielen darauf ab, den inhärent linearen Verlauf der Zustandsraumrekursion beizubehalten und eine effiziente Ausführung entweder durch Faltung oder einen selektiven Scanprozess zu ermöglichen. Dieser Ansatz verringert die Leistungsunterschiede zu Transformatoren in praktischen Anwendungen erheblich. Unter diesen Fortschritten sticht Mamba als Zustandsraummodell hervor, das die Ziele früherer Forschung widerspiegelt und beeindruckende Leistungsniveaus zeigt, die mit Transformatoren auf Skalen von bis zu 2.8 Milliarden Parametern vergleichbar sind. Dies wird erreicht, indem eingabeabhängiges Gating auf die Eingaben der Rekursion des Zustandsraummodells (SSM) angewendet wird und gleichzeitig eine effiziente Berechnung durch die Verwendung maßgeschneiderter selektiver Scan-Kernel gewährleistet wird.

Mischung aus Expertenmodellen

Mixture of Expert (MoE)-Modelle erreichen eine Trennung zwischen den Inferenzkosten und der Gesamtparameteranzahl durch selektive Aktivierung von Parametern während des Vorwärtsdurchlaufs. Anstatt alle Parameter zu verwenden, leiten diese Modelle Token an bestimmte Experten für Multilayer Perceptron (MLP) weiter. Im Idealfall ist jeder Experte darauf zugeschnitten, eine bestimmte Art von Eingabe zu verarbeiten, wobei ein Routing-Mechanismus, im Wesentlichen ein kompaktes neuronales Netzwerk, den am besten geeigneten Experten für jedes Token ermittelt. Dieser Ansatz zielt darauf ab, die umfassende Ausdruckskraft eines Modells mit einer äquivalenten Anzahl von Parametern in einer dichteren Konfiguration zu bewahren, jedoch mit erheblich geringerem Rechenaufwand. Typischerweise ist der Router eine Abbildung der linearen Schichten von Token auf Expertenindizes, wobei jeder Experte einfach ein Standardtransformator-Mehrschicht-Perzeptron ist. Entwickler müssen jedoch noch die optimale Trainingsmethode für den Router herausfinden, da das Expertenzuweisungsproblem nicht differenzierbar ist und Mixture-of-Expert-Modelle häufig Probleme mit der Lastverteilung und Trainingsstabilität zwischen verschiedenen Experten im Hinblick auf die Hardwareeffizienz haben. 

Architektur

Im Kern verwendet BlackMamba ein Standardtransformatormodell, das aus verschachtelten MLP-Blöcken und Aufmerksamkeitsblöcken besteht, die nacheinander entlang eines Reststroms hinzugefügt werden. Mittlerweile ersetzen die meisten Mixture-of-Expert-Modelle einfach die mehrschichtigen Perzeptronblöcke durch eine geroutete Expertenschicht. Andererseits ersetzt das BlackMamba-Framework nicht nur den mehrschichtigen Perzeptronblock im Transformator durch eine geroutete Expertenschicht, sondern auch die Aufmerksamkeitsschicht durch eine Mamba State Space Model-Schicht. Die Architektur des BlackMamba-Frameworks wird in der folgenden Abbildung veranschaulicht. 

Training und Datensatz

Das BlackMamba-Modell wird auf über 300 Milliarden Token in einem benutzerdefinierten Datensatz trainiert und verwendet die SwiGLU-Aktivierungsfunktion für die Experten-Mehrschicht-Perzeptrone. Das Framework trainiert mit 8 Experten, eine Zahl, die die Entwickler als die richtige Balance und den richtigen Kompromiss zwischen dem Speicherbedarf und den Inferenzkosten des Modells erachteten. Der benutzerdefinierte Datensatz, der zum Trainieren des BlackMamba-Frameworks verwendet wird, besteht aus einer Mischung bereits vorhandener Open-Source-Datensätze, darunter Starcoder, SlimPajama, Pile und mehr. Die folgende Tabelle zeigt die Gewichtungen der einzelnen Datensätze, die zum Training des BlackMamba-Frameworks verwendet werden. Insgesamt enthält der Datensatz 1.8 Billionen Token. 

BlackMamba: Ergebnisse

Um einen fairen Vergleich zwischen Mamba und BlackMamba zu gewährleisten, haben die Entwickler beide Modelle mit denselben Trainingsparametern auf denselben Trainingsdaten trainiert. Das BlackMamba-Framework ist in der Lage, sowohl Mamba- als auch Transformer-Modelle zu übertreffen, da es zum Inferenzzeitpunkt die gleiche Größe des Vorwärtsdurchlaufmodells aufweist und Gleitkommaoperationen pro Sekunde trainiert. Die folgende Abbildung zeigt die Zeit, die zum autoregressiven Generieren einer Sequenz einer bestimmten Länge aus einer anfänglichen Ein-Token-Eingabeaufforderung als Funktion der Sequenzlänge benötigt wird. 

Darüber hinaus werden die Latenzvorteile der Mixture of Expert- und Mamba-Modelle im BlackMamba-Framework kombiniert, was zu deutlich schnelleren Inferenzzeiten im Vergleich zu Transformer-Modellen, reinen Mamba-Modellen und MoE-Modellen führt. Darüber hinaus ist der Inferenzvorteil des BlackMamba-Frameworks direkt proportional zu den Sequenzlängen, wodurch BlackMamba bei der Generierung langer Sequenzen äußerst effektiv ist. Im weiteren Verlauf zeigt die folgende Abbildung die Anzahl der den BlackMamba-Modellen zugewiesenen Token mit 340 Millionen bzw. 640 Millionen Parametern. Wie man sehen kann, weist ein Großteil der Schichten aufgrund des verbesserten Sinkhorn-Algorithmus, der von den BlackMamba-Modellen implementiert wird, ein hohes Maß an Expertenbalance auf. 

Die folgende Tabelle enthält die Bewertungsergebnisse des BlackMamba-Frameworks im Vergleich zu einer Reihe vorab trainierter Open-Source-Sprachmodelle. Wie man beobachten kann, ist das BlackMamba-Framework in der Lage, mit den meisten Frameworks aller Baselines zu konkurrieren und diese zu übertreffen. Darüber hinaus ist anzumerken, dass die Modelle, die BlackMamba übertreffen, eine wesentlich höhere Anzahl von Parametern aufweisen und der Leistungsunterschied minimal ist, was auf die Leistungsfähigkeit des BlackMamba-Frameworks mit weniger Parametern hinweist. 

Abschließende Überlegungen

In diesem Artikel haben wir über BlackMamba gesprochen, eine neuartige Architektur, die das Mamba State Space Model mit Mixture of Expert-Modellen kombiniert, um die Vorteile beider Frameworks zu nutzen. Experimente mit BlackMamba haben gezeigt, dass es das bestehende Mamba-Framework und die Transformer-Baselines sowohl bei Trainings-FLOPs als auch bei Inferenz übertrifft. Die außergewöhnliche Leistung des BlackMamba-Frameworks zeigt, dass es die Fähigkeiten der Mamba- und MoE-Frameworks außergewöhnlich gut erben und kombinieren kann, da es die kostengünstige und schnelle Inferenz von MoE mit der linearen Komplexitätsgenerierung von Mamba kombiniert. Wir haben darüber gesprochen, wie die Architektur des BlackMamba-Frameworks in der Lage ist, stark trainierte Large Language Models, das bestehende Mamba-Framework und Mixture of Expert-Modelle in Bezug auf Trainings-FLOPs und Inferenzkosten zu übertreffen. Darüber hinaus erbt das BlackMamba-Framework gleichzeitig die Generierungs-FLOPs und das reduzierte Training sowohl von Mixture of Expert-Modellen als auch vom Mamba-Framework.