Title: Understanding and Enhancing Mask-Based Pretraining towards Universal Representations
Venue: NeurIPS 2025
Accepted: yes
Keywords: Mask-based pretraining

self-supervised learning

high-dimensional statistics

computational biology
Abstract: 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²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²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 [this URL](https://github.com/MingzeDong/r2mae).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the theoretical foundation and empirical behavior of mask-based pretraining, which is widely used in self-supervised learning for language, vision, and biological data. The authors propose a unified theoretical framework based on high-dimensional linear regression (including spiked covariance models) to explain why and when mask-based pretraining is effective, and how the optimal mask ratio depends on data and model size. Building on these insights, they introduce a simple pretraining method called R²MAE (Randomly Random Mask AutoEncoding), where the mask ratio is randomly sampled for each training batch. The method is evaluated on DNA sequence and single-cell gene expression tasks, showing modest improvements over traditional approaches. While the theoretical analysis is mathematically solid within linear settings, its extension to deep nonlinear models remains questionable, and the empirical gains are relatively small.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- The approach is easy to implement and performs well on some biological datasets, showing minor improvements over standard random masking.\n- The paper attempts to provide a theoretical explanation for the mask-based pretraining phenomenon, which is a meaningful direction.\n\n## Weaknesses\n\n1. The main idea is just randomizing the mask ratio, which is very similar to dynamic masking and scheduled masking from earlier works [1,2,3,4,5]. The paper doesn’t fairly compare or cite these, and the novelty is overclaimed.\n2. All the theoretical analysis is based on linear models, while real masked autoencoders are highly nonlinear deep networks. There’s a big gap between the theory and practice, so the “universality” claim isn’t convincing [1,6].\n3. Experiments are only on a few models and datasets, with no evidence for challenging or large-scale biological tasks. The generality claim is not well supported [1,7].\n4. The biological results show only marginal improvements, and there’s almost no discussion on whether these gains are meaningful for real science or medicine [7].\n5. The Introduction is poorly constructed—it promises universal advances for language and vision, but the paper only does very limited, niche biological experiments (DNA sequence and single-cell gene expression) or toy datasets like MNIST, with absolutely no evaluation on general language or vision tasks. There’s no bio-specific design either; the method is generic and not tailored at all.\n6. The writing is lengthy and dense, with key points buried in technical details and formulas.\n7. There are issues in the theoretical proofs. For example, in the spiked covariance analysis, the authors make strong assumptions (like high-dimensional limits and feature alignment) without really checking if they hold in practice. Many derivations just assume ideal Gaussian data and feature structures, but small deviations in real data can break the conclusions. The claim that an optimal mask ratio always exists under spiked covariance only holds under strict conditions, which are not verified in their experiments. Much of the analysis is quite similar to recent works like MaskTwins [8], which also use linear model or spiked covariance for mask theory, and while those papers have their own limitations, they actually discuss model assumptions and limitations more carefully. This paper overstates its theoretical contribution and doesn't clearly position itself with respect to these existing analyses.  \n\n### References\n\n1. Ankner et al., Dynamic masking rate schedules for MLM pretraining, arXiv:2305.15096 (2023). [link](https://arxiv.org/abs/2305.15096)\n2. Yang et al., Learning better masking for better language model pre-training, arXiv:2208.10806 (2022). [link](https://arxiv.org/abs/2208.10806)\n3. Wettig et al., Should you mask 15% in masked language modeling?, arXiv:2202.08005 (2022). [link](https://arxiv.org/abs/2202.08005)\n4. Bao et al., BEiT: BERT Pre-Training of Image Transformers, arXiv:2106.08254 (2021). [link](https://arxiv.org/abs/2106.08254)\n5. He et al., Masked Autoencoders Are Scalable Vision Learners, CVPR 2022. [link](https://arxiv.org/abs/2111.06377)\n6. Chen et al., TokenUnify: Scalable Autoregressive Visual Pre-training with Mixture Token Prediction, arXiv:2405.16847 (2024). [link](https://arxiv.org/abs/2405.16847)\n7. Dalla-Torre et al., Nucleotide transformer: building and evaluating robust foundation models for human genomics, Nature Methods, 2025. [link](https://www.nature.com/articles/s41592-024-02523-z)\n8. Wang et al., MaskTwins: Dual-form Complementary Masking for Domain-Adaptive Image Segmentation, ICML 2025. [link](https://openreview.net/pdf?id=9CpeZ8BzPO)\n\nQUESTIONS:\nSee in Weaknesses\n\nLIMITATIONS:\nSee in Weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a theoretical framework for understanding mask-based pre-training through the lens of high-dimensional linear regression. The authors demonstrate that masked autoencoding can be characterized by the test risk of ridgeless regression with masked covariates. Following these insights, they propose R2MAE (Randomly Random Mask AutoEncoding), which samples masking ratios uniformly during training. Their method achieves state-of-the-art results on DNA sequence and single-cell gene expression tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The reduction of masked pretraining to high-dimensional linear regression provides tractable closed-form analysis while reproducing key empirical phenomena like optimal masking ratios and plateau behaviors.\n- Theoretical predictions are validated across multiple architectures (MLPs, CNNs, Transformers) and datasets, confirming that benefits only emerge in overparameterized regimes.\n- Method achieves consistent improvements over existing methods on DNA and single-cell tasks\n- Rigorous mathematical analysis\n\nWeaknesses:\n- The proposed method (R^2MAE) is only evaluated on biological data despite claims of universality, with no results on standard vision or language benchmarks.\n- The linear model ignores optimization dynamics, architectural inductive biases, and nonlinear feature interactions that may be crucial to masked pretraining\n\nQUESTIONS:\nWhat is the relationship between this method and autoregressive modeling? Fully causal masks in fact also expose a uniformly sample masking ratio to every token in the sequence (first token has full masking, last token has almost no masking). Connections to this paradigm would be very valuable for further understanding.\n\nHow does R2MAE's performance scale with dataset size and diversity? A test on an even more studied domain with clearly highly non-homogenous data such as images would be even more convincing and would make me upgrade my score to a 6. If the method fails to improve performance on image encoding (beat MAE) why is that the case?\n\nHow does your framework relate to recent work on spatially-structured masking strategies in vision, such as SiamMAE (Gupta et al., 2023) or CWM (Unifying Machine Vision ... Bear et al., 2023) which utilize specifically structured masking strategies to obtain specific behaviors from the model (such as separation of appearance and motion information) which are not present in uniform, non biased masking approaches?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel theoretical framework to understand mask-based pretraining through the lens of high-dimensional ridgeless linear regression. The authors show that both qualitative and quantitative behaviors of masked autoencoding across domains (vision, language, biology) can be recapitulated through this simplified yet analytically tractable model. Based on insights from this theory, they propose a simple yet effective pretraining strategy called R2MAE (Randomly Random Mask AutoEncoding), which applies uniformly sampled masking ratios during pretraining. Empirical results on DNA sequence and single-cell gene expression data demonstrate that R2MAE outperforms existing state-of-the-art approaches across multiple metrics.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. Theoretical Contributions: The use of high-dimensional linear models to explain masking behaviors is theoretically grounded. The analysis reveals key phenomena such as optimal masking ratios, phase transitions, and performance plateaus. Closed-form expressions for test risk under various covariance structures (isotropic, spiked, general) are novel and informative.\n\n2. Clear Empirical Validation: The authors validate theoretical predictions on diverse architectures (MLPs, CNNs, Transformers) and domains (vision, language, biology). Experiments on real biological data (DNA, single-cell) show that theoretical predictions align well with empirical trends.\n\n3. Practical Impact via R2MAE: Despite its simplicity, R2MAE consistently improves performance over standard and dynamic masking schemes in biological settings. The strategy is easy to implement and generalizable.\n\nWeaknesses:\n1. Theoretical-Algorithmic Gap: While this work provides novel insights into masked pretraining, some of the underlying technical assumptions are quite strong. In particular, the reliance on a linear model and simplified data distribution assumptions may limit the direct applicability of the theoretical findings to practical deep learning settings involving nonlinear architectures and complex data structures.\n\n2.  Simplified Feature Independence in $\\beta$ Estimation:  The assumption that each regression target ($\\beta_i$) is estimated independently or with minimal interaction may ignore feature co-adaptation, which is critical in practice.\n\n3. Minimal Comparison on General Datasets: Most empirical validations beyond theory focus on DNA and single-cell data. It would strengthen the paper to include results on standard vision/language benchmarks with R2MAE (e.g., ImageNet, GLUE).\n\nQUESTIONS:\n1. Have you considered incorporating R2MAE into more general-purpose pretrained models (e.g., ViT or LLaMA) to evaluate generalization beyond biology?\n\n2. My understanding is that R²MAE randomly samples a different masking ratio for each input from a uniform distribution—is this correct? If so, based on your theoretical analysis, is it possible to determine an optimal masking ratio for each individual sample, rather than relying on random sampling?\n\nLIMITATIONS:\nThe main limitation lies in the strong assumption of a linear model and the minimal interaction between features, which may oversimplify the complex dependencies present in real-world deep networks.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nHere, the authors propose a theoretical analysis that leads to a masking strategy for masked language models called \"R^2MAE\", and demonstrate its effectiveness in self-supervised representation learning for DNA and single-cell.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper presents a generally strong theoretical analysis, which essentially argues that the success of the masked-based pretraining objective can be explained as a trade-off between minimizing bias and optimizing risk (i.e. ensuring that for each feature, the model learns which other features can be  used to reconstruct the feature, while minimizing overfitting). They do convincing empirical experiments to validate the key predictions of their theoretical analysis, showing that the masked-based pretraining is only beneficial in the overparameterized regime, and that optimal masking ratio differs by evaluation task.\n\nBased upon these insights, the authors propose a new masking strategy (R^2MAE), which contrary to previous adaptive masking proposals, simply samples masking ratio during pre-training at uniform, and demonstrate the effectiveness of this strategy on two different biological domains, DNA sequences and single-cell gene expression. These empirical results are strong, demonstrating that their proposed method can outperform even SOTA proposals on at least some kinds of tasks in their respective domains, on standardized benchmarks at least, and critically, frequently over fixed/dynamic masking ratio proposals.\n\nEven as the empirical result is strong, I'm struggling to connect their masking strategy proposal to the theory presented - I understand that different masking ratios have different bias-risk trade-offs, but how does sampling masking ratio at uniform necessarily lead to a better final model? By intuition, I'd expect this to lead to a model that's an \"expectation\" over all of these masking ratios (including ones with high bias/risk), so I'm not fully understanding why this is effective.\n\nQUESTIONS:\n(See questions in strengths/weaknesses)\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-04"}