Understanding and Enhancing Mask-Based Pretraining for Universal Biological Representations
Mingze Dong, Leda Wang, Yuval Kluger
Mask-based pretraining has become a cornerstone of modern large-scale models across natural language processing, computer 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**), designed for learning universal representations capturing features of all scales from data. We implement **R$^2$MAE** in DNA sequence and single-cell self-supervised models, where it improves upon state-of-the-art models in terms of various performance metrics for MLP and transformer-based architectures, outperforming standard and more complicated mask pretraining schemes.