Переосмысление зависимости патчей в маскированных автокодерах

Rethinking Patch Dependence for Masked Autoencoders

January 25, 2024

Авторы: Letian Fu, Long Lian, Renhao Wang, Baifeng Shi, Xudong Wang, Adam Yala, Trevor Darrell, Alexei A. Efros, Ken Goldberg

cs.AI

Аннотация

В данной работе мы пересматриваем межфрагментные зависимости в механизме декодирования маскированных автоэнкодеров (MAE). Мы разлагаем этот механизм декодирования для восстановления маскированных фрагментов в MAE на самовнимание и перекрестное внимание. Наши исследования показывают, что самовнимание между маскированными фрагментами не является необходимым для обучения качественных представлений. В связи с этим мы предлагаем новую структуру предварительного обучения: Маскированные Автоэнкодеры с Перекрестным Вниманием (CrossMAE). Декодер CrossMAE использует только перекрестное внимание между маскированными и видимыми токенами, без ухудшения производительности на последующих задачах. Этот дизайн также позволяет декодировать только небольшое подмножество маскированных токенов, повышая эффективность. Кроме того, каждый блок декодера теперь может использовать различные признаки энкодера, что приводит к улучшению обучения представлений. CrossMAE демонстрирует сопоставимую с MAE производительность при 2,5–3,7-кратном сокращении вычислительных затрат на декодирование. Он также превосходит MAE в классификации на ImageNet и сегментации экземпляров на COCO при одинаковых вычислительных ресурсах. Код и модели: https://crossmae.github.io.

English

In this work, we re-examine inter-patch dependencies in the decoding mechanism of masked autoencoders (MAE). We decompose this decoding mechanism for masked patch reconstruction in MAE into self-attention and cross-attention. Our investigations suggest that self-attention between mask patches is not essential for learning good representations. To this end, we propose a novel pretraining framework: Cross-Attention Masked Autoencoders (CrossMAE). CrossMAE's decoder leverages only cross-attention between masked and visible tokens, with no degradation in downstream performance. This design also enables decoding only a small subset of mask tokens, boosting efficiency. Furthermore, each decoder block can now leverage different encoder features, resulting in improved representation learning. CrossMAE matches MAE in performance with 2.5 to 3.7times less decoding compute. It also surpasses MAE on ImageNet classification and COCO instance segmentation under the same compute. Code and models: https://crossmae.github.io