Understanding and Enhancing Mask-Based Pretraining towards Universal Representations

Thank you very much for your comments. Our rebuttal is organized as follows. We will first present new results that address general feedback from all reviewers. They include new experiments on vision and language models and a strengthened connection between R$^2$MAE and our theoretical framework. We will next provide a point-by-point response to specific concerns.

New Results 1: Experiments of R$^2$MAE on Vision and Language Models

In general, reviewers appreciated our empirical study but expressed curiosity about whether R$^2$MAE would achieve SOTA performance in more established domains like vision and language modeling. We were fortunate to gain access to additional computational resources during the rebuttal period to perform these experiments.

Vision: Our implementations closely follow established practices from the vision MAE work [1]. We implemented different mask ratio settings on the ViT-base MAE model using their official Pytorch codebase. The settings include: 1. Default MAE (constant 0.75 mask ratio, MR); 2. R$^2$MAE (p ~ U(0.6, 0.9)); 3. Dynamic MR [3] (linearly decreasing MR from 0.9 to 0.6); 4. High (0.9) and low (0.5) fixed MR baselines. Given the time limit, all models were trained for 150 epochs. While this is shorter than the 800 epochs in [1], their analysis shows predictable improvements with longer training schedules. Therefore, while our absolute accuracies may be suboptimal, we expect the relative performance across settings to be comparable and the conclusion to hold if trained longer. All other hyperparameters follow the defaults in [1]. Models were then fine-tuned for classification for 100 epochs, following [1].

Our results below show R$^2$MAE achieves the best performance in both top-1 and top-5 accuracy. We found that R$^2$MAE introduces fluctuations in the pretraining loss, likely due to variable sequence lengths since MAE only encodes unmasked tokens. This presents a potential area for even further improvement.

Language: We trained RoBERTa-base and RoBERTa-medium models on the FineWeb dataset for 10B tokens (max seq length 128, comparable to [2]) and fine-tuned them on GLUE benchmarks (MNLI, QQP, SST-2, QNLI). We evaluated: 1. Default MLM (MR 0.15); 2. R$^2$MAE (p ~ U(0.15, 0.4)); 3. Dynamic MR [3] (0.4 to 0.15); 4. MLM with a fixed 0.4 MR. Fine-tuning accuracies are comparable to those in [2]. The results show R$^2$MAE achieves the best overall rank, with a more pronounced advantage in the larger RoBERTa-Base model.

Note: In these fine-tuning evaluations, the advantage of R$^2$MAE appears smaller than in our linear probing or zero-shot evaluations. This is expected, as the identical fine-tuning pipeline can reduce differences from pre-training. Nevertheless, the consistent advantage of R$^2$MAE across tasks and domains strongly supports its effectiveness and generalizability.

New Results 2: Explicit Connection between R$^2$MAE and the High-Dimensional Linear Regression Framework

Inspired by reviewer comments, we evaluated R$^2$MAE within our linear model framework. Specifically, we simulated the latent space model from Fig. 1E. Intriguingly, we observed that for features of different strengths, R$^2$MAE's risk tends to align with the optimal risk, with its "effective masking ratio" decreasing as feature strength decays.

Our new results support the intuition that R$^2$MAE achieves a “feature-adaptive” masking ratio that approximates dataset-specific optimal masking ratios, allowing it to outperform any fixed ratio in settings with multi-scale features. This is validated by real data (Tables 4 and 5), where R$^2$MAE achieves near-optimal reconstruction across its entire masking range (0.1-0.5), while fixed-rate settings only well cover a smaller range. This result aligns with the robust advantage of R$^2$MAE and strengthens the value of our work.

Point-by-point Response to Specific Concerns

W1. Relevant works: Thank you for the comment. We would like to note that we indeed cited all relevant literature mentioned in the review. Furthermore, we have a dedicated section (Appendix A.1) discussing previous works, and our benchmarks included several of these approaches. While R$^2$MAE is simple, it remains uncovered by previous work. Our theoretical framework is also novel in clarifying unexplained phenomena in mask-based pre-training. Given this, we respectfully disagree with the judgment that our work “does not fairly compare or cite these, and the novelty is overclaimed.” We are happy to address more specific comments if provided.

W2. Gap between the theory and practice: We discussed this gap in the “Relation with real network optimization” part of Section 3.1. Despite the gap, using linear models to explain behaviors in neural networks is a well-established approach in the literature [4, 5]. To verify our findings generalize, we performed extensive evaluations on MLPs, CNNs, and Transformers (Figs. 1-2). Importantly, these evaluations highlighted several phenomena that are explained by our overparametrized linear model but not by alternative frameworks, suggesting our model serves as a valuable “working theory” of masked pre-training.

W3-W5. On the limitations of experimental results:

1. On "marginal" improvement: We respectfully disagree. R$^2$MAE provides consistent and, in several cases, substantial performance improvements. For instance, we improved the PRAUC for the OMIM regulatory variant prediction task by 16.5% over the SOTA model reported in [6]. In brain gene expression data, we improved cell-level age prediction Spearman's r by 10.9% over an optimal MAE. While improvements are smaller in a number of cases, they are consistent across metrics, tasks, and data domains, including the new vision/language results. This improvement is a solid advance for a pre-training scheme that does not alter model architecture.

2. On task relevance and scientific discovery: Our tasks are well-established benchmarks. The DNA variant prediction task is widely used for evaluating DNA models (e.g., Nucleotide Transformer [NT] cited by the reviewer). The OMIM benchmark itself is a large-scale task evaluating over two million DNA sequences. Our experiments span a range of evaluation scenarios from zero-shot (DNA) and linear probing (gene expression) to fine-tuning (ViT, RoBERTa).

3. On the lack of vision/language results: We agree this is an important point and have now provided these results in the first section of our response.

W5-W6. On paper writing: We appreciate the feedback. We are unclear how to address general criticisms like “lengthy and dense writing” but would be happy to modify our text if more specific suggestions are provided. Our introduction aims to motivate a general framework, and we did not find sentences that “promise universal advances for language and vision” as criticized. We believe these comments may stem from a misunderstanding.

W7. On theory assumptions and the MaskTwins paper: Our theoretical work does not aim to prove properties of neural networks directly. Instead, we define a tractable linear model, analyze its properties, and then empirically evaluate if these properties hold in real networks. We empirically verified our findings across multiple data models (Fig. 1E) and network architectures (Fig. 2), in order to support our framework's validity.

We were not aware of MaskTwins as it was published after our submission deadline, but we will cite it in the next version. We have reviewed the paper and find the technical approach to be very different. MaskTwins uses compression sensing theory to derive bounds for signal discovery in a signal+noise model. Our work uses random matrix theory to characterize test risk in an overparameterized regression setting. The goal is also distinct. We would be happy to revise any part of our manuscript that the reviewer feels is an overstatement.

Thanks again for your thorough evaluation. Please let us know if there are any additional concerns.

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." CVPR. 2022. [2] Wettig, Alexander, et al. "Should You Mask 15% in Masked Language Modeling?." EACL. 2023. [3] Ankner, Zachary, et al. "Dynamic Masking Rate Schedules for MLM Pretraining." EACL. 2024. [4] Hastie, Trevor, et al. "Surprises in high-dimensional ridgeless least squares interpolation." Annals of statistics. 2022. [5] Bahri, Yasaman, et al. "Explaining neural scaling laws." PNAS. 2024. [6] Benegas, Gonzalo, et al. "A DNA language model based on multispecies alignment predicts the effects of genome-wide variants." Nature Biotechnology, 2025.

Thank you very much for your comments. Our rebuttal is organized as follows. We will first present new results that address general feedback from all reviewers. They include new experiments on vision and language models and a strengthened connection between R$^2$MAE and our theoretical framework. We will next provide a point-by-point response to specific concerns.

New Results 1: Experiments of R$^2$MAE on Vision and Language Models

In general, reviewers appreciated our empirical study but expressed curiosity about whether R$^2$MAE would achieve SOTA performance in more established domains like vision and language modeling. We were fortunate to gain access to additional computational resources during the rebuttal period to perform these experiments.

Vision: Our implementations closely follow established practices from the vision MAE work [1]. We implemented different mask ratio settings on the ViT-base MAE model using their official Pytorch codebase. The settings include: 1. Default MAE (constant 0.75 mask ratio, MR); 2. R$^2$MAE (p ~ U(0.6, 0.9)); 3. Dynamic MR [3] (linearly decreasing MR from 0.9 to 0.6); 4. High (0.9) and low (0.5) fixed MR baselines. Given the time limit, all models were trained for 150 epochs. While this is shorter than the 800 epochs in [1], their analysis shows predictable improvements with longer training schedules. Therefore, while our absolute accuracies may be suboptimal, we expect the relative performance across settings to be comparable and the conclusion to hold if trained longer. All other hyperparameters follow the defaults in [1]. Models were then fine-tuned for classification for 100 epochs, following [1].

Our results below show R$^2$MAE achieves the best performance in both top-1 and top-5 accuracy. We found that R$^2$MAE introduces fluctuations in the pretraining loss, likely due to variable sequence lengths since MAE only encodes unmasked tokens. This presents a potential area for even further improvement.

Language: We trained RoBERTa-base and RoBERTa-medium models on the FineWeb dataset for 10B tokens (max seq length 128, comparable to [2]) and fine-tuned them on GLUE benchmarks (MNLI, QQP, SST-2, QNLI). We evaluated: 1. Default MLM (MR 0.15); 2. R$^2$MAE (p ~ U(0.15, 0.4)); 3. Dynamic MR [3] (0.4 to 0.15); 4. MLM with a fixed 0.4 MR. Fine-tuning accuracies are comparable to those in [2]. The results show R$^2$MAE achieves the best overall rank, with a more pronounced advantage in the larger RoBERTa-Base model.

Note: In these fine-tuning evaluations, the advantage of R$^2$MAE appears smaller than in our linear probing or zero-shot evaluations. This is expected, as the identical fine-tuning pipeline can reduce differences from pre-training. Nevertheless, the consistent advantage of R$^2$MAE across tasks and domains strongly supports its effectiveness and generalizability.

New Results 2: Explicit Connection between R$^2$MAE and the High-Dimensional Linear Regression Framework

Inspired by reviewer comments, we evaluated R$^2$MAE within our linear model framework. Specifically, we simulated the latent space model from Fig. 1E. Intriguingly, we observed that for features of different strengths, R$^2$MAE's risk tends to align with the optimal risk, with its "effective masking ratio" decreasing as feature strength decays.

Our new results support the intuition that R$^2$MAE achieves a “feature-adaptive” masking ratio that approximates dataset-specific optimal masking ratios, allowing it to outperform any fixed ratio in settings with multi-scale features. This is validated by real data (Tables 4 and 5), where R$^2$MAE achieves near-optimal reconstruction across its entire masking range (0.1-0.5), while fixed-rate settings only well cover a smaller range. This result aligns with the robust advantage of R$^2$MAE and strengthens the value of our work.

Point-by-point response to other specific concerns

In response to the reviewer’s comments, we have performed extensive experiments on vision and language masked pretraining models. Our responses to each question are provided as follows.

Q1. Relation to next-token prediction. Thank you for this very insightful question. Despite the apparent relevance, we believe next-token prediction cannot be adequately addressed by our current theoretical framework. This is because next-token prediction is a token-wise, rather than sample-wise, prediction task. These prediction tasks within an input sequence are highly dependent. This constitutes a clear gap with our current setup, and addressing it will require future research. This is an important point, and we will discuss it in the next version of our paper.

Q2. R²MAE's performance on larger and more diverse datasets. Our results on vision and language modeling show that R²MAE provides consistent improvements (albeit sometimes marginal), even for larger models trained on diverse datasets across various data domains. Our evaluations on two configurations of RoBERTa models show that R²MAE provides consistent improvements across model configurations, whereas the performance of other settings varies. These findings align with our results on biological data and strengthen our conclusions.

Q3. Comparison with structured masking strategies. This is another great point. The cited structured masking strategies are designed to enforce the model to learn specific patterns in the data of interest. Therefore, they may not be directly comparable to ordinary MAE (or R²MAE) pretraining. Nevertheless, the masking procedure in the cited works is an ordinary random mask with a fixed rate, a setting where incorporating R²MAE could potentially improve performance.

Meanwhile, a variety of structured masking strategies have been proposed to directly improve general pretraining models, as discussed in the literature and in Appendix A.1 of our paper. The RoBERTa masking study [2] compared several masking strategies and concluded that these strategies do not outperform simply tuning the masking rate. Our evaluations of learnable masking strategies for biological data, as well as a parallel study on biological data [4], reached the same conclusion.

Thank you again for your thorough evaluation. Please let us know if our responses have fully addressed your concerns or if any additional information is needed.

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." CVPR. 2022. [2] Wettig, Alexander, et al. "Should You Mask 15% in Masked Language Modeling?." EACL. 2023. [3] Ankner, Zachary, et al. "Dynamic Masking Rate Schedules for MLM Pretraining." EACL. 2024. [4] Richter, Till, et al. "Delineating the effective use of self-supervised learning in single-cell genomics." Nature Machine Intelligence. 2025.

Thank you very much for your comments. Our rebuttal is organized as follows. We will first present new results that address general feedback from all reviewers. They include new experiments on vision and language models and a strengthened connection between R$^2$MAE and our theoretical framework. We will next provide a point-by-point response to specific concerns.

New Results 1: Experiments of R$^2$MAE on Vision and Language Models

In general, reviewers appreciated our empirical study but expressed curiosity about whether R$^2$MAE would achieve SOTA performance in more established domains like vision and language modeling. We were fortunate to gain access to additional computational resources during the rebuttal period to perform these experiments.

Vision: Our implementations closely follow established practices from the vision MAE work [1]. We implemented different mask ratio settings on the ViT-base MAE model using their official Pytorch codebase. The settings include: 1. Default MAE (constant 0.75 mask ratio, MR); 2. R$^2$MAE (p ~ U(0.6, 0.9)); 3. Dynamic MR [3] (linearly decreasing MR from 0.9 to 0.6); 4. High (0.9) and low (0.5) fixed MR baselines. Given the time limit, all models were trained for 150 epochs. While this is shorter than the 800 epochs in [1], their analysis shows predictable improvements with longer training schedules. Therefore, while our absolute accuracies may be suboptimal, we expect the relative performance across settings to be comparable and the conclusion to hold if trained longer. All other hyperparameters follow the defaults in [1]. Models were then fine-tuned for classification for 100 epochs, following [1].

Our results below show R$^2$MAE achieves the best performance in both top-1 and top-5 accuracy. We found that R$^2$MAE introduces fluctuations in the pretraining loss, likely due to variable sequence lengths since MAE only encodes unmasked tokens. This presents a potential area for even further improvement.

Language: We trained RoBERTa-base and RoBERTa-medium models on the FineWeb dataset for 10B tokens (max seq length 128, comparable to [2]) and fine-tuned them on GLUE benchmarks (MNLI, QQP, SST-2, QNLI). We evaluated: 1. Default MLM (MR 0.15); 2. R$^2$MAE (p ~ U(0.15, 0.4)); 3. Dynamic MR [3] (0.4 to 0.15); 4. MLM with a fixed 0.4 MR. Fine-tuning accuracies are comparable to those in [2]. The results show R$^2$MAE achieves the best overall rank, with a more pronounced advantage in the larger RoBERTa-Base model.

Note: In these fine-tuning evaluations, the advantage of R$^2$MAE appears smaller than in our linear probing or zero-shot evaluations. This is expected, as the identical fine-tuning pipeline can reduce differences from pre-training. Nevertheless, the consistent advantage of R$^2$MAE across tasks and domains strongly supports its effectiveness and generalizability.

New Results 2: Explicit Connection between R$^2$MAE and the High-Dimensional Linear Regression Framework

Inspired by reviewer comments, we evaluated R$^2$MAE within our linear model framework. Specifically, we simulated the latent space model from Fig. 1E. Intriguingly, we observed that for features of different strengths, R$^2$MAE's risk tends to align with the optimal risk, with its "effective masking ratio" decreasing as feature strength decays.

Our new results support the intuition that R$^2$MAE achieves a “feature-adaptive” masking ratio that approximates dataset-specific optimal masking ratios, allowing it to outperform any fixed ratio in settings with multi-scale features. This is validated by real data (Tables 4 and 5), where R$^2$MAE achieves near-optimal reconstruction across its entire masking range (0.1-0.5), while fixed-rate settings only well cover a smaller range. This result aligns with the robust advantage of R$^2$MAE and strengthens the value of our work.

Point-by-point response to other specific concerns

In response to the reviewer’s comments, we have performed extensive experiments on vision and language mask pretraining models. Our response to each question is provided as follows.

W1. Theoretical-Algorithmic Gap: This is a good point. We acknowledge that our theoretical results are based on linear models that deviate from real neural networks, as mentioned in the subsection “Relation with Real Network Optimization” in Section 3.1. Nevertheless, a rich body of literature has applied ridgeless regression models and random matrix theory to successfully recapitulate neural network behaviors like double descent and the scaling law [4,5]. Despite the deviations, we believe that our extensive evaluations of neural network and linear model behaviors convincingly support the validity of our framework as a “working theory.”

W2. β Estimation: Indeed, the β estimation procedure is a simplification that enables feature-wise analysis. Nevertheless, we argue that the “feature” here does not necessarily mean individual pixels or tokens. Intuitively, the MSE loss of a sample can be approximated by the reconstruction loss in a latent feature space, linked by a decoding transformation. This insight suggests that we can investigate mask reconstruction behavior by evaluating different eigenvectors of the design matrix covariance as β. Fully addressing the underlying non-linearity remains an open problem and a promising direction for future research.

W3, Q1. Comparison on General Datasets: Thank you for raising this excellent point. We agree that incorporating results on image and language modeling would greatly improve our work. Please refer to the earlier section for these results. Overall, they demonstrate the superior performance of R²MAE, which aligns well with the results on biological data.

Q2. The possibility of determining an optimal masking ratio per sample: This is a good question. First, as we demonstrated in Section 3.5 of the paper, the optimal masking ratio varies for different downstream tasks. Therefore, in the pretraining phase, finding a single optimal masking ratio may not be well-defined, as it defeats the purpose of training a universal model. Second, a real-world downstream task (e.g., image classification, GLUE tasks) can be highly complex and require learning features at different scales. In contrast, R²MAE learns a spectrum of features during pretraining through a feature-adaptive masking mechanism, a point we clarified earlier with our new results. We will incorporate relevant discussions on this topic in future versions of our work to make our presentation more self-contained. Please let us know if there are any additional concerns.

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." CVPR. 2022. [2] Wettig, Alexander, et al. "Should You Mask 15% in Masked Language Modeling?." EACL. 2023. [3] Ankner, Zachary, et al. "Dynamic Masking Rate Schedules for MLM Pretraining." EACL. 2024. [4] Hastie, Trevor, et al. "Surprises in high-dimensional ridgeless least squares interpolation." Annals of statistics. 2022. [5] Bahri, Yasaman, et al. "Explaining neural scaling laws." Proceedings of the National Academy of Sciences. 2024.

Thank you very much for your comments. Our rebuttal is organized as follows. We will first present new results that address general feedback from all reviewers. They include new experiments on vision and language models and a strengthened connection between R$^2$MAE and our theoretical framework. We will next provide a point-by-point response to specific concerns.

New Results 1: Experiments of R$^2$MAE on Vision and Language Models

In general, reviewers appreciated our empirical study but expressed curiosity about whether R$^2$MAE would achieve SOTA performance in more established domains like vision and language modeling. We were fortunate to gain access to additional computational resources during the rebuttal period to perform these experiments.

Vision: Our implementations closely follow established practices from the vision MAE work [1]. We implemented different mask ratio settings on the ViT-base MAE model using their official Pytorch codebase. The settings include: 1. Default MAE (constant 0.75 mask ratio, MR); 2. R$^2$MAE (p ~ U(0.6, 0.9)); 3. Dynamic MR [3] (linearly decreasing MR from 0.9 to 0.6); 4. High (0.9) and low (0.5) fixed MR baselines. Given the time limit, all models were trained for 150 epochs. While this is shorter than the 800 epochs in [1], their analysis shows predictable improvements with longer training schedules. Therefore, while our absolute accuracies may be suboptimal, we expect the relative performance across settings to be comparable and the conclusion to hold if trained longer. All other hyperparameters follow the defaults in [1]. Models were then fine-tuned for classification for 100 epochs, following [1].

Our results below show R$^2$MAE achieves the best performance in both top-1 and top-5 accuracy. We found that R$^2$MAE introduces fluctuations in the pretraining loss, likely due to variable sequence lengths since MAE only encodes unmasked tokens. This presents a potential area for even further improvement.

Language: We trained RoBERTa-base and RoBERTa-medium models on the FineWeb dataset for 10B tokens (max seq length 128, comparable to [2]) and fine-tuned them on GLUE benchmarks (MNLI, QQP, SST-2, QNLI). We evaluated: 1. Default MLM (MR 0.15); 2. R$^2$MAE (p ~ U(0.15, 0.4)); 3. Dynamic MR [3] (0.4 to 0.15); 4. MLM with a fixed 0.4 MR. Fine-tuning accuracies are comparable to those in [2]. The results show R$^2$MAE achieves the best overall rank, with a more pronounced advantage in the larger RoBERTa-Base model.

Note: In these fine-tuning evaluations, the advantage of R$^2$MAE appears smaller than in our linear probing or zero-shot evaluations. This is expected, as the identical fine-tuning pipeline can reduce differences from pre-training. Nevertheless, the consistent advantage of R$^2$MAE across tasks and domains strongly supports its effectiveness and generalizability.

New Results 2: Explicit Connection between R$^2$MAE and the High-Dimensional Linear Regression Framework

Inspired by reviewer comments, we evaluated R$^2$MAE within our linear model framework. Specifically, we simulated the latent space model from Fig. 1E. Intriguingly, we observed that for features of different strengths, R$^2$MAE's risk tends to align with the optimal risk, with its "effective masking ratio" decreasing as feature strength decays.

Our new results support the intuition that R$^2$MAE achieves a “feature-adaptive” masking ratio that approximates dataset-specific optimal masking ratios, allowing it to outperform any fixed ratio in settings with multi-scale features. This is validated by real data (Tables 4 and 5), where R$^2$MAE achieves near-optimal reconstruction across its entire masking range (0.1-0.5), while fixed-rate settings only well cover a smaller range. This result aligns with the robust advantage of R$^2$MAE and strengthens the value of our work.

Point-by-point response to other specific concerns

The only concern from the reviewer is the connection between R²MAE and the linear model. We are grateful for the feedback which inspired the second part of our new results. Here, we highlight the difference between the R²MAE setting and the fixed mask ratio setting in the linear model framework. In the linear model, R²MAE translates to first sampling a row-wise masking ratio $p_i \sim U(p_{\text{min}}, p_{\text{max}})$ and then using it to sample each row of the masking matrix $Z$. This introduces a subtle difference in the effective sample matrix $X_{sub}$, as the sampling probability for each row is no longer a constant ($p$) but depends on the specific value of $p_i$. Due to this difference, deriving the closed-form solution for the test risk becomes very challenging and remains a topic for future research. At the very least, this complexity suggests that the test risk of R²MAE is not identical to that of any fixed masking ratio. We further empirically validated this, as shown in the second part of our new results. We believe these results effectively address the concern. Please also refer to the first part of new results which contains additional validation in vision and language models that improve our work. Please let us know if there are any additional concerns.

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." CVPR. 2022. [2] Wettig, Alexander, et al. "Should You Mask 15% in Masked Language Modeling?." EACL. 2023. [3] Ankner, Zachary, et al. "Dynamic Masking Rate Schedules for MLM Pretraining." EACL. 2024.