GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathy Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V Le, Yonghui Wu, Zhifeng Chen, Claire Cui

Статья: https://arxiv.org/abs/2112.06905

В конце прошлого года буквально через неделю после статьи про Gopher от DeepMind (https://t.me/gonzo_ML/742) вышла статья от Гугла про большие языковые модели, про которую перед НГ мы не успели написать. Исправляемся. И пользуясь случаем передаю привет Диме Лепихину, одному из соавторов :)

Работа интересна по крайней мере несколькими вещами. Во-первых, модель свежая, обучалась в августе-сентябре 2021 (в отличие от Гофера, которого держали чуть ли не год). Во-вторых, это разреженная модель с MoE и она является логическим продолжением историй про GShard и Switch Transformer (https://t.me/gonzo_ML/472). В третьих, в работе явно сравнили dense и sparse модели сопоставимых масштабов (по числу активных параметров при обработке одного токена) и показали разницу.

Обо всём по порядку.

MoE (Mixture-of-Experts) становится популярной темой. За прошлый год MoE пришёл в трансформеры для изображений (Vision MoE, V-MoE https://arxiv.org/abs/2106.05974, где натренировали модель на 15B параметров), в сентябре MS написал про мультизадачную и мультиязычную модель с MoE под названием Z-code M3, что значит Multitask, Multilingual и MoE (на 20B параметров, https://arxiv.org/abs/2109.10465), и под конец года уже после GLaM появились языковые модели с MoE от FB/Meta (в похожей по сути работе работе с лишь немного меньшей, чем GLaM на 1.2T параметров, моделью в 1.1T параметров, https://arxiv.org/abs/2112.10684), а в начале 2022 вышла новая статья от MS с модификацией MoE под названием Pyramid-ResidualMoE, PR-MoE, и с дистиллированной Mixture-of-Students, MoS (https://arxiv.org/abs/2201.05596). В общем, тема горячая, а главное, судя по цифрам полезная.

Авторы текущей работы разработали семейство моделей, которые они назвали GLaM (Generalist Language Model, не путать с Хинтоновским GLOM, https://arxiv.org/abs/2102.12627). Это модели на основе декодера трансформера как и GPT-3, но не как BERT (который энкодер), и не как T5 или, кстати, предыдущие основанные на MoE GShard-M4 или Switch Transformer (которые полноценные encoder-decoder).

Обучено семейство моделей с dense бэйзлайнами на 0.1B, 1.7B, 8B параметров и разным количеством экспертов в MoE (32, 64, 128, 256), а также большие 137B и 64B+64E (эксперта).

Для обучения собрали свой высококачественный фильтрованный датасет на 1.6T токенов (вроде как чисто англоязычный, хотя явно не сказано).

В модели в каждом втором трансформерном слое стандартный MLP заменяется на MoE. В MoE есть сколько-то экспертов (например, 64), то есть независимых полносвязных сетей (тех самых MLP, вероятно такой же архитектуры как обычно, где сначала размерность увеличивают, потом возвращают к обычной embedding dimension для слоя трансформера). Специальная функция выбора (gating function) через softmax выбирает двух экспертов (в switch transformer был один эксперт), выходом этого слоя является взвешенная комбинация выходов этих двух экспертов (также не совсем понял, какие именно это веса -- результаты softmax или что).

Также стандартные позиционные эмбеддинги трансформера заменены на относительные, а в обычных (не MoE) трансформерных MLP/FF слоях первая линейная проекция и функция активации заменены на GLU (Gated Linear Unit, который делает покомпонентное умножение двух линейных преобразований входа, по сути обучаемый гейтинг) + GELU (Gaussian Error Linear Unit, который довольно часто используется как функция активации в трансформерах). Это, наверное, чтобы даже если кто-то путает GLU и GELU, то было бы не страшно :) LayerNorm заменён на RMSNorm.

При обучении всё это распределяется по кластеру TPUv4 (для самой большой модели использовались 1024 TPU чипа), эксперты с одинаковыми индексами в разных слоях живут на одних и тех же устройствах, чтобы граф вычислений был одинаковым для каждого MoE слоя. В шардинге полагаются на компилятор GSPMD (https://arxiv.org/abs/2105.04663). Используется дополнительный MoE loss из работы про GShard, чтобы более равномерно раскидывались токены по экспертам. Авторы добились стабильного обучения сетей разного размера с одним набором гиперпараметров.

Для оценки моделей отобрали 29 датасетов (8 для задач NLG, 21 для NLU). Это оригинальный набор из 42 датасетов, на котором оценивали GPT-3 (https://t.me/gonzo_ML/305), с отброшенными 7 синтетическими задачами на арифметику и т.п., а также 6 датасетами для машинного перевода. Сравнивают на zero-shot и one-shot.

Сравнивают разные модели с упором на одинаковое количество вычислений или количество активных параметров на токен.

В целом по 5 из 7 категорий задач GLaM заметно превзошла GPT-3. При этом GLaM 64B/64E (самая большая) с 96.6B активных весов на токен (у самой большой GPT-3 их 175B) во время инференса требует примерно в два раза меньше вычислений для того же входа.

На TriviaQA побили старую SoTA. На ANLI дают качество выше, чем у Megatron-NLG с 530B параметров, при этом требуя менее 20% вычислений последнего.

Авторы также поисследовали bias’ы модели, показали, что она установила SoTA на WinoGender, превзойдя GPT-3.

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

И с ростом сети (и соответственно ростом размера отдельного эксперта), и с ростом числа экспертов, качество увеличивается.

По сравнению с dense моделями у MoE видны два преимущества.

Во-первых, MoE оказываются более эффективны по данным — им требуется меньше данных, чем плотным моделям с сопоставимыми вычислительными требованиями, чтобы достичь сравнимого перформанса на zero-shot/one-shot.

Во-вторых, по количеству вычислений и потребляемой энергии MoE эффективнее плотных сетей. А если сравниваться с GPT-3, то самый большой GLaM потребовал примерно ⅓ энергии от обучения GPT-3.

Правда, тут же мы сравниваем обучение GPT-3 на GPU, которыми были V100, и GLaM на TPUv4, что сильно разное. По цифрам у железа вроде разница не безумная: у одного чипа TPUv3 было что-то в районе 120 TFLOPS/450 Вт, у TPUv4 грубо в районе 250 TFLOPS на фиг знает сколько ватт, возможно 450-500, у V100 было в районе 125 TFLOPS/300 Вт для тензорных ядер и FP16, у A100 было бы 312 TFLOPS/300-400 Вт. Тут конечно много дьявола в деталях, да и на пиковые значения плохо ориентироваться, но разница в железе кажется тоже должна была повлиять и не особо корректно всю разницу в потреблении приписывать архитектуре MoE.

Ещё из любопытных цифр на злобу дня — это carbon emissions. Полное обучение GLaM на 600B токенов оценено в 456 МВт*ч и 40.2 тонны выбросов CO2. Для GPT-3 приводятся цифры в 1297 МВт*ч (в другой таблице 1287, что кажется более правильно и соответствует другому источнику — https://arxiv.org/abs/2104.10350) и 552 тонны CO2. Для сравнения, в статье Паттерсона по упомянутой только что ссылке полёт целого самолёта в обе стороны SFO-JFK даёт в среднем 180.4 тонн. На этом фоне всё же и GPT-3 выглядела не так страшно (ну да, целых 3 таких полёта, что мне кажется несопоставимо с пользой модели), а GLaM так вообще круто.

Жаль вот только, модель не опубликована и будет ли как-то доступна, неясно.

Оригинал: https://t.me/gonzo_ML/840