Paper: Representation Deficiency in Masked Language Modeling (ICLR-2024)

Reviewer comment (Theory):
The interplay between pretraining and fine-tuning is very complicated and not fully understood yet, so the argument in the introduction that dimensions exclusively used for [MASK] tokens have not been pretrained to represent real tokens — thus requiring training from scratch on downstream data (raising the risk of overfitting) or becoming unused — is not substantiated. The cited works by Hendrycks et al. (2019) and Kumar et al. (2022) are not clearly shown to support this argument. Moreover, Figure (a) in the paper shows that inputs without a mask token have higher-rank representations, which could indicate that having mask tokens during pretraining does not impact the rank of real tokens when mask tokens are not used.

Author response:
**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.

Reviewer comment (Theory):
Theorem 2.2 appears to be based on some unwritten assumptions, such as the model needing to be an attention-only model without MLP and residual layers.

Author response:
**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.

Reviewer comment (Theory):
There is a mismatch between the setup for Lemma 2.1 and Theorem 2.2. Theorem 2.2 is based on some unwritten and unrealistic assumptions, while Lemma 2.1 is based on empirical results showing that (full) transformers can fit real-world high-rank distributions. According to Dong et al. (2021), under the assumptions of Theorem 2.2, the rank of the representation converges to 1, implying that Lemma 2.1 would not hold under these conditions. Therefore, it is inappropriate to use Theorem 2.2 together with Lemma 2.1 to derive the paper’s conclusion.

Author response:
**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.

Reviewer comment (Theory):
The theoretical results have weak implications because they only suggest that the representation of real tokens cannot be full-rank but do not characterize how far the representation is from being full-rank. It is possible that the rank could be $d - 1$, in which case the “representation deficiency” problem may not be a significant issue.

Author response:
**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.

Reviewer comment (Theory):
In Lemma 2.1, the rank of the representation matrices seems to refer to the representation matrix of the whole corpus. However, in practice, models are fed much shorter sequences (e.g., 512 tokens), so the rank of the masked token distribution is at most 512 × 15% in expectation, meaning mask tokens cannot occupy too many dimensions. This may not be a major issue, but the arguments starting from Theorem 2.2 should be rewritten to reflect that, because Theorem 2.2 concerns sequences fed into a transformer model, the representation matrix should be of an example rather than the whole corpus, and therefore Lemma 2.1 cannot be used directly.

Author response:
**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.

Reviewer comment (Novelty):
This work has similarities with Dong et al. (2021), as Dong et al. also plotted the representation rank of transformer models, and Theorem 2.2 in this paper is largely based on a proof from Dong et al. (2021). The authors should give more explicit credit to Dong et al.

Author response:
**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.

Reviewer comment (Experiments):
Although the paper mentions in Footnote 2 that some MLM training settings (like Google's original BERT) keep 10% of [MASK] tokens as original, and that subsequent studies like Wettig et al. (2022) found this trick unnecessary, an exploration of the effective rank of models using this 80:10:10 technique, similar to Figure 1, would be beneficial.

Author response:
**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.

Reviewer comment (Novelty):
Missing reference: [this paper](https://proceedings.mlr.press/v97/gong19a.html) looked into attention patterns in pretrained BERT-like LMs, uncovering patterns related to low-rankness. This study could provide additional insights into the topic at hand.

Author response:
**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.

Reviewer comment (Novelty):
In the analysis module, the use of effective rank to represent MLM's representative capacity overlaps with section 3.1 of ISOTROPY (https://openreview.net/pdf?id=xYGNO86OWDH), where effective rank is used to denote the isotropy of representations. The differences between these two approaches are minor.

Author response:
**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.

Reviewer comment (Novelty):
In the experimental module, the model structure used in this paper is identical to the one used in Mask Later (https://arxiv.org/pdf/2211.04898.pdf).

Author response:
**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.

Reviewer comment (Novelty):
The authors should explicitly define how this work differs from existing works, otherwise, it appears as if the paper merely uses the theoretical tool of ISOTROPY to analyze why Mask Later is effective. Without a distinct contribution, the overall impact of this work seems limited.

Author response:
**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.