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).

От свёрток MLP ощутимо отличается по локальности рецептивного поля. Полносвязные слои в итоге всё равно выучивают какие-то локальные фичи, это было видно и по MLP-Mixer, и по ResMLP, и по gMLP. Правда на более глубоких слоях MLP это обычно теряется (но может просто не хватает данных?). Также любопытно, что у двух очень похожих архитектур (MLP-Mixer и ResMLP) фичи получаются ощутимо разные, то есть это не какие-то универсальные визуальные фичи? Требует разбирательства.

От трансформеров с весами, определяемыми данными, все остальные заметно отличаются тем, что у остальных веса выучиваются в процессе обучения и далее от данных не зависят (хотя, конечно, есть более продвинутые динамические свёртки, например, в знаковой работе “Pay Less Attention with Lightweight and Dynamic Convolutions” https://arxiv.org/abs/1901.10430).

При этом у MLP есть несколько боттлнеков.

Первый очевидный — это переобучение, когда данных мало, так что надо готовить датасеты размера JFT-300M и больше. Второй — это вычислительная сложность, которая как и у стандартных трансформеров обычно квадратичная от размера входа (хотя MLP-Mixer вроде как военной смекалкой обходил это и был линейным). Третья проблема — заточенность MLP на конкретные входные разрешения, ибо размеры матриц весов полносвязных слоёв прошиты внутри и размерность входа в них обычно фигурирует. Ну и ещё одна потенциальная проблема в том, что в дефолтной композиции в виде пачки однотипных блоков, модель заточена на одно конкретное разрешение, а для некоторых задач хорошо бы иметь пирамиды фич (та же проблема есть и в дефолтном трансформере, поэтому там появились и рулят Swin Transformers и т.п.).

После MLP-Mixer и ко пошёл вал работ, устраняющих то или иное ограничение. Если посмотреть на Fig.6 из работы, то видно, что всё это так или иначе вариации на тему со сходной структурой и все кроме gMLP содержат тандем spatial MLP и channel MLP. Значимая часть улучшений при этом фокусируется на spatial MLP, чтобы сделать его нечувствительным к разрешению. Прям генеалогию можно простраивать, а если мы ещё придём к кодированию структуры всех сетей через какой-то универсальный язык (aka геном), то прям и биоинформатические методы можно было бы сюда применить.

Авторы поделили все варианты на три большие группы. Группа 1 содержит полные пространственные и канальные маппинги, группа 2 содержит аксиальные маппинги по пространственным измерениям плюс канальные, и группа 3 только канальные маппинги.

Аксиальные маппинги — это оптимизация обычных полных, чтобы иметь меньше параметров, но при этом всё равно захватывать осмысленные дальние взаимодействия, только теперь они должны быть строго по вертикальной и горизонтальной осям, что конечно уменьшает выразительность (в работе есть пример картинки со вратарём в прыжке за мячом, где мяч в стороне по диагонали и такое бы не поймалось), но возможно всё равно good enough. Vision Permutator (ViP) — пример такой модели, в его Permute-MLP три параллельные ветви: для горизонтальных, вертикальных и канальных проекций. В Sparse MLP аналогично. В RaftMLP горизонтальная и вертикальная ветви работают последовательно.

Модели на основе только проекций каналов заменяют пространственные полносвязные слои на канальные проекции (или свёртки 1x1), что приводит к тому, что патчи более не взаимодействуют между собой и они для восстановления рецептивного поля прибегают к различным сдвигам или перемещениям карт фич, чтобы у разных позиций была возможность взаимодействовать через канальное измерение. Выглядит, честно говоря, как “купи козу” - “продай козу”. Среди таких моделей S2MLP, S2MLPv2, AS-MLP (Axial Shifted) или Cycle MLP (эти прибегают к deformable convolution). Только HireMLP предлагает что-то ощутимо новое по сравнению с соседями, используя inner-region rearrangements (разбивают карты фич на группы и группируют вдоль пространственных измерений, это даёт локальные взаимодействия) и cross-region rearrangements (сдвигают токены в каком-то направлении на заданный шаг, это даёт более глобальные взаимодействия). Ещё есть ConvMLP, но он так вообще скорее CNN.

Модели с полными пространственными проекциями скорее пытаются улучшить результативность базовых вариантов. Уже описанный gMLP (https://t.me/gonzo_ML/787) как раз отсюда. Также в группе Sparse-MLP (который не то же самое, что упомянутый выше Sparse MLP, а модель на основе MoE, Mixture-of-Experts, https://t.me/gonzo_ML/472) или CCS (Circulant Channel-Specific Token Mixing MLP, решивший проблему чувствительности MLP к переносу фичи в пространстве через хитрую структуру весов).

Ну и кроме модификаций самих блоков внутри MLP есть ещё направление про общую структуру сети. Разные авторы пытаются отойти от изотропного дизайна MLP-Mixer/ViT, когда все уровни модели работают со входами одинакового размера и дают single-stage структуру. Здесь появляются так называемые two-stage альтернативы, когда начальная часть сети работает на более мелких патчах (7x7), а затем сливает их в более крупные (14x14). Это делают ViP и S2MLPv2. А другие (Sparse MLP, HireMLP, AS-MLP, CycleMLP) применяют пирамидальную структуры, где последовательно патчи укрупняются (в этих конкретных получается four-stage). И пирамидальные модели обычно дают лучший перформанс.

В целом качество современных моделей достойное, MLP сравнимы с современными трансформерами и CNN, причём и по объёмам вычислений тоже. Особенно те MLP, что всё-таки включают в себя что-то близкое к CNN. Отсюда авторы делают вывод, что текущие объёмы данных недостаточны для того, чтобы чистая MLP модель выучила фичи с нуля, и подходящий inductive bias рулит. Я в целом согласен, что разумен скорее срединный путь, чем полный отказ от inductive biases. Фундаментально это конечно интересное направление и его стоит копать дальше, и для новых типов данных оно нужно, и для старых вдруг мы найдём что-то интересное, или выяснится, что наши biases таки suboptimal (что в общем запросто, но опять же для наших целей оно может быть good enough). Эволюция на Земле тоже вон миллиарды лет крутилась прежде чем сформировала всё то, что у нас сейчас есть, и все имеющиеся у нас inductive biases дорогого стоят (в том числе и в терминах вычислений, коих за это время было проведено немеряно).

Из другого интересного, разные CNN были неплохо подвержены adversarial attacks. По имеющимся данным Vision трансформеры (ViT) вроде как более устойчивы к ним, а MLP где-то посередине. И с corruption robustness (https://arxiv.org/abs/2106.13122) тоже так. Но тут ещё есть что исследовать.

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

В общем, итоговые выводы авторов в том, что пока ещё слишком много компромиссов, чтобы считать приход MLP сменой парадигмы. Они пока ещё работают с патчами, а не пикселями (как правило вычислительной эффективности ради). Также довольно много движения к CNN и ради той же вычислительной эффективности, и ради независимости от разрешения входа. И как уже упоминалось, inductive biases ещё не хватает для независимого полёта.

Но зато фронт работ ого-го! Тут же до кучи и интерпретация этих моделей, и возможно более специальные оптимизаторы, и железо, и более продвинутые дизайны архитектур, и более подходящие варианты self-supervised learning.

Следующий год будет весёлым!