KI-Modelle und Plattformen
Jamba: AI21 Labs’ Neues Hybrid-Transformer-Mamba-Sprachmodell
Unite.AI zu deinen bevorzugten Quellen auf Google hinzufügen
Sprachmodelle haben rasante Fortschritte erlebt, wobei Transformer-basierte Architekturen die Entwicklung in der Verarbeitung von natürlichen Sprachen angeführt haben. Allerdings haben sich mit der Skalierung der Modelle die Herausforderungen bei der Handhabung langer Kontexte, der Speichereffizienz und der Durchsatzleistung verschärft.
AI21 Labs hat mit Jamba eine Lösung vorgestellt, einem state-of-the-art großen Sprachmodell (LLM), das die Stärken von Transformer- und Mamba-Architekturen in einem hybriden Framework kombiniert. Dieser Artikel beschreibt Jambas Architektur, Leistung und mögliche Anwendungen.
Überblick über Jamba
Jamba ist ein hybrides großes Sprachmodell, das von AI21 Labs entwickelt wurde und eine Kombination von Transformer-Schichten und Mamba-Schichten mit einem Mixture-of-Experts (MoE)-Modul integriert. Diese Architektur ermöglicht es Jamba, den Speicherbedarf, die Durchsatzleistung und die Leistung auszugleichen, was es zu einem leistungsstarken Werkzeug für eine breite Palette von NLP-Aufgaben macht. Das Modell ist so konzipiert, dass es in einer einzigen 80-GB-GPU passt und bietet eine hohe Durchsatzleistung und einen kleinen Speicherbedarf, während es gleichzeitig eine Spitzenleistung auf verschiedenen Benchmarks erzielt.
Die Architektur von Jamba
Jambas Architektur ist der Eckpfeiler seiner Fähigkeiten. Sie basiert auf einem neuartigen hybriden Design, das Transformer-Schichten mit Mamba-Schichten verbindet und MoE-Module integriert, um die Kapazität des Modells zu erhöhen, ohne die Rechenanforderungen signifikant zu erhöhen.
1. Transformer-Schichten
Die Transformer-Architektur ist zum Standard für moderne LLMs geworden, da sie effizient parallele Verarbeitung und die Erfassung von Langstreckenabhängigkeiten in Texten ermöglicht. Allerdings ist ihre Leistung oft durch hohe Speicher- und Rechenanforderungen begrenzt, insbesondere bei der Verarbeitung langer Kontexte. Jamba behebt diese Einschränkungen, indem es Mamba-Schichten integriert, die wir im nächsten Abschnitt betrachten werden.
2. Mamba-Schichten
Mamba ist ein neues State-Space-Modell (SSM), das darauf ausgelegt ist, Langstreckenbeziehungen in Sequenzen effizienter zu handhaben als herkömmliche RNNs oder sogar Transformer. Mamba-Schichten sind besonders effektiv bei der Reduzierung des Speicherbedarfs, der mit der Speicherung von Schlüssel-Wert-Caches in Transformern verbunden ist. Durch die Verbindung von Mamba-Schichten mit Transformer-Schichten reduziert Jamba den gesamten Speicherbedarf, während es gleichzeitig eine hohe Leistung erzielt, insbesondere bei Aufgaben, die die Handhabung langer Kontexte erfordern.
3. Mixture-of-Experts (MoE)-Module
Das MoE-Modul in Jamba bietet einen flexiblen Ansatz, um die Kapazität des Modells zu skalieren. MoE ermöglicht es dem Modell, die Anzahl der verfügbaren Parameter zu erhöhen, ohne die aktiven Parameter während der Inferenz proportional zu erhöhen. In Jamba wird MoE auf einige der MLP-Schichten angewendet, wobei der Router-Mechanismus die Top-Experten auswählt, die für jeden Token aktiviert werden sollen. Diese selektive Aktivierung ermöglicht es Jamba, eine hohe Effizienz zu erhalten, während es komplexe Aufgaben handhabt.
Das folgende Bild zeigt die Funktionalität eines Induktionskopfes in einem hybriden Attention-Mamba-Modell, einem wichtigen Merkmal von Jamba. In diesem Beispiel ist der Attention-Kopf für die Vorhersage von Labels wie “Positiv” oder “Negativ” bei der Sentiment-Analyse verantwortlich. Die hervorgehobenen Wörter zeigen, wie die Aufmerksamkeit des Modells stark auf Label-Tokens aus den Few-Shot-Beispielen fokussiert ist, insbesondere im kritischen Moment vor der Vorhersage des endgültigen Labels. Dieser Aufmerksamkeitsmechanismus spielt eine entscheidende Rolle bei der Fähigkeit des Modells, in-Kontext-Lernen durchzuführen, bei dem das Modell das geeignete Label basierend auf dem gegebenen Kontext und den Few-Shot-Beispielen ableiten muss.
Die Leistungsverbesserungen, die durch die Integration von Mixture-of-Experts (MoE) mit der Attention-Mamba-Hybridarchitektur erzielt werden, werden in der Tabelle hervorgehoben. Durch die Verwendung von MoE erhöht Jamba seine Kapazität, ohne die Rechenkosten proportional zu erhöhen. Dies ist insbesondere in der signifikanten Leistungssteigerung bei verschiedenen Benchmarks wie HellaSwag, WinoGrande und Natural Questions (NQ) erkennbar. Das Modell mit MoE erzielt nicht nur eine höhere Genauigkeit (z. B. 66,0 % bei WinoGrande im Vergleich zu 62,5 % ohne MoE), sondern zeigt auch verbesserte Log-Wahrscheinlichkeiten in verschiedenen Domänen (z. B. -0,534 bei C4).
Schlüsselmerkmale der Architektur
- Schichten-Zusammensetzung: Jambas Architektur besteht aus Blöcken, die Mamba- und Transformer-Schichten in einem bestimmten Verhältnis (z. B. 1:7, was bedeutet, dass für jede sieben Mamba-Schichten eine Transformer-Schicht vorhanden ist) kombinieren. Dieses Verhältnis wird für optimale Leistung und Effizienz abgestimmt.
- MoE-Integration: Die MoE-Schichten werden alle paar Schichten angewendet, wobei 16 Experten verfügbar sind und die Top-2-Experten pro Token aktiviert werden. Diese Konfiguration ermöglicht es Jamba, effektiv zu skalieren, während es die Kompromisse zwischen Speicherbedarf und Recheneffizienz managt.
- Normalisierung und Stabilität: Um die Stabilität während des Trainings zu gewährleisten, integriert Jamba RMSNorm in den Mamba-Schichten, was hilft, Probleme wie große Aktivierungsspitzen zu vermeiden, die bei der Skalierung auftreten können.
Jambas Leistung und Benchmarking
Jamba wurde umfassend auf eine breite Palette von Benchmarks getestet und zeigte eine wettbewerbsfähige Leistung auf dem gesamten Gebiet. Die folgenden Abschnitte heben einige der wichtigsten Benchmarks hervor, bei denen Jamba hervorragend abgeschnitten hat, und zeigen seine Stärken in allgemeinen NLP-Aufgaben und Langkontext-Szenarien.
1. Übliche NLP-Benchmarks
Jamba wurde auf mehrere akademische Benchmarks ausgewertet, darunter:
- HellaSwag (10-Shot): Eine Aufgabe zum gemeinsamen Verständnis, bei der Jamba eine Leistung von 87,1 % erzielte und viele konkurrierende Modelle übertraf.
- WinoGrande (5-Shot): Eine weitere Aufgabe zum Verständnis, bei der Jamba 82,5 % erreichte und seine Fähigkeit unter Beweis stellte, komplexe sprachliche Argumentationen zu handhaben.
- ARC-Challenge (25-Shot): Jamba zeigte eine starke Leistung mit einem Ergebnis von 64,4 %, was seine Fähigkeit widerspiegelt, schwierige Multiple-Choice-Fragen zu meistern.
In aggregierten Benchmarks wie MMLU (5-Shot) erzielte Jamba einen Wert von 67,4 %, was seine Robustheit bei verschiedenen Aufgaben zeigt.
2. Langkontext-Bewertungen
Ein herausragendes Merkmal von Jamba ist seine Fähigkeit, extrem lange Kontexte zu handhaben. Das Modell unterstützt eine Kontextlänge von bis zu 256K Token, die längste unter den öffentlich verfügbaren Modellen. Diese Fähigkeit wurde mit dem Needle-in-a-Haystack-Benchmark getestet, bei dem Jamba eine außergewöhnliche Abrufgenauigkeit bei verschiedenen Kontextlängen zeigte, einschließlich bis zu 256K Token.
3. Durchsatz und Effizienz
Jambas hybride Architektur verbessert den Durchsatz erheblich, insbesondere bei langen Sequenzen.
In Tests, die den Durchsatz (Token pro Sekunde) bei verschiedenen Modellen verglichen, übertraf Jamba seine Mitbewerber konstant, insbesondere in Szenarien mit großen Batch-Größen und langen Kontexten. Beispielsweise erzielte Jamba bei einem Kontext von 128K Token den dreifachen Durchsatz von Mixtral, einem vergleichbaren Modell.
Verwendung von Jamba: Python
Für Entwickler und Forscher, die mit Jamba experimentieren möchten, hat AI21 Labs das Modell auf Plattformen wie Hugging Face bereitgestellt, um es für eine breite Palette von Anwendungen zugänglich zu machen. Der folgende Code-Ausschnitt zeigt, wie man Jamba lädt und Text generiert:
<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p> <p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1") tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p> <p>input_ids = tokenizer("In der jüngsten Super-Bowl-LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p> <p>outputs = model.generate(input_ids, max_new_tokens=216)</p> print(tokenizer.batch_decode(outputs))
Dieses einfache Skript lädt das Jamba-Modell und den Tokenizer, generiert Text basierend auf einem gegebenen Eingabe-Prompt und gibt die generierte Ausgabe aus.
Fine-Tuning von Jamba
Jamba ist als Basis-Modell konzipiert, was bedeutet, dass es für spezifische Aufgaben oder Anwendungen fein abgestimmt werden kann. Das Feinabstimmen ermöglicht es Benutzern, das Modell an Nischen-Domänen anzupassen und die Leistung bei spezialisierten Aufgaben zu verbessern. Der folgende Code-Ausschnitt zeigt, wie man Jamba mit der PEFT-Bibliothek fein abstimmt:
import torch from datasets import load_dataset from trl import SFTTrainer, SFTConfig from peft import LoraConfig from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments <p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1") model = AutoModelForCausalLM.from_pretrained( "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p> <p>lora_config = LoraConfig(r=8, target_modules=[ "embed_tokens","x_proj", "in_proj", "out_proj", # mamba "gate_proj", "up_proj", "down_proj", # mlp "q_proj", "k_proj", "v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p> <p>dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote") trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset, ) trainer.train()
Dieser Code-Ausschnitt fein abstimmt Jamba auf einem Datensatz von englischen Zitaten, indem er die Parameter des Modells an die spezifische Aufgabe der Textgenerierung in einer spezialisierten Domäne anpasst.
Bereitstellung und Integration
AI21 Labs hat die Jamba-Familie auf verschiedenen Plattformen und Bereitstellungsvarianten verfügbar gemacht:
- Cloud-Plattformen:
- Verfügbar auf großen Cloud-Anbietern wie Google Cloud Vertex AI, Microsoft Azure und NVIDIA NIM (NVDA ).
- Bald verfügbar auf Amazon Bedrock, Databricks Marketplace und Snowflake Cortex.
- AI-Entwicklungsframeworks:
- Integration mit beliebten Frameworks wie LangChain und LlamaIndex (in Kürze).
- AI21 Studio:
- Direkter Zugriff über AI21s eigene Entwicklungsplattform.
- Hugging Face:
- Modelle zum Download und zur Experimentierung verfügbar.
- On-Premises-Bereitstellung:
- Optionen für private, vor-Ort-Bereitstellung für Organisationen mit spezifischen Sicherheits- oder Compliance-Anforderungen.
- Benutzerdefinierte Lösungen:
- AI21 bietet maßgeschneiderte Modellanpassung und Feinabstimmungsdienste für Unternehmenskunden an.
Entwicklerfreundliche Funktionen
Jamba-Modelle verfügen über mehrere integrierte Funktionen, die sie besonders attraktiv für Entwickler machen:
- Funktionsaufruf: Integrieren Sie leicht externe Tools und APIs in Ihre AI-Workflows.
- Strukturierte JSON-Ausgabe: Generieren Sie saubere, parsbare Datenstrukturen direkt aus natürlichen Spracheneingaben.
- Dokument-Objekt-Verdauung: Verarbeiten Sie effizient komplexe Dokumentstrukturen.
- RAG-Optimierungen: Integrierte Funktionen zur Verbesserung von Retrieval-augmented-Generation-Pipelines.
Diese Funktionen, kombiniert mit der Fähigkeit des Modells, lange Kontexte zu handhaben und effizient zu verarbeiten, machen Jamba zu einem vielseitigen Werkzeug für eine breite Palette von Entwicklungszenarien.
Ethische Überlegungen und verantwortungsvolle KI
Während die Fähigkeiten von Jamba beeindruckend sind, ist es wichtig, seine Verwendung mit einer verantwortungsvollen KI-Einstellung anzugehen. AI21 Labs betont mehrere wichtige Punkte:
- Basis-Modell-Natur: Jamba 1.5-Modelle sind vorgefertigte Basis-Modelle ohne spezifische Ausrichtung oder Anpassung.
- Fehlen von eingebauten Sicherheitsvorkehrungen: Die Modelle verfügen nicht über inhärente Moderationsmechanismen.
- Sorgfältige Bereitstellung: Zusätzliche Anpassungen und Sicherheitsvorkehrungen sollten implementiert werden, bevor Jamba in Produktionsumgebungen oder mit Endbenutzern verwendet wird.
- Datenschutz: Bei der Verwendung von Cloud-basierten Bereitstellungen sollten Sie sich der Datenverarbeitung und der Einhaltung von Vorschriften bewusst sein.
- Bewusstsein für Vorurteile: Wie alle großen Sprachmodelle kann Jamba Vorurteile widerspiegeln, die in seinen Trainingsdaten vorhanden sind. Benutzer sollten sich dieser Vorurteile bewusst sein und geeignete Gegenmaßnahmen ergreifen.
Indem Sie diese Faktoren berücksichtigen, können Entwickler und Organisationen Jambas Fähigkeiten verantwortungsvoll und ethisch einsetzen.
Ein neues Kapitel in der KI-Entwicklung?
Die Einführung der Jamba-Familie durch AI21 Labs markiert einen bedeutenden Meilenstein in der Evolution von großen Sprachmodellen. Durch die Kombination der Stärken von Transformern und State-Space-Modellen, die Integration von Mixture-of-Experts-Techniken und die Erweiterung der Kontextlänge und der Verarbeitungsgeschwindigkeit eröffnet Jamba neue Möglichkeiten für KI-Anwendungen in verschiedenen Branchen.
Wenn die KI-Gemeinschaft weiterhin auf dieser innovativen Architektur aufbaut, können wir weitere Fortschritte in der Modell-Effizienz, der Langkontext-Verständnis und der praktischen KI-Bereitstellung erwarten. Die Jamba-Familie repräsentiert nicht nur eine neue Reihe von Modellen, sondern einen möglichen Paradigmenwechsel in der Entwicklung und Implementierung von großen KI-Systemen.
