Jamba: AI21 Labs nya hybridmodell för transformer-mambaspråk

Språkmodeller har sett en snabb utveckling, med transformer-baserade arkitekturer i spetsen för naturlig språkbehandling. Men när modellerna skalar upp, har utmaningarna med att hantera långa sammanhang, minneseffektivitet och genomströmning blivit mer påtagliga.

AI21 Labs har introducerat en ny lösning med Jamba, en state-of-the-art stor språkmodell (LLM) som kombinerar styrkorna hos både transformer- och Mamba-arkitekturer i en hybridram. Den här artikeln beskriver Jambas arkitektur, prestanda och potentiella tillämpningar.

Översikt av Jamba

Jamba är en hybrid stor språkmodell utvecklad av AI21 Labs, som utnyttjar en kombination av transformerlager och Mambalager, integrerat med ett Mixture-of-Experts (MoE)-modul. Denna arkitektur tillåter Jamba att balansera minnesanvändning, genomströmning och prestanda, vilket gör den till ett kraftfullt verktyg för en mängd olika NLP-uppgifter. Modellen är utformad för att passa inom en enda 80GB GPU, med hög genomströmning och liten minnesavtryck samtidigt som den upprätthåller state-of-the-art-prestanda på olika benchmarktest.

Jambas arkitektur

Jambas arkitektur är grunden för dess förmågor. Den byggs på en ny hybriddesign som växlar transformerlager med Mambalager, med MoE-moduler för att förbättra modellens kapacitet utan att signifikant öka beräkningskraven.

1. Transformerlager

Transformerarkitekturen har blivit standarden för moderna LLM:er på grund av dess förmåga att hantera parallell bearbetning effektivt och fånga långdistansberoenden i text. Men dess prestanda begränsas ofta av höga minnes- och beräkningskrav, särskilt när det gäller att bearbeta långa sammanhang. Jamba hanterar dessa begränsningar genom att integrera Mambalager, som vi kommer att undersöka nästa.

2. Mambalager

Mamba är en nyligen utvecklad state-space-modell (SSM) som är utformad för att hantera långdistansrelationer i sekvenser mer effektivt än traditionella RNN:er eller till och med transformer. Mambalager är särskilt effektiva för att minska minnesavtrycket som är förknippat med att lagra nyckel-värde-cacheminne i transformer. Genom att växla Mambalager med transformerlager minskar Jamba den totala minnesanvändningen samtidigt som den upprätthåller hög prestanda, särskilt i uppgifter som kräver långsammanhangshantering.

3. Mixture-of-Experts (MoE)-moduler

MoE-modulen i Jamba introducerar en flexibel metod för att skala modellens kapacitet. MoE tillåter modellen att öka antalet tillgängliga parametrar utan att proportionellt öka de aktiva parametrarna under inferens. I Jamba appliceras MoE på vissa av MLP-lagren, med routningsmekanismen som väljer de bästa experterna att aktivera för varje token. Denna selektiva aktivering möjliggör för Jamba att upprätthålla hög effektivitet samtidigt som den hanterar komplexa uppgifter.

Bilden nedan demonstrerar funktionen hos en induktionshuvud i en hybrid Attention-Mambamodell, en viktig funktion i Jamba. I det här exemplet är uppmärksamhetsrubriken ansvarig för att förutsäga etiketter som “Positiv” eller “Negativ” som svar på sentimentanalysuppgifter. De markerade orden visar hur modellens uppmärksamhet är starkt fokuserad på etiketttoken från de fåskottsexemplen, särskilt i det kritiska ögonblicket innan förutsägelsen av den slutliga etiketten. Denna uppmärksamhetsmekanism spelar en avgörande roll i modellens förmåga att utföra in-context-lärande, där modellen måste härleda den lämpliga etiketten baserat på den givna kontexten och fåskottsexemplen.

Prestandaförbättringarna som erbjuds genom att integrera Mixture-of-Experts (MoE) med hybrid Attention-Mambamodellen framhävs i tabellen. Genom att använda MoE ökar Jamba sin kapacitet utan att proportionellt öka beräkningskostnaderna. Detta är särskilt tydligt i den betydande prestandaförbättringen över olika benchmarktest som HellaSwag, WinoGrande och Natural Questions (NQ). Modellen med MoE uppnår inte bara högre noggrannhet (t.ex. 66,0% på WinoGrande jämfört med 62,5% utan MoE) utan visar också förbättrade log-sannolikheter över olika domäner (t.ex. -0,534 på C4).

Nyckelarkitekturfunktioner

• Lagerkomposition: Jambas arkitektur består av block som kombinerar Mamba och transformerlager i ett visst förhållande (t.ex. 1:7, vilket innebär ett transformerlager för varje sju Mambalager). Detta förhållande är justerat för optimal prestanda och effektivitet.
• MoE-integrering: MoE-lagren appliceras varannan lager, med 16 experter tillgängliga och de två bästa experterna aktiverade per token. Denna konfiguration möjliggör för Jamba att skala effektivt samtidigt som den hanterar avvägningarna mellan minnesanvändning och beräkningseffektivitet.
• Normalisering och stabilitet: För att säkerställa stabilitet under utbildning, inkorporerar Jamba RMSNorm i Mambalagren, vilket hjälper till att mildra problem som stora aktiveringspikar som kan uppstå i skala.

Jambas prestanda och benchmarking

Jamba har testats rigoröst mot en mängd olika benchmarktest, med konkurrenskraftig prestanda över hela linjen. Följande avsnitt framhäver några av de viktigaste benchmarktesten där Jamba har utmärkt sig, vilket visar dess styrkor i både allmänna NLP-uppgifter och långsammanhangsscenarier.

1. Vanliga NLP-benchmarktest

Jamba har utvärderats på flera akademiska benchmarktest, inklusive:

• HellaSwag (10-skott): En uppgift för sunt förnuft som Jamba uppnådde en prestandapoäng på 87,1%, och överträffade många konkurrerande modeller.
• WinoGrande (5-skott): En annan resonemangs-uppgift där Jamba uppnådde en poäng på 82,5%, och visade sin förmåga att hantera komplex lingvistisk resonemang.
• ARC-utmaning (25-skott): Jamba visade stark prestanda med en poäng på 64,4%, vilket reflekterar dess förmåga att hantera utmanande flervalsfrågor.

I aggregerade benchmarktest som MMLU (5-skott) uppnådde Jamba en poäng på 67,4%, vilket indikerar dess robusthet över olika uppgifter.

2. Långsammanhangsbedömningar

En av Jambas utmärkande funktioner är dess förmåga att hantera extremt långa sammanhang. Modellen stöder en sammanhangslängd på upp till 256K token, den längsta bland offentligt tillgängliga modeller. Denna funktion testades med hjälp av Needle-in-a-Haystack-benchmarktestet, där Jamba visade exceptionell återhämtningsnoggrannhet över varierande sammanhangslängder, inklusive upp till 256K token.

3. Genomströmning och effektivitet

Jambas hybridarkitektur förbättrar genomströmningen avsevärt, särskilt med långa sekvenser.

I tester som jämförde genomströmning (token per sekund) över olika modeller, presterade Jamba konsekvent bättre än sina konkurrenter, särskilt i scenarier som involverade stora batchstorlekar och långa sammanhang. Till exempel, med en sammanhang på 128K token, uppnådde Jamba tre gånger högre genomströmning än Mixtral, en jämförbar modell.

Användning av Jamba: Python

För utvecklare och forskare som är angelägna om att experimentera med Jamba, har AI21 Labs tillhandahållit modellen på plattformar som Hugging Face, vilket gör den tillgänglig för en mängd olika tillämpningar. Följande kodsnutt demonstrerar hur man kan ladda och generera text med Jamba:

<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p>

<p>model = AutoModelForCausalLM.from_pretrained(&quot;ai21labs/Jamba-v0.1&quot;)
tokenizer = AutoTokenizer.from_pretrained(&quot;ai21labs/Jamba-v0.1&quot;)</p>

<p>input_ids = tokenizer(&quot;I den senaste Super Bowl LVIII,&quot;, return_tensors=&#039;pt&#039;).to(model.device)[&quot;input_ids&quot;]</p>

<p>outputs = model.generate(input_ids, max_new_tokens=216)</p>

print(tokenizer.batch_decode(outputs))

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(&quot;ai21labs/Jamba-v0.1&quot;)
model = AutoModelForCausalLM.from_pretrained(
&quot;ai21labs/Jamba-v0.1&quot;, device_map=&#039;auto&#039;, torch_dtype=torch.bfloat16)</p>

<p>lora_config = LoraConfig(r=8,
target_modules=[
&quot;embed_tokens&quot;,&quot;x_proj&quot;, &quot;in_proj&quot;, &quot;out_proj&quot;, # mamba
&quot;gate_proj&quot;, &quot;up_proj&quot;, &quot;down_proj&quot;, # mlp
&quot;q_proj&quot;, &quot;k_proj&quot;, &quot;v_proj&quot;
# attention],
task_type=&quot;CAUSAL_LM&quot;, bias=&quot;none&quot;)</p>

<p>dataset = load_dataset(&quot;Abirate/english_quotes&quot;, split=&quot;train&quot;)
training_args = SFTConfig(output_dir=&quot;./results&quot;,
num_train_epochs=2,
per_device_train_batch_size=4,
logging_dir=&#039;./logs&#039;,
logging_steps=10, learning_rate=1e-5, dataset_text_field=&quot;quote&quot;)
trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args,
peft_config=lora_config, train_dataset=dataset,
)
trainer.train()