Jamba (language model) (PL)

Jamba — to rodzina dużych modeli językowych (LLM), opracowana przez izraelską firmę badawczą AI21 Labs. Jamba stanowi pierwszą w swoim rodzaju hybrydową architekturę, która łączy kluczowe elementy dwóch dominujących podejść w tworzeniu AI: transformerów oraz modeli przestrzeni stanów (State Space Models, SSM), w szczególności architektury Mamba.

Głównym celem Jamba jest rozwiązanie fundamentalnego kompromisu współczesnych LLM: wysokiej jakości i wydajności (właściwych transformerom) w stosunku do efektywności i zdolności przetwarzania bardzo długich kontekstów (właściwych SSM). Łącząc te podejścia i dodając do nich rzadkość przy pomocy Mixture-of-Experts (MoE), Jamba oferuje model, który jednocześnie jest wydajny, efektywny i zdolny do pracy z ogromnymi ilościami tekstu w jednym zapytaniu.

Architektura Jamba w szczegółach

Jamba nie tylko naprzemiennie stosuje warstwy transformera i Mamba. Wykorzystuje starannie zaprojektowaną strukturę blokową, gdzie każdy blok składa się z ośmiu warstw.

Struktura jednego bloku Jamba:

• Jedna warstwa Transformera: Ta warstwa odpowiada za „głębokie" rozumienie i złożone wnioskowanie. W tę warstwę wbudowana jest architektura Mixture-of-Experts (MoE).
• Siedem warstw Mamba: Te warstwy następują po warstwie transformera i odpowiadają za efektywne przetwarzanie sekwencji oraz „przepuszczanie" informacji przez długi kontekst.

Taka asymetryczna struktura pozwala modelowi efektywnie zarządzać zasobami obliczeniowymi: ciężkie, lecz wydajne operacje transformera wykonywane są rzadziej, natomiast lekkie i szybkie operacje Mamba wykonywane są częściej.

Integracja Mixture-of-Experts (MoE)

W Jamba stosowana jest architektura MoE w celu dalszego zwiększenia efektywności.

• MoE stosowany jest wyłącznie do bloków w pełni połączonych (FFN) wewnątrz warstw-transformerów. Warstwy Mamba pozostają gęste.
• W pierwszym modelu Jamba było 16 ekspertów.
• Dla każdego tokenu sieć-router wybiera 2 najlepszych ekspertów (Top-2 gating).

Oznacza to, że choć całkowita liczba parametrów modelu jest duża (52 mld), na każdym kroku przetwarzania tokenu w warstwie-transformerze aktywne są tylko 2 z 16 ekspertów, co sprawia, że obliczenia są bardzo szybkie.

Ewolucja modeli Jamba

Jamba-v0.1 (marzec 2024)

Pierwszy model przedstawiony w ramach tej rodziny posiada następujące charakterystyki:

Specyfikacje techniczne Jamba-v0.1

Charakterystyka	Wartość
Całkowita liczba parametrów	52 miliardy
Aktywne parametry	~12 miliardów
Liczba ekspertów (MoE)	16 (2 aktywnych)
Okno kontekstowe	256 000 tokenów
Licencja	Apache 2.0

Dzięki swojej hybrydowej architekturze, Jamba-1 jest w stanie przetwarzać kontekst o długości 256 000 tokenów, co odpowiada mniej więcej powieści liczącej 400 stron, i może być wdrożona na jednym konsumenckim GPU z pamięcią 80 GB.

Jamba-1.5 (2024)

W 2024 roku AI21 Labs przedstawiła zaktualizowaną rodzinę modeli Jamba 1.5, obejmującą dwie wersje: Jamba 1.5 Mini (12B aktywnych parametrów z 52B łącznych) oraz Jamba 1.5 Large (94B aktywnych parametrów z 398B łącznych). Modele te wykazują istotne ulepszenia w zakresie wydajności:

• Do 2,5 raza szybszy wnioskowanie na długich kontekstach w porównaniu z konkurentami.
• Obsługę dziewięciu języków, w tym angielskiego, hiszpańskiego, francuskiego i arabskiego.

Kluczowe zalety i wydajność

• Ogromne okno kontekstowe: 256 000 tokenów — jedno z największych okien spośród wszystkich dostępnych (w tym własnościowych) modeli w momencie jej premiery. Sprawia to, że Jamba jest idealna do zadań wymagających analizy dużych dokumentów: umów prawnych, prac naukowych, całych baz kodu lub długich dialogów.
• Wysoka wydajność i efektywność: W testach Jamba demonstruje wydajność porównywalną lub przewyższającą czołowe otwarte modele podobnej wielkości, takie jak Llama i Mixtral, przy jednoczesnym wykazywaniu 3-krotnie wyższej przepustowości na długich kontekstach.
• Otwartość i dostępność: Jamba dystrybuowana jest na permisywnej licencji Apache 2.0, co umożliwia jej swobodne wykorzystanie w celach komercyjnych i badawczych. Wagi modelu dostępne są na platformie Hugging Face.

Wyniki na benchmarkach

Jamba 1.5 osiąga konkurencyjne wyniki na różnych benchmarkach:

• Jamba 1.5 Mini uzyskała 46,1 punktu w Arena Hard, co czyni ją wiodącym publicznym modelem w swojej kategorii.
• Jamba 1.5 Large uzyskała 65,4 punktu w Arena Hard, przewyższając Llama 3.1 70B i 405B.

Zastosowanie i dostępność

Jamba jest zoptymalizowana pod kątem zastosowań biznesowych i obsługuje takie możliwości jak wywoływanie funkcji, strukturyzowany output w JSON oraz przetwarzanie dokumentów. Model dostępny jest na wielu platformach, w tym:

• Hugging Face
• Google Cloud Vertex AI
• Microsoft Azure
• NVIDIA API catalog
• Amazon Bedrock
• AI21 Studio

W celu wspierania ekonomicznie efektywnego wnioskowania AI21 Labs przedstawiła ExpertsInt8 — nową technikę kwantyzacji, pozwalającą na uruchomienie Jamba 1.5 Large na maszynie z 8 GPU po 80 GB bez utraty jakości przy przetwarzaniu kontekstu o rozmiarze 256K tokenów.

Literatura

• Lieber, O.; et al. (2024). Jamba: A Hybrid Transformer‑Mamba Language Model. arXiv:2403.19887.
• Lieber, O.; et al. (2024). Jamba‑1.5 Models and ExpertsInt8 Quantization. OpenReview JFPaD7lpBD.
• Gu, A.; Dao, T. (2023). Mamba: Linear‑Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
• Gu, A.; et al. (2021). S4: Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396.
• Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
• Yun, L.; et al. (2024). Toward Inference‑Optimal Mixture‑of‑Expert Large Language Models. arXiv:2404.02852.
• Liu, J.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
• Gupta, V.; et al. (2024). Lynx: Enabling Efficient MoE Inference through Dynamic Batch‑Aware Expert Selection. arXiv:2411.08982.
• Liu, J.; et al. (2024). A Survey on Inference Optimization Techniques for Mixture of Experts Models. arXiv:2412.14219.
• Hsieh, C.‑P.; et al. (2024). RULER: What's the Real Context Size of Your Long‑Context Language Models?. arXiv:2404.06654.