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

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

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

BlackMamba: Введение в MoE для государственных космических моделей

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

За последние несколько лет были предприняты значительные усилия в попытке преодолеть эти ограничения, и внимание было переключено на разработку архитектурных альтернатив каноническим моделям преобразователей плотного внимания, причем модели SSM и MoE являются наиболее многообещающими кандидатами на архитектуру. Ключевым преимуществом, которое дает предпочтение моделям пространства состояний по сравнению с моделями архитектуры преобразователя, является линейная вычислительная сложность относительно длины входной последовательности, обеспечиваемая SSM, в отличие от квадратичной сложности, предлагаемой преобразователями. Теоретически, линейная вычислительная сложность относительно длины входной последовательности позволяет моделям пространства состояний обрабатывать более крупные последовательности, чем модели с архитектурой преобразователя, для заданного бюджета FLOPS или операций с плавающей запятой в секунду, а также отображать константу авторегрессионной генерации при вычислениях без KV-кеша. Недавно разработанные модели пространства состояний, включая Mamba, RetNet и некоторые другие, продемонстрировали эффективный вывод и обучение длинных последовательностей, а также конкурентоспособную производительность задач моделирования языка для преобразователей с аналогичными свойствами масштабирования. С другой стороны, архитектура «Смесь экспертных моделей» набирает популярность в качестве альтернативы плотным преобразователям, поскольку она способствует значительному сокращению количества логических выводов и флопс обучения, необходимых для достижения качества, сравнимого с плотной моделью. Модели MoE (Mixture of Experts) работают, активируя только редкий выбор общих параметров во время одного прямого прохода. Они используют функцию маршрутизации, чтобы определить, какие «эксперты» будут задействованы в заданном контексте. Этот подход создает разделение между вычислительными затратами на вывод и общим количеством параметров, позволяя повысить производительность в рамках фиксированного бюджета вывода, хотя и с увеличенным количеством параметров и большими требованиями к памяти.

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

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

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

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

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

Государственные космические модели

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

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

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

Смесь экспертных моделей

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

Архитектура

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

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

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

БлэкМамба: Результаты

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

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

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

Заключение

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