Jamba: AI21 Labs’ ny hybrid Transformer-Mamba sprogmodel

Sprogmodeller har oplevet hurtige fremskridt, og Transformer-baserede arkitekturer har ført an i naturlig sprogbehandling. Men når modellerne skalerer, er udfordringerne vedrørende håndtering af lange kontekster, hukommelseseffektivitet og gennemstrømning blevet mere udtalt.

AI21 Labs har introduceret en ny løsning med Jamba, et state-of-the-art stort sprogmodel (LLM), der kombinerer styrkerne fra både Transformer og Mamba-arkitekturer i en hybrid ramme. Denne artikel beskriver Jambas arkitektur, præstation og potentielle anvendelser.

Overblik over Jamba

Jamba er et hybrid stort sprogmodel udviklet af AI21 Labs, der udnytter en kombination af Transformer-lag og Mamba-lag, integreret med en Mixture-of-Experts (MoE)-modul. Denne arkitektur giver Jamba mulighed for at balancere hukommelsesbrug, gennemstrømning og præstation, hvilket gør det til et kraftfuldt værktøj til en bred vifte af NLP-opgaver. Modellen er designet til at kunne passe inden for en enkelt 80GB GPU, hvilket giver høj gennemstrømning og en lille hukommelsesaftryk, samtidig med at den opretholder state-of-the-art-præstation på forskellige benchmarks.

Jambas arkitektur

Jambas arkitektur er hjørnestenen i dets evner. Den er bygget på en ny hybrid design, der vekselvirker Transformer-lag med Mamba-lag, og inkorporerer MoE-moduler for at forbedre modellens kapacitet uden at øge beregningskravene væsentligt.

1. Transformer-lag

Transformer-arkitekturen er blevet standarden for moderne LLM’er på grund af deres evne til at håndtere parallel procesering effektivt og fange lange afhængigheder i tekst. Men deres præstation er ofte begrænset af høje hukommelses- og beregningskrav, især når der håndteres lange kontekster. Jamba løser disse begrænsninger ved at integrere Mamba-lag, som vi vil udforske næste.

2. Mamba-lag

Mamba er et nyt state-space-model (SSM) designet til at håndtere lange afstande i sekvenser mere effektivt end traditionelle RNN’er eller selv Transformers. Mamba-lag er særligt effektive til at reducere den hukommelsesaftryk, der er forbundet med lagring af nøgle-værdi (KV)-caches i Transformers. Ved at vekselvirke Mamba-lag med Transformer-lag reducerer Jamba den samlede hukommelsesbrug, samtidig med at den opretholder høj præstation, især i opgaver, der kræver lange kontekster.

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

MoE-modulen i Jamba introducerer en fleksibel tilgang til at skala modellens kapacitet. MoE giver modellen mulighed for at øge antallet af tilgængelige parametre uden at øge de aktive parametre proportionalt under inferens. I Jamba anvendes MoE på nogle af MLP-lagene, og router-mekanismen vælger de top-eksperter, der skal aktiveres for hvert token. Denne selektive aktivering giver Jamba mulighed for at opretholde høj effektivitet, samtidig med at den håndterer komplekse opgaver.

Billedet nedenfor demonstrerer funktionen af en induktionshoved i en hybrid Attention-Mamba-model, en nøglefunktion i Jamba. I dette eksempel er attention-hovedet ansvarligt for at forudsige mærker som “Positiv” eller “Negativ” i svar til sentiment-analyseopgaver. De markerede ord illustrerer, hvordan modellens attention er stærkt fokuseret på mærke-token fra de få-shot-eksempler, især lige før forudsigelsen af det endelige mærke. Denne attention-mekanisme spiller en afgørende rolle i modellens evne til at udføre in-context-læring, hvor modellen skal slutte sig til det passende mærke baseret på den givne kontekst og få-shot-eksempler.

Præstationsforbedringerne, der tilbydes ved at integrere Mixture-of-Experts (MoE) med den hybrid Attention-Mamba-arkitektur, er fremhævet i tabel. Ved at anvende MoE øger Jamba sin kapacitet uden at øge beregningsomkostningerne proportionelt. Dette er særligt tydeligt i den betydelige forbedring af præstation på forskellige benchmarks som HellaSwag, WinoGrande og Natural Questions (NQ). Modellen med MoE opnår ikke kun højere nøjagtighed (f.eks. 66,0% på WinoGrande i forhold til 62,5% uden MoE), men viser også forbedrede log-sandsynligheder på tværs af forskellige domæner (f.eks. -0,534 på C4).

Nøglearkitektoniske funktioner

• Lag-sammensætning: Jambas arkitektur består af blokke, der kombinerer Mamba og Transformer-lag i en bestemt forhold (f.eks. 1:7, hvilket betyder en Transformer-lag for hver syv Mamba-lag). Dette forhold er justeret for optimal præstation og effektivitet.
• MoE-integration: MoE-lagene anvendes hver få lag, med 16 eksperter til rådighed og de top-2 eksperter aktiveres per token. Denne konfiguration giver Jamba mulighed for at skala effektivt, samtidig med at den håndterer kompromiserne mellem hukommelsesbrug og beregnings-effektivitet.
• Normalisering og stabilitet: For at sikre stabilitet under træning inkorporerer Jamba RMSNorm i Mamba-lagene, hvilket hjælper med at afværge problemer som store aktiverings-spike, der kan opstå i størrelsesorden.

Jambas præstation og benchmarking

Jamba er blevet grundigt testet mod en bred vifte af benchmarks, og har demonstreret konkurrencedygtig præstation på tværs af brættet. Følgende afsnit fremhæver nogle af de nøglebenchmarks, hvor Jamba har excelleret, og viser dens styrker i både generelle NLP-opgaver og lange kontekster.

1. Almindelige NLP-benchmarks

Jamba er blevet evaluueret på flere akademiske benchmarks, herunder:

• HellaSwag (10-shot): En fælles fornuftig grundighed-opgave, hvor Jamba opnåede en præstations-score på 87,1%, og overgik mange konkurrerende modeller.
• WinoGrande (5-shot): En anden grundighed-opgave, hvor Jamba scorede 82,5%, og viste igen sin evne til at håndtere komplekse lingvistiske grundighed.
• ARC-Challenge (25-shot): Jamba demonstrerede stærk præstation med en score på 64,4%, og reflekterede sin evne til at håndtere udfordrende multiple-choice-spørgsmål.

I samlede benchmarks som MMLU (5-shot) opnåede Jamba en score på 67,4%, og indikerede sin robusthed på tværs af forskellige opgaver.

2. Lange kontekst-evalueringer

En af Jambas fremhævede funktioner er dens evne til at håndtere ekstremt lange kontekster. Modellen understøtter en kontekstlængde på op til 256K token, den længste blandt offentligt tilgængelige modeller. Denne funktion blev testet ved hjælp af Needle-in-a-Haystack-benchmark, hvor Jamba viste exceptionel udlejningsnøjagtighed på tværs af varierende kontekstlængder, herunder op til 256K token.

3. Gennemstrømning og effektivitet

Jambas hybrid-arkitektur forbedrer gennemstrømningen, især med lange sekvenser.

I tests, der sammenlignede gennemstrømning (token per sekund) på tværs af forskellige modeller, opførte Jamba sig konsekvent bedre end sine ligemænd, især i scenarier med store batch-størrelser og lange kontekster. For eksempel opnåede Jamba en gennemstrømning, der var tre gange højere end Mixtral, en sammenlignelig model, med en kontekst på 128K token.

Brug af Jamba: Python

For udviklere og forskere, der er ivrige efter at eksperimentere med Jamba, har AI21 Labs gjort modellen tilgængelig på platforme som Hugging Face, hvilket gør det muligt at anvende den i en bred vifte af anvendelser. Følgende kode-eksempel demonstrerer, hvordan man kan indlæse og generere tekst ved hjælp af 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 seneste 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()