MAGMA: Manifold Regularization for MAEs
This work addresses a performance bottleneck in MAEs for computer vision researchers, offering an incremental improvement through regularization.
The paper tackles the lack of regularization in Transformer-based Masked Autoencoders (MAEs) for self-supervised learning by introducing MAGMA, a batch-wide layer-wise regularization loss, which significantly improves MAE performance and also benefits other SSL methods like VICReg and SimCLR.
Masked Autoencoders (MAEs) are an important divide in self-supervised learning (SSL) due to their independence from augmentation techniques for generating positive (and/or negative) pairs as in contrastive frameworks. Their masking and reconstruction strategy also nicely aligns with SSL approaches in natural language processing. Most MAEs are built upon Transformer-based architectures where visual features are not regularized as opposed to their convolutional neural network (CNN) based counterparts, which can potentially hinder their performance. To address this, we introduce MAGMA, a novel batch-wide layer-wise regularization loss applied to representations of different Transformer layers. We demonstrate that by plugging in the proposed regularization loss, one can significantly improve the performance of MAE-based models. We further demonstrate the impact of the proposed loss on optimizing other generic SSL approaches (such as VICReg and SimCLR), broadening the impact of the proposed approach. Our code base can be found at https://github.com/adondera/magma.