Repensando a Dependência de Patches para Autoencoders Mascarados

Rethinking Patch Dependence for Masked Autoencoders

January 25, 2024

Autores: Letian Fu, Long Lian, Renhao Wang, Baifeng Shi, Xudong Wang, Adam Yala, Trevor Darrell, Alexei A. Efros, Ken Goldberg

cs.AI

Resumo

Neste trabalho, reexaminamos as dependências entre patches no mecanismo de decodificação de autoencoders mascarados (MAE). Decompomos esse mecanismo de decodificação para reconstrução de patches mascarados no MAE em auto-atenção e atenção cruzada. Nossas investigações sugerem que a auto-atenção entre patches mascarados não é essencial para aprender boas representações. Para tanto, propomos um novo framework de pré-treinamento: Cross-Attention Masked Autoencoders (CrossMAE). O decodificador do CrossMAE utiliza apenas atenção cruzada entre tokens mascarados e visíveis, sem degradação no desempenho em tarefas subsequentes. Esse design também permite decodificar apenas um pequeno subconjunto de tokens mascarados, aumentando a eficiência. Além disso, cada bloco do decodificador pode agora aproveitar diferentes características do codificador, resultando em uma melhoria na aprendizagem de representações. O CrossMAE iguala o desempenho do MAE com 2,5 a 3,7 vezes menos computação de decodificação. Ele também supera o MAE na classificação do ImageNet e na segmentação de instâncias do COCO sob a mesma quantidade de computação. Código e modelos: 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