Возрождение М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 году. Эта работа ввела понятие разреженно-шлюзированного слоя MoE, который использовал механизм шлюза, добавляющий разреженность и шум к процессу выбора экспертов, гарантируя, что только подмножество экспертов активируется для каждой входной последовательности.

С тех пор несколько других работ еще больше продвинули применение 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 известны своей большей склонностью к нестабильности обучения по сравнению с их плотными аналогами. Эта проблема возникает из-за разреженного и условного характера активации экспертов, что может привести к трудностям в распространении градиентов и сходимости. Техники, такие как потеря шлюза (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 (GOOGL ) 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, является активной областью исследований. Специализированные ускорители и распределенные фреймворки обучения, предназначенные для эффективного xửления разреженных и условных вычислительных моделей MoE, могут еще больше повысить их производительность и масштабируемость.

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

Заключение

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

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