Hierarchical Perceiver

Joao Carreira, Skanda Koppula, Daniel Zoran, Adria Recasens, Catalin Ionescu, Olivier Henaff, Evan Shelhamer, Relja Arandjelovic, Matt Botvinick, Oriol Vinyals, Karen Simonyan, Andrew Zisserman, Andrew Jaegle

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

Код: пока нет (https://github.com/deepmind/deepmind-research/issues/332)

Дальнейшее развитие Perceiver (https://t.me/gonzo_ML/545) и Perceiver IO (https://t.me/gonzo_ML/700).

Если проследить развитие универсальных (general-purpose) архитектур в контексте трансформеров, то уже сам трансформер был большим достижением в этом плане, потому что не делал сильных предположений про свойства данных и имел меньшее количество implicit biases по сравнению, например, со свёрточными сетями. Мы много писали по этой теме, когда недавно разбирали возрождение MLP (https://t.me/gonzo_ML/800). Работа с разными модальностями обычно была устроена через специализированный препроцессинг входов разных модальностей (например, нарезание картинок на патчи в случае ViT). Заодно это также решало проблему больших входов, с которыми у классических трансформеров были проблемы из-за квадратичной сложности механизма внимания.

Perceiver был шагом вперёд, потому что был способен работать с различными модальностями и с размерами входов в сотни тысяч элементов за счёт не зависящих от размера входа латентных переменных. Perceiver IO добавил такие же свойства для выходов. Оба использовали глобальную операцию внимания, это было эффективнее оригинальных трансформеров, но всё равно глобальность мешала скейлить персиверы дальше и, например, обрабатывать картинки и видео высокого разрешения. Различная дополнительная информация, например, пространственные координаты, конкатенировались со входом в виде специальных фич, например, Фурье для картинок и аудио. Это требовало доменно-специфичной предобработки.

В текущей работе про иерархический персивер (Hierarchival Perceiver, HiP) авторы делают персиверы ещё более быстрыми, масштабируемыми на ещё больше входов и выходов, а также уменьшают необходимость feature engineering’а для входов. Заодно модели получаются и более точными.

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

Как это всё работает?

На вход модели подаются тензоры X размерности M*C, где M — число токенов, а C — каналов. В случае разных модальностей и вообще данных различный препроцессинг с flattening и padding позволяет свести всё к единой форме.

Perceiver вводил новую сущность, обучаемые латентные векторы Z (latents) размерности K*D, где К — число латентных векторов, а D — их размерность. Эти векторы обновлялись через cross-attention ко входу (Q из латентной части, K+V их входной).

После cross-attention был обычный self-attention над этими латентными векторами Z. Собственно так и снимаем квадратичную сложность от размера входа. Она либо O(N*M), либо O(N^2), но не O(M^2).

Далее в HiP используется тот факт, что после применения flattening к многоразмерному входу, какая-то локальная структура в сигнале всё равно сохраняется. Например, если вы сделали flattening для картинки по строкам, то внутри строки структура осталась. Но между различными строками, конечно, потерялась.

HiP пытается использовать эту сохранившуюся локальную структуру с помощью групп. Вход каждого блока HiP разбивается на G групп, по M/G токенов каждая. Каждая из групп обрабатывается независимо последовательность операций Perciever: сначала cross-attention каждой группы с набором K латентных векторов, затем self-attention + MLP, как в обычном трансформере. В каждой группе это происходит независимо, но параметры расшарены, кроме латентных векторов, которые в каждой группе свои. Выходы всех групп затем объединяются и на выходе блока имеем K*G токенов. Затем это подаётся в следующий аналогичный блок и процедура повторяется.

Блоки организованы иерархически, например, в энкодере сначала 4 группы, потом 2, потом 1 (бутылочное горлышко), а затем в декодере делаем зеркально 2 группы и на выходе 4 (такая схема в работе обозначается как 4-2-1-2-4). Но вы не обязаны делать симметрично, можно сделать как угодно, например, 1-8-2-4. При этом с уменьшением количества групп обычно увеличивается число каналов (это тоже отличается от обычного трансформера, где в каждом слое одинаковое число каналов). Это если нужна модель для генерации dense outputs. Если нужна лишь классификация, то можно остановится на bottleneck и брать Q из обучаемого латентного вектора, а фичи (K,V) из боттлнека. Также в модели есть skip-connections между энкодером и декодером по аналогии с U-Net.

Модели типа ViT выучивали позиционные эмбеддинги для токенов и добавляли их к репрезентации токена. Это удобно и общо, и хорошо работает для ViT с небольшим числом входов (256 например), но плохо работает со входами сильно большей размерности (а в персивере это на порядки больше). В персиверах хорошо работали Фурье-фичи, которые конкатенировались с каждым входом. Но с этими фичами было две проблемы: 1) их надо дизайнить вручную для каждой новой модальности, 2) они большие и занимают много памяти для входов большого разрешения. Для видео из 32 кадров размера 224*224 одни только эти фичи занимают 13Gb (в примере с 64 частотными полосами и 258 фичами Фурье для случае ImageNet).

В данной работе авторы переключаются на обучаемые низкоразмерные (размерности 16 или 32) позиционные эмбеддинги. Секретным соусом для успешного обучения оказался маскирующий автоэнкодер (MAE), когда случайно и равномерно маскируется фиксированная доля входного тензора. Причём доля эта большая (например, 85-90%), а маскируемые токены не заменяются на какой-то спец.токен типа <MASK>, а удаляются из входа. Чтобы восстановить их на выходе, на последний слой иерархического декодера подаются выученные позиционные эмбеддинги замаскированных (выброшенных) токенов. Этот подход аналогичен описанному в интересной работе из FAIR под названием “Masked Autoencoders Are Scalable Vision Learners” (https://arxiv.org/abs/2111.06377), про которую я так и не успел написать.

В работе предлагают две конфигурации HiP: HiP-16 и HiP-256. HiP-16 имеет структуру 16-4-1-1. HiP-256 глубже и у него в начале и конце сети есть по два дополнительных блока с 256 и 64 группами, все остальные блоки у них одинаковые, получается структура 256-64-16-4-1-1.

Ещё в HiP-16 обучаемые позиционные эмбеддинги размерности 32, а в HiP-256 они размерности 16.

Все эксперименты в два этапа: 1) предобучение через MAE, 2) файнтюнинг на supervised разметке.

На ImageNet HiP везде лучше и быстрее Perceiver IO. SoTA не бьёт, но всё равно достойно. Предобучение через MAE очень помогает. Локальная структура видимо таки используется, потому что на перемешанных пикселях работает хуже.

На AudioSet (клип из 8 кадров 224*224 + 2.56 секунды аудио в 48КГц, и всего 401К входных сэмплов для видео и 123К для аудио) в multilabel классификации или бьёт (на аудио) или сравним (на аудио+видео) с обычным персивером, но зато работает на сильно более длинных входах.

На PASCAL VOC и семантической сегментации HiP-16 лучше Perceiver IO и примерно как ResNet-50, но немного помедленнее.

На классификации объектов по облакам точек из ModelNet40 HiP-16 без MAE хуже, чем Perceiver IO, а с MAE лучше, но непонятно почему, позиционные эмбеддинги выглядят похоже, возможно дело в регуляризации, которую даёт предобучение.

В общем ещё один шаг в сторону развития персиверов и заодно к универсальности (меньше дизайнить ручками). А заодно ещё один плюс в копилку полезности self-supervised pre-training и dense masked auto-encoding. В работе намечен большой план на будущее, это и другие подходы к предобучению (возможно, контрастному), и что будет при множестве маленьких входов, а не малом числе больших, а также можно ли отказаться от аугментации и перейти к anymodal models в автоматическом режиме.

Видимо, ждём в скором будущем новой версии персивера.

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