Jamba: Nowy hybrydowy model językowy Transformer-Mamba AI21 Labs

Modele językowe przechodzą przez szybki postęp, a architektury oparte na Transformerach prowadzą w przetwarzaniu języka naturalnego. Jednak wraz ze skalowaniem modeli, wyzwania związane z obsługą długich kontekstów, wydajnością pamięci i przepływnością stają się bardziej wyraźne.

AI21 Labs wprowadziło nowe rozwiązanie z Jamba, model językowy o najwyższej jakości (LLM), który łączy zalety architektur Transformer i Mamba w ramach hybrydowego frameworku. Ten artykuł szczegółowo opisuje architekturę Jamba, wydajność i potencjalne zastosowania.

Przegląd Jamba

Jamba to hybrydowy duży model językowy opracowany przez AI21 Labs, wykorzystujący połączenie warstw Transformer i warstw Mamba, zintegrowanych z modułem Mixture-of-Experts (MoE). Ta architektura pozwala Jamba na równowagę pomiędzy użyciem pamięci, przepływnością i wydajnością, czyniąc go potężnym narzędziem dla szerokiego zakresu zadań NLP. Model został zaprojektowany, aby zmieścić się w jednej 80GB GPU, oferując wysoką przepływność i mały ślad pamięci, jednocześnie utrzymując najwyższą wydajność na różnych benchmarkach.

Architektura Jamba

Architektura Jamba jest kamieniem węgielnym jego możliwości. Została zbudowana na nowatorskim projekcie hybrydowym, który przeplata warstwy Transformer z warstwami Mamba, włączając moduły MoE w celu zwiększenia pojemności modelu bez znacznego zwiększania wymagań obliczeniowych.

1. Warstwy Transformer

Architektura Transformer stała się standardem dla nowoczesnych LLM ze względu na jej zdolność do efektywnego przetwarzania równoległego i przechwytywania dalekosiężnych zależności w tekście. Jednak jej wydajność jest często ograniczona przez wysokie wymagania pamięci i obliczeniowe, szczególnie podczas przetwarzania długich kontekstów. Jamba rozwiązuje te ograniczenia, integrując warstwy Mamba, które zostaną omówione w dalszej części.

2. Warstwy Mamba

Mamba to niedawny model przestrzeni stanu (SSM) zaprojektowany do obsługi dalekosiężnych relacji w sekwencjach w sposób bardziej wydajny niż tradycyjne RNN lub nawet Transformatory. Warstwy Mamba są szczególnie skuteczne w redukowaniu śladu pamięci związanego z przechowywaniem pamięci podręcznej klucz-wartość (KV) w Transformatorkach. Przeplatając warstwy Mamba z warstwami Transformer, Jamba redukuje ogólny zużycie pamięci, jednocześnie utrzymując wysoką wydajność, szczególnie w zadaniach wymagających długich kontekstów.

3. Mixture-of-Experts (MoE) Moduły

Moduł MoE w Jamba wprowadza elastyczne podejście do skalowania pojemności modelu. MoE pozwala modelowi zwiększyć liczbę dostępnych parametrów bez proporcjonalnego zwiększania aktywnych parametrów podczas inferencji. W Jamba MoE jest stosowany do niektórych warstw MLP, z mechanizmem routera wybierającym najlepszych ekspertów do aktywacji dla każdego tokenu. Ta selektywna aktywacja umożliwia Jamba utrzymanie wysokiej wydajności, jednocześnie radząc sobie z złożonymi zadaniami.

Poniższy obraz ilustruje funkcjonalność głowy indukcyjnej w hybrydowym modelu Attention-Mamba, kluczowej cechy Jamba. W tym przykładzie głowa uwagi jest odpowiedzialna za przewidywanie etykiet, takich jak “Pozytywny” lub “Negatywny” w odpowiedzi na zadania analizy sentymentu. Podświetlone słowa pokazują, jak uwaga modelu jest silnie skupiona na tokenach etykiet z przykładów few-shot, szczególnie w krytycznym momencie przed przewidywaniem ostatecznej etykiety. Ten mechanizm uwagi odgrywa kluczową rolę w zdolności modelu do wykonywania uczenia w kontekście, gdzie model musi wnioskować o odpowiedniej etykiecie na podstawie danego kontekstu i przykładów few-shot.

Poprawa wydajności oferowana przez integrację Mixture-of-Experts (MoE) z hybrydową architekturą Attention-Mamba jest przedstawiona w tabeli. Używając MoE, Jamba zwiększa swoją pojemność bez proporcjonalnego zwiększania kosztów obliczeniowych. Jest to szczególnie widoczne w znacznym wzroście wydajności w różnych benchmarkach, takich jak HellaSwag, WinoGrande i Natural Questions (NQ). Model z MoE nie tylko osiąga wyższą dokładność (np. 66,0% w WinoGrande w porównaniu z 62,5% bez MoE), ale także wykazuje poprawione logarytmiczne prawdopodobieństwa w różnych dziedzinach (np. -0,534 w C4).

Kluczowe cechy architektury

• Skład warstw: Architektura Jamba składa się z bloków, które łączą Mamba i warstwy Transformer w określonym stosunku (np. 1:7, co oznacza jedną warstwę Transformer dla każdych siedmiu warstw Mamba). Ten stosunek jest dostosowany do optymalnej wydajności i efektywności.
• Integracja MoE: Warstwy MoE są stosowane co kilka warstw, z 16 dostępnymi ekspertami i aktywowanymi 2 ekspertami na token. Ta konfiguracja pozwala Jamba na efektywne skalowanie, jednocześnie zarządzając kompromisami pomiędzy użyciem pamięci a efektywnością obliczeniową.
• Normalizacja i stabilność: Aby zapewnić stabilność podczas szkolenia, Jamba włącza RMSNorm w warstwach Mamba, co pomaga złagodzić problemy, takie jak duże szczyty aktywacji, które mogą wystąpić w skali.

Wydajność i benchmarking Jamba

Jamba został rygorystycznie przetestowany na szerokim zakresie benchmarków, wykazując konkurencyjną wydajność we wszystkich obszarach. Poniższe sekcje podkreślają niektóre z kluczowych benchmarków, na których Jamba wyróżnił się, prezentując swoje atuty w zadaniach NLP ogólnych i długich kontekstach.

1. Typowe benchmarki NLP

Jamba został oceniony w kilku akademickich benchmarkach, w tym:

• HellaSwag (10-shot): Zadanie rozumienia zdrowego rozsądku, w którym Jamba osiągnął wynik 87,1%, przewyższając wiele konkurujących modeli.
• WinoGrande (5-shot): Inne zadanie rozumienia, w którym Jamba uzyskał wynik 82,5%, ponownie pokazując swoją zdolność do radzenia sobie z złożonymi lingwistycznymi rozumowaniami.
• ARC-Challenge (25-shot): Jamba wykazał się silną wydajnością z wynikiem 64,4%, odzwierciedlając swoją zdolność do zarządzania wymagającymi pytaniami wielokrotnego wyboru.

W agregowanych benchmarkach, takich jak MMLU (5-shot), Jamba osiągnął wynik 67,4%, wskazując na swoją wytrzymałość w różnorodnych zadaniach.

2. Oceny długich kontekstów

Jedną z wyróżniających się cech Jamba jest jego zdolność do obsługi niezwykle długich kontekstów. Model obsługuje długość kontekstu do 256K tokenów, co jest najdłuższe wśród publicznie dostępnych modeli. Ta zdolność została przetestowana przy użyciu benchmarku Needle-in-a-Haystack, gdzie Jamba wykazał się wyjątkową dokładnością pobierania w różnych długościach kontekstu, w tym do 256K tokenów.

3. Przepływność i efektywność

Hybrydowa architektura Jamba znacznie poprawia przepływność, szczególnie w przypadku długich sekwencji.

W testach porównujących przepływność (tokeny na sekundę) w różnych modelach, Jamba konsekwentnie przewyższał swoich rywali, szczególnie w scenariuszach z dużymi rozmiarami partii i długimi kontekstami. Na przykład, z kontekstem 128K tokenów, Jamba osiągnął 3-krotnie większą przepływność niż Mixtral, porównywalny model.

Używanie Jamba: Python

Dla deweloperów i badaczy, którzy chcą eksperymentować z Jamba, AI21 Labs udostępnił model na platformach takich jak Hugging Face, czyniąc go dostępnym dla szerokiego zakresu zastosowań. Poniższy fragment kodu pokazuje, jak załadować i wygenerować tekst przy użyciu Jamba:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

input_ids = tokenizer("W ostatnim Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=216)

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

tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
model = AutoModelForCausalLM.from_pretrained(
"ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)

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")

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()