AI21 Labs усиливает генеративные AI трансформеры благодаря сотрудничеству с Jamba

С момента публикации революционной научной статьи "Внимание — это всё, что вам нужно" в 2017 году трансформеры заняли центральное место в области генеративного ИИ. Однако transformers — не единственный подход к генеративному ИИ. AI21 Labs представила новую структуру под названием "Jamba", которая стремится выйти за рамки традиционных трансформеров.

Jamba объединяет модель Mamba, основанную на модели структурированного состояния (SSM), с архитектурой трансформеров для создания оптимизированного решения в области генеративного ИИ. Название "Jamba" расшифровывается как Joint Attention and Mamba Architecture, что позволяет использовать сильные стороны как SSM, так и трансформеров. Эта модель выпущена в открытый доступ под лицензией Apache 2.0.

Хотя Jamba не предназначена для замены существующих крупных языковых моделей (LLM) на основе трансформеров, ожидается, что она станет ценным дополнением в определенных приложениях. Представители AI21 Labs утверждают, что Jamba может surpass традиционные трансформерные модели в задачах генеративного рассуждения, что подтверждают такие бенчмарки, как HellaSwag. Однако на критических бенчмарках, таких как Massive Multitask Language Understanding (MMLU), подход Jamba пока не превосходит модели на основе трансформеров, которые оценивают способности решения задач.

AI21 Labs специализируется на генеративном ИИ для корпоративных приложений и недавно привлекла 155 миллионов долларов в августе 2023 года для развития своих инициатив. В числе ее корпоративных предложений есть Wordtune, инструмент, предназначенный для помощи организациям в создании контента, соответствующего их тону и брендингу. В 2023 году компания сообщила, что успешно конкурировала с гигантом генеративного ИИ OpenAI за привлечение корпоративных клиентов.

Традиционно технологии LLM от AI21 Labs использовали архитектуру трансформеров, включая семью LLM Jurassic-2, которая является частью платформы обработки естественного языка AI21 Studio и доступна через API для интеграции в бизнес. Однако Jamba представляет собой переход к гибридной модели SSM и трансформера.

Несмотря на выдающуюся роль трансформеров в генеративном ИИ, у них есть определенные ограничения. Существенной проблемой является то, что с увеличением окон контекста время вывода замедляется. Как объясняют исследователи AI21 Labs, механизм внимания трансформера масштабируется с длиной последовательности, что приводит к снижению пропускной способности, поскольку каждый токен зависит от всей предыдущей последовательности. Это делает применение к длинным контекстам неэффективным.

Другой проблемой является значительный объем памяти, необходимый для масштабирования трансформеров. Их потребности в памяти растут с увеличением длины контекста, что затрудняет обработку длинных контекстов или нескольких параллельных задач без значительных аппаратных ресурсов. Подход SSM стремится решить эти проблемы с контекстом и памятью.

Архитектура Mamba SSM, изначально разработанная исследователями Университета Карнеги-Меллон и Принстона, требует меньше памяти и использует другой механизм внимания для управления большими окнами контекста. Тем не менее, она не может обеспечить такое же качество вывода, как модели трансформеров. Гибридный подход Jamba сочетает оптимизацию ресурсов и контекста SSM с возможностями вывода трансформеров.

AI21 Labs утверждает, что Jamba имеет окно контекста в 256K и предлагает в три раза большую пропускную способность для длинных контекстов по сравнению с Mixtral 8x7B. Примечательно, что Jamba позиционируется как единственная модель в своем классе, способная обрабатывать до 140K контекста на одном GPU.

Подобно Mixtral, Jamba включает модель Mixture of Experts (MoE). Однако Jamba использует MoE в своей гибридной архитектуре SSM и трансформеров, что позволяет достигать более высокого уровня оптимизации. В частности, во время вывода активируются лишь 12 миллиардов из доступных 52 миллиардов параметров Jamba, что делает ее более эффективной, чем модель только на трансформерах сопоставимого размера, согласно AI21 Labs.

На данный момент Jamba находится на ранней стадии разработки и пока не является частью корпоративных предложений AI21 Labs, хотя компания планирует в ближайшее время выпустить обучающую версию на платформе AI21 в бета-версии.