Jamba (language model) (CS)

Jamba — je rodina velkých jazykových modelů (LLM) vyvinutá izraelskou výzkumnou společností AI21 Labs. Jamba představuje svého druhu první hybridní architekturu, která spojuje klíčové prvky dvou dominantních přístupů ve vývoji AI: transformerů a modelů stavového prostoru (State Space Models, SSM), konkrétně architektury Mamba.

Hlavním cílem Jamba je vyřešit zásadní kompromis současných LLM: vysoká kvalita a výkon (vlastní transformerům) versus efektivita a schopnost zpracovávat velmi dlouhé kontexty (vlastní SSM). Spojením těchto přístupů a přidáním řídkosti pomocí Mixture-of-Experts (MoE) nabízí Jamba model, který je zároveň výkonný, efektivní a schopný pracovat s obrovskými objemy textu v jednom dotazu.

Architektura Jamba v detailech

Jamba neprosté střídá vrstvy transformeru a Mamba. Využívá pečlivě navrženou blokovou strukturu, kde každý blok se skládá z osmi vrstev.

Struktura jednoho bloku Jamba:

• Jedna vrstva Transformeru: Tato vrstva zajišťuje „hluboké" porozumění a složité uvažování. Do této vrstvy je zabudována architektura Mixture-of-Experts (MoE).
• Sedm vrstev Mamba: Tyto vrstvy následují za vrstvou transformeru a zajišťují efektivní zpracování sekvencí a „protahování" informací přes dlouhý kontext.

Tato asymetrická struktura umožňuje modelu efektivně řídit výpočetní zdroje: náročné, ale výkonné operace transformeru se provádějí méně často, zatímco lehké a rychlé operace Mamba se provádějí častěji.

Integrace Mixture-of-Experts (MoE)

Jamba využívá architekturu MoE pro další zvýšení efektivity.

• MoE se aplikuje pouze na plně propojené bloky (FFN) uvnitř vrstev transformeru. Vrstvy Mamba zůstávají husté.
• V prvním modelu Jamba bylo 16 expertů.
• Pro každý token vybírá síť-směrovač 2 nejlepší experty (Top-2 gating).

To znamená, že ačkoli je celkový počet parametrů modelu velký (52 miliard), při každém kroku zpracování tokenu ve vrstvě transformeru jsou aktivní pouze 2 z 16 expertů, což výpočty značně urychluje.

Evoluce modelů Jamba

Jamba-v0.1 (březen 2024)

První model představený v rámci této rodiny má následující charakteristiky:

Technické charakteristiky Jamba-v0.1

Charakteristika	Hodnota
Celkový počet parametrů	52 miliard
Aktivní parametry	~12 miliard
Počet expertů (MoE)	16 (2 aktivní)
Kontextové okno	256 000 tokenů
Licence	Apache 2.0

Díky své hybridní architektuře je Jamba-1 schopna zpracovat kontext o délce 256 000 tokenů, což odpovídá přibližně románu o 400 stranách, a lze ji nasadit na jedné spotřebitelské GPU s 80 GB paměti.

Jamba-1.5 (2024)

V roce 2024 představila AI21 Labs aktualizovanou rodinu modelů Jamba 1.5, která zahrnuje dvě verze: Jamba 1.5 Mini (12 miliard aktivních parametrů z celkových 52 miliard) a Jamba 1.5 Large (94 miliard aktivních parametrů z celkových 398 miliard). Tyto modely vykazují podstatná zlepšení výkonu:

• Až 2,5krát rychlejší inference na dlouhých kontextech ve srovnání s konkurencí.
• Podporu devíti jazyků, včetně angličtiny, španělštiny, francouzštiny a arabštiny.

Klíčové výhody a výkon

• Obrovské kontextové okno: 256 000 tokenů — jedno z největších oken mezi všemi dostupnými (včetně proprietárních) modely v době jeho vydání. To činí Jamba ideální pro úlohy vyžadující analýzu rozsáhlých dokumentů: právních smluv, vědeckých prací, celých kódových bází nebo dlouhých dialogů.
• Vysoký výkon a efektivita: V testech Jamba vykazuje výkon srovnatelný s předními otevřenými modely podobné velikosti, jako jsou Llama a Mixtral, nebo je předčí, přičemž dosahuje 3krát vyšší propustnosti na dlouhých kontextech.
• Otevřenost a dostupnost: Jamba je šířena pod permisivní licencí Apache 2.0, která umožňuje její volné využití pro komerční i výzkumné účely. Váhy modelu jsou dostupné na platformě Hugging Face.

Výsledky na benchmarcích

Jamba 1.5 dosahuje konkurenceschopných výsledků na různých benchmarcích:

• Jamba 1.5 Mini získala 46,1 bodu v Arena Hard, což z ní činí vedoucí veřejný model ve své kategorii.
• Jamba 1.5 Large získala 65,4 bodu v Arena Hard, čímž překonává Llama 3.1 70B a 405B.

Využití a dostupnost

Jamba je optimalizována pro podnikové aplikace a podporuje funkce jako volání funkcí, strukturovaný výstup v JSON a zpracování dokumentů. Model je dostupný na mnoha platformách, včetně:

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

Pro podporu ekonomicky efektivní inference představila AI21 Labs ExpertsInt8 — novou techniku kvantizace, která umožňuje nasadit Jamba 1.5 Large na stroji s 8 GPU po 80 GB bez ztráty kvality při zpracování kontextu o délce 256 tisíc tokenů.

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.