[Перевод] Mamba‑3B‑SlimPJ: Модель, которая бросает вызов лучшим Transformer-ам

14 декабря 2023 года · Альберт Гу, Три Дао

https://cartesia.ai/blog/mamba-3b-slimpj

Мы рады представить вам самую мощную на сегодняшний день языковую модель Mamba — Mamba-3B-SlimPJ. И это не просто очередной релиз — это настоящий прорыв! В сотрудничестве с Cartesia и Together мы выпустили эту модель под лицензией Apache 2.0. Она обучена на 600 миллиардов токенов и, внимание, показывает такую же производительность, как и топовые 3B-модели на базе Transformer, но сэкономив 17% вычислительных ресурсов! Если хотите узнать все детали — читайте статью на arXiv и смотрите открытый код на GitHub.

А если вы уже хотите попробовать модель в деле, вот ссылка на веса на HuggingFace.

Немного о том, что такое Mamba

Архитектура Mamba — это результат долгих экспериментов с моделями на основе State-Space (например, S4) и суперэффективными алгоритмами (как FlashAttention). В итоге у нас получилась модель, которая достойно конкурирует с Transformer-ами, но при этом — вот где магия — она имеет линейную масштабируемость по длине последовательности и работает быстрее. В рамках сотрудничества с Cartesia и Together мы выпустили Mamba-модель с 2.8 миллиарда параметров, обученную на 600 миллиардов токенов из набора данных SlimPajama и под лицензией Apache 2.0.

Да, вдумайтесь: 600 миллиардов токенов! И эта версия Mamba-3B-SlimPJ по качеству не уступает таким сильным моделям, как BTLM-3B-8K, но использует на 17% меньше вычислительных операций (FLOPs). Модель BTLM-3B-8K, кстати, основана на мощнейшей архитектуре Transformer и обучена на тех же 600 миллиардов токенов. Это подтверждает, что Mamba — это не просто альтернатива, а реальный претендент на лидирующие позиции среди базовых моделей.

Как мы обучали Mamba

Мы обучали Mamba-3B-SlimPJ на тех самых 600 миллиардов токенов с контекстом длиной 2048, используя те же гиперпараметры, что и для Mamba-3B, которая обучалась на данных Pile (300 миллиардов токенов). Единственное отличие — более плавное снижение скорости обучения, чтобы учесть больший объём данных. Для токенизации использовали GPT-NeoX токенизатор, а набор данных SlimPajama — это, в сущности, "почищенная" и более аккуратная версия RedPajama.

За что мы любим open-source? За то, что разные команды могут строить свои модели на основе работы коллег, улучшая и данные, и саму модель.

Оценка производительности

Mamba-3B-SlimPJ показывает результаты на уровне очень сильных моделей Transformer (например, BTLM-3B-8K), при этом экономя 17% вычислительных ресурсов. Обычно, чем больше данных и вычислительных мощностей, тем лучше модель. Например, модель StableLM-3B-4E1T, которая обучалась на 7 раз большем количестве токенов, всё ещё показывает лучшие результаты, чем Mamba-3B-SlimPJ или BTLM-3B-8K.

Мы протестировали Mamba-3B-SlimPJ на 10 задачах (SIQA и RACE-middle пока не доступны для оценки через lm-evaluation-harness). Это такие задачи, как BoolQ, PIQA, HellaSwag, WinoGrande, ARC easy, ARC challenge, OpenBookQA, RACE-high, TruthfulQA и MMLU. Все тесты были в формате zero-shot (за исключением MMLU, где использовались 5 примеров). Мы приводим нормализованные точности для PIQA, HellaSwag, ARC-e, ARC-c, OpenBookQA и MMLU, а также обычные точности для BoolQ, WinoGrande, RACE-high и TruthfulQA.

Что дальше?

Transformer-ы вроде BTLM-3B-8K могут использовать более продвинутые техники, такие как обучение с переменной длиной последовательностей и максимальная параметризация обновлений. Мы с нетерпением ждём возможности применить эти техники и к обучению Mamba в будущем.

Мы невероятно рады видеть, как люди интересуются моделями State-Space и альтернативами Transformer-ам в принципе, не говоря уже о Mamba. Один из наших мотивов для этого релиза — дать возможность исследователям и энтузиастам провести эксперименты с этой мощной архитектурой, а также использовать её для создания чат-ботов и моделей, обученных на инструкциях. Мы верим, что Mamba может стать отличной базовой архитектурой для моделей, работающих с текстом, аудио и видео.

Немного о Cartesia

Cartesia — это компания, которая строит базовые модели нового поколения на архитектурах будущего, таких как модели на основе State-Space. Мы мега-рады, что у нас в команде ведущим учёным работает Альберт Гу. Хотите следить за нашими разработками и получить ранний доступ? Подписывайтесь на нас!

Благодарности

Мы хотим выразить благодарность Cerebras за набор данных SlimPajama и за BTLM-3B-8K. Также огромное спасибо EleutherAI за набор данных Pile и lm-evaluation-harness.