Jamba
Jamba — это семейство больших языковых моделей (LLM), разработанное израильской исследовательской компанией AI21 Labs. Jamba представляет собой первую в своём роде гибридную архитектуру, которая объединяет ключевые элементы двух доминирующих подходов в разработке ИИ: трансформеров и моделей пространства состояний (State Space Models, SSM), в частности, архитектуры Mamba[1].
Основная цель Jamba — решить фундаментальный компромисс современных LLM: высокое качество и производительность (свойственные трансформерам) против эффективности и способности обрабатывать сверхдлинные контексты (свойственные SSM). Объединив эти подходы и добавив к ним разреженность с помощью Mixture-of-Experts (MoE), Jamba предлагает модель, которая одновременно является мощной, эффективной и способной работать с огромными объёмами текста в одном запросе.
Архитектура Jamba в деталях
Jamba не просто чередует слои трансформера и Mamba. Она использует тщательно спроектированную блочную структуру, где каждый блок состоит из восьми слоёв.
Структура одного блока Jamba:
- Один слой Трансформера: Этот слой отвечает за "глубокое" понимание и сложные рассуждения. В этот слой встроена архитектура Mixture-of-Experts (MoE).
- Семь слоёв Mamba: Эти слои следуют за слоем трансформера и отвечают за эффективную обработку последовательности и "протаскивание" информации через длинный контекст[2].
Такая асимметричная структура позволяет модели эффективно управлять вычислительными ресурсами: тяжёлые, но мощные операции трансформера выполняются реже, в то время как лёгкие и быстрые операции Mamba выполняются чаще.
Интеграция Mixture-of-Experts (MoE)
В Jamba используется архитектура MoE для дальнейшего повышения эффективности.
- MoE применяется только к полносвязным блокам (FFN) внутри слоёв-трансформеров[3]. Слои Mamba остаются плотными.
- В первой модели Jamba было 16 экспертов.
- Для каждого токена сеть-маршрутизатор выбирает 2 лучших эксперта (Top-2 gating).
Это означает, что хотя общее число параметров модели велико (52 млрд), на каждом шаге обработки токена в слое-трансформере активны только 2 из 16 экспертов, что делает вычисления очень быстрыми.
Эволюция моделей Jamba
Jamba-v0.1 (март 2024)
Первая модель, представленная в рамках этого семейства, имеет следующие характеристики:
| Характеристика | Значение |
|---|---|
| Общее число параметров | 52 миллиарда |
| Активные параметры | ~12 миллиардов |
| Число экспертов (MoE) | 16 (2 активных) |
| Контекстное окно | 256 000 токенов |
| Лицензия | Apache 2.0[4] |
Благодаря своей гибридной архитектуре, Jamba-1 способна обрабатывать контекст длиной 256 000 токенов, что эквивалентно примерно 400-страничному роману, и может быть развёрнута на одной потребительской GPU с 80 ГБ памяти[5].
Jamba-1.5 (2024)
В 2024 году AI21 Labs представила обновлённое семейство моделей Jamba 1.5, включающее две версии: Jamba 1.5 Mini (12B активных параметров из 52B общих) и Jamba 1.5 Large (94B активных параметров из 398B общих)[6]. Эти модели демонстрируют существенные улучшения в производительности:
- До 2.5 раз более быстрый вывод на длинных контекстах по сравнению с конкурентами.
- Поддержку девяти языков, включая английский, испанский, французский и арабский[7].
Ключевые преимущества и производительность
- Огромное контекстное окно: 256 000 токенов — одно из самых больших окон среди всех доступных (в том числе и проприетарных) моделей на момент её выпуска. Это делает Jamba идеальной для задач, требующих анализа больших документов: юридических контрактов, научных работ, целых кодовых баз или длинных диалогов.
- Высокая производительность и эффективность: В тестах Jamba демонстрирует производительность, сопоставимую или превосходящую ведущие открытые модели аналогичного размера, такие как Llama и Mixtral, при этом показывая в 3 раза более высокую пропускную способность на длинных контекстах[8].
- Открытость и доступность: Jamba распространяется под разрешительной лицензией Apache 2.0, что позволяет её свободное использование в коммерческих и исследовательских целях. Веса модели доступны на платформе Hugging Face.
Результаты на бенчмарках
Jamba 1.5 показывает конкурентоспособные результаты на различных бенчмарках[9]:
- Jamba 1.5 Mini набрала 46.1 балл в Arena Hard, что делает её лидирующей публичной моделью в своей категории[10].
- Jamba 1.5 Large набрала 65.4 балла в Arena Hard, превосходя Llama 3.1 70B и 405B.
Применение и доступность
Jamba оптимизирована для бизнес-применений и поддерживает такие возможности, как вызов функций, структурированный вывод в JSON и обработку документов. Модель доступна на множестве платформ, включая:
- Hugging Face
- Google Cloud Vertex AI
- Microsoft Azure
- NVIDIA API catalog
- Amazon Bedrock[9]
- AI21 Studio
Для поддержки экономически эффективного вывода AI21 Labs представила ExpertsInt8 — новую технику квантизации, позволяющую разместить Jamba 1.5 Large на машине с 8 GPU по 80GB без потери качества при обработке контекста в 256K токенов[11].
Литература
- 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.
Примечания
- ↑ «Announcing Jamba: AI21's Groundbreaking SSM-Transformer Model». AI21 Labs Blog. [1]
- ↑ Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
- ↑ «Jamba Documentation». Hugging Face Transformers. [2]
- ↑ «ai21labs/Jamba-v0.1». Hugging Face. [3]
- ↑ «AI21 Labs' Jamba: A New Hybrid LLM Architecture». Gradient Flow. [4]
- ↑ «Announcing the Jamba-1.5 model family». AI21 Labs Blog. [5]
- ↑ «ai21labs/AI21-Jamba-Large-1.5». Hugging Face. [6]
- ↑ «AI21 Labs разбивает новые барьеры с помощью Jamba». ITinAI. [7]
- ↑ 9,0 9,1 «Вышла Jamba 1.5: гибридная модель от AI21 Labs». Дзен. [8]
- ↑ «Jamba-1.5 family of models by AI21 Labs is now available in Amazon Bedrock». AWS What's New. [9]
- ↑ «ExpertsInt8: A new paradigm for efficient inference of MoE-based LLMs». OpenReview. [10]