From Pixels to Components: Eigenvector Masking for Visual Representation Learning

TL;DR

It is posited that predicting masked from visible components involves more high-level features, allowing the masking strategy to extract more useful representations, and is corroborated by the empirical findings which demonstrate improved image classification performance for component over pixel masking.

Abstract

Predicting masked from visible parts of an image is a powerful self-supervised approach for visual representation learning. However, the common practice of masking random patches of pixels exhibits certain failure modes, which can prevent learning meaningful high-level features, as required for downstream tasks. We propose an alternative masking strategy that operates on a suitable transformation of the data rather than on the raw pixels. Specifically, we perform principal component analysis and then randomly mask a subset of components, which accounts for a fixed ratio of the data variance. The learning task then amounts to reconstructing the masked components from the visible ones. Compared to local patches of pixels, the principal components of images carry more global information. We thus posit that predicting masked from visible components involves more high-level features, allowing our masking strategy to extract more useful representations. This is corroborated by our empirical findings which demonstrate improved image classification performance for component over pixel masking. Our method thus constitutes a simple and robust data-driven alternative to traditional masked image modeling approaches.

Figures (15)

• Figure 1: From Pixels to Components. Masked image modeling involves reconstructing masked-out patches of pixels from visible ones. Instead of masking in pixel space (top), we propose applying random masks to a transformed version of the image, specifically to its principal component representation (bottom). Two disjoint sets of components are used as input and reconstruction target.
• Figure 2: Overview of the Principal Masked Autoencoder. A principal masked autoencoder (PMAE) differs from a vanilla MAE he_masked_2021 by performing the masking in the space of principal components ${\mathbf{x}}_\text{PC}\!=\!\text{PCA}({\mathbf{x}})$ rather than in pixel space. The visible principal components $\bm {\mathbf{m}} \odot {\mathbf{x}}_{\text{PC}}$ are then projected back into the observation space and serve as the input for a ViT encoder-decoder architecture. Masked principal components, $(\mathbf{1}\!-\!{\mathbf{m}}) \odot {\mathbf{x}}_{\text{PC}}$, serve as the reconstruction target.
• Figure 3: Masking in pixel space. Image (left) with a random spatial mask partially removing relevant information (middle) and removing all semantic information (right). The latter is an example in which MIM would likely fail to learn useful representations.
• Figure 4: Mask Design in PMAE. 1. Perform PCA 2. For each batch, randomly shuffle the principal components and select a subset to construct the input (light blue), while the remaining components are used to create the reconstruction target (dark blue). In PMAE$_{\text{ocl}}$, the input components are chosen to explain $100\!\times\!(1\!-\!r)\%$ of the data's variance with $r$ optimized for downstream performance (here $R_{\text{opt}}\!=\!0.15$). In PMAE$_{\text{rd}}$, the input explains between 10% and 90% of the variance with $r$ sampled independently and uniformly in $[0.1,0.9]$ for each batch.
• Figure 5: Impact of the Masking Ratio. MAE (top) and PMAE (bottom) linear probing accuracy for varying masking ratios. The masking ratio is a sensitive and data-dependent hyper-parameter. While for MAE a clear masking guideline is hard to extract, for PMAE we observe close-to-optimal performance across datasets for $20\%$ of the data variance masked. (Right) Learning curves. Linear probe accuracy for CIFAR10 classification across training epochs. PMAE outperforms MAE's final performance after 200 epochs.