[10/🧵] This work is the result of an amazing team effort w/ Julius von Kügelgen, Alain Ryser, Thomas Sutter, Bernhard Schölkopf, Julia Vogt
📜 arXiv: arxiv.org/abs/2502.06314
👩💻 Code: github.com/alicebizeul/...
[9/🧵] As a result, PMAE’s masking ratio becomes a more interpretable and robust hyperparameter!
Unlike MAEs—where the optimal ratio varies across datasets—we show that masking PCs that account for 20% of the data variance consistently yields near-optimal performance.
[8/🧵] What about the masking ratio?
In MAEs, this ratio represents the proportion of masked-out pixels.
In PMAE, we make the masking ratio more data-driven by leveraging PCA. The masking ratio now reflects the proportion of data variance captured by the set of masked PCs.
[7/🧵] We show that PMAE outperforms MAEs in downstream image classification on CIFAR10, TinyImageNet and MedMNIST datasets.
Using a ViT-Tiny, we observe an average 38% improvement in linear probing performance compared to MAEs with the standard 75% masking ratio.
[6/🧵] However, instead of working with a subset of pixels, the ViT processes the original image with a subset of its principal components (PCs) masked out. The model is then trained to output images that, when projected onto the masked PCs, match the ground truth.
[5/🧵] Our approach, Principal Masked Autoencoders (PMAE), closely follows the design of the Masked Autoencoder (MAE): a Vision Transformer (ViT) encoder-decoder is trained to reconstruct missing information from the visible parts.
[4/🧵] We posit that this reduces the redundancy between visible and masked-out information and ensures the visible information is predictive of masked-out components.
[3/🧵] Need a refresher on PCA?
For natural images, projecting data into its principal components partitions the information into a set of global features.
By masking principal components instead of raw pixels, we effectively mask more global rather than local features.
[2/🧵] What if, instead of masking pixels, we mask information in a more meaningful space using off-the-shelf image transformations?
We keep it simple: we consider the space of principal components and reconstruct masked-out principal components instead of raw pixels.
[1/🧵] Unlike text, images are not compact representations. Masking and reconstructing 75% of raw pixels—a common practice in MIM—can thus lead to failure cases:
❌ Visible pixels may be redundant with the masked ones.
❌ Visible pixels may be unpredictive of the masked regions.
✨New Preprint ✨ Ever thought that reconstructing masked pixels for image representation learning seems sub-optimal?
In our new preprint, we show how masking principal components—rather than raw pixel patches— improves Masked Image Modelling (MIM).
Find out more below 🧵