AI21 Labs stärkt Generative AI Transformers durch Zusammenarbeit mit Jamba

Seit der Veröffentlichung des bahnbrechenden Forschungsartikels "Attention is All You Need" im Jahr 2017 haben Transformer eine zentrale Rolle im Bereich der generativen KI eingenommen. Allerdings sind Transformer nicht der einzige brauchbare Ansatz für generative KI. AI21 Labs hat ein neuartiges Framework namens „Jamba“ vorgestellt, das über herkömmliche Transformer hinausgehen möchte.

Jamba kombiniert das Mamba-Modell, das auf dem Structured State Space Model (SSM) basiert, mit der Transformer-Architektur, um eine optimierte generative KI-Lösung zu schaffen. Der Begriff „Jamba“ steht für Joint Attention and Mamba Architecture und zielt darauf ab, die Stärken von SSM und Transformern zu nutzen. Dieses Modell wird als Open Source unter der Apache 2.0-Lizenz veröffentlicht.

Obwohl Jamba nicht als Ersatz für bestehende Transformer-basierte große Sprachmodelle (LLMs) gedacht ist, wird erwartet, dass es in bestimmten Anwendungen wertvolle Ergänzungen bietet. AI21 Labs gibt an, dass Jamba in generativen Schlußfolgerungsaufgaben besseren Nachwuchs als herkömmliche Transformer-Modelle zeigt, was durch Benchmarks wie HellaSwag belegt wird. Allerdings übertrifft es noch nicht die Transformer-Modelle bei kritischen Benchmarks wie dem Massive Multitask Language Understanding (MMLU), der die Problemlösungsfähigkeiten bewertet.

AI21 Labs hat sich auf generative KI für Unternehmensanwendungen spezialisiert und hat im August 2023 155 Millionen US-Dollar gesammelt, um seine Initiativen weiter voranzutreiben. Zu den Unternehmensangeboten gehört Wordtune, ein Tool, das Organisationen unterstützt, Inhalte zu generieren, die mit ihrem Ton und ihrer Marke übereinstimmen. Im Jahr 2023 berichtete das Unternehmen, dass es erfolgreich gegen den generativen KI-Riesen OpenAI bei der Gewinnung von Unternehmenskunden konkurriert hat.

Traditionell hat die LLM-Technologie von AI21 Labs die Transformer-Architektur verwendet, einschließlich der Jurassic-2 LLM-Familie, die Teil der AI21 Studio-Plattform für natürliche Sprachverarbeitung (NLP) ist und über APIs zur Unternehmensintegration verfügbar ist. Jamba stellt jedoch einen Wendepunkt in Richtung eines hybriden SSM- und Transformer-Modells dar.

Trotz der herausragenden Rolle von Transformern in der generativen KI haben sie bestimmte Einschränkungen. Ein erhebliches Problem besteht darin, dass die Inferenz bei größeren Kontextfenstern langsamer wird. Wie die Forscher von AI21 Labs erklären, skaliert der Aufmerksamkeitsmechanismus eines Transformers mit der Sequenzlänge, was zu einer verringerten Durchsatzrate führt, da jedes Token auf die gesamte vorherige Sequenz angewiesen ist. Dies macht Anwendungen mit langen Kontexten ineffizient.

Ein weiteres Problem betrifft den hohen Speicherbedarf, der für das Skalieren von Transformern erforderlich ist. Ihre Speicheranforderungen steigen mit der Kontextlänge, was die Verarbeitung langer Kontexte oder mehrerer paralleler Aufgaben ohne erhebliche Hardware-Ressourcen erschwert. Der SSM-Ansatz zielt darauf ab, diese Kontext- und Speicherprobleme zu lösen.

Die Mamba-SSM-Architektur, die ursprünglich von Forschern der Carnegie Mellon- und Princeton-Universitäten entwickelt wurde, benötigt weniger Speicher und verwendet einen anderen Aufmerksamkeitsmechanismus zur Verwaltung großer Kontextfenster. Allerdings hat sie Schwierigkeiten, die gleiche Ausgabequalität wie Transformer-Modelle zu erzielen. Der hybride Ansatz von Jamba kombiniert die Ressourcen- und Kontextoptimierung von SSM mit den Ausgabefähigkeiten von Transformern.

AI21 Labs behauptet, dass das Jamba-Modell ein Kontextfenster von 256K bietet und für lange Kontexte dreimal so viel Durchsatz wie Mixtral 8x7B erzielt. Bemerkenswert ist, dass Jamba als das einzige Modell seiner Größenklasse positioniert ist, das bis zu 140K Kontext auf einer einzigen GPU verarbeiten kann.

Ähnlich wie Mixtral beinhaltet Jamba ein Mixture of Experts (MoE)-Modell. Jamba nutzt jedoch MoE innerhalb seines hybriden SSM-Transformer-Frameworks, was höhere Optimierungsstufen ermöglicht. Insbesondere aktivieren die MoE-Schichten von Jamba während der Inferenz nur 12 Milliarden seiner verfügbaren 52 Milliarden Parameter, was es effizienter macht als ein reines Transformer-Modell in vergleichbarer Größe, so AI21 Labs.

Derzeit befindet sich Jamba noch in der frühen Phase und gehört noch nicht zu den Unternehmensangeboten von AI21 Labs, jedoch plant das Unternehmen, bald eine Schulungsversion auf der AI21-Plattform in Beta einzuführen.