Modele i platformy AI
Jamba: nowy hybrydowy model językowy AI21 Labs łączący architektury Transformer i Mamba
Dodaj Unite.AI do preferowanych źródeł w Google
Modele językowe przechodzą przez szybki rozwój, a architektury oparte na Transformerze 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 przepustowością stają się coraz bardziej wyraźne.
AI21 Labs wprowadziło nowe rozwiązanie z Jambą, najnowocześniejszym dużym modelem językowym (LLM), który łączy zalety architektur Transformer i Mamba w ramach hybrydowej struktury. Artykuł ten szczegółowo opisuje architekturę Jambi, jej wydajność i potencjalne zastosowania.
Przegląd Jambi
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 Jambie na równowagę pomiędzy użyciem pamięci, przepustowością i wydajnością, czyniąc ją potężnym narzędziem dla szerokiego zakresu zadań NLP. Model został zaprojektowany, aby zmieścić się w jednej karcie graficznej 80GB, oferując wysoką przepustowość i niewielką stopę pamięci, jednocześnie utrzymując najnowocześniejszą wydajność na różnych benchmarkach.
Architektura Jambi
Architektura Jambi jest kamieniem węgielnym jej możliwości. Została zbudowana na nowatorskiej hybrydowej konstrukcji, która 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 przy przetwarzaniu 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 nowy model przestrzeni stanu (SSM) zaprojektowany do efektywniejszego radzenia sobie z dalekosiężnymi relacjami w sekwencjach niż tradycyjne RNN lub nawet Transformatory. Warstwy Mamba są szczególnie skuteczne w redukowaniu śladu pamięci związanego z przechowywaniem pamięci klucz-wartość (KV) w Transformatorkach. Przeplatając warstwy Mamba z warstwami Transformer, Jamba zmniejsza ogólne użycie pamięci, jednocześnie utrzymując wysoką wydajność, szczególnie w zadaniach wymagających długiego kontekstu.
3. Moduł Mixture-of-Experts (MoE)
Moduł MoE w Jambie wprowadza elastyczną strategię skalowania pojemności modelu. MoE pozwala modelowi zwiększyć liczbę dostępnych parametrów bez proporcjonalnego zwiększania aktywnych parametrów podczas inferencji. W Jambie 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 Jambie utrzymanie wysokiej wydajności przy jednoczesnym radzeniu sobie z złożonymi zadaniami.
Poniższy obraz demonstruje funkcjonalność głowy indukcyjnej w hybrydowym modelu Attention-Mamba, kluczowej cechy Jambi. W tym przykładzie głowa uwagi jest odpowiedzialna za przewidywanie etykiet, takich jak “Pozytywny” lub “Negatywny” w odpowiedzi na zadania analizy nastrojów. Podświetlone słowa ilustrują, 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ć odpowiednią etykietę na podstawie danego kontekstu i przykładów few-shot.
Poprawy wydajności oferowane przez integrację Mixture-of-Experts (MoE) z hybrydową architekturą Attention-Mamba są podkreślone 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 Jambi składa się z bloków, które łączą warstwy 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 dwoma najlepszymi ekspertami na token. Ta konfiguracja pozwala Jambie skalować efektywnie, jednocześnie zarządzając kompromisami pomiędzy użyciem pamięci a efektywnością obliczeniową.
- Normalizacja i stabilność: Aby zapewnić stabilność podczas treningu, Jamba wprowadza RMSNorm w warstwach Mamba, co pomaga złagodzić problemy, takie jak duże skoki aktywacji, które mogą wystąpić w skali.
Wydajność i benchmarking Jambi
Jamba została poddana rygorystycznym testom 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 się wyróżnia, prezentując jej siłę w zadaniach NLP ogólnych i długich kontekstów.
1. Typowe benchmarki NLP
Jamba została oceniona na kilku akademickich benchmarkach, w tym:
- HellaSwag (10-shot): Zadanie rozumienia zdrowego rozsądku, w którym Jamba osiągnęła wynik 87,1%, przewyższając wiele konkurujących modeli.
- WinoGrande (5-shot): Inne zadanie rozumienia, w którym Jamba uzyskała wynik 82,5%, ponownie pokazując swoją zdolność do radzenia sobie z złożonymi lingwistycznymi rozumowaniami.
- ARC-Challenge (25-shot): Jamba wykazała silną wydajność z wynikiem 64,4%, odzwierciedlając jej zdolność do zarządzania wymagającymi pytaniami wielokrotnego wyboru.
W agregowanych benchmarkach, takich jak MMLU (5-shot), Jamba osiągnęła wynik 67,4%, wskazując na jej solidność w różnorodnych zadaniach.
2. Oceny długiego kontekstu
Jedną z wyróżniających się cech Jambi jest jej zdolność do radzenia sobie z ekstremalnie długimi kontekstami. 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, w którym Jamba wykazała wyjątkową dokładność odzyskiwania w różnych długościach kontekstu, w tym do 256K tokenów.
3. Przepustowość i efektywność
Hybrydowa architektura Jambi znacząco poprawia przepustowość, szczególnie w przypadku długich sekwencji.
W testach porównujących przepustowość (tokeny na sekundę) w różnych modelach, Jamba konsekwentnie przewyższała 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ęła trzykrotnie większą przepustowość niż Mixtral, porównywalny model.
Użycie Jambi: Python
Dla deweloperów i badaczy, którzy chcą eksperymentować z Jambą, AI21 Labs udostępnił model na platformach takich jak Hugging Face, co umożliwia dostęp do szerokiego zakresu zastosowań. Poniższy fragment kodu demonstruje, jak załadować i wygenerować tekst przy użyciu Jambi:
<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("W ostatnim 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))
Ten prosty skrypt ładuje model Jambi i tokenizator, generuje tekst na podstawie danego wejścia i drukuje wygenerowany wynik.
Dokształcanie Jambi
Jamba została zaprojektowana jako model podstawowy, co oznacza, że może być dostosowywana do konkretnych zadań lub aplikacji. Dokształcanie pozwala użytkownikom adaptować model do niszowych dziedzin, poprawiając wydajność w specjalistycznych zadaniach. Poniższy przykład pokazuje, jak dokształcić Jambę przy użyciu biblioteki PEFT:
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()
Ten fragment kodu dokształca Jambę na zbiorze danych angielskich cytatów, dostosowując parametry modelu do lepszego dopasowania do zadania generowania tekstu w specjalistycznej dziedzinie.
Wdrożenie i integracja
AI21 Labs udostępniła rodzinę Jambi na różnych platformach i opcjach wdrożeniowych:
- Platformy chmurowe:
- Dostępna na głównych dostawcach chmury, w tym Google Cloud Vertex AI, Microsoft Azure i NVIDIA NIM (NVDA ).
- Wkrótce dostępna na Amazon Bedrock, Databricks Marketplace i Snowflake Cortex.
- Ramowe deweloperskie AI:
- Integracja z popularnymi ramami, takimi jak LangChain i LlamaIndex (wkrótce).
- AI21 Studio:
- Bezpośredni dostęp przez platformę deweloperską AI21.
- Hugging Face:
- Modele dostępne do pobrania i eksperymentowania.
- Wdrożenie na miejscu:
- Opcje wdrożenia prywatnego dla organizacji z określonymi wymogami bezpieczeństwa lub zgodności.
- Rozwiązania niestandardowe:
- AI21 oferuje usługi dostosowywania i dokształcania modelu dla klientów przedsiębiorstw.
Cechy przyjazne dla deweloperów
Modele Jambi posiadają kilka wbudowanych funkcji, które sprawiają, że są one szczególnie atrakcyjne dla deweloperów:
- Wywoływanie funkcji: Łatwo integruj zewnętrzne narzędzia i API z przepływami AI.
- Wydajność JSON: Generuj czyste, parsowalne struktury danych bezpośrednio z wejść języka naturalnego.
- Przetwarzanie dokumentów: Efektywnie przetwarzaj i rozumiej złożone struktury dokumentów.
- Optymalizacje RAG: Wbudowane funkcje do poprawy potoku generacji z uzupełnieniem.
Te funkcje, w połączeniu z długim oknem kontekstu i efektywnym przetwarzaniem, sprawiają, że Jamba jest wszechstronnym narzędziem dla szerokiego zakresu scenariuszy deweloperskich.
Zagadnienia etyczne i odpowiedzialne AI
Podczas gdy możliwości Jambi są imponujące, ważne jest, aby podejść do jej użycia z odpowiedzialnym podejściem do AI. AI21 Labs podkreśla kilka ważnych punktów:
- Natura modelu podstawowego: Modele Jambi 1.5 są wstępnie trenowanymi modelami podstawowymi bez konkretnego wyrównania lub dostosowania.
- Brak wbudowanych zabezpieczeń: Modele nie posiadają wewnętrznych mechanizmów moderacji.
- Ostrozne wdrożenie: Dodatkowe adaptacje i zabezpieczenia powinny być wdrożone przed użyciem Jambi w środowiskach produkcyjnych lub z użytkownikami końcowymi.
- Prywatność danych: Podczas korzystania z wdrożeń opartych na chmurze, należy być świadomym obsługi i wymogów zgodności danych.
- Świadomość biasu: Jak wszystkie duże modele językowe, Jamba może odzwierciedlać biasy obecne w danych szkoleniowych. Użytkownicy powinni być tego świadomi i wdrożyć odpowiednie środki zaradcze.
Poprzez uwzględnienie tych czynników, deweloperzy i organizacje mogą wykorzystywać możliwości Jambi w sposób odpowiedzialny i etyczny.
Nowy rozdział w rozwoju AI?
Wprowadzenie rodziny Jambi przez AI21 Labs oznacza znaczący kamień milowy w ewolucji dużych modeli językowych. Łącząc zalety Transfomerów i modeli przestrzeni stanu, integrując techniki Mixture-of-Experts i poszerzając granice długości kontekstu i szybkości przetwarzania, Jamba otwiera nowe możliwości dla aplikacji AI w różnych branżach.
Podczas gdy społeczność AI kontynuuje eksplorację i rozwój tej innowacyjnej architektury, możemy oczekiwać dalszych postępów w efektywności modelu, zrozumieniu długich kontekstów i praktycznym wdrożeniu AI. Rodzina Jambi reprezentuje nie tylko nowy zestaw modeli, ale potencjalną zmianę w podejściu do projektowania i wdrażania dużych systemów AI.
