Jamba: нова гібридна мова модель AI21 Labs, що поєднує архітектури Transformer і Mamba

Моделі мови пройшли швидкий розвиток, а архітектури на основі Transformer стали стандартом для сучасних моделей обробки природної мови. Однак при збільшенні масштабу моделей виникли проблеми з обробкою довгих контекстів, ефективністю пам’яті та пропускною здатністю.

AI21 Labs представила нове рішення з Jamba, сучасною великомасштабною мовною моделлю (LLM), що поєднує сильні сторони архітектур Transformer і Mamba в гібридній рамці. Ця стаття описує архітектуру Jamba, її продуктивність та потенційні застосування.

Огляд Jamba

Jamba – це гібридна великомасштабна мова модель, розроблена AI21 Labs, що використовує комбінацію шарів Transformer і Mamba, інтегрованих з модулем Mixture-of-Experts (MoE). Ця архітектура дозволяє Jamba балансувати використання пам’яті, пропускну здатність і продуктивність, роблячи її потужним інструментом для широкого спектра завдань обробки природної мови. Модель розроблена для роботи в рамках одного 80GB GPU, забезпечуючи високу пропускну здатність і малий відбиток пам’яті, зберігаючи при цьому найвищу продуктивність на різних бенчмарках.

Архітектура Jamba

Архітектура Jamba є основою її можливостей. Вона побудована на новій гібридній конструкції, що чередує шари Transformer з шарами Mamba, включно з модулями MoE для підвищення здатності моделі без суттєвого збільшення обчислювальних вимог.

1. Шари Transformer

Архітектура Transformer стала стандартом для сучасних великомасштабних мовних моделей завдяки своїй здатності ефективно обробляти паралельну обробку і захоплювати довгострокові залежності в тексті. Однак її продуктивність часто обмежена високими вимогами до пам’яті і обчислювальних ресурсів, особливо при обробці довгих контекстів. Jamba вирішує ці обмеження, інтегруючи шари Mamba, які ми розглянемо далі.

2. Шари Mamba

Mamba – це недавно розроблена модель простору стану (SSM), призначена для більш ефективної обробки довгострокових відносин в послідовностях порівняно з традиційними РНН або навіть Transformer. Шари Mamba особливо ефективні при зменшенні відбитку пам’яті, пов’язаної з зберіганням кешів ключ-значення (KV) в Transformer. Чередуючи шари Mamba з шарами Transformer, Jamba зменшує загальне використання пам’яті, зберігаючи при цьому високу продуктивність, особливо в завданнях, що вимагають обробки довгих контекстів.

3. Модуль Mixture-of-Experts (MoE)

Модуль MoE в Jamba вводить гнучкий підхід до масштабування здатності моделі. MoE дозволяє моделі збільшувати кількість доступних параметрів без пропорційного збільшення активних параметрів під час висновку. У Jamba MoE застосовується до деяких шарів MLP, з механізмом роутера, який вибирає верхніх експертів для активації для кожного токена. Ця селективна активація дозволяє Jamba підтримувати високу ефективність при обробці складних завдань.

Нижче наведено зображення, що демонструє функціональність індукційного голови в гібридній моделі Attention-Mamba, ключовій функції Jamba. У цьому прикладі голова уваги відповідає за передбачення міток, таких як “Позитивний” або “Негативний”, у відповідь на завдання аналізу настрою. Виділені слова показують, як увага моделі сильно зосереджена на мітках токенів з прикладів few-shot, особливо в критичний момент перед передбаченням кінцевої мітки. Цей механізм уваги грає важливу роль у здатності моделі виконувати навчання в контексті, де модель повинна вивести відповідну мітку на основі заданого контексту і прикладів few-shot.

Покращення продуктивності, забезпечені інтеграцією Mixture-of-Experts (MoE) з гібридною архітектурою Attention-Mamba, виділені в таблиці. Використання MoE дозволяє Jamba збільшити свою здатність без пропорційного збільшення обчислювальних витрат. Це особливо очевидно в суттєвому підвищенні продуктивності на різних бенчмарках, таких як HellaSwag, WinoGrande і Natural Questions (NQ). Модель з MoE не тільки досягає вищої точності (наприклад, 66,0% на WinoGrande порівняно з 62,5% без MoE), але також демонструє покращені логарифмічні ймовірності в різних доменах (наприклад, -0,534 на C4).

Ключові архітектурні особливості

• Композиція шарів: Архітектура Jamba складається з блоків, що поєднують шари Mamba і Transformer у певній пропорції (наприклад, 1:7, що означає один шар Transformer на кожні сім шарів Mamba). Ця пропорція налаштовується для оптимальної продуктивності і ефективності.
• Інтеграція MoE: Шари MoE застосовуються кожні кілька шарів, з 16 експертами, доступними для активації двох найкращих експертів на токен. Ця конфігурація дозволяє Jamba масштабуватися ефективно, керуючи компромісами між використанням пам’яті та обчислювальною ефективністю.
• Нормалізація і стабільність: Для забезпечення стабільності під час навчання Jamba включає RMSNorm у шарах Mamba, що допомагає мінімізувати проблеми, такі як великі сплески активації, які можуть виникнути при масштабуванні.

Продуктивність і бенчмаркінг Jamba

Jamba пройшла суворе тестування на широкому спектрі бенчмарків, демонструючи конкурентоспроможну продуктивність по всьому спектру. Наступні розділи виділяють деякі ключові бенчмарки, на яких Jamba продемонструвала свої сильні сторони як у загальних завданнях обробки природної мови, так і в сценаріях довгого контексту.

1. Загальні бенчмарки обробки природної мови

Jamba була оцінена на декількох академічних бенчмарках, включаючи:

• HellaSwag (10-shot): Завдання зі спільним розумінням, де Jamba досягла продуктивності 87,1%, перевершивши багато конкуруючих моделей.
• WinoGrande (5-shot): Інше завдання зі розумінням, де Jamba набрала 82,5%, знову демонструючи свою здатність обробляти складні лінгвістичні завдання.
• ARC-Challenge (25-shot): Jamba продемонструвала сильну продуктивність з результатом 64,4%, відображаючи її здатність керувати складними завданнями з多 вибором.

У агрегатних бенчмарках, таких як MMLU (5-shot), Jamba досягла результату 67,4%, вказуючи на її стабільність у різних завданнях.

2. Оцінки довгого контексту

Однією з видатних особливостей Jamba є її здатність обробляти надзвичайно довгі контексти. Модель підтримує довжину контексту до 256К токенів, що є найдовшою серед публічно доступних моделей. Ця здатність була протестована за допомогою бенчмарка Needle-in-a-Haystack, де Jamba показала виняткову точність пошуку по різним довжинам контексту, включаючи до 256К токенів.

3. Пропускна здатність і ефективність

Гібридна архітектура Jamba суттєво покращує пропускну здатність, особливо при обробці довгих послідовностей.

У тестах, що порівнюють пропускну здатність (токенів за секунду) серед різних моделей, Jamba послідовно перевершувала своїх конкурентів, особливо в сценаріях, що включають великі розміри пакетів і довгі контексти. Наприклад, з контекстом 128К токенів, Jamba досягла тричі вищої пропускної здатності порівняно з Mixtral, порівнянною моделлю.

Використання Jamba: Python

Для розробників і дослідників, які хочуть експериментувати з Jamba, AI21 Labs надала модель на платформах, таких як Hugging Face, зробивши її доступною для широкого спектра застосунків. Наступний фрагмент коду демонструє, як завантажити і згенерувати текст за допомогою Jamba:

<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("In the recent Super Bowl 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))

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", # mamba
"gate_proj", "up_proj", "down_proj", # mlp
"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()