MAGMA: Manifold Regularization for MAEs

TL;DR

MAGMA addresses the lack of regularization in transformer-based Masked Autoencoders by introducing a batch-wide, layer-wise manifold regularization that enforces cross-layer representational consistency. Formulated as $L = L_{SSL} + \lambda L_{Reg}$ with $L_{Reg}$ computed via a normalized Laplacian of a batch-wide RBF affinity on a reference layer, MAGMA pulls together target-layer representations when they are similar in the reference layer but divergent in the target layer. Empirically, MAGMA yields significant gains on MAE and U-MAE across CIFAR-100, STL-10, Tiny-ImageNet, and ImageNet-100, and provides improvements to other SSL baselines like SimCLR and VICReg, supported by training-dynamics analyses and qualitative visualizations. The approach is efficient, scalable, and broadly applicable to SSL with Transformer backbones, offering a practical path to stronger self-supervised representations with minimal architectural changes.

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.

Figures (6)

• Figure 1: Visualization for the proposed regularization loss MAGMA with MAE: MAGMA penalizes representations that are close in the latent space of intermediate layer $k$ but far apart in layer $l$ latent space. This induces a regularization effect across different layers that preserves inter-sample and intra-batch relationships thus enforcing consistency in the latent representation space. Note that we demonstrate this for MAE based pre-training with a transformer encoder-decoder architecture such as ViT.
• Figure 2: (a) The regularization loss showed for MAE and M-MAE. For MAE we calculate the loss without backpropagating. For M-MAE, we apply the loss after 10 warmup epochs, and take it out after 100 epochs. (b) The online accuracy was obtained by training a linear layer on the representations produced by the encoder throughout pretraining. The accuracy slightly drops for M-MAE when the regularization kicks in but increases at a significantly higher rate compared to MAE.
• Figure 3: Effect of regularization. Implication: if the representations from any two layers are close, then the output representation will also be close.
• Figure 4: PaCMAP plots for MAE-based methods. Applying MAGMA on top of U-MAE leads to compact and well-defined clusters.
• Figure 5: Visualization of PCA's leading component for features extracted from different layers of a ViT-B pretrained using MAE, M-MAE (ours), U-MAE, and MU-MAE (ours).