Jamba

Материал из Systems analysis wiki
Перейти к навигации Перейти к поиску

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)

Первая модель, представленная в рамках этого семейства, имеет следующие характеристики:

Технические характеристики Jamba-v0.1
Характеристика Значение
Общее число параметров 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.

Примечания

  1. «Announcing Jamba: AI21's Groundbreaking SSM-Transformer Model». AI21 Labs Blog. [1]
  2. Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
  3. «Jamba Documentation». Hugging Face Transformers. [2]
  4. «ai21labs/Jamba-v0.1». Hugging Face. [3]
  5. «AI21 Labs' Jamba: A New Hybrid LLM Architecture». Gradient Flow. [4]
  6. «Announcing the Jamba-1.5 model family». AI21 Labs Blog. [5]
  7. «ai21labs/AI21-Jamba-Large-1.5». Hugging Face. [6]
  8. «AI21 Labs разбивает новые барьеры с помощью Jamba». ITinAI. [7]
  9. 9,0 9,1 «Вышла Jamba 1.5: гибридная модель от AI21 Labs». Дзен. [8]
  10. «Jamba-1.5 family of models by AI21 Labs is now available in Amazon Bedrock». AWS What's New. [9]
  11. «ExpertsInt8: A new paradigm for efficient inference of MoE-based LLMs». OpenReview. [10]