MAGMA: Manifold Regularization for MAEs

TL;DR

MAGMA introduces a new regularization loss for Transformer-based Masked Autoencoders, significantly enhancing their performance and also benefiting other SSL methods like VICReg and SimCLR.

Contribution

The paper proposes a novel batch-wide layer-wise regularization loss for Transformer-based MAEs, improving their effectiveness and generalizing to other SSL approaches.

Findings

Regularization improves MAE performance.

The loss benefits other SSL methods.

Code available online.

Abstract

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.