Coordination Among Neural Modules Through a Shared Global Workspace
Anirudh Goyal, Aniket Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Mozer, Yoshua Bengio
Статья: https://arxiv.org/abs/2103.01197

И ещё одна релевантная статья в значительной степени от тех же авторов, что и про RIM/BRIM. На этот раз суперсвежая, от 1 марта.

В современном DL есть большое движение в сторону более богатых и модульных репрезентаций. С RIM это вообще очевидно, но RIM пока не массовая архитектура. Есть активное движение к архитектуре Mixture of Experts (MoE), применявшейся в свежих Switch Transformer или переводчике M2M-100. Но даже обычные трансформеры по сути сегментируют вход по позициям, где у каждой позиции есть своя репрезентация.

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

И тут мы снова возвращаемся к упомянутой в прошлый раз теории глобального рабочего пространства (Global Workspace Theory, GWT), где вводится ограниченный канал (то самое глобальное рабочее пространство), через который взаимодействуют различные модули-специалисты, при этом соревнуясь за доступ к нему.

Данная работа вдохновлена GWT и авторы утверждают, что подобная архитектура даёт гибкость и генерализацию.

В целом, система составленная из модулей-специалистов с разреженными взаимодействиями обладают хорошими свойствами типа 1) возможности масштабирования (легко добавить новых специалистов); 2) лучшей устойчивости (выбивание части специалистов может не сильно повредить системе); 3) вычислительной эффективности (вычисления производятся в основном локально, что хорошо). Но синхронизация всё равно нужна.

GWT предлагает простой и понятный способ это устроить. Есть глобальная общая репрезентация, называющаяся рабочим пространством или доской. Она может быть модифицирована теоретически любым специалистом и далее рассылается всем специалистам в системе. При этом, чтобы в неё что-то записать, специалист соревнуется с остальными и по факту возникает некая приоритизация сигналов специалистов.

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

Для работы с общим пространством (SW, shared workspace) используется механизм внимания (это обычное multihead dot-product attention как в трансформерах), где элементы этого пространства, как и специалисты, рассматриваются как неупорядоченное множество. SW — это в каком-то смысле память. Скорее рабочая память.

Данная архитектура накручивается как на RIM, так и на трансформеры. В RIM всё и так уже модуляризовано (каждый модуль — это специалист), а в трансформерах специалистом считается каждая отдельная позиция.

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

Процесс работы устроен следующим образом:

1. Вход обрабатывается, чтобы получить его репрезентацию каждым специалистом. По сути здесь мы получаем какие-то начальные представления, с которыми дальше работать.

2. Специалисты соревнуются за право записи в глобальное пространство. Этот механизм по сути обеспечивает реагирование только на критически важные сигналы, так чтобы всякая ерунда не захламляла GWT. Соревнования сделано через внимание, Q поступает из памяти (то самое GWT), а K и V от специалистов. Память по результату обновляется либо через soft-attention от всех специалистов (но по факту мы знает, что обычно там доминирует их небольшое число, ибо софтмакс), либо вообще через top-k softmax. Первые модели дальше будут называться SSW (shared workspace with soft competition), вторые HSW (hard attention). Содержимое памяти обновляется через хитрый механизм с гейтами как в RMC (https://arxiv.org/abs/1806.01822).

3. Содержимое памяти бродкастится всем специалистам. Специалисты обновляют своё состояние снова через механизм внимания (генерят свои Q, получают K,V из памяти, прибавляют к своим состояниям взвешенную таким образом через внимание память). После обновления состояния с учётом памяти специалисты вычисляют свои обычные функции (LSTM/GRU или FFN для трансформера).

Заодно как бонус получаем линейную вычислительную сложность от числа специалистов.

Неожиданная приятность другого рода в том, что эта работа ссылается на наш с Мишей Бурцевым Memory Transformer (https://arxiv.org/abs/2006.11527). У нас тоже в одном из вариантов в трансформере был боттлнек в виде памяти небольшого размера.

И дальше собственно эксперименты.

На задачке определения равносторонних треугольников (когда на картинке есть три кластера точек) трансформеры с общим рабочим пространством (особенно HSW) сходятся быстрее обычных и дают более высокое качество. Sparse transformer, кстати, оказывается здесь особенно плох.

Задачки на object tracking из CATER dataset тоже показывают, что трансформеры с SW лучше обычных.

На задачах relational reasoning из CLEVR, где по картинке и текстовому вопросу надо ответить про определённые свойства объектов и их отношения с другими, трансформеры с SW сходятся сильно быстрее обычных и дают более высокое итоговое качество.

На задачах physical reasoning и всё тех же отскакивающих мячиков, что и раньше (надо предсказать будущие кадры видео), RIM + SW лучше всех и это преимущество растёт и с ростом горизонта предсказания, и с ростом числа мячиков.

Была ещё задачка на RL из Starcraft 2, где надо предсказывать будущие наблюдения агентов, RIM + SW тоже бьём и обычный LSTM, и RIM. Что интересно, обычный RIM с попарными взаимодействиями между специалистами, здесь даже хуже LSTM.

В общем хорошая тема, ждём развитий.

Perceiver: General Perception with Iterative Attention
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, Joao Carreira
Статья: https://arxiv.org/abs/2103.03206
Код (неавторский): https://github.com/lucidrains/perceiver-pytorch

Работа, нацеленная на мультимодальность, чтобы одна универсальная архитектура на основе трансформера могла работать с данными разной природы.

Сейчас работа с различными модальностями часто основана на выборе подходящих задаче inductive biases, учитывающих знание о том, как соотносятся между собой элементы входа. Например, свёрточные сетки полагаются на локальность фич в изображениях и имеют подходящий для этого bias. С текстами хорошо работает трансформер. Ну и так далее. Смешивать разные модальности в одной модели в целом непросто.

Хочется заменить это на одну универсальную архитектуру, делающую по возможности меньше таких предположений о структуре входа. Ну то есть давайте заложим поменьше biases, сделаем универсальную архитектуру, и позволим ей выучить всё из данных. Если получится, не придётся затачиваться на отдельные модальности и можно будет использовать одну общую архитектуру для смешанных модальностей.

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

Как оно работает?

Основная идея подхода — завести маленький набор латентных переменных (latent array). Эти юниты будут через механизм внимания обращаться ко входу (потенциально большому) и таким образом окажутся боттлнеком, который должен отсеять важные части входа. И делать это perceiver может итеративно, потому что таких блоков в сети будет несколько. Каждый раз, соответственно, он может обращаться к разным частям входа, в зависимости от того, что он “узнал” на предыдущих шагах.

Авторы предлагают посмотреть на эту модель как на выполняющую кластеризацию входа end-to-end, где скрытые юниты являются центрами кластеров. Не знаю, мне сходу неочевидно, почему это так.

Поскольку временная/пространственная информация всё-таки важна, а модель на структуру входа напрямую не закладывается, авторы добавляют к каждому входному элементу позиционные энкодинги на основе Фурье (синусы и косинусы от логарифмического банка частот, а координата внутри измерения масштабируется в диапазон [-1,1]). Эмбеддинги не добавляются ко входу (как в оригинальном трансформере), а конкатенируются с ним.

Авторы считают, что эта история с позиционными эмбеддингами не дискредитирует их подход c уменьшением implicit biases. Типа, мы даём сети фичи, а как их использовать её дело, сама выучит и решит как лучше. К тому же их легче адаптировать к новым доменам, чем собирать новую архитектуру. Ну и вообще мультимодальные модели так легче собирать.

Внутри модели повторяются два основных блока: 1) кросс-внимание (cross-attention), где Q приходит из низкоразмерной латентной части, а K/V из входа (большого), и 2) обычное внимание (self-attention трансформера), трансформирующее эти latent’ы. Эти блоки могут повторяться много раз, и даже иметь расшаренные веса между этими повторениями. Тогда получается что-то типа развёрнутой по глубине RNN или универсального трансформера.

Стандартное внимание по-прежнему остаётся квадратичным, но оно здесь от низкоразмерного latent’а (<=1024), так что не так страшно. Можно заодно глубины добавить, чтобы модель посложнее была (на ImageNet сделали 48 таких блоков). У latent’ов есть обучаемые позиционные эмбеддинги. Сам этот latent array, так понимаю, выучивается, а не зависит от входа или какого-то рандома.

Эксперименты показали, что на ImageNet персивер в ~44M параметров бьёт заточенный на это ResNet-50 (у него правда обычно ~25M), сильно бьёт обычный трансформер (но там пришлось уменьшать размер входа до 64x64 чтобы он смог работать) и в общем сравним с недавним ViT (https://t.me/gonzo_ML/434).

На датасете AudioSet, где надо классифицировать звуковые события в видео, моделька дала SoTA на чистой видео модальности, а также лучшая и на чистом аудио, и на audio+video, хоть там разница и поменьше.

На датасете с облаками точек ModelNet40 не SoTA, но весьма достойно, особенно учитывая, что SoTA весьма сложна.

В общем, модель работает. Что особенно прекрасно, она хорошо контрибьютит в большое свежее движение трансформеров от текстов к картинкам и к другим модальностям, а тут и к мультимодальности.