Возрождение Мixture-of-Experts для Эффективных Больших Моделей Языка

В мире обработки естественного языка (NLP) стремление к созданию более крупных и способных языковых моделей было движущей силой многих недавних достижений. Однако, поскольку эти модели растут в размерах, вычислительные требования для обучения и вывода становятся все более требовательными, что приводит к пределам доступных аппаратных ресурсов.

Вступает Мixture-of-Experts (MoE), техника, которая обещает облегчить эту вычислительную нагрузку, позволяя обучать более крупные и мощные языковые модели. Ниже мы обсудим MoE, изучим его происхождение, внутреннюю структуру и его применения в моделях языка на основе трансформеров.

Происхождение Мixture-of-Experts

Концепция Мixture-of-Experts (MoE) может быть отнесена к началу 1990-х годов, когда исследователи изучали идею условной вычисления, где части нейронной сети выбираются на основе входных данных. Одним из пионерских работ в этой области была статья “Адаптивная смесь локальных экспертов” Джейкобса и др. в 1991 году, которая предложила рамки обучения с учителем для ансамбля нейронных сетей, каждая из которых специализируется на разных областях входного пространства.

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

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

Мixture-of-Experts в Трансформерах

Мixture of Experts

Хотя идея MoE существует уже несколько десятилетий, ее применение к моделям языка на основе трансформеров относительно недавнее. Трансформеры, которые стали де-факто стандартом для моделей языка уровня государства, состоят из нескольких слоев, каждый из которых содержит механизм само-внимания и нейронную сеть прямого распространения (FFN).

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

Одной из первых работ, которая продемонстрировала потенциал MoE в трансформерах, была статья “Непомерно большие нейронные сети: Слои Мixture-of-Experts с разреженным управлением” Шейзера и др. в 2017 году. Эта работа ввела понятие слоя Мixture-of-Experts с разреженным управлением, который использовал механизм управления, добавляющий разреженность и шум к процессу выбора экспертов, гарантируя, что только подмножество экспертов активируется для каждого входа.

С тех пор несколько других работ еще больше продвинули применение MoE к трансформерам, решая проблемы, такие как нестабильность обучения, балансировка нагрузки и эффективный вывод. Заметными примерами являются Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) и GLaM (Du et al., 2022).

Преимущества Мixture-of-Experts для Моделей Языка

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

Например, рассмотрим модель языка с плотным слоем FFN из 7 миллиардов параметров. Если мы заменим этот слой на слой MoE, состоящий из восьми экспертов, каждый из которых имеет 7 миллиардов параметров, общее количество параметров увеличивается до 56 миллиардов. Однако во время вывода, если мы активируем только два эксперта на токен, вычислительная стоимость эквивалентна модели с 14 миллиардами параметров, поскольку она вычисляет два умножения матриц по 7 миллиардам параметров.

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

Проблемы и Рассмотрения

Хотя модели MoE предлагают убедительные преимущества, их принятие и развертывание также сопряжены с несколькими проблемами и рассмотрениями:

1. Нестабильность Обучения: Модели MoE известны тем, что более склонны к нестабильности обучения по сравнению с их плотными аналогами. Эта проблема возникает из-за разреженной и условной природы активации экспертов, что может привести к трудностям в распространении градиентов и сходимости. Были предложены методы, такие как потеря роутера z (Zoph et al., 2022), для смягчения этих нестабильностей, но дальнейшие исследования все еще необходимы.
2. Настройка и Переобучение: Модели MoE склонны к переобучению более легко во время настройки, особенно когда задача внизу имеет относительно небольшой набор данных. Это поведение объясняется повышенной емкостью и разреженностью моделей MoE, что может привести к чрезмерной специализации на обучающих данных. Требуются тщательные регуляризация и стратегии настройки для смягчения этой проблемы.
3. Требования к Памяти: Хотя модели MoE могут снизить вычислительные затраты во время вывода, они часто имеют более высокие требования к памяти по сравнению с плотными моделями подобного размера. Это связано с тем, что все веса экспертов необходимо загрузить в память, даже если только подмножество активируется для каждого входа. Ограничения памяти могут ограничить масштабируемость моделей MoE на устройствах с ограниченными ресурсами.
4. Балансировка Нагрузки: Для достижения оптимальной вычислительной эффективности важно сбалансировать нагрузку между экспертами, гарантируя, что ни один эксперт не перегружен, в то время как другие остаются недоиспользованными. Эта балансировка нагрузки обычно достигается с помощью вспомогательных потерь во время обучения и тщательной настройки коэффициента емкости, который определяет максимальное количество токенов, которое можно назначить каждому эксперту.
5. Передача Данных: В распределенных сценариях обучения и вывода модели MoE могут ввести дополнительную передачу данных из-за необходимости обмена информацией об активации и градиентах между экспертами, расположенными на разных устройствах или ускорителях. Эффективные стратегии передачи данных и проектирование модели, учитывающее аппаратное обеспечение, являются важными для смягчения этой передачи.

Несмотря на эти проблемы, потенциальные преимущества моделей MoE в ermögлении более крупных и способных языковых моделей спровоцировали значительные исследовательские усилия для решения и смягчения этих проблем.

Пример: Mixtral 8x7B и GLaM

Чтобы проиллюстрировать практическое применение MoE в моделях языка, рассмотрим два заметных примера: Mixtral 8x7B и GLaM.

Mixtral 8x7B – это вариант MoE модели языка Mistral, разработанный Anthropic. Он состоит из восьми экспертов, каждый из которых имеет 7 миллиардов параметров, в результате чего общее количество параметров составляет 56 миллиардов. Однако во время вывода активируется только два эксперта на токен, эффективно снижая вычислительную стоимость до уровня модели с 14 миллиардами параметров.

Mixtral 8x7B продемонстрировал впечатляющие результаты, превзойдя модель Llama с 70 миллиардами параметров, предлагая намного более быстрые времена вывода. Была выпущена также версия Mixtral 8x7B, настроенная на инструкции, называемая Mixtral-8x7B-Instruct-v0.1, которая еще больше расширила его возможности в следовании естественным языковым инструкциям.

Другим заметным примером является GLaM (Google Language Model), крупномасштабная модель MoE, разработанная Google. GLaM использует архитектуру декодера-только трансформера и была обучена на огромном наборе данных из 1,6 триллиона токенов. Модель достигает впечатляющих результатов на оценках с несколькими выстрелами и одним выстрелом, соответствуя качеству GPT-3, используя только одну треть энергии, необходимой для обучения GPT-3.

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

Архитектура Grok-1

GROK MIXTURE OF EXPERT

Grok-1 – это модель MoE на основе трансформера с уникальной архитектурой, предназначенной для максимизации эффективности и производительности. Давайте рассмотрим ключевые характеристики:

1. Параметры: С огромными 314 миллиардами параметров Grok-1 является самой большой открытой LLM на сегодняшний день. Однако благодаря архитектуре MoE только 25% весов (приблизительно 86 миллиардов параметров) активны в любой момент времени, повышая возможности обработки.
2. Архитектура: Grok-1 использует архитектуру Мixture-of-8-Experts, где каждый токен обрабатывается двумя экспертами во время вывода.
3. Слои: Модель состоит из 64 слоев трансформера, каждый из которых включает механизм само-внимания и плотный блок.
4. Токенизация: Grok-1 использует токенизатор SentencePiece с размером словаря 131 072 токенов.
5. Вложения и Позиционная Кодировка: Модель имеет 6 144-мерные вложения и использует роторную позиционную кодировку, позволяя более динамичную интерпретацию данных по сравнению с традиционной фиксированной позиционной кодировкой.
6. Внимание: Grok-1 использует 48 головок внимания для запросов и 8 головок внимания для ключей и значений, каждая из которых имеет размер 128.
7. Длина Контекста: Модель может обрабатывать последовательности длиной до 8 192 токенов, используя точность bfloat16 для эффективных вычислений.

Производительность и Детали Реализации

Grok-1 продемонстрировал впечатляющие результаты, превзойдя LLaMa 2 70B и Mixtral 8x7B с оценкой MMLU 73%, демонстрируя свою эффективность и точность на различных тестах.

Однако важно отметить, что Grok-1 требует значительных ресурсов GPU из-за своего огромного размера. Текущая реализация в открытом выпуске фокусируется на проверке правильности модели и использует неэффективную реализацию слоя MoE, чтобы избежать необходимости в пользовательских ядрах.

Тем не менее, модель поддерживает шардирование активации и квантование 8-бит, которые могут оптимизировать производительность и снизить требования к памяти.

В замечательном шаге xAI выпустил Grok-1 под лицензией Apache 2.0, сделав его веса и архитектуру доступными для глобального сообщества для использования и вклада.

Открытый выпуск включает репозиторий примеров кода JAX, который демонстрирует, как загрузить и запустить модель Grok-1. Пользователи могут скачать веса контрольных точек, используя клиент торрента или напрямую через HuggingFace Hub, облегчая доступ к этой новаторской модели.

Будущее Мixture-of-Experts в Моделях Языка

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

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

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

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

Заключение

Техника Мixture-of-Experts возникла как мощный инструмент в стремлении к более крупным и способным языковым моделям. Выбирая активацию экспертов на основе входных данных, модели MoE предлагают перспективное решение вычислительным проблемам, связанным с масштабированием плотных моделей. Хотя еще остаются проблемы, которые необходимо решить, такие как нестабильность обучения, переобучение и требования к памяти, потенциальные преимущества моделей MoE в плане вычислительной эффективности, масштабируемости и экологической устойчивости делают их интересной областью исследований и разработок.

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