Die Entwicklung von Large Language Models (LLMs), die auf decoder-only-Transformer-Modellen basieren, hat eine entscheidende Rolle bei der Transformation des Natural Language Processing (NLP)-Bereichs gespielt und verschiedene Deep-Learning-Anwendungen wie Reinforcement Learning, Zeitreihenanalyse, Bildverarbeitung und vieles mehr vorangetrieben. Dennoch stellen diese decoder-only-Transformer-Modelle trotz ihrer Skalierbarkeit und starken Leistung erhebliche Nachteile dar. Obwohl sie ausdrucksstark sind, erfordert der Aufmerksamkeitsmechanismus in Transformer-Modellen hohe Rechenressourcen während des Trainings und der Inferenz, was eine erhebliche Menge an Speicher für die Sequenzlänge und quadratische FLOPs erfordert. Diese hohen Rechenanforderungen begrenzen die Kontextlänge der Transformer-Modelle, machen autoregressive Generierungsaufgaben proportional teuer mit der Skalierung und behindern die Fähigkeit der Modelle, aus kontinuierlichen Datenströmen zu lernen und Sequenzen unbegrenzter Länge effizient zu verarbeiten.

In jüngster Zeit haben State Space Models (SSMs) bemerkenswerte Fähigkeiten und Leistungen gezeigt und konkurrieren mit Transformer-Architektur-Modellen in großen Modellierungsbenchmarks, während sie eine lineare Speicherkomplexität in Bezug auf die Sequenzlänge erreichen, im Gegensatz zu der quadratischen Komplexität von Transformern. Darüber hinaus hat Mamba, ein kürzlich veröffentlichtes State Space Model, herausragende Leistungen in einer Reihe von Sprachmodellierungsaufgaben und der Verarbeitung langer Sequenzen gezeigt. Gleichzeitig haben Mixture-of-Expert-Modelle (MoE) beeindruckende Leistungen gezeigt, während sie die Latenz und die Rechenkosten der Inferenz erheblich reduzierten, allerdings auf Kosten eines größeren Speicherbedarfs. Basierend auf Mamba- und MoE-Modellen wird in diesem Artikel BlackMamba diskutiert, eine neue Architektur, die das Mamba-State-Space-Modell mit MoE-Modellen kombiniert, um die Vorteile beider Frameworks zu nutzen. Experimente mit BlackMamba haben gezeigt, dass es in der Lage ist, das bestehende Mamba-Framework und Transformer-Baselines in Bezug auf Trainings-FLOPs und Inferenz zu übertreffen. Die außergewöhnliche Leistung des BlackMamba-Frameworks zeigt, dass es die Fähigkeiten des Mamba- und MoE-Frameworks effektiv kombinieren kann, indem es die schnelle und kostengünstige Inferenz von MoE mit linearer Komplexität von Mamba verbindet.

Dieser Artikel zielt darauf ab, das BlackMamba-Framework in-depth zu beleuchten. Wir erkunden den Mechanismus, die Methodik und die Architektur des Frameworks, sowie seinen Vergleich mit state-of-the-art-Bild- und Video-Generierungsframeworks. Lassen Sie uns beginnen.

BlackMamba: Eine Einführung in MoE für State Space Modelle

Die Entwicklung von Large Language Models (LLMs), insbesondere solchen, die auf decoder-only-Transformer-Architekturen basieren, hat den Natural Language Processing (NLP)-Bereich und verschiedene Deep-Learning-Anwendungen wie Reinforcement Learning, Zeitreihenanalyse, Bildverarbeitung und vieles mehr vorangetrieben. Dennoch stellen diese decoder-only-Transformer-Modelle trotz ihrer Skalierbarkeit und starken Leistung erhebliche Nachteile dar. Der Aufmerksamkeitsmechanismus, ein wichtiger Bestandteil von Transformer-Modellen, erfordert umfangreiche Rechenressourcen für das Training und die Inferenz. Dies erfordert eine erhebliche Menge an Speicher, die mit der Sequenzlänge und quadratischen FLOPs anwächst. Diese hohen Rechenanforderungen begrenzen die Kontextlänge der Transformer-Modelle, machen autoregressive Generierungsaufgaben proportional teuer mit der Skalierung und behindern die Fähigkeit der Modelle, aus kontinuierlichen Datenströmen zu lernen und 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 alternativer Architekturen zu den kanonischen dichten Aufmerksamkeits-Transformer-Modellen gerichtet, wobei State Space Modelle und MoE-Modelle die vielversprechendsten Kandidaten-Architekturen sind. Der Hauptvorteil, der durch die Verwendung von State Space Modellen anstelle von Transformer-Architektur-Modellen erzielt wird, ist die lineare Rechenkomplexität in Bezug auf die Sequenzlänge, die von State Space Modellen im Gegensatz zu der quadratischen Komplexität von Transformern angeboten wird. Theoretisch ermöglicht die lineare Rechenkomplexität in Bezug auf die Sequenzlänge es State Space Modellen, größere Sequenzen als Transformer-Architektur-Modelle zu verarbeiten, und rendert autoregressive Generierung konstant in der Berechnung ohne einen KV-Cache. Kürzlich entwickelte State Space Modelle, einschließlich Mamba, RetNet und einigen anderen, haben effiziente lange Sequenzinferenz und Training sowie wettbewerbsfähige Sprachmodellierungsaufgaben im Vergleich zu Transformern mit ähnlichen Skalierungseigenschaften gezeigt. Andererseits haben Mixture-of-Expert-Modelle-Architekturen an Popularität gewonnen, da sie eine erhebliche Reduzierung der Inferenz- und Trainings-FLOPs ermöglichen, die für die Erzielung einer vergleichbaren Qualität zu einem dichten Modell erforderlich sind. MoE-Modelle (Mixture of Experts) operieren, indem sie nur eine sparse Auswahl der Gesamtparameter während eines einzelnen Vorwärtsdurchlaufs aktivieren. Sie verwenden eine Routing-Funktion, um zu bestimmen, welche “Experten” basierend auf dem gegebenen Kontext aktiviert werden. Dieser Ansatz schafft eine Trennung zwischen den Rechenkosten der Inferenz und der Gesamtzahl der Parameter, was eine verbesserte Leistung innerhalb eines festen Inferenzbudgets ermöglicht, allerdings mit einem erhöhten Speicherbedarf.

Diese Entwicklung in der Architektur bietet erhebliche Vorteile gegenüber traditionellen Transformern und stellt eine spannende Richtung für weitere Entwicklungen dar. Wir gehen davon aus, dass die Integration dieser Verbesserungen in ein kombiniertes Mamba-MoE-Modell die Sprachmodellierungsfähigkeiten und -effizienz erheblich über die von Standard-Transformer-Modellen hinaus steigern könnte. Die erwarteten Vorteile einer Mamba-MoE-Architektur im Vergleich zu einem traditionellen dichten Transformer-Modell umfassen:

Mamba: Erreicht lineare Rechenkomplexität in Bezug auf die Sequenzlänge für beide Phasen, Training und Inferenz. Es ermöglicht autoregressive Generierung in konstanter Zeit und mit konstantem Speicherbedarf.

MoE: Bietet Inferenzgeschwindigkeit und TrainingsrechenEffizienz, die mit einem kleineren, dichten Baseline-Modell vergleichbar sind, während es eine Modellqualität beibehält, die mit der eines Modells mit einer äquivalenten Anzahl von Parametern wie das dichtere Modell vergleichbar ist.

Es ist jedoch wichtig zu betonen, dass Transformer-Architektur-Modelle immer noch state-of-the-art sind und konsistente und bemerkenswerte starke Leistungen bei Sprachmodellierungsaufgaben und Sequenzverarbeitungsaufgaben gezeigt haben. Im Kern verwendet die Transformer-Architektur Selbstaufmerksamkeit, die eine quadratische all-to-all-Vergleich der Dot-Produkt-Ähnlichkeiten zwischen den Einbettungen verschiedener Token in einer Sequenz durchführt und eine lineare Abbildung auf einen Ausgabevector durchführt. Das Transformer-Modell besteht aus Selbstaufmerksamkeitsblöcken, die zwischen MLP- oder Multi-Layer-Perceptron-Blöcken gestapelt sind, die weiterhin aus einem zweischichtigen MLP mit einer gegebenen Aktivierungsfunktion bestehen.

BlackMamba: Architektur und Methodik

State Space Modelle

State Space Modelle gehören zu der Gruppe von Sequenzmodellen mit linearer Komplexität in Bezug auf die Länge der Eingabesequenz. Die Architektur von State Space Modellen stimmt mehr mit rekurrenten neuronalen Netzen und Convolutional Neural Networks überein als mit Aufmerksamkeits-basierten Architekturen und ist von einem kontinuierlichen dynamischen System inspiriert, das eine 1-dimensionale Funktion durch einen impliziten latenten Raum abbildet. Ein lineares dynamisches System ermöglicht effiziente parallele Berechnungen mithilfe eines assoziativen oder eines Convolution-Scans. In der Praxis war die rekurrente Natur von State Space Modellen der Grund, warum sie noch nicht auf hochparallelen AI-Hardware wie GPUs übernommen wurden. Die Entwicklung von State Space Modellen wie RWKV und Mamba hat jedoch parallele Scan-Kernel verwendet, um rekurrente Operationen effizient auf GPUs abzubilden, wodurch das Training neuer Architekturen mit einer Effizienz ermöglicht wurde, die mit der von Transformer-Modellen vergleichbar ist.

Die inhärente quadratische Komplexität in Bezug auf die Sequenzlänge in Transformern ist eine bekannte Einschränkung, die die Vernunft und das Verständnis über sehr lange Kontexte behindert. Kürzliche Innovationen haben die Idee der Erweiterung der Kontextlänge eingeführt, um es Transformern zu ermöglichen, auf einer machbaren Skala trainiert zu werden, bevor sie auf viel längere Kontexte während der Inferenz angewendet werden. Trotz dieser Fortschritte erfordert der Inferenzprozess immer noch eine erhebliche Menge an Rechenressourcen und Speicher, insbesondere für die Aufrechterhaltung des Key-Value-Caches, was es zu einem ressourcenintensiven Unterfangen macht. Kürzliche Forschungsbemühungen haben sich auf die Verbesserung der ausdrucksstarken Fähigkeiten von State Space Modellen konzentriert, indem sie eingegebene Abhängigkeitsmechanismen ähnlich den Query-, Key-, Value-Matrizen in Aufmerksamkeitsmechanismen einbeziehen.

Diese Bemühungen zielen darauf ab, die inhärente lineare Progression der State Space Rekursion zu erhalten, was eine effiziente Ausführung durch entweder Convolution oder einen selektiven Scan-Prozess ermöglicht. Dieser Ansatz verringert die Leistungsunterschiede mit Transformern in praktischen Anwendungen erheblich. Unter diesen Fortschritten ragt Mamba als ein State Space Modell hervor, das die Ziele vorheriger Forschung widerspiegelt und beeindruckende Leistungen zeigt, die mit denen von Transformern bei Skalen von bis zu 2,8 Milliarden Parametern vergleichbar sind. Es erreicht dies, indem es eingegebene Abhängigkeitsmechanismen auf die Eingaben der State Space Modell-Rekursion anwendet, während es gleichzeitig eine effiziente Berechnung durch den Einsatz von speziellen selektiven Scan-Kernels ermöglicht.

Mixture of Expert Modelle

Mixture-of-Expert-Modelle (MoE) erreichen eine Trennung zwischen den Inferenzkosten und der Gesamtzahl der Parameter, indem sie Parameter während des Vorwärtsdurchlaufs selektiv aktivieren. Anstatt alle Parameter zu verwenden, leiten diese Modelle Token an bestimmte Multilayer-Perceptron-Experten weiter. Idealisiert ist jeder Experte auf die Verarbeitung einer bestimmten Art von Eingabe ausgerichtet, wobei ein Routing-Mechanismus, im Wesentlichen ein kompaktes neuronales Netzwerk, den am besten geeigneten Experten für jedes Token bestimmt. Dieser Ansatz zielt darauf ab, die umfassende ausdrucksstarke Kraft eines Modells mit einer äquivalenten Anzahl von Parametern wie das dichtere Modell zu bewahren, jedoch mit erheblich reduzierten Rechenanforderungen. Typischerweise ist der Router eine Abbildung der linearen Schichten von Token zu Experten-Indizes, wobei jeder Experte einfach ein Standard-Transformer-Multilayer-Perceptron ist. Entwickler müssen jedoch noch die optimale Trainingsmethode für den Router finden, da das Experten-Zuweisungsproblem nicht differenzierbar ist und Mixture-of-Expert-Modelle oft mit Lastausgleichs- und Trainingsstabilitätsproblemen bei der Hardware-Effizienz zu kämpfen haben.

Architektur

Im Kern verwendet BlackMamba ein Standard-Transformer-Modell, das aus abwechselnden MLP-Blöcken und Aufmerksamkeitsblöcken besteht, die in einer Residual-Stream-Sequenz angeordnet sind. Die meisten Mixture-of-Expert-Modelle ersetzen einfach die Multilayer-Perceptron-Blöcke durch eine geroutete Expertenschicht. Andererseits ersetzt das BlackMamba-Framework nicht nur den Multilayer-Perceptron-Block im Transformer durch eine geroutete Expertenschicht, sondern ersetzt auch die Aufmerksamkeitsschicht durch eine Mamba-State-Space-Modell-Schicht. Die Architektur des BlackMamba-Frameworks wird in der folgenden Abbildung gezeigt.

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-Multilayer-Perceptrons. Das Framework trainiert mit 8 Experten, eine Zahl, die die Entwickler als richtige Balance und Kompromiss zwischen dem Speicherbedarf und den Inferenzkosten des Modells fanden. Der benutzerdefinierte Datensatz, der zum Trainieren des BlackMamba-Frameworks verwendet wird, besteht aus einer Mischung aus bereits existierenden Open-Source-Datensätzen, einschließlich Starcoder, SlimPajama, Pile und mehr. Die folgende Tabelle zeigt die Gewichte jedes Datensatzes, der zum Trainieren des BlackMamba-Frameworks verwendet wird. Insgesamt gibt es 1,8 Billionen Token im Datensatz.

BlackMamba: Ergebnisse

Um einen fairen Vergleich zwischen Mamba und BlackMamba zu ermöglichen, haben die Entwickler beide Modelle mit denselben Trainingsparametern auf denselben Trainingsdaten trainiert. Das BlackMamba-Framework kann sowohl Mamba als auch Transformer-Modelle bei identischer Modellgröße während der Inferenz und des Trainings in Bezug auf Floating-point-Operationen pro Sekunde übertreffen. Die folgende Abbildung zeigt die Zeit, die benötigt wird, um eine Sequenz einer bestimmten Länge autoregressiv aus einem initialen Ein-Token-Prompt zu generieren, als Funktion der Sequenzlänge.

Darüber hinaus werden die Latenzvorteile sowohl von Mixture-of-Expert- als auch von Mamba-Modellen im BlackMamba-Framework kombiniert, was zu erheblich 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, was BlackMamba bei der Generierung langer Sequenzen extrem effektiv macht. Wenn man weitermacht, zeigt die folgende Abbildung die Anzahl der Token, die den BlackMamba-Modellen mit 340 Millionen und 640 Millionen Parametern zugewiesen werden. Wie man sehen kann, zeigen die meisten Schichten aufgrund des verbesserten Sinkhorn-Algorithmus, der von den BlackMamba-Modellen implementiert wird, ein hohes Maß an Expertenbalance.

Die folgende Tabelle zeigt die Bewertungsergebnisse des BlackMamba-Frameworks im Vergleich zu einer Reihe von Open-Source-Vorabtrainierten Sprachmodellen. Wie man sehen kann, kann das BlackMamba-Framework mit der Mehrheit der Frameworks über alle Baselines konkurrieren und übertreffen. Darüber hinaus ist es erwähnenswert, dass die Modelle, die BlackMamba übertreffen, erheblich mehr Parameter haben und die Leistungsunterschiede minimal sind, was die Fähigkeit des BlackMamba-Frameworks mit weniger Parametern zeigt.

Schlussgedanken

In diesem Artikel haben wir über BlackMamba gesprochen, eine neue Architektur, die das Mamba-State-Space-Modell mit Mixture-of-Expert-Modellen kombiniert, um die Vorteile beider Frameworks zu nutzen. Experimente mit BlackMamba haben gezeigt, dass es in der Lage ist, das bestehende Mamba-Framework und Transformer-Baselines in Bezug auf Trainings-FLOPs und Inferenz zu übertreffen. Die außergewöhnliche Leistung des BlackMamba-Frameworks zeigt, dass es die Fähigkeiten des Mamba- und MoE-Frameworks effektiv kombinieren kann, indem es die schnelle und kostengünstige Inferenz von MoE mit linearer Komplexität von Mamba verbindet. Wir haben darüber gesprochen, wie die Architektur des BlackMamba-Frameworks in der Lage ist, starke trainierte Large Language Modelle, bestehende Mamba-Frameworks und Mixture-of-Expert-Modelle in Bezug auf Trainings-FLOPs und Inferenzkosten zu übertreffen. Darüber hinaus erbt das BlackMamba-Framework die Generierungs-FLOPs und das reduzierte Training von beiden Mixture-of-Expert-Modellen und Mamba-Frameworks gleichzeitig.