Разработка больших языковых моделей (LLM) на основе декодерных трансформерных моделей сыграла решающую роль в трансформации области обработки естественного языка (NLP), а также в продвижении различных приложений глубокого обучения, включая обучение с подкреплением, анализ временных рядов, обработку изображений и многое другое. Однако, несмотря на их масштабируемость и сильную производительность, LLM на основе декодерных трансформерных моделей все еще сталкиваются с существенными недостатками. Хотя они являются выразительными, механизм внимания в трансформерных моделях требует больших вычислительных ресурсов как во время вывода, так и во время обучения, что требует значительной памяти для длины последовательности и квадратичных операций (FLOPs). Это высокое вычислительное требование ограничивает длину контекста моделей, увеличивает стоимость автoreгенеративных задач при масштабировании и препятствует способности моделей учиться на непрерывных потоках данных и обрабатывать последовательности неограниченной длины эффективно.

В последнее время модели пространства состояний (SSM) продемонстрировали замечательные возможности и производительность, конкурируя с моделями на основе трансформерной архитектуры в крупномасштабных моделях, а также достигая линейной сложности по отношению к длине последовательности и линейному времени. Более того, Mamba, недавно выпущенная модель пространства состояний, показала исключительную производительность в ряде задач моделирования языка и обработки длинных последовательностей. Одновременно с этим модели смеси экспертов (MoE) также показали впечатляющую производительность, значительно снижая задержку и вычислительные затраты на вывод, хотя и за счет увеличения памяти. Основываясь на моделях Mamba и MoE, эта статья будет обсуждать BlackMamba, новую архитектуру, которая сочетает модель пространства состояний Mamba с моделями смеси экспертов, чтобы использовать преимущества, предлагаемые обеими框ами. Эксперименты на BlackMamba продемонстрировали ее способность превосходить существующую框у Mamba и базовые модели трансформеров как в обучении FLOPs, так и в выводе. Исключительная производительность BlackMamba показывает, что она может эффективно сочетать возможности Mamba и MoE, предлагая быстрый и экономически эффективный вывод от MoE с линейной сложностью генерации от Mamba.

Эта статья направлена на углубленное освещение BlackMamba. Мы исследуем механизм, методологию и архитектуру этой рамки, а также ее сравнение с передовыми моделями генерации изображений и видео. Давайте начнем.

BlackMamba: Введение в MoE для моделей пространства состояний

Развитие больших языковых моделей (LLM), в частности тех, которые основаны на декодерных трансформерных архитектурах, существенно повлияло на область обработки естественного языка (NLP) и расширилось на различные приложения глубокого обучения, включая обучение с подкреплением, анализ временных рядов, обработку изображений и многое другое. Однако, несмотря на их масштабируемость и сильную производительность, эти декодерные трансформерные LLM сталкиваются с существенными проблемами. Механизм внимания, ключевая особенность трансформерных моделей, требует больших вычислительных ресурсов как во время вывода, так и во время обучения. Это включает необходимость памяти, которая растет с длиной последовательности и операций (FLOPs), которые увеличиваются квадратично. Такие интенсивные вычислительные потребности ограничивают длину контекста моделей, увеличивают стоимость автoreгенеративных задач при масштабировании и препятствуют способности моделей учиться на непрерывных потоках данных или обрабатывать последовательности неограниченной длины эффективно.

Были предприняты значительные усилия в последние годы, чтобы преодолеть эти ограничения, и внимание было сосредоточено на разработке альтернативных архитектур каноническим плотным моделям внимания трансформеров с помощью моделей пространства состояний и моделей смеси экспертов, являющихся наиболее перспективными кандидатами. Основным преимуществом, полученным за счет предпочтения моделей пространства состояний над моделями на основе трансформерной архитектуры, является линейная вычислительная сложность по отношению к длине входной последовательности, предлагаемая моделями пространства состояний, в отличие от квадратичной сложности, предлагаемой трансформерами. Теоретически, линейная вычислительная сложность по отношению к длине входной последовательности позволяет моделям пространства состояний обрабатывать более длинные последовательности, чем модели на основе трансформерной архитектуры, для заданного бюджета FLOPS или операций с плавающей запятой в секунду, и обеспечивает автoreгенеративную генерацию постоянной вычислительной сложности без кэша KV. Недавно разработанные модели пространства состояний, включая RWKV и Mamba, продемонстрировали эффективное обучение и вывод длинных последовательностей, а также конкурентную производительность задач моделирования языка с трансформерами аналогичных масштабов. С другой стороны, модели смеси экспертов набирают популярность как альтернатива плотным трансформерам, поскольку они обеспечивают значительное снижение операций FLOPs и затрат на вывод, необходимых для достижения качества, сравнимого с плотной моделью. Модели MoE (смесь экспертов) работают путем активации только разреженного подмножества общего количества параметров во время одного прямого прохода. Они используют функцию маршрутизации для определения, какие “эксперты” вызываются на основе заданного контекста. Этот подход создает разделение между вычислительными затратами на вывод и общим количеством параметров, позволяя добиться лучшей производительности в рамках фиксированного бюджета вывода, хотя и за счет увеличения количества параметров и большей потребности в памяти.

Этот прогресс в архитектуре предлагает существенные преимущества над традиционными трансформерами и представляет собой интересное направление для дальнейшего развития. Мы полагаем, что интеграция этих улучшений в объединенную модель Mamba-MoE может существенно ускорить возможности моделирования языка и эффективность за пределами стандартных моделей трансформеров. Ожидаемые преимущества архитектуры Mamba-MoE по сравнению с традиционной плотной моделью трансформера включают:

Mamba: Достигает линейной вычислительной сложности относительно длины входной последовательности как для фазы обучения, так и для фазы вывода. Это позволяет осуществлять автoreгенеративную генерацию в постоянном времени и с постоянным использованием памяти.

MoE: Предлагает скорость вывода и эффективность обучения, сравнимые с меньшей плотной базовой моделью, сохраняя при этом уровень качества модели, сопоставимый с моделью, имеющей эквивалентное количество параметров, как и в более плотной конфигурации.

С учетом этого, важно отметить, что модели трансформерной архитектуры все еще являются лучшими и демонстрируют последовательную и замечательную производительность на задачах моделирования языка и обработки последовательностей. В своей основе архитектура трансформера использует само-внимание, которое выполняет квадратное все-к- всем сравнение скалярных произведений между вложениями различных токенов в последовательности и выполняет линейное отображение на выходной вектор. Модель трансформера состоит из блоков само-внимания, сложенных между блоками MLP (многослойный перцептрон), которые, в свою очередь, состоят из двухслойного MLP с заданной функцией активации.

BlackMamba: Архитектура и методология

Модели пространства состояний

Модели пространства состояний принадлежат к группе моделей последовательностей с линейной сложностью по отношению к длине входной последовательности. Архитектура моделей пространства состояний более соответствует рекуррентным нейронным сетям и свёрточным нейронным сетям, чем архитектуре, основанной на внимании, и вдохновлена непрерывной динамической системой, которая отображает одномерную функцию через隐латентное пространство. Линейная динамическая система делает параллельные вычисления эффективными с помощью либо ассоциативного, либо свёрточного сканирования. На практике, рекуррентная природа моделей пространства состояний была причиной, по которой они еще не были приняты на высокопараллельном оборудовании AI, таком как GPU. Однако появление моделей пространства состояний, таких как RWKV и Mamba, использовало параллельные ядра сканирования для эффективного отображения рекуррентных операций на GPU, тем самым облегчив обучение новых архитектур с эффективностью, сравнимой с той, которую достигают модели трансформера.

Известная квадратичная сложность по отношению к длине последовательности внутри трансформеров является хорошо известным ограничением, которое препятствует рассуждениям и пониманию на очень длинных контекстах. Недавние инновации ввели идею расширения длины контекста, позволяя трансформерам обучаться на осуществимом масштабе, прежде чем применяться к гораздо более длинным контекстам во время вывода. Несмотря на эти достижения, процесс вывода все еще требует значительных вычислительных ресурсов и памяти, особенно для поддержания кэша Key-Value (KV), что делает его ресурсоемким мероприятием. Недавние исследования были сосредоточены на улучшении выразительных возможностей моделей пространства состояний путем включения входозависимых механизмов шлюзования, аналогичных матрицам Query, Key, Value (QKV), найденным в механизмах внимания.

Эти усилия направлены на сохранение внутренней линейной прогрессии рекурсии пространства состояний, позволяя выполнять эффективные вычисления либо через свёрточное, либо через селективное сканирование. Этот подход существенно сужает разрыв в производительности с трансформерами в практических приложениях. Среди этих достижений Mamba выделяется как модель пространства состояний, отражающая цели предыдущих исследований, демонстрирующая впечатляющие уровни производительности, сравнимые с трансформерами на масштабах до 2,8 миллиарда параметров. Она достигает этого, применяя входозависимое шлюзование к входам рекурсии модели пространства состояний (SSM), обеспечивая при этом эффективные вычисления с помощью специальных селективных ядер сканирования.

Модели смеси экспертов

Модели смеси экспертов (MoE) достигают разделения между стоимостью вывода и общим количеством параметров путем селективной активации параметров во время прямого прохода. Вместо использования всех параметров, эти модели направляют токены в конкретные эксперты MLP. Идеально, каждый эксперт предназначен для обработки определенного типа входных данных, с механизмом маршрутизации, по сути компактной нейронной сетью, определяющей наиболее подходящего эксперта для каждого токена. Этот подход направлен на сохранение полной выразительной силы модели с эквивалентным количеством параметров в более плотной конфигурации, но с существенно сниженными вычислительными требованиями. Обычно маршрутизатор является отображением линейных слоев от токенов до индексов экспертов, с каждым экспертом, являющимся стандартным многослойным перцептроном. Однако разработчикам еще предстоит выяснить оптимальный метод обучения маршрутизатора, поскольку задача назначения экспертов является недифференцируемой, и модели смеси экспертов часто испытывают трудности с балансировкой нагрузки и стабильностью обучения для эффективности оборудования.

Архитектура

В своей основе BlackMamba использует стандартную модель трансформера, состоящую из чередующихся блоков MLP и блоков внимания, добавленных последовательно вдоль остаточного потока. Теперь большинство моделей смеси экспертов просто заменяют блоки многослойного перцептрона на маршрутизированный слой экспертов. С другой стороны, BlackMamba не только заменяет блок многослойного перцептрона в трансформере на маршрутизированный слой экспертов, но также заменяет слой внимания на слой модели пространства состояний Mamba. Архитектура BlackMamba демонстрируется на следующей схеме.

Обучение и набор данных

Модель BlackMamba обучена на более чем 300 миллиардах токенов в пользовательском наборе данных и использует функцию активации SwiGLU для многослойных перцептронов экспертов. Рамка обучается с 8 экспертами, количество, которое разработчики сочли правильным балансом между памятью и стоимостью вывода модели. Пользовательский набор данных, использованный для обучения BlackMamba, состоит из смеси существующих открытых наборов данных, включая Starcoder, SlimPajama, Pile и многое другое. Следующая таблица демонстрирует веса каждого набора данных, использованного для обучения BlackMamba. Всего в наборе данных 1,8 триллиона токенов.

BlackMamba: Результаты

Чтобы обеспечить справедливое сравнение между Mamba и BlackMamba, разработчики обучили обе модели с одинаковыми параметрами обучения на одних и тех же данных. BlackMamba способна превосходить как Mamba, так и модели трансформеров для одинакового размера прямого прохода модели во время вывода, а также обучения FLOPs. Следующая схема демонстрирует время, необходимое для генерации последовательности заданной длины автoreгенеративно из начального однотокенного сигнала в зависимости от длины последовательности.

Более того, преимущества задержки как моделей смеси экспертов, так и Mamba объединены в BlackMamba, что приводит к существенно более быстрым временам вывода по сравнению с моделями трансформеров, чистыми моделями Mamba и моделями MoE. Кроме того, преимущество вывода BlackMamba прямо пропорционально длине последовательности, что делает BlackMamba чрезвычайно эффективной для генерации длинных последовательностей. Двигаясь дальше, следующая схема иллюстрирует количество токенов, назначенных моделям BlackMamba с 340 миллионами и 640 миллионами параметров соответственно. Как можно увидеть, большинство слоев демонстрируют высокий уровень баланса экспертов в результате улучшенного алгоритма Sinkhorn, реализованного в BlackMamba.

Следующая таблица охватывает оценочные баллы BlackMamba по сравнению с рядом открытых предварительно обученных языковых моделей. Как можно наблюдать, BlackMamba способна конкурировать и превосходить большинство рамок по всем базовым показателям. Кроме того, стоит отметить, что модели, превосходящие BlackMamba, имеют значительно большее количество параметров, и разрыв в производительности минимальный, что указывает на способность BlackMamba с меньшим количеством параметров.

Окончательные мысли

В этой статье мы говорили о BlackMamba, новой архитектуре, которая сочетает модель пространства состояний Mamba с моделями смеси экспертов, чтобы использовать преимущества, предлагаемые обеими рамками. Эксперименты на BlackMamba продемонстрировали ее способность превосходить существующую рамку Mamba и базовые модели трансформеров как в обучении FLOPs, так и в выводе. Исключительная производительность BlackMamba показывает, что она может эффективно сочетать возможности Mamba и MoE, предлагая быстрый и экономически эффективный вывод от MoE с линейной сложностью генерации от Mamba. Мы говорили о том, как архитектура BlackMamba способна превосходить сильные обученные большие языковые модели, существующую рамку Mamba и модели смеси экспертов в плане обучения FLOPs и стоимости вывода. Кроме того, BlackMamba также наследует генерацию FLOPs и снижение обучения как от моделей смеси экспертов, так и от рамки Mamba одновременно.