Mixtral of Experts

Albert Q. Jiang; Alexandre Sablayrolles; Antoine Roux; Arthur Mensch; Blanche Savary; Chris Bamford; Devendra Singh Chaplot; Diego de las Casas; Emma Bou Hanna; Florian Bressand; Gianna Lengyel; Guillaume Bour; Guillaume Lample; Lélio Renard Lavaud; Lucile Saulnier; Marie-Anne Lachaux; Pierre Stock; Sandeep Subramanian; Sophia Yang; Szymon Antoniak; Teven Le Scao; Théophile Gervet; Thibaut Lavril; Thomas Wang; Timothée Lacroix; William El Sayed

Статьи · arXiv 2401.04088

Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour et al.

Опубликовано 2024-01-08cs.LGcs.CL arXiv PDF

Читать с переводом и ИИ

Кратко в трёх строках

🎯 Масштабирование емкости и качества трансформеров без чрезмерного роста задержек инференса и вычислительных затрат.
🔧 Архитектура декодера на основе Sparse Mixture of Experts (SMoE), заменяющая стандартные FFN-блоки на 8 экспертов SwiGLU и выбирающая 2 эксперта на токен с помощью обучаемого маршрутизатора.
📈 Доступ к 47B параметров в сумме при активации лишь 13B на токен, с превосходством над Llama 2 70B на MMLU (70,6% против 69,9%) и GSM8K (74,4% против 69,6%).

Аннотация

Мы представляем Mixtral 8x7B — языковую модель на основе разреженной смеси экспертов (Sparse Mixture of Experts, SMoE). Mixtral имеет ту же архитектуру, что и Mistral 7B, с тем отличием, что каждый слой состоит из 8 блоков прямого распространения (feedforward, т. е. экспертов). Для каждого токена на каждом слое сеть-маршрутизатор выбирает двух экспертов для обработки текущего состояния и объединения их выходов. Несмотря на то что каждый токен обрабатывается только двумя экспертами, выбор экспертов может меняться на каждом временном шаге. В результате каждый токен имеет доступ к 47B параметров, но во время инференса использует лишь 13B активных параметров. Mixtral обучалась с контекстным окном в 32k токенов и превосходит либо не уступает Llama 2 70B и GPT-3.5 на всех оцениваемых бенчмарках. В частности, Mixtral значительно опережает Llama 2 70B в задачах по математике, генерации кода и многоязычных тестах. Мы также представляем модель, дообученную следовать инструкциям, — Mixtral 8x7B - Instruct, которая превосходит GPT-3.5 Turbo, Claude-2.1, Gemini Pro и модель Llama 2 70B - chat на бенчмарках с участием людей-оценщиков. Обе модели, базовая и instruct, выпущены под лицензией Apache 2.0.

Оригинальная аннотация (англ.)

We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine-tuned to follow instructions, Mixtral 8x7B - Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B - chat model on human benchmarks. Both the base and instruct models are released under the Apache 2.0 license.

Разбор статьи

Проблема и мотивация

Плотные (dense) авторегрессионные языковые модели требуют вычислительных затрат, прямо пропорциональных общему числу параметров для каждого обрабатываемого токена. По мере масштабирования моделей до 70B параметров и выше их инференс с низкой задержкой и высокой пропускной способностью становится вычислительно дорогим. Авторы стремились масштабировать емкость модели, сохраняя стоимость инференса на токен на уровне значительно более компактных моделей, и создать модель с открытыми весами, способную сравниться или превзойти Llama 2 70B и закрытые бейзлайны вроде GPT-3.5.

Ключевые идеи

Sparse Mixture of Experts (SMoE): Замена плотных блоков сетей прямого распространения (FFN) во всех 32 слоях трансформера на 8 отдельных сетей-экспертов с маршрутизацией каждого токена только в топ-2 эксперта.
Разделение общего и активного числа параметров: Архитектура содержит 47B параметров в общей сложности, но для прямого прохода каждого токена активирует лишь 13B параметров, что соответствует вычислительным затратам плотной модели среднего размера.
Полноценное контекстное окно 32k: Нативная поддержка последовательностей длиной до 32 768 токенов со 100% точностью извлечения данных в бенчмарках passkey.
Синтаксическая, а не тематическая маршрутизация: Эмпирический анализ маршрутизации токенов показал, что выбор экспертов определяется в первую очередь синтаксической функцией и временной локальностью, а не высокоуровневыми тематическими доменами.

Метод (с ключевыми формулами)

Mixtral модифицирует архитектуру Mistral 7B ( $32$ слоя, скрытая размерность $4096$ , $32$ головы внимания, $8$ голов key-value), заменяя каждый FFN-подблок на MoE-модуль, состоящий из $n = 8$ экспертов ( $E_{0}, \dots, E_{n - 1}$ ), каждый из которых реализован как сеть SwiGLU с промежуточной размерностью $14336$ .

Для входного представления токена $x$ сеть-маршрутизатор вычисляет разреженное распределение гейтинга по экспертам, используя топ- $K$ логитов ( $K = 2$ ): $G (x) := Softmax (TopK (x \cdot W_{g}))$ где оператор Top-K присваивает $- \infty$ всем координатам вне $K$ наибольших значений логитов $ℓ = x \cdot W_{g} \in R^{n}$ : $(TopK (ℓ))_{i} := {ℓ_{i} - \infty если ℓ_{i} входит в топ - K координат ℓ, иначе .$

Итоговый выход слоя $y$ вычисляется как взвешенная гейтингом сумма выходов экспертов: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ Эффективное выполнение достигается с помощью вычислительных ядер, таких как MegaBlocks, которые представляют MoE-вычисления в виде умножения разреженных матриц, а также за счет Expert Parallelism на нескольких GPU.

Результаты (с важнейшими числами)

Общие бенчмарки и рассуждения: На MMLU (5-shot) Mixtral 8x7B набирает 70,6% против 69,9% у Llama 2 70B и 70,0% у GPT-3.5. На GSM8K (8-shot, maj@8) результат составляет 74,4% против 69,6% у Llama 2 70B. В бенчмарке MATH (4-shot, maj@4) Mixtral достигает 28,4% против 13,8% у Llama 2 70B.
Генерация кода: Mixtral существенно опережает Llama 2 70B на HumanEval (40,2% против 29,3%) и MBPP (60,7% против 49,8%).
Многоязычные задачи: Повышение доли многоязычных данных позволило Mixtral обойти Llama 2 70B на французском, немецком, испанском и итальянском языках в тестах ARC-Challenge, HellaSwag и MMLU (например, французский MMLU: 70,9% против 64,3%).
Instruction Fine-Tuning и Arena Elo: Mixtral 8x7B – Instruct (обученная с помощью SFT и Direct Preference Optimization) набирает 8,30 балла на MT-Bench и достигает рейтинга 1121 на LMSys Chatbot Arena Elo, превосходя Claude-2.1 (1117), GPT-3.5-Turbo (1117) и Llama-2-70b-chat (1077).
Извлечение из контекста: Демонстрирует 100% точность в синтетической задаче извлечения ключа (passkey retrieval) на всем контекстном окне в 32k токенов.

Ограничения и открытые вопросы

Требования к памяти против вычислений: Хотя вычислительная сложность масштабируется исходя из 13B активных параметров, для обслуживания модели требуется объем VRAM, достаточный для размещения всех 47B параметров, что требует систем с несколькими GPU или применения техник offloading.
Узкие места Expert Parallelism: В работе выявлена высокая временная локальность — соседние токены выбирают одних и тех же экспертов значительно чаще, чем при случайном распределении (до 28,4% повторений первого выбора на 15-м слое против 12,5% при случайном выборе). Это может приводить к дисбалансу нагрузки и перегрузке отдельных GPU при использовании Expert Parallelism.
Отсутствие тематической специализации: Маршрутизатор естественным образом не разделяет токены по концептуальным доменам (например, биология против философии), оставляя открытым вопрос о том, как проектировать гейтинг, если требуется явная предметная специализация экспертов.

Ключевые термины

Sparse Mixture of Experts (SMoE): Архитектура нейронных сетей, в которой условные вычисления направляют каждый входной токен в разреженное подмножество подсетей прямого распространения (экспертов) вместо активации всех параметров модели.
Active Parameter Count: Количество параметров, фактически загружаемых и вычисляемых во время инференса для одного токена; у Mixtral это 13B из общих 47B разреженных параметров.
SwiGLU: Вариант функции активации Gated Linear Unit, сочетающий Swish и поэлементное умножение, используемый во внутренней архитектуре FFN-блока каждого эксперта.
TopK Gating: Механизм маршрутизации, применяющий линейную проекцию к представлению токена, сохраняющий топ-K наибольших значений, маскирующий остальные в минус бесконечность и нормализующий их с помощью softmax.
Expert Parallelism (EP): Стратегия распределенного параллелизма, при которой разные эксперты размещаются на отдельных GPU, а активации токенов пересылаются между устройствами к назначенным им экспертам.
Direct Preference Optimization (DPO): Алгоритм выравнивания (alignment), оптимизирующий параметры языковой модели напрямую на парах предпочтений без необходимости обучения отдельной модели вознаграждения в рамках RL.
Passkey Retrieval: Синтетический тест оценки способности модели находить и возвращать короткий числовой ключ, случайно размещенный внутри длинного контекста длиной до 32k токенов.
Positional Locality: Наблюдаемая тенденция, при которой последовательные токены в последовательности с высокой вероятностью (намного превышающей равномерное распределение) назначаются одному и тому же эксперту.

Вопросы

Почему Mixtral называется 8x7B, если у нее всего 47B параметров вместо 56B?

Несмотря на наличие 8 экспертов на слой и общую основу в виде архитектуры Mistral 7B, веса механизма внимания (self-attention) и эмбеддингов разделяются между всеми экспертами. Восемь раз дублируются только подблоки прямого распространения (FFN), что дает 47B суммарных параметров вместо 8 × 7B = 56B.

Сколько экспертов активируется для каждого токена при инференсе?

Маршрутизатор выбирает ровно 2 экспертов ( $K = 2$ ) из 8 доступных на каждом слое для каждого токена, что дает 13B активных параметров на один прямой проход.

Специализируются ли разные эксперты на конкретных научных или предметных темах?

Явной доменной специализации обнаружено не было. Анализ на датасете The Pile показал практически идентичное распределение выбора экспертов по таким темам, как ArXiv, PubMed и PhilPapers; логика маршрутизации гораздо сильнее коррелирует с синтаксисом, конкретным токеном и позиционной локальностью.

Под какой лицензией выпущены модели Mixtral 8x7B?

Как базовая модель Mixtral 8x7B, так и Mixtral 8x7B – Instruct выпущены под свободной лицензией Apache 2.0 для академического и коммерческого использования.

Обзор на этой странице написан ИИ по тексту статьи и может содержать ошибки. Подробности смотрите в самой статье.

Читать с переводом и ИИ