CoMA: Complementary Masking and Hierarchical Dynamic Multi-Window Self-Attention in a Unified Pre-training Framework

TL;DR

CoMA introduces a complementary masking strategy and a hierarchical vision transformer with dynamic multi-window self-attention, significantly improving pre-training efficiency and adaptability in image representation learning.

Contribution

The paper proposes CoMA with a novel complementary masking approach and DyViT with dynamic multi-window self-attention, enhancing learning efficiency and model adaptability over existing MAE methods.

Findings

Pre-trained CoMA-DyViT matches MAE performance with only 12% of pre-training epochs.

DyViT reduces parameters and FLOPs while improving fine-grained feature learning.

Pre-training time per epoch is reduced by 10% with DyViT.

Abstract

Masked Autoencoders (MAE) achieve self-supervised learning of image representations by randomly removing a portion of visual tokens and reconstructing the original image as a pretext task, thereby significantly enhancing pretraining efficiency and yielding excellent adaptability across downstream tasks. However, MAE and other MAE-style paradigms that adopt random masking generally require more pre-training epochs to maintain adaptability. Meanwhile, ViT in MAE suffers from inefficient parameter use due to fixed spatial resolution across layers. To overcome these limitations, we propose the Complementary Masked Autoencoders (CoMA), which employ a complementary masking strategy to ensure uniform sampling across all pixels, thereby improving effective learning of all features and enhancing the model's adaptability. Furthermore, we introduce DyViT, a hierarchical vision transformer that employs a Dynamic Multi-Window Self-Attention (DM-MSA), significantly reducing the parameters and FLOPs while improving fine-grained feature learning. Pre-trained on ImageNet-1K with CoMA, DyViT matches the downstream performance of MAE using only 12% of the pre-training epochs, demonstrating more effective learning. It also attains a 10% reduction in pre-training time per epoch, further underscoring its superior pre-training efficiency.