Perceiver IO: A General Architecture for Structured Inputs & Outputs
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, João Carreira
Статья: https://arxiv.org/abs/2107.14795
Пост: https://deepmind.com/blog/article/building-architectures-that-can-handle-the-worlds-data
Код (на JAX/Haiku): https://github.com/deepmind/deepmind-research/tree/master/perceiver
Код (неавторский, на PyTorch): https://github.com/krasserm/perceiver-io

В марте мы писали про архитектуру для работы с мультимодальными данными от DeepMind под названием Perceiver (https://t.me/gonzo_ML/545).

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

Но была у персивера серьёзная проблема. Хоть он и скейлился на разные входы, выходы его были просты — он годился для классификации по заданному числу классов, но не годился для генерации сложных выходов произвольного размера (картинок, текстов, множества символов и т.д.).

В июле 2021 вышел Perceiver IO, который устраняет этот недостаток.

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

На входе происходит трансформация входных данных в латентный массив (query здесь обучаемый и его размер не зависит от размера входа, всё как в стандартном Perceiver), далее трансформер обрабатывает латенты (эта часть называется processor), и на выходе происходит трансформация латентов в выход нужной структуры (которая определяется query с количеством элементов равным размерности выхода). Сами эти query обучаются, могут быть созданы вручную или же являться функциями от входа.

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

В качестве выходов и задач используются:
1) тексты (размером 512-2048 токенов или UTF-8 байт) и MLM-предобучение + задачи из GLUE (выходные размерности от 8 до 2048);
2) видео (размером 22816) и задача предсказания оптического потока (с выходом размера 11408);
3) видео+аудио+метка в задаче автокодирования видео из датасета Kinetics-700-2020 (размера 50657x704 на входе и 803297x512 на выходе; разница, видимо, из-за того, что на входе обрабатывают патчи 4x4, а выход генерят попиксельно)
4) множество юнитов StarCraft и задача выбора юнитов для замены трансформера в AlphaStar (система DeepMind, игравшая в StarCraft).
5) картинки (50176 пикселей) и классификация на 1000 классов.

Получаются вполне достойные результаты.

На языковых задачах получается обучить модель без токенизатора (на уровне UTF-8), соответствующую аналогичной по FLOPS бейзлайн-модели с SentencePiece токенизатором.

На определении оптического потока модель вполне конкурентоспособна в сравнении со специально созданными для этого методами (PWCNet, RAFT), а на одной из задач даже получила SoTA, что классно с учётом общности модели и отсутствия существенной кастомизации под задачу.

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

В AlphaStar получилось заменить оригинальный трансформер Perceiver’ом IO с 32 латентами (на входе и выходе он работает с 512 юнитами).