Modality-Agnostic Self-Supervised Learning with Meta-Learned Masked Auto-Encoder

TL;DR

This paper introduces MetaMAE, a modality-agnostic self-supervised learning framework that leverages meta-learning to improve masked auto-encoder performance across diverse data modalities.

Contribution

It develops a unified MAE framework enhanced with meta-learning techniques, enabling effective SSL across multiple modalities, which is a novel approach in the field.

Findings

MetaMAE significantly outperforms prior baselines on the DABS benchmark.

The integration of meta-learning improves the reconstruction quality in SSL.

MetaMAE demonstrates strong generalization across different data modalities.

Abstract

Despite its practical importance across a wide range of modalities, recent advances in self-supervised learning (SSL) have been primarily focused on a few well-curated domains, e.g., vision and language, often relying on their domain-specific knowledge. For example, Masked Auto-Encoder (MAE) has become one of the popular architectures in these domains, but less has explored its potential in other modalities. In this paper, we develop MAE as a unified, modality-agnostic SSL framework. In turn, we argue meta-learning as a key to interpreting MAE as a modality-agnostic learner, and propose enhancements to MAE from the motivation to jointly improve its SSL across diverse modalities, coined MetaMAE as a result. Our key idea is to view the mask reconstruction of MAE as a meta-learning task: masked tokens are predicted by adapting the Transformer meta-learner through the amortization of unmasked tokens. Based on this novel interpretation, we propose to integrate two advanced meta-learning techniques. First, we adapt the amortized latent of the Transformer encoder using gradient-based meta-learning to enhance the reconstruction. Then, we maximize the alignment between amortized and adapted latents through task contrastive learning which guides the Transformer encoder to better encode the task-specific knowledge. Our experiment demonstrates the superiority of MetaMAE in the modality-agnostic SSL benchmark (called DABS), significantly outperforming prior baselines. Code is available at https://github.com/alinlab/MetaMAE.