Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

Статьи · arXiv 1701.06538

Невероятно большие нейронные сети: слой Mixture-of-Experts с разреженным гейтингом

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

Опубликовано 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

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

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

🎯 Плотные нейросети сталкиваются с квадратичным ростом затрат при расширении емкости модели для обработки огромных датасетов.
🔧 Представлен слой Sparsely-Gated Mixture-of-Experts (MoE) с гейтингом Noisy Top-K и вспомогательными лоссами балансировки нагрузки между шардами GPU.
📈 Модели языка и перевода масштабированы до 137 млрд параметров, превзойдя лучшую опубликованную перплексию LM на 24–39% при меньших затратах.

Аннотация

Способность нейронной сети усваивать информацию ограничена числом её параметров. Условные вычисления (conditional computation), при которых части сети активируются индивидуально для каждого примера, теоретически предлагались как способ кардинального увеличения емкости модели без пропорционального роста вычислительных затрат. Однако на практике возникают серьезные алгоритмические трудности и проблемы с производительностью. В этой работе мы решаем данные проблемы и реализуем потенциал условных вычислений, добиваясь более чем 1000-кратного увеличения емкости модели при незначительных потерях вычислительной эффективности на современных GPU-кластерах. Мы представляем слой Sparsely-Gated Mixture-of-Experts (MoE), состоящий из тысяч полносвязных подсетей. Обучаемая сеть гейтинга определяет разреженную комбинацию этих экспертов для каждого входного примера. Мы применяем MoE к задачам языкового моделирования и машинного перевода, где емкость модели критически важна для поглощения огромных объемов знаний из обучающих корпусов. Мы представляем архитектуры моделей, в которых MoE объемом до 137 миллиардов параметров применяется сверточно между слоями stacked LSTM. На крупных бенчмарках языкового моделирования и машинного перевода эти модели достигают существенно лучших результатов по сравнению с state-of-the-art при более низких вычислительных затратах.

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

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

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

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

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

Хотя условные вычисления (активация лишь подмножеств модели для каждого примера) предлагались теоретически, ранние реализации сталкивались с трудностями на современных ускорителях из-за:

Ориентации GPU на однородные арифметические операции, а не на ветвление.
«Проблемы сжимающегося батча» (shrinking batch problem), когда распределение подмножеств примеров оставляет отдельным подсетям крошечные и неэффективные размеры батчей.
Узких мест в пропускной способности межсоединений между устройствами.
Несбалансированной маршрутизации, когда несколько подсетей перегружены, а остальные простаивают.

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

Слой Sparsely-Gated MoE: Полносвязные подсети («эксперты») размещаются между рекуррентными слоями (такими как stacked LSTM). Для каждой позиции вычисляется только разреженное top- $k$ подмножество экспертов.
Noisy Top-K Gating: Настраиваемый гауссов шум добавляется к логитам гейтинга перед выбором top- $k$ и softmax для обеспечения непрерывных, дифференцируемых градиентов и стимулирования сбалансированного исследования экспертов.
Комбинация параллелизма по данным и по модели: Стандартные слои и гейтинг выполняются с параллелизмом по данным (data-parallel) на разных устройствах, тогда как эксперты разделяются между GPU с параллелизмом по модели (model-parallel). Батчи со всех устройств объединяются, чтобы эксперты получали достаточно большие агрегированные мини-батчи.
Ограничения балансировки: Явные вспомогательные функции потерь для важности экспертов (importance) и их нагрузки (load) предотвращают коллапс сети гейтинга к вырожденным решениям, перегружающим лишь малую часть экспертов.

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

Слой MoE состоит из $n$ полносвязных сетей-экспертов $E_{1}, \dots, E_{n}$ и сети гейтинга $G (x)$ . Выход слоя равен: $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ Вычисления пропускаются во всех случаях, когда $G (x)_{i} = 0$ .

Noisy Top-K Gating

Чтобы объединить разреженность с дифференцируемым гейтингом, в линейные проекции добавляется шум: $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty если v_{i} входит в число первых k элементов v в противном случае$ $G (x) = Softmax (KeepTopK (H (x), k))$

Лоссы для балансировки нагрузки

Чтобы обеспечить равный суммарный вес гейтинга, определяется лосс важности (importance loss) на батче $X$ с использованием коэффициента вариации ( $C V$ ): $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ Чтобы гарантировать равное количество примеров на эксперта, гладкая оценка нагрузки использует функцию распределения стандартного нормального закона $Φ$ : $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

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

1 Billion Word Benchmark: Модель MoE-4096-h с сопоставимой вычислительной сложностью достигла тестовой перплексии 34,1 (снижение на 24% по сравнению с эквивалентными baseline LSTM). Высокобюджетная модель MoE достигла тестовой перплексии 28,0 за 47 часов на 32 GPU K40, превзойдя предыдущий state-of-the-art результат 34,7 (на который требовалось 59 часов на 32 GPU K40).
Корпус 100 Billion Word: Масштабирование до 65 536 экспертов (68 млрд параметров) снизило тестовую перплексию на 39% по сравнению с плотным baseline (28,9 против 47,0), сохранив производительность 0,72 TFLOPS/GPU.
Машинный перевод: MoE-модифицированная GNMT достигла 40,56 BLEU на WMT'14 En $\to$ Fr (+1,34 BLEU по сравнению с GNMT без RL) и 26,03 BLEU на WMT'14 En $\to$ De (+1,12 BLEU), требуя менее половины операций на временной шаг.
Многоязычный перевод: Единая многоязычная модель MoE превзошла мультиязычный baseline GNMT на 11 из 12 языковых пар (до +5,84 BLEU на паре корейский $\to$ английский) и снизила перплексию на валидации на 19%.

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

Экстремальная разреженность может приводить к убывающей отдаче или падению качества: перплексия ухудшилась при масштабировании с 65 536 до 131 072 экспертов на датасете из 100 млрд слов.
Накладные расходы на распределенные коммуникации по-прежнему зависят от соотношения размера скрытого слоя к размерностям входа/выхода, из-за чего эксперты малой емкости менее эффективно используют пропускную способность сети.
Применение MoE рекуррентно внутри шагов рекуррентности RNN ломает оптимизацию развертывания батча (batch-unrolling) по позициям последовательности, требуя схем пересчета активаций (activation recomputation).

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

Sparsely-Gated Mixture-of-Experts Layer (MoE): Слой нейронной сети, содержащий множество параллельных подсетей-экспертов, в котором обучаемая сеть гейтинга выбирает лишь разреженное подмножество экспертов для каждого входного примера.
Conditional Computation: Архитектурная парадигма, при которой части нейронной сети динамически активируются или отключаются для каждого примера индивидуально с целью увеличения емкости параметров без пропорционального роста вычислений.
Noisy Top-K Gating: Механизм гейтинга, который добавляет настраиваемый гауссов шум к логитам перед выбором k наибольших значений и применением softmax, что обеспечивает обратное распространение ошибки и предотвращает коллапс маршрутизации.
Shrinking Batch Problem: Уменьшение эффективного размера батча, поступающего в отдельные подсети, возникающее при разреженном распределении входов по множеству экспертов, что снижает арифметическую эффективность оборудования.
Importance Loss: Вспомогательная функция потерь, штрафующая квадрат коэффициента вариации суммы значений гейтов для каждого эксперта по батчу с целью поощрения равного использования экспертов.
Load Loss: Вспомогательная функция потерь, минимизирующая дисперсию расчетного вероятностно-взвешенного количества токенов, назначенного экспертам, для балансировки аппаратной нагрузки.
Hierarchical Mixture of Experts: Двухуровневая структура MoE, в которой первичная сеть гейтинга направляет входы во вторичные кластеры MoE для уменьшения коэффициента ветвления и накладных расходов на распределенную передачу данных.

Вопросы

Как сеть гейтинга обеспечивает обратное распространение ошибки при дискретной маршрутизации top-k?

Поскольку выбор top-k применяется до softmax, ненулевые значения гейтов для выбранных экспертов сохраняют непрерывные производные относительно весов входа, что позволяет градиентам протекать обратно через сеть гейтинга.

Как решается проблема сжимающегося батча (shrinking batch problem) на распределенных устройствах?

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

Зачем в сети гейтинга необходим шум?

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

Что происходит при масштабировании модели до 131 072 экспертов?

Перплексия незначительно ухудшилась с 28,9 (при 65 536 экспертах) до 29,2 (при 131 072 экспертах), а вычислительная эффективность снизилась до 0,30 TFLOPS/GPU, вероятно, из-за экстремальной разреженности и фиксированного, немасштабированного размера батча.

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

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