Pay Attention to MLPs
Hanxiao Liu, Zihang Dai, David R. So, Quoc V. Le
Статья: https://arxiv.org/abs/2105.08050
Код (неофициальный): https://github.com/lucidrains/g-mlp-pytorch  

В области в последние полгода происходит большое бурление. По инерции всё ещё кажется, что фронтир многих исследований в computer vision и nlp уже (или соответственно всё ещё) находится в трансформерах. И там действительно много всего происходит, так что следить за всем потоком публикаций уже просто нереально и надо держать небольшой НИИ (а лучше автоматического агента), чтобы разбирать только свежие работы. Но это всё же не совсем уже фронтир. А вот старые добрые многослойные персептроны таки фронтир!

Буквально через пару недель после MLP-Mixer (https://t.me/gonzo_ML/776) и кучки других практически одновременно вышедших похожих работ появилась новая интересная работа на аналогичную тему.

Авторы снова подходят со стороны замены трансформера на варианты полносвязных сетей, экспериментируют с разными вариантами MLP и выбирают одну наиболее эффективную конфигурацию. Полученная модель называется gMLP и её проверяют на задачах классификации картинок ImageNet и на обучении BERT’а.

Собственно авторы сделали следующее. Как и трансформер, сеть состоит из набора одинаковых блоков, на вход которых прилетают эмбеддинги из предыдущего слоя (или входные), а на выходе вылетают обработанные эмбеддинги той же размерности. Размерность входа n*d, где n — число токенов (или длина последовательности), а d — размер эмбеддинга (ну плюс ещё подразумеваемое измерение для батча).

Модель по входам и выходам совместима с BERT/ViT.

Внутри нет никаких блоков self-attention, а только нормализация, channel projection, активация (GELU), блок spatial projection и снова channel projection на выходе

Channel projection — это обычные линейные проекции как в FFN-слоях трансформеров, в конфигурации аналогичной BERT_base это 768x3072 и 3072x768 (на входе и выходе соответственно, причём на выходе скорее первое измерение в два раза меньше указанного, потому что там хитрый сплит по каналам и гейтинг, о которых ниже).

Spatial projection — это самое мясо gMLP, слой, который устраивает взаимодействие между различными токенами (а не каналами). Если этот слой заменить на identity, то из блока получится обычная FFN, где каждый токен обрабатывается независимо и друг с другом они не взаимодействуют.

Целью работы было найти такой вариант этого блока, который позволит заложить в систему сложные взаимодействия между токенами. В самом эффективном из вариантов, который далее в работе и используется, этот блок сначала делает разделение каналов  (split) на два потока одинакового размера, которые в конце поэлементно перемножатся. Первый поток без изменений отправляется к выходу, а второй сначала прогоняется через нормализацию, а затем линейную проекцию размерности n*n (где как раз для каждого элемента можно определить влияющие на него другие элементы), и далее эти два потока поэлементно (по измерению каналов) перемножаются, то есть имеем некий gating, когда один вход выбирает, что возьмём из другого. Чтобы при обучении это не взорвалось, данную функцию инициализируют так, чтобы на старте это была практически identity трансформация. Хороший разбор кода есть тут https://nn.labml.ai/transformers/gmlp/index.html.

Данный блок называется Spatial Gating Unit (SGU). Это напоминает гейтинг из LSTM/GRU или скорее даже гейтинг из Highway Networks (https://arxiv.org/abs/1505.00387) или GLU (https://arxiv.org/abs/1612.08083), только в отличие от последних он вычисляется не по проекции скрытого измерения (которое здесь называется измерением каналов), а по проекции пространственного (которое отвечает за cross-token interactions). По мне так ещё одно существенное отличие от гейтинга в том, что обычно выход гейта нормализован (например, через сигмоиду) так, чтобы он был в диапазоне [0,1], здесь же, хоть на входе нормализация и есть, далее работает фактически полноценная multiplicative interaction, где в принципе итоговый вес может после линейного преобразования быть любым, и ещё вопрос, кто там в итоге кого гейтит.

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

Полученную модель проверяют на картинках и текстах. На картинках решают задачу классификации ImageNet, аналогичным ViT образом, когда картинка конвертируется в набор патчей 16x16 пикселей. Авторы обнаружили, что gMLP любит переобучаться, поэтому в него добавили разных регуляризаций по аналогии с более производительным DeiT (улучшенный ViT, если кто не следил https://arxiv.org/abs/2012.12877). Модели подобрали примерно соответствующие по числу параметров ViT/DeiT.

Результаты хороши. SoTA для картинок ожидаемо не побита (там, конечно, есть свои зверски затюненные герои из свёрточного мира), но gMLP ощутимо бьёт ViT и примерно соответствует DeiT. Ну то есть модель без self-attention прекрасно справляется с изображениями. А также gMLP бьёт недавно описанные MLP-Mixer/ResMLP (https://t.me/gonzo_ML/776).

Предыдущие MLP пытались заменить ViT, но не шли в NLP, а авторы текущей работы идут. На задаче обучения BERT-подобной модели попробовали разные варианты блоков MLP и собственно пришли к варианту multiplicative + split. Модель получилась сравнимая с BERT_base, и намного лучше его же с отключенным механизмом внимания (что, видимо, заслуга SGU). Также gMLP ощутимо лучше MLP-Mixer.

Для BERT-подобной модели сравниваются по трём метрикам: perplexity самой предобученной модели, а также результаты файнтюнинга модели на две задачи из GLUE — SST-2 и MNLI.

На предобучении gMLP сначала отстаёт от трансформера по perplexity, но потом при увеличении размера модели начинает обгонять. Степенной закон в целом довольно близкий к трансформеру.

Веса spatial projections повизуализировали и нашли разные красивые выученные фильтры. Также для задачи MLM (masked language model) обнаружили, что gMLP выучивает Тёплицевы матрицы в качестве весов линейной проекции в SGU, что является выученной из данных инвариантностью к перемещению (неважно, в какой позиции стоял токен, который надо заполнить). Это позволяет также более экономно хранить эти матрицы в модели.

С файнтюнингом интересно, на SST-2 gMLP всегда обгоняет трансформер, а на MNLI всегда отстаёт. Авторы делают вывод, что для задач MNLI у трансформера более подходящий inductive bias, и так как в этих задачах необходимо работать с двумя предложениями (вместо одного как в SST-2), то видимо gMLP не хватает способности глядеть на соседние предложения.

Поэтому сделали ещё один вариант модели под названием aMLP (“a” потому что attention), где в блок SGU добавили одну единственную голову self-attention размера 64. Её выход плюсуется к выходу spatial projection до перемножения с другим стримом внутри блока.

aMLP берёт лучшее из двух миров и бьёт трансформер везде. Получается, что inductive bias механизма внимания и spatial gating несколько разные и дополняют друг друга.

Короче, получается, что в NLP можно в общем-то и без трансформеров, старыми добрыми MLP. Намёки на это были и за год до описанных работ (например, https://arxiv.org/abs/2005.13895). Были, правда, движения и в обратную сторону, чтобы наоборот все FFN заменить вниманием (https://arxiv.org/abs/1907.01470).  

Официального кода нет, но есть неофициальный. Из интересного, есть реализация GPT на базе gMLP (PyTorch: https://github.com/lucidrains/g-mlp-gpt, JAX: https://github.com/lucidrains/mlp-gpt-jax).

Are we ready for a new paradigm shift? A Survey on Visual Deep MLP
Ruiyang Liu, Yinghui Li, Linmi Tao, Dun Liang, Shi-Min Hu, Hai-Tao Zheng
Статья: https://arxiv.org/abs/2111.04060

И снова про MLP. Как мы уже писали, примерно полгода назад начался и до сих пор продолжается кембрийский взрыв разнообразия полносвязных сетей (или многослойных персептронов, MLP, как их часто называют), связанный с появлением работ, показавших, что и на MLP можно достигать качества трансформеров при работе с картинками и текстом. Главное, данных подгонять побольше, в объёмах существенно больших размера ImageNet, который сам когда-то казался гигантским. Либо добавлять новые inductive biases взамен утраченных (внимание или свёртки).

Началось это с MLP-Mixer (https://t.me/gonzo_ML/776) и нескольких других местами крайне похожих работ, вышедших практически одновременно с ним — явно из какой-то общей ноосферы все набирались.

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

Здесь авторы фокусируются только на задачах компзрения и задаются вопросом, станут ли MLP — наиболее простая архитектура с глобальным рецептивным полем, но без механизмов внимания — новой парадигмой в компьютерном зрении?

В целом про MLP уже давно известно, что оные являются универсальными аппроксиматорами, то есть по идее могут приблизить более-менее любые практически полезные непрерывные функции, так что и разные задачи компзрения и не только должны бы уметь решать. Но чего-то всё время для этого не хватало. То градиенты затухали (это, кстати, можно хорошо пронаблюдать даже в Tensorflow Playground, http://playground.tensorflow.org/, всего лишь на максимально доступных там шести слоях) и глубокие сети не обучались (и тут много всего придумали, от жадного послойного обучения ещё во времена Ивахненко, до residual connections и хитрых регуляризаций и нормализаций), то железо слабое и никто дождаться не может (тогда появились свёрточные сети с вычислительно более эффективными biases, а за ними после успеха и железо), то, как выяснилось, данных не хватало, чтобы вменяемые и полезные фичи выучить (как недавно выяснилось сначала для трансформеров на ViT, https://t.me/gonzo_ML/434, и затем и на MLP, https://t.me/gonzo_ML/776).

Переход к CNN, а затем и к трансформерам (ViT и последовавшими за ним), были как раз отдельными сменами парадигм в комп.зрении, и что важно, каждая была построена на базисе достаточно большого датасета (сначала ImageNet, потом более крупные коллекции типа JFT-300M), плюс производительного железа (сначала игровые GPU, теперь много всего, включая свежие продвинутые GPU и TPU, по этой теме я много писал тут https://blog.inten.to/hardware-for-deep-learning-part-4-asic-96a542fe6a81, пора дополнять и обновлять, особенно про GPU). В принципе, и то и другое на новом уровне есть и сейчас, осталось понять, не отсутствует ли что-то ещё? О чём и статья.

Для этого авторы пытаются разобраться и сравнить свойства разных вариантов свёрток/трансформеров/MLP, а также прослеживают, что произошло начиная с появления MLP-Mixer и куда ушли архитектуры за эти полгода.

Между разными архитектурами и моделями есть довольно большое пересечение и до некоторой степени эквивалентность в отдельных частях. Например, channel mixing в MLP-Mixer а также MLP в ViT это по сути то же, что и свёртки 1x1, и вопрос тогда, что считать свёрточной архитектурой? Да и механизмы внимания неоднократно пытались сводить к различным свёрткам. Но в целом всё равно есть несколько осей, по которым разные модели разделяются: обучаемые/необучаемые (как вариант — data driven) веса, рецептивное поле (глобальное/локальное), чувствительность к разрешению, чувствительность к позиции фичи, чувствительность к канальному измерению, а также следующие из архитектурных решений важные практические характеристики в виде числа параметров и количества вычислений (FLOPS).