Jamba: AI21 Labs’ ny hybrid Transformer-Mamba språkmodell

Språkmodellene har vært gjennom raske fremgang, med Transformer-baserte arkitekturer i spissen for naturlig språkbehandling. Imidlertid, når modellene skalerer, har utfordringene med å håndtere lange kontekster, minnehåndtering og gjennomstrømming blitt mer fremtredende.

AI21 Labs har introdusert en ny løsning med Jamba, en state-of-the-art stor språkmodell (LLM) som kombinerer styrkene til både Transformer og Mamba-arkitekturer i en hybrid ramme. Denne artikkelen detaljerer Jambas arkitektur, ytelse og potensielle anvendelser.

Overblik over Jamba

Jamba er en hybrid stor språkmodell utviklet av AI21 Labs, som utnytter en kombinasjon av Transformer-lag og Mamba-lag, integrert med en Mixture-of-Experts (MoE)-modul. Denne arkitekturen tillater Jamba å balansere minnebruk, gjennomstrømming og ytelse, og gjør den til et kraftig verktøy for en rekke NLP-oppdrag. Modellen er designet for å passe innenfor en enkelt 80GB GPU, og tilbyr høy gjennomstrømming og en liten minneavtrykk samtidig som den opprettholder state-of-the-art-ytelse på ulike benchmark.

Jambas arkitektur

Jambas arkitektur er hjørnestenen til dens evner. Den er bygget på en ny hybrid design som veksler Transformer-lag med Mamba-lag, og inkorporerer MoE-moduler for å forbedre modellens kapasitet uten å øke beregningskostnadene betydelig.

1. Transformer-lag

Transformer-arkitekturen har blitt standarden for moderne LLM-er på grunn av dens evne til å håndtere parallell prosessering effektivt og fange lange avhengigheter i tekst. Imidlertid er dens ytelse ofte begrenset av høye minne- og beregningskrav, spesielt når det gjelder å prosessere lange kontekster. Jamba løser disse begrensningene ved å integrere Mamba-lag, som vi skal se nærmere på neste.

2. Mamba-lag

Mamba er en ny state-space-modell (SSM) designet for å håndtere lange avhengigheter i sekvenser mer effektivt enn tradisjonelle RNN-er eller selv Transformers. Mamba-lag er spesielt effektive til å redusere minneavtrykket forbundet med lagring av nøkkel-verdi (KV)-cacher i Transformers. Ved å veksle Mamba-lag med Transformer-lag, reduserer Jamba det totale minnebruket samtidig som den opprettholder høy ytelse, spesielt i oppdrag som krever håndtering av lange kontekster.

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

MoE-modulen i Jamba introduserer en fleksibel tilnærming til å skalerer modellkapasitet. MoE tillater modellen å øke antallet tilgjengelige parametre uten å øke de aktive parameterne proporsjonalt under inferens. I Jamba brukes MoE på noen av MLP-lagene, med router-mekanismen som velger de beste ekspertene å aktivere for hver token. Denne selektive aktivering ermögner Jamba å opprettholde høy effektivitet samtidig som den håndterer komplekse oppdrag.

Bildet under demonstrerer funksjonaliteten til en induksjonshead i en hybrid Attention-Mamba-modell, en nøkkeltilnærming i Jamba. I dette eksempelet er attention-headet ansvarlig for å forutsi etiketter som “Positiv” eller “Negativ” i respons til sentiment-analyseoppdrag. De highlightede ordene viser hvordan modellens attention er sterkt fokusert på etikett-tokens fra de få-shot-eksemplene, spesielt i det kritiske øyeblikket før forutsigelsen av den endelige etiketten. Denne attention-mekanismen spiller en avgjørende rolle i modellens evne til å utføre in-context-læring, hvor modellen må slutte seg til den riktige etiketten basert på den gitt konteksten og få-shot-eksemplene.

Ytelsesforbedringene som tilbys ved å integrere Mixture-of-Experts (MoE) med den hybrid Attention-Mamba-arkitekturen er highlightet i tabellen. Ved å bruke MoE, øker Jamba sin kapasitet uten å øke beregningskostnadene proporsjonalt. Dette er spesielt tydelig i den betydelige ytelsesøkningen på ulike benchmark, som HellaSwag, WinoGrande og Natural Questions (NQ). Modellen med MoE oppnår ikke bare høyere nøyaktighet (f.eks. 66,0% på WinoGrande sammenlignet med 62,5% uten MoE), men viser også forbedrede log-sannsynligheter på ulike domener (f.eks. -0,534 på C4).

Nøkkelarkitektoniske trekk

• Lagkomposisjon: Jambas arkitektur består av blokker som kombinerer Mamba og Transformer-lag i et bestemt forhold (f.eks. 1:7, som betyr ett Transformer-lag for hver syv Mamba-lag). Dette forholdet er justert for optimal ytelse og effektivitet.
• MoE-integrasjon: MoE-lagene brukes hver få lag, med 16 eksperter tilgjengelige og de to beste eksperter aktivert per token. Denne konfigurasjonen tillater Jamba å skalerer effektivt samtidig som den håndterer avveiingene mellom minnebruk og beregnings-effektivitet.
• Normalisering og stabilitet: For å sikre stabilitet under trening, inkorporerer Jamba RMSNorm i Mamba-lagene, som hjelper til å mildne problemer som store aktiverings-spike som kan oppstå i skala.

Jambas ytelse og benchmarking

Jamba er blitt grundig testet mot en rekke benchmark, og viser konkurrerende ytelse over hele linjen. Følgende avsnitt highlighter noen av de viktigste benchmarkene hvor Jamba har utmerket seg, og viser dens styrker i både generelle NLP-oppdrag og lange-kontekst-scenarier.

1. Vanlige NLP-benchmark

Jamba er blitt evaluert på flere akademiske benchmark, inkludert:

• HellaSwag (10-shot): En common sense reasoning-oppdrag hvor Jamba oppnådde en ytelsespoeng på 87,1%, og overgikk mange konkurrerende modeller.
• WinoGrande (5-shot): Et annet reasoning-oppdrag hvor Jamba scoret 82,5%, og viste igjen dens evne til å håndtere komplekse lingvistiske resonneringer.
• ARC-Challenge (25-shot): Jamba viste sterk ytelse med en poengsum på 64,4%, og reflekterte dens evne til å håndtere utfordrende multiple-choice-spørsmål.

I aggregat-benchmark som MMLU (5-shot) oppnådde Jamba en poengsum på 67,4%, og indikerte dens robusthet over ulike oppdrag.

2. Lang-kontekst-evalueringer

En av Jambas utmerkede trekk er dens evne til å håndtere ekstremt lange kontekster. Modellen støtter en kontekstlengde på opptil 256K tokens, den lengste blant offentlig tilgjengelige modeller. Denne evnen ble testet ved å bruke Needle-in-a-Haystack-benchmark, hvor Jamba viste unik gjenfinningsnøyaktighet over ulike kontekstlengder, inkludert opptil 256K tokens.

3. Gjennomstrømming og effektivitet

Jambas hybrid-arkitektur forbedrer gjennomstrømmingen betydelig, spesielt med lange sekvenser.

I tester som sammenligner gjennomstrømming (tokens per sekund) over ulike modeller, utmerket Jamba seg konsekvent, spesielt i scenarier med store batch-størrelser og lange kontekster. For eksempel, med en kontekst på 128K tokens, oppnådde Jamba 3 ganger gjennomstrømmingen til Mixtral, en sammenlignbar modell.

Bruke Jamba: Python

For utviklere og forskere som ønsker å eksperimentere med Jamba, har AI21 Labs gjort modellen tilgjengelig på plattformer som Hugging Face, og gjort den tilgjengelig for en rekke anvendelser. Følgende kode-utdrag demonstrerer hvordan du kan laste og generere tekst 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;In the recent 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()