Модели и платформы ИИ
Джамба: Новый Гибридный Трансформер-Мамба Языковая Модель от AI21 Labs
Добавьте Unite.AI в избранные источники в Google
Языковые модели пережили быстрый прогресс, и архитектуры на основе трансформеров стали стандартом для современной обработки естественного языка. Однако по мере увеличения масштаба моделей возросли проблемы с обработкой длинных контекстов, эффективностью памяти и пропускной способностью.
AI21 Labs представила новое решение с Джамбой, современной большой языковой моделью (LLM), которая сочетает сильные стороны как трансформерных, так и мамба-архитектур в гибридной структуре. Эта статья описывает архитектуру Джамбы, ее производительность и потенциальные применения.
Обзор Джамбы
Джамба – это гибридная большая языковая модель, разработанная AI21 Labs, которая сочетает трансформерные слои и мамба-слои, интегрированные с модулем Mixture-of-Experts (MoE). Эта архитектура позволяет Джамбе сбалансировать использование памяти, пропускную способность и производительность, что делает ее мощным инструментом для широкого спектра задач обработки естественного языка. Модель разработана для работы в рамках одного 80ГБ GPU, обеспечивая высокую пропускную способность и небольшой объем памяти при сохранении передового уровня производительности на различных бенчмарках.
Архитектура Джамбы
Архитектура Джамбы является основой ее возможностей. Она построена на новой гибридной конструкции, которая чередует трансформерные слои с мамба-слоями, включая модули MoE для повышения емкости модели без значительного увеличения вычислительных требований.
1. Трансформерные Слои
Архитектура трансформера стала стандартом для современных больших языковых моделей благодаря своей способности эффективно обрабатывать параллельную обработку и захватывать длинные зависимости в тексте. Однако ее производительность часто ограничена высокими требованиями к памяти и вычислениям, особенно при обработке длинных контекстов. Джамба решает эти ограничения, интегрируя мамба-слои, которые мы рассмотрим далее.
2. Мамба-Слои
Мамба – это недавно разработанная модель состояния-пространства (SSM), предназначенная для более эффективной обработки длинных отношений в последовательностях по сравнению с традиционными РНН или даже трансформерами. Мамба-слои особенно эффективны в снижении объема памяти, связанного с хранением кэшей ключ-значение (KV) в трансформерах. Чередуя мамба-слои с трансформерными слоями, Джамба снижает общее использование памяти, сохраняя высокую производительность, особенно в задачах, требующих обработки длинных контекстов.
3. Модули Mixture-of-Experts (MoE)
Модуль MoE в Джамбе вводит гибкий подход к масштабированию емкости модели. MoE позволяет модели увеличивать количество доступных параметров без пропорционального увеличения активных параметров во время вывода. В Джамбе MoE применяется к некоторым слоям MLP, с механизмом маршрутизации, выбирающим лучших экспертов для активации для каждого токена. Этот избирательный механизм активации позволяет Джамбе поддерживать высокую эффективность при обработке сложных задач.
Ниже приведено изображение, демонстрирующее функциональность индукционного заголовка в гибридной модели Attention-Mamba, ключевой особенности Джамбы. В этом примере заголовок внимания отвечает за предсказание меток, таких как “Положительный” или “Отрицательный”, в ответ на задачи анализа настроений. Выделенные слова иллюстрируют, как внимание модели сильно сосредоточено на токенах меток из примеров с несколькими выстрелами, особенно в критический момент перед предсказанием окончательной метки. Этот механизм внимания играет решающую роль в способности модели выполнять контекстное обучение, когда модель должна выводить соответствующую метку на основе заданного контекста и примеров с несколькими выстрелами.
Улучшения производительности, обеспечиваемые интеграцией Mixture-of-Experts (MoE) с гибридной архитектурой Attention-Mamba, выделены в таблице. Используя MoE, Джамба увеличивает свою емкость без пропорционального увеличения вычислительных затрат. Это особенно очевидно в значительном увеличении производительности на различных бенчмарках, таких как HellaSwag, WinoGrande и Natural Questions (NQ). Модель с MoE не только достигает более высокой точности (например, 66,0% на WinoGrande по сравнению с 62,5% без MoE), но также демонстрирует улучшенные логарифмические вероятности в различных доменах (например, -0,534 на C4).
Ключевые Архитектурные Особенности
- Состав Слоев: Архитектура Джамбы состоит из блоков, которые сочетают мамба и трансформерные слои в определенном соотношении (например, 1:7, что означает один трансформерный слой для каждого семи мамба-слоев). Это соотношение настроено для оптимальной производительности и эффективности.
- Интеграция MoE: Слои MoE применяются каждые несколько слоев, с 16 доступными экспертами и двумя лучшими экспертами, активированными на каждый токен. Эта конфигурация позволяет Джамбе эффективно масштабироваться, управляя компромиссами между использованием памяти и вычислительной эффективностью.
- Нормализация и Стабильность: Для обеспечения стабильности во время обучения Джамба включает RMSNorm в мамба-слоях, что помогает смягчить проблемы, такие как большие всплески активации, которые могут возникнуть при масштабировании.
Производительность и Бенчмаркинг Джамбы
Джамба была тщательно протестирована на широком спектре бенчмарков, демонстрируя конкурентоспособную производительность во всех аспектах. В следующих разделах выделены некоторые ключевые бенчмарки, где Джамба продемонстрировала свои сильные стороны как в общих задачах обработки естественного языка, так и в сценариях с длинными контекстами.
1. Общие Бенчмарки Обработки Естественного Языка
Джамба была оценена на нескольких академических бенчмарках, включая:
- HellaSwag (10-шот): Задача рассуждения на основе здравого смысла, где Джамба достигла производительности 87,1%, превзойдя многие конкурирующие модели.
- WinoGrande (5-шот): Другая задача рассуждения, где Джамба набрала 82,5%, еще раз демонстрируя свою способность обрабатывать сложные лингвистические рассуждения.
- ARC-Challenge (25-шот): Джамба продемонстрировала сильную производительность с результатом 64,4%, отражая ее способность управлять сложными вопросами с несколькими вариантами ответов.
В агрегированных бенчмарках, таких как MMLU (5-шот), Джамба достигла результата 67,4%, указывая на ее прочность в различных задачах.
2. Оценки с Длинными Контекстами
Одной из выдающихся особенностей Джамбы является ее способность обрабатывать чрезвычайно длинные контексты. Модель поддерживает длину контекста до 256К токенов, что является самым длинным среди публично доступных моделей. Эта способность была протестирована с помощью бенчмарка Needle-in-a-Haystack, где Джамба показала исключительную точность извлечения в различных контекстных длинах, включая до 256К токенов.
3. Пропускная Способность и Эффективность
Гибридная архитектура Джамбы значительно улучшает пропускную способность, особенно при работе с длинными последовательностями.
В тестах, сравнивающих пропускную способность (токенов в секунду) между различными моделями, Джамба последовательно превосходила своих конкурентов, особенно в сценариях с большими размерами пакетов и длинными контекстами. Например, с контекстом из 128К токенов Джамба достигла в 3 раза большей пропускной способности, чем Mixtral, сравнимая модель.
Использование Джамбы: Python
Для разработчиков и исследователей, желающих экспериментировать с Джамбой, AI21 Labs предоставила модель на платформах, таких как Hugging Face, что делает ее доступной для широкого спектра приложений. Ниже приведен фрагмент кода, демонстрирующий, как загрузить и сгенерировать текст с помощью Джамбы:
<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p> <p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1") tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p> <p>input_ids = tokenizer("В недавнем Супербоуле LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p> <p>outputs = model.generate(input_ids, max_new_tokens=216)</p> print(tokenizer.batch_decode(outputs))
Этот простой скрипт загружает модель Джамбы и токенизатор, генерирует текст на основе заданного входного запроса и выводит сгенерированный вывод.
Настройка Джамбы
Джамба разработана как базовая модель, что означает, что ее можно настроить для конкретных задач или приложений. Настройка позволяет пользователям адаптировать модель к нишевым областям, улучшая производительность в специализированных задачах. Ниже приведен пример того, как настроить Джамбу с помощью библиотеки PEFT:
import torch from datasets import load_dataset from trl import SFTTrainer, SFTConfig from peft import LoraConfig from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments <p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1") model = AutoModelForCausalLM.from_pretrained( "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p> <p>lora_config = LoraConfig(r=8, target_modules=[ "embed_tokens","x_proj", "in_proj", "out_proj", # мамба "gate_proj", "up_proj", "down_proj", # мlp "q_proj", "k_proj", "v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p> <p>dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote") trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset, ) trainer.train()
Этот фрагмент кода настраивает Джамбу на наборе данных английских цитат, регулируя параметры модели для лучшего соответствия конкретной задаче генерации текста в специализированной области.
Развертывание и Интеграция
AI21 Labs сделала семейство Джамбы широко доступным через различные платформы и варианты развертывания:
- Облачные Платформы:
- Доступна на основных облачных провайдерах, включая Google Cloud Vertex AI, Microsoft Azure и NVIDIA NIM (NVDA ).
- Скоро будет доступна на Amazon Bedrock, Databricks Marketplace и Snowflake Cortex.
- Фреймворки Разработки ИИ:
- Интеграция с популярными фреймворками, такими как LangChain и LlamaIndex (планируется).
- AI21 Studio:
- Прямой доступ через платформу разработки AI21.
- Hugging Face:
- Модели доступны для скачивания и экспериментов.
- Развертывание на Премиссе:
- Варианты для частного развертывания на месте для организаций с особыми требованиями безопасности или соответствия.
- Пользовательские Решения:
- AI21 предлагает индивидуальную настройку модели и услуги по тонкой настройке для корпоративных клиентов.
Возможности для Разработчиков
Модели Джамбы поставляются с несколькими встроенными возможностями, которые делают их особенно привлекательными для разработчиков:
- Вызов Функций: Легко интегрируйте внешние инструменты и API в ваши рабочие процессы ИИ.
- Структурированный JSON-Вывод: Генерируйте чистые, парсируемые структуры данных напрямую из естественного языка.
- Документный Объектный Дайджест: Эффективно обрабатывайте и понимайте сложные структуры документов.
- Оптимизации RAG: Встроенные функции для улучшения конвейеров генерации с извлечением.
Эти возможности, в сочетании с длинным контекстным окном модели и эффективной обработкой, делают Джамбу универсальным инструментом для широкого спектра сценариев разработки.
Этические Рассмотрения и Ответственный ИИ
Хотя возможности Джамбы впечатляющие, важно подойти к ее использованию с ответственным подходом к ИИ. AI21 Labs подчеркивает несколько важных моментов:
- Базовая Модель: Модели Джамбы 1.5 – это базовые модели без конкретной настройки или инструкций.
- Отсутствие Встроенных Защитных Механизмов: Модели не имеют встроенных механизмов модерации.
- Осторожное Развертывание: Дополнительная адаптация и защитные меры должны быть реализованы перед использованием Джамбы в производственных средах или с конечными пользователями.
- Защита Данных: При использовании облачных развертываний будьте внимательны к обработке и требованиям соответствия данных.
- Осведомленность о Предвзятости: Как и все большие языковые модели, Джамба может отражать предвзятости, присутствующие в ее обучающих данных. Пользователи должны быть осведомлены об этом и реализовать соответствующие меры.
Следуя этим рекомендациям, разработчики и организации могут использовать возможности Джамбы ответственно и этично.
Новая Глава в Разработке ИИ?
Введение семейства Джамбы от AI21 Labs знаменует собой значительный этап в эволюции больших языковых моделей. Объединяя сильные стороны трансформеров и моделей состояния-пространства, интегрируя техники mixture-of-experts и расширяя границы длины контекста и скорости обработки, Джамба открывает новые возможности для приложений ИИ в различных отраслях.
По мере того, как сообщество ИИ продолжает исследовать и строить на основе этой инновационной архитектуры, мы можем ожидать дальнейших достижений в эффективности модели, понимании длинных контекстов и практическом развертывании ИИ. Семейство Джамбы представляет не только новую серию моделей, но и потенциальный сдвиг в том, как мы подходим к проектированию и реализации крупномасштабных систем ИИ.
