MLP-Mixer: An all-MLP Architecture for Vision
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy
Статья: https://arxiv.org/abs/2105.01601
Код: https://github.com/google-research/vision_transformer
Модели: https://github.com/google-research/vision_transformer#available-mixer-models
Что мы всё про трансформеры, да про трансформеры. За пределами трансформеров тоже много жизни и она не менее интересная. И сами любители трансформеров тоже так считают. Вот, например, вышедшая в мае 2021 года работа от авторов ViT (https://t.me/gonzo_ML/434), где они собрали модель на основе полносвязных сетей (или MLP, multi-layer perceptron) и показали, что полносвязные сети неожиданно тоже могут давать очень высокое качество на изображениях. Модель назвали MLP-Mixer.
Всё это движение, от CNN к ViT и теперь к MLP — это движение ко всё большему устранению inductive biases из модели и переход к выучиванию нужных biases из данных. В каком-то смысле это продолжает движение от классического ML к DL, когда мы отказались от предрассчитанных фич и стали выучивать их из данных. Но отказались мы тогда не полностью, потому что, например, в картинках, где оно и выстрелило первоначально, мы заложились на встроенные в архитектуру решения, те самые inductive biases, которые в случае свёрточных сетей (CNN) были про локальность фич и некоторую инвариантность к перемещению. Trade-off здесь везде понятен -- для выучивания чего-то из данных, нужно больше данных, поэтому переход к новым моделям с меньшим числом (или с менее специфичными) biases требует больше данных для достижения того же качества. Что мы и наблюдали в случае ViT — он начинал круто работать на датасетах, существенно больше ImageNet’а. В случае MLP-Mixer этот тренд логично продолжается.
Разберём, что у неё внутри.
Как и ViT, модель принимает на вход последовательность патчей изображения, те же самые неперекрывающиеся кусочки, скажем 16x16 пикселей, они же токены, которые на входе в модель линейно преобразуются в эмбеддинги фиксированного размера (C). К каждому токену применяется одна и та же матрица преобразований и далее весь вход может быть представлен как матрица размером патчи*каналы (S*C), и эта размерность сохраняется внутри модели.
Далее модель состоит из одинаковых блоков (Mixer layer), которые сначала применяют token-mixing MLP, а потом channel-mixing MLP. Первый работает на столбцах матрицы (и как бы перемешивает все токены в пределах одного канала), а второй на строках (и перемешивает каналы в пределах токена, то есть позиции). Матрицы общие, в первом случае на столбцы, а во втором на строки.
На выходе модели global average pooling и обычная классификационная полносвязная голова.
Каждый из этих MLP состоит в свою очередь из двух полносвязных слоёв, между которыми стоит нелинейность GELU. Скрытые размерности внутри этих MLP (D_s и D_c соответственно) задаются независимо (в работе D_c это 4xC, а D_s это C/2), и ничего из этого не зависит от размера последовательности, то есть сложность модели линейная относительно входа, что было не так для ViT с его квадратичной сложностью.
Также в модели есть skip-connections и используется LayerNorm. Никаких позиционных эмбеддингов как в ViT нету, так как token-mixing MLP и так чувствителен к порядку входных элементов.
Модель также как и ViT написана на JAX/Flax, и выложена в том же репозитории.
В экспериментах обучили семейство моделей с аналогичным ViT именованием. Например, модель B/16 — это модель Base (12 слоёв) с размером патча 16x16 пикселей. Самая тяжёлая модель — H/14 в 431M параметров.
Модели предобучали на датасетах разного размера (ILSVRC2021 ImageNet, ImageNet-21k, а также на проприетарном JFT-300M) в разрешении 224^2 с различными аугментациями и регуляризациями (RandAugment, mixup, dropout, stochastic depth). Затем модели файнтюнили на задачи классификации (ImageNet оригинальный и чищенный ReaL, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102, VTAB-1k).
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy
Статья: https://arxiv.org/abs/2105.01601
Код: https://github.com/google-research/vision_transformer
Модели: https://github.com/google-research/vision_transformer#available-mixer-models
Что мы всё про трансформеры, да про трансформеры. За пределами трансформеров тоже много жизни и она не менее интересная. И сами любители трансформеров тоже так считают. Вот, например, вышедшая в мае 2021 года работа от авторов ViT (https://t.me/gonzo_ML/434), где они собрали модель на основе полносвязных сетей (или MLP, multi-layer perceptron) и показали, что полносвязные сети неожиданно тоже могут давать очень высокое качество на изображениях. Модель назвали MLP-Mixer.
Всё это движение, от CNN к ViT и теперь к MLP — это движение ко всё большему устранению inductive biases из модели и переход к выучиванию нужных biases из данных. В каком-то смысле это продолжает движение от классического ML к DL, когда мы отказались от предрассчитанных фич и стали выучивать их из данных. Но отказались мы тогда не полностью, потому что, например, в картинках, где оно и выстрелило первоначально, мы заложились на встроенные в архитектуру решения, те самые inductive biases, которые в случае свёрточных сетей (CNN) были про локальность фич и некоторую инвариантность к перемещению. Trade-off здесь везде понятен -- для выучивания чего-то из данных, нужно больше данных, поэтому переход к новым моделям с меньшим числом (или с менее специфичными) biases требует больше данных для достижения того же качества. Что мы и наблюдали в случае ViT — он начинал круто работать на датасетах, существенно больше ImageNet’а. В случае MLP-Mixer этот тренд логично продолжается.
Разберём, что у неё внутри.
Как и ViT, модель принимает на вход последовательность патчей изображения, те же самые неперекрывающиеся кусочки, скажем 16x16 пикселей, они же токены, которые на входе в модель линейно преобразуются в эмбеддинги фиксированного размера (C). К каждому токену применяется одна и та же матрица преобразований и далее весь вход может быть представлен как матрица размером патчи*каналы (S*C), и эта размерность сохраняется внутри модели.
Далее модель состоит из одинаковых блоков (Mixer layer), которые сначала применяют token-mixing MLP, а потом channel-mixing MLP. Первый работает на столбцах матрицы (и как бы перемешивает все токены в пределах одного канала), а второй на строках (и перемешивает каналы в пределах токена, то есть позиции). Матрицы общие, в первом случае на столбцы, а во втором на строки.
На выходе модели global average pooling и обычная классификационная полносвязная голова.
Каждый из этих MLP состоит в свою очередь из двух полносвязных слоёв, между которыми стоит нелинейность GELU. Скрытые размерности внутри этих MLP (D_s и D_c соответственно) задаются независимо (в работе D_c это 4xC, а D_s это C/2), и ничего из этого не зависит от размера последовательности, то есть сложность модели линейная относительно входа, что было не так для ViT с его квадратичной сложностью.
Также в модели есть skip-connections и используется LayerNorm. Никаких позиционных эмбеддингов как в ViT нету, так как token-mixing MLP и так чувствителен к порядку входных элементов.
Модель также как и ViT написана на JAX/Flax, и выложена в том же репозитории.
В экспериментах обучили семейство моделей с аналогичным ViT именованием. Например, модель B/16 — это модель Base (12 слоёв) с размером патча 16x16 пикселей. Самая тяжёлая модель — H/14 в 431M параметров.
Модели предобучали на датасетах разного размера (ILSVRC2021 ImageNet, ImageNet-21k, а также на проприетарном JFT-300M) в разрешении 224^2 с различными аугментациями и регуляризациями (RandAugment, mixup, dropout, stochastic depth). Затем модели файнтюнили на задачи классификации (ImageNet оригинальный и чищенный ReaL, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102, VTAB-1k).
GitHub
GitHub - google-research/vision_transformer
Contribute to google-research/vision_transformer development by creating an account on GitHub.
Меряли качество классификации, пропускную способность (картинок в секунду) и общее время обучения на TPUv3.
В результате получили очень высокий перформанс. Чуть хуже SoTA, но всё равно весьма сильно, что было неожиданно для MLP. Модель хорошо оверфитится, поэтому регуляризация очень важна.
По скорости Mixer-H/14 в 2.5 раза быстрее ViT-H/14 и почти в два раза быстрее BiT-R152x4.
В целом по качеству MLP-Mixer сидит на Парето фронте и он чуть хуже ViT/BiT на малых размерах датасетов и малом вычислительном бюджете, но выходит вперёд на больших числах. Снова The Bitter Lesson.
Модель проверяют на перестановках пикселей, глобальной (по всему изображению, общая перестановка для всех изображений) и патчи+пиксели внутри патча (снова общая для всех). Ожидаемо, на патчи+пиксели Mixer вообще не деградирует, но что интересно, при глобальной перестановке Mixer деградирует меньше (на 45%), чем ResNet (75%).
Авторы визуализировали фичи, среди них есть как глобальные на всё изображение, так и вполне локальные, в том числе одни фичи в разных фазах.
Интересно, что буквально в то же время (через 1-3 дня) на arxiv запостили ещё несколько очень близких работ, например “Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks” от Tsinghua University (https://arxiv.org/abs/2105.02358), “Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet” из Оксфорда (https://arxiv.org/abs/2105.02723) и “ResMLP: Feedforward networks for image classification with data-efficient training” из FB/Сорбонны/INRIA (https://arxiv.org/abs/2105.03404). Другая китайская работа конца мая 2021 “Can Attention Enable MLPs To Catch Up With CNNs?” (https://arxiv.org/abs/2105.15078) кратко анализирует эти разные решения.
Оксфордская работа практически такая же, только порядок слоёв обратный, сначала фичи, потом патчи. В ResMLP в первом блоке про token-mixing один слой вместо двух и везде афинные преобразования вместо LayerNorm. В китайской работе self-attention вроде как заменяется на external attention, реализуемый через простые линейные слои.
Такие вот дела. Не совсем понимаю, откуда такая точная синхронизация разных работ по времени.
В результате получили очень высокий перформанс. Чуть хуже SoTA, но всё равно весьма сильно, что было неожиданно для MLP. Модель хорошо оверфитится, поэтому регуляризация очень важна.
По скорости Mixer-H/14 в 2.5 раза быстрее ViT-H/14 и почти в два раза быстрее BiT-R152x4.
В целом по качеству MLP-Mixer сидит на Парето фронте и он чуть хуже ViT/BiT на малых размерах датасетов и малом вычислительном бюджете, но выходит вперёд на больших числах. Снова The Bitter Lesson.
Модель проверяют на перестановках пикселей, глобальной (по всему изображению, общая перестановка для всех изображений) и патчи+пиксели внутри патча (снова общая для всех). Ожидаемо, на патчи+пиксели Mixer вообще не деградирует, но что интересно, при глобальной перестановке Mixer деградирует меньше (на 45%), чем ResNet (75%).
Авторы визуализировали фичи, среди них есть как глобальные на всё изображение, так и вполне локальные, в том числе одни фичи в разных фазах.
Интересно, что буквально в то же время (через 1-3 дня) на arxiv запостили ещё несколько очень близких работ, например “Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks” от Tsinghua University (https://arxiv.org/abs/2105.02358), “Do You Even Need Attention? A Stack of Feed-Forward Layers Does Surprisingly Well on ImageNet” из Оксфорда (https://arxiv.org/abs/2105.02723) и “ResMLP: Feedforward networks for image classification with data-efficient training” из FB/Сорбонны/INRIA (https://arxiv.org/abs/2105.03404). Другая китайская работа конца мая 2021 “Can Attention Enable MLPs To Catch Up With CNNs?” (https://arxiv.org/abs/2105.15078) кратко анализирует эти разные решения.
Оксфордская работа практически такая же, только порядок слоёв обратный, сначала фичи, потом патчи. В ResMLP в первом блоке про token-mixing один слой вместо двух и везде афинные преобразования вместо LayerNorm. В китайской работе self-attention вроде как заменяется на external attention, реализуемый через простые линейные слои.
Такие вот дела. Не совсем понимаю, откуда такая точная синхронизация разных работ по времени.