Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, Vedant Misra
Статья: https://arxiv.org/abs/2201.02177

Недавняя (на самом деле уже не совсем, она была на ICLR 2021 на 1st Mathematical Reasoning in General Artificial Intelligence Workshop) прикольная работа от OpenAI из серии про природу вещей, вернее про природу обучения и генерализации в нейросетях.

Авторы продемонстрировали прикольный феномен, который они назвали “grokking” (“врубание” или “схватывание”? Кстати, я не знал, что это слово придумал Хайнлайн, https://en.wikipedia.org/wiki/Grok), когда нейросеть резко переходит от качества случайного угадывания к идеальному качеству, причём случается это сильно после точки оверфиттинга (когда на обучающем датасете уже идеально, а на тестовом всё фигово).

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

Что делали?

Авторы взяли алгоритмически сгенерированный датасет с бинарными операциями вида a∘b=c, где все эти “a”, “b”, “c”, “=” и “∘” — это просто токены, а сами операции — это различные варианты бинарных операций в модулярной арифметике (по модулю 97, но не уверен, только или ещё и по другим модулям тоже) типа сложения/вычитания, умножения/деления, суммы квадратов и т.п.

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

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

В трансформере всего 2 слоя, 4 головы внимания, размер эмбеддинга 128, это даёт около 400K параметров без учёта эмбеддингов. По дефолту обучают AdamW.

Собственно главный феномен в том, что точность на обучающей выборке довольно быстро добирается до близкой к идеальной (например, за 1K шагов оптимизации), но на тестовой выборке генерализации не видно до 100K шагов, а в районе 1М точность на тестовом датасете наконец добирается до уровня точности на обучающем.

Этот эффект проверили на разных моделях, оптимизаторах и размерах датасета и более-менее везде на маленьких датасетах он воспроизводится. На более крупных датасетах кривые качества на трейне и тесте следуют друг другу более тесно, а с уменьшением датасета время оптимизации до достижения высокого уровня качества быстро растёт. Для некоторых операций (например, x^3 + xy^2 + y (mod 97)) генерализации в пределах заданного вычислительного бюджета не случается на любом размере датасета вплоть до 95%.

В абляциях пробовали разные интервенции: Adam в различных вариантах — обычный и с полными батчами (точный градиент на всём датасете), с полным батчем и градиентным шумом, с шумом для весов модели, с низким или высоким learning rate, с residual dropout, а также AdamW с двумя разными вариантами weight decay — к началу координат или к инициализации.

Большой эффект на data efficiency оказывает weight decay (L2-лосс на веса модели), особенно wd к началу координат (это и стало потом дефолтом во всех экспериментах), сокращая необходимое число сэмплов более чем в два раза по сравнению с другими изменениями.

Я не до конца понял, опускается ли лосс на трейне прямо до нуля, или он какой-то ненулевой и этого хватает, чтобы понемногу куда-то выбраться по ландшафту приспособленности и перестроить веса сети. Или он прям нулевой и они на одном weight decay добираются до нужной точки. Поскольку это работает (хоть и сильно хуже) и без weight decay, и даже на full-batch Adam (где градиенты типа точные), то я склоняюсь к первому варианту.

Это как-то перекликается с историей про двойной спуск (double descent, про него, кстати, у Коли Михайловского недавно был хороший семинар с первооткрывателем, Михаилом Белкиным, https://ntr.ai/webinar/nauchno-tehnicheskij-vebinar-chemu-uchit-glubokoe-obuchenie/) и особенно с историей про двойной спуск не относительно сложности модели (про что оригинальная работа, https://www.pnas.org/content/116/32/15849.short), а про двойной спуск относительно количества итераций обучения (тоже, кстати, работа OpenAI, https://arxiv.org/abs/1912.02292).

Для тех, кто не следил, в двух словах. Двойной спуск по сложности модели даёт неожиданную картинку поведения validation loss в зависимости от сложности модели — увеличивая сложность модели мы не уходим навсегда в область переобучения (как учит нас классическая теория), а быстро проходим её и начинаем получать loss ещё более низкий, чем в лучшей точке “классического режима”. То есть в большие модели и нулевой лосс на трейне вкладываться полезно, нет от них адского переобучения. Двойной спуск по числу итераций показывает, что и тупо по времени обучения похожая картинка тоже может возникать — обучаете модель дольше, когда валидационный лосс начал уже возрастать, и в какой-то момент он снова начинает уменьшаться (привет early stopping’у, не позволяющему многим дойти дотуда). Про это, кстати, у Коли тоже был вебинар с Дмитрием Ветровым, рекомендую (https://ntr.ai/webinar/nauchno-tehnicheskij-vebinar-neobychnye-svojstva-funkczii-poter-v-glubinnom-obuchenii/).

Всё это контринтуитивные истории, теоретическая база под которые ещё не подведена. Вот с гроккингом тоже что-то похожее. И более того, двойной спуск по числу итераций в данной работе также показали.

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

Такие вот дела. Всем успешного гроккинга!

Double Descent

Любителям судоку посвящается

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://yangx.top/gonzo_ML/742) вышла статья от Гугла про большие языковые модели, про которую перед НГ мы не успели написать. Исправляемся. И пользуясь случаем передаю привет Диме Лепихину, одному из соавторов :)

Работа интересна по крайней мере несколькими вещами. Во-первых, модель свежая, обучалась в августе-сентябре 2021 (в отличие от Гофера, которого держали чуть ли не год). Во-вторых, это разреженная модель с MoE и она является логическим продолжением историй про GShard и Switch Transformer (https://yangx.top/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://yangx.top/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 так вообще круто.

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