MAGMA: Manifold Regularization for MAEs

Masked Autoencoders (MAEs) are an important di-vide in self-supervised learning (SSL) due to their in-dependence from augmentation techniques for generating positive (and/or negative) pairs as in contrastive frame-works. Their masking and reconstruction strategy also nicely aligns with SSL approaches in natural language pro-cessing. Most MAEs are built upon Transformer-based ar-chitectures where visual features are not regularized as op-posed 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 rep-resentations of different Transformer layers. We demon-strate 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 here: https://github.com/adondera/magma