Replanteando la Dependencia de Parches para Autoencoders Enmascarados
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
Resumen
En este trabajo, reexaminamos las dependencias entre parches en el mecanismo de decodificación de los autoencoders enmascarados (MAE). Descomponemos este mecanismo de decodificación para la reconstrucción de parches enmascarados en MAE en auto-atención y atención cruzada. Nuestras investigaciones sugieren que la auto-atención entre parches enmascarados no es esencial para aprender buenas representaciones. Con este fin, proponemos un nuevo marco de preentrenamiento: Cross-Attention Masked Autoencoders (CrossMAE). El decodificador de CrossMAE aprovecha únicamente la atención cruzada entre tokens enmascarados y visibles, sin degradación en el rendimiento en tareas posteriores. Este diseño también permite decodificar solo un pequeño subconjunto de tokens enmascarados, mejorando la eficiencia. Además, cada bloque del decodificador puede ahora aprovechar diferentes características del codificador, lo que resulta en un aprendizaje de representaciones mejorado. CrossMAE iguala el rendimiento de MAE con un cálculo de decodificación entre 2.5 y 3.7 veces menor. También supera a MAE en clasificación de ImageNet y segmentación de instancias en COCO bajo el mismo cómputo. Código y 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