Representation Deficiency in Masked Language Modeling

Thank you very much for your thoughtful feedback! We address your raised points as follows.

The interplay between pretraining and fine-tuning is very complicated and not fully understood yet: While we agree that theoretical understanding of the relationship between pretraining and fine-tuning is still lacking, the general consensus in the literature supports the pivotal role of pretraining in enhancing both the effectiveness (Kumar et al.) and robustness (Hendrycks et al.) upon fine-tuning. Therefore, analogous to why a non-pretrained model is inferior to a pretrained counterpart, the model dimensions that are not pretrained to represent real tokens (due to the deficiency issue caused by [MASK] tokens) will be less effective on downstream data. Furthermore, we have shown in Figure 5 (b) that the gap in effective rank persists in fine-tuning, confirming the lasting impact of pretraining.

Figure (a) in this paper shows that inputs without a mask token have higher-rank representations: Sorry for the confusion, but our Figure 1 (a) shows the opposite – inputs without [MASK] suffer from lower-rank representations.

The unwritten assumptions of Theorem 2.2: In our original submission, we explicitly stated the assumption and rationale for only analyzing the rank change induced by self-attention at the beginning of our proof in Appendix A. In the updated paper version, we have moved them to the main paper for better clarity. In summary, only the rank changes induced by self-attention reflect the extent of contextualization in token representations, and the effectiveness of text encoders is typically attributed to the contextualized representations. While MLPs/residual connections can also change the rank of representations, they do not provide each token with new information from other tokens. While we do not account for MLPs and residual connections, our analysis validates that the rank deficiency is caused by the self-attention mechanism, and in practice, MLPs and residual connections do not prevent the issue from happening.

The mismatch between the setup for Lemma 2.1 and the setup for Theorem 2.2: We’d like to clarify that Lemma 2.1 does not impose any architectural assumptions and it holds for any sequence modeling architecture. The empirical studies in Dong et al. (Figure 2 in their paper) are under a different setup than our Lemma 2.1: According to Section 4.1 of Dong et al., the pure attention architecture is not the one directly trained for the MLM task, whereas in our case the model is directly trained with MLM.

The weak implication of the theoretical results; it is possible that the representation of the real tokens has rank $d-1$: Our primary goal is to utilize theoretical analysis to substantiate the presence of the representation deficiency issue. We have shown empirically that the representation deficiency issue results in a nontrivial gap in effective rank as large as ~50 dimensions, especially in deeper layers. Furthermore, we demonstrate that addressing this issue via our proposed MAE-LM method yields consistent performance improvements across various downstream tasks. Therefore, the identified issue is a crucial problem in MLM pretraining that warrants attention and remediation.

The representation matrix of each example v.s. the whole corpus: Thanks for the question. We’d like to clarify that while the analysis in Lemma 2.1 is conducted at the corpus level, the conclusion of Lemma 2.1 can extend seamlessly to individual sequences. Specifically, if the effective rank of [MASK] token increases throughout Transformer layers when considering the entire corpus, it logically follows that this phenomenon holds true by expectation for the majority, if not all, individual samples in the corpus.

The contribution of this work and Dong et al.: While Dong et al. also plot the representation rank, their focus is not on investigating the issue related to [MASK] tokens in MLM pretraining. Regarding our theoretical analysis, we indeed draw inspiration from Dong et al., and we explicitly mentioned this in the full proof in our original submission (Appendix A). In the updated paper, we have added a paragraph at the beginning of our full proof in Appendix A to give Dong et al. more credits.

Thanks again for your thoughtful comments. We have incorporated the revision accordingly in the updated paper. Please let us know if there are any further questions.

Thank you very much for your thoughtful feedback, and we appreciate your positive acknowledgment of our work! We address your raised points as follows.

Using 80:10:10 masking technique: In our preliminary experiments, we found that the 80:10:10 masking technique did not provide notable performance gain compared to directly masking out 15% of the tokens, and thus we used the latter in our evaluations. That said, we agree that validating our findings under the 80:10:10 masking technique is beneficial. We have added the effective rank plot in Appendix E, Figure 6 of our updated paper. The results are largely similar to Figure 1, which confirms that randomly replacing a small percentage of [MASK] tokens with real ones does not effectively address the representation deficiency issue, as the ratio of [MASK] tokens in pretraining is still high.

Missing reference: Thank you for pointing out the reference, and we have added the discussions of this work to both Section 5 and Appendix F. We agree that our findings in this paper are related to the paper you mentioned: The rank deficiency problem might have caused de-contextualization in token representations, which is reflected by the localized self-attention patterns as discovered in that study.

Thanks again for your review! Please let us know if you have any further questions.

Thank you very much for your thoughtful feedback, and we appreciate your positive acknowledgment of our work! We discuss the implications of our work on autoregressive LMs as follows.

While autoregressive LM pretraining generally does not introduce artificial symbols such as [MASK], our analyses can be extended to show that the representation deficiency issue can also arise in autoregressive pretraining when certain real tokens exist exclusively in the pretraining data but are either absent or occur infrequently in downstream data. Similar to the impact of [MASK] tokens, these tokens occupy dimensions during pretraining that may not be effectively utilized in downstream tasks. Consequently, it is desirable to maximize the vocabulary overlap between pretraining data and downstream data, which can be realized via pretraining data selection, training corpora pre-processing, and vocabulary pruning. We leave these explorations as future work, and we have added the discussions to Appendix F in the updated paper.

Thanks again for your review! Please let us know if you have any further questions.

Thank you very much for your thoughtful feedback! We address your raised points as follows.

The use of effective rank to represent MLM's representative capacity overlaps with section 3.1 of ISOTROPY: We’d like to clarify that we follow the definition of effective rank in ISOTROPY (Cai et al.) to only compute the rank and empirically showcase the identified representation deficiency issue, and we have not used it in our theoretical analyses. Our major contributions are (1) investigating the discrepancy between pretraining and fine-tuning of MLM models caused by [MASK] tokens from the perspective of representation rank, and (2) demonstrating that a simple architectural modification in MLM—excluding [MASK] tokens from the encoder’s input—can effectively address the representation deficiency and result in improved downstream performance. These contributions are orthogonal to Cai et al. We have updated our paper to provide clarifications in Section 2.2.

The model structure used in this paper is identical to the one used in Mask Later: Compared to the 3ML approach proposed by Liao et al., our proposed MAE-LM shares the same masked autoencoder architecture. However, the fundamental motivations, focuses, results and insights of the two works are different:

• Focuses: Liao et al. focus on an efficiency aspect, by excluding [MASK] tokens to reduce the sequence length. In contrast, our primary goal is to use principled empirical and theoretical analyses to illuminate an overlooked issue in MLM—specifically, the impact of [MASK] tokens on the model's representation power. The MAE-LM framework is introduced as a simple solution to rectify the rank deficiency issue arising from [MASK] tokens, aiming for the understanding of its influence on the learned representations and the downstream task performance.
• Results: In contrast to Liao et al.'s efficiency-driven emphasis, our experimental results substantiate that MAE-LM effectively mitigates the representation gap induced by [MASK] tokens, consistently outperforming standard MLM across various tasks and pretraining settings. To our knowledge, such a representation learning benefit of the masked autoencoder architecture has not been explored before.
• Insights: Importantly, our contributions may extend beyond the specific model architecture and offer insights for future developments in MLM pretraining. For example, one may regularize the [MASK] token representations to reside in the same space of real token representations as another potential solution to the representation deficiency problem. Our analyses also underscore the importance of considering real tokens that exist solely in pretraining data, yet are absent or occur rarely in downstream data. Such tokens, akin to [MASK] tokens, have the potential to compromise model representation power.

We have incorporated a brief discussion in Section 3 of our revised paper to emphasize the distinctiveness of our contribution.

Thanks again for your review! Please let us know if you have any further questions.

Thank you for your prompt response.

1. Analysis moudle. After carefully reviewing your reply and Appendix A, I have reconsidered the theoretical analysis section of the paper. I recognize that both lemmas in the paper are non-trivial, and the proof methods used differ from ISOTROPY. So my concerns about the analysis module have been well addressed.

2. Experiment moudle. Nevertheless, my concerns persist regarding the experiment module. Despite different motivations, Liao et al.'s results have already demonstrated that MAE-LM can achieve performance comparable to, or even superior to, traditional MLM (as shown in Table 3 of Liao et al.'s paper). In other words, we already hold a prior expectation that, regardless of the correctness of the theoretical analysis in this paper, the MAE-LM structure will work.

However, I fully agree that the theoretical analysis of this paper "may extend beyond the specific model architecture and offer insights for future developments in MLM pretraining." Consequently, I have raised the score to 6.