Understanding and Enhancing Mask-Based Pretraining towards Universal Representations
This work addresses the problem of improving universal representations for researchers and practitioners in machine learning, offering a novel theoretical framework and method that enhances state-of-the-art models, though it builds incrementally on existing mask-based techniques.
The paper tackled the unclear role and limits of mask-based pretraining in learning data representations by characterizing its behavior through test risk in high-dimensional linear regression, and proposed a new pretraining scheme, R^2MAE, which outperformed standard masking schemes across vision, language, DNA sequence, and single-cell models.
Mask-based pretraining has become a cornerstone of modern large-scale models across language, vision, and recently biology. Despite its empirical success, its role and limits in learning data representations have been unclear. In this work, we show that the behavior of mask-based pretraining can be directly characterized by test risk in high-dimensional minimum-norm ("ridge-less") linear regression, without relying on further model specifications. Further analysis of linear models uncovers several novel aspects of mask-based pretraining. The theoretical framework and its implications have been validated across diverse neural architectures (including MLPs, CNNs, and Transformers) applied to both vision and language tasks. Guided by our theory, we propose an embarrassingly simple yet overlooked pretraining scheme named Randomly Random Mask AutoEncoding (R$^2$MAE), which enforces capturing multi-scale features from data and is able to outperform optimal fixed mask ratio settings in our linear model framework. We implement R$^2$MAE in vision, language, DNA sequence, and single-cell models, where it consistently outperforms standard and more complicated masking schemes, leading to improvements for state-of-the-art models. Our code is available at: https://github.com/MingzeDong/r2mae