Blessing of Class Diversity in Pre-training

Abstract: This paper presents a new statistical analysis aiming to explain the recent superior achievements of the pre-training techniques in natural language processing (NLP).
We prove that when the classes of the pre-training task (e.g., different words in masked language model task) are sufficiently diverse, in the sense that the least singular value of the last linear layer in pre-training is large, then pre-training can significantly improve the sample efficiency of downstream tasks. Inspired by our theory, we propose a new regularization technique that targets the multi-class pre-training: a \emph{diversity regularizer only to the last linear layer} in the pre-training phase.
Our empirical results show that this technique consistently boosts the performance of the pre-trained BERT model on different downstream tasks.

Reviews and author responses:

 Thank you for your reply!

Regarding BART, to clarify, the decoder can be viewed as solving a structured multi-class prediction problem (tokens are predicted auto-regressively) whereas masked language modeling can be viewed as an unstructured multiclass prediction problem.

 We thank you for your suggestions! And we are glad the reviewer appreciates our theoretical results. However, we largely disagree with your comments.

1. First, we disagree with your general philosophy on evaluating theoretical papers, especially when you said

> this paper is a theory-only one, which cannot meet the standard of ICLR for sure.

- There are hundreds of papers (maybe more) in ICML / NeurIPS / ICLR that are purely theoretical (no experiment) and do not have a direct consequence in practical improvements. **The  "empirical usefulness" of these theoretical papers is in the long run and may not be obvious at the time the papers appear.**

- As a concrete example, deep learning is mostly empirical, say 5 years ago. However, in recent years, there are hundreds of papers on deep learning theory (many of which are without any experiment) that help us gain a deeper understanding of deep learning. 


2. We disagree with your comments about the 
community":

> However, in an engineering area like NLP or the representation learning of the ICLR, the community still eventually expects to see better performance.

- We agree that the community wants to see better performance for sure. However, we believe there are also a number of researchers who want to understand why these pre-training models work so well, *at a quantitative level like what our theorems state, which is more precise than the intuitive level "large enough data and diverse enough representation may help"*.

- We believe **characterizing NLP and representation learning as engineering areas is not appropriate and a bit too narrow.** There are many researchers who want to study the *scientific aspect* of NLP and representation learning. Furthermore, the community is evolving: a paper should not be rejected because the paper does present what the current community wants to see. 

- Lastly, our paper is not restricted to one community. Our paper can also benefit, e.g., the statistical learning theory community.



 I read this paper again after I also check the others' review and all  your response.
I should appreciate what you did in the theoretical part, I don't doubt the correctness of the proposed theory with my rough check, as I believe it righ follows the intuition like large enough data and diverse enough representation may help. 
I am considering what we expect from this work and what the community may benefit from the theoretical or empirical results of this paper. Unfortunately, my conclusion is negative.

It is indeed nice to well formulize the empircal or intuitive observation in a work like pre-trained language models and their usage. This is the contribution of this paper as the authors claims. However, in an engineering area like NLP or the representaiton learning of the ICLR, the community still eventually expects to see better performance. Meanwhile, the empirical results are a good indicator to show how strong the theoretical results are. If, the empirical enhancement derived the corresponding theory is not so significant (as I argued in previous comments), then, I have to speculate the theory itself is trivial derivation or formulization. On the basis of such consideration, I will keep my current evaluation of this paper, though I would like to give suggestions to make this paper better in the future.

(1) In the theoretical part, the authors may consider to better address the non-trivialness of  their error bound estimation. As I say the trivialness, it means any further improvment over the trivial esitmation. For example, the trivial estimation on errors may be 100\%, then any percent below 100\% may be non-trivial.

(2) The authors may consider improving or altering their enhancement method in pre-trained language model. The current experimental results is problematic at least if not inconvincible (unless you now report significant improvement on large setting of pre-trained LM with your current regularizer). With this invalid part, this paper is a theory-only one, which cannot meet the standard of ICLR for sure.

(3) The derivation of the theoretical results, as the authors claims, it is more than suitable for pre-trained language model. I agree. This is a general result on any transfer learning but not limited to the pre-training/fine-tuning mode. I am afraid pre-training language models are not a good enough testbed for the proposed theoretical results as they are too large to imply onefold evaluation scenario.




 Thanks for your reply to my questions. I'm still skeptical on how Assumption 1 holds, even approximately, although I'm a bit surprised that prior papers have also adopted this assumption.

Regarding BART, I'm not clear why it will not fall in the same framework though. While being an encoder-decoder model, the decoder head can still be argued to performing multi-class classification at every step of auto-regressive language prediction. Its good to thus separate out the claim to be specific to masked language models.

As another reviewer pointed out, the paper is indeed a bit weak empirically - however the strong theoretical framework kind of balances it. Thus, I'm not increasing my score, however I would like to see this paper to be presented in some form in the conference.

 We believe the reviewer misunderstood the theme of this paper. 
We have stated in our paper and in the rebuttal that our **main contribution is the theory, not the regularizer**.  We propose the new regularization technique only to showcase that our theory can inspire us to design stronger regularizers. 

We have stated that our experiments are preliminary. If the reviewer thinks the sentence “boosts the performance of the pre-trained BERT model on different downstream tasks” is not appropriate, we are happy to change it to a softer tone, such as "shows the potential to boost the performance of pre-trained BERT model on downstream classification tasks".

By the way, can you provide the reference for [1]? We believe you missed it in your reply. Thank you.

 Thanks for the feedback and clarification on the motivation of this work.

My concerns are kept on the insufficient and maybe inconvincible experiments.

1.	The experiments in the current format are not enough. I do not get the idea of “boosts the performance of the pre-trained BERT model on different downstream tasks”. The paper actually conducts experiments on GLUE, 8 sentence classification tasks. More experiments on diverse tasks may shed the full light. Classification task cannot be seen as a typical one among NLP. Besides, only BERT-base model is experimented while I really expect to see results on bigger model. My experience show that big and base models may show quite different characteristics in pre-training of language models. Though I understand the computational resource limitation, however, no results on big model keep all the conclusions of this work much less inconvincible.

For detailed results, I just notice that the dev results on QNLI and RTE dev much lower than those in other literature [1]. The lower reported baselines make the claimed performance boost further fragile. 

2.	For the most contribution of this work, the regularizer, which slightly boosts the performances on BERT-base model. However, I regard there are no direct ablation experiments which can be seen to demonstrate its effect right from the regularizer. Is it better than tuning different dropout probabilities or weight decay rates?

So far, I cannot vote for accepting this paper in current stage due to the poor empirical experiments. The authors are encouraged to conduct further experiments to facilitate readers to understand the proposed regularizer properly.


 We thank you for your valuable feedback.

First of all, we want to emphasize that our paper is a theoretical paper that aims to explain the benefits of pre-training. Our paper is NOT a methodology / empirical paper that aims to improve existing methods and achieve state-of-the-art performances. Our experiments are preliminary and only showcase that our theory can inspire us to design stronger regularizers. 


Below we address your comments on our work.
1. “multi-task pre-training and multi-class pre-training” (to point 1 and point 4 in the review): We believe the reviewer has a misunderstanding on our paper. In the fourth paragraph of Section 1, we argue that the theory of multi-task pre-training CANNOT explain NLP pre-training. Hence, in the sixth paragraph, we claim that we should study class diversity instead of task diversity and discuss MLM pre-training as an example. 
2. “The theoretical part is relatively small (to point 2 in the review)”: we note that (Du et al., 2021, Remark 6.2) only focused on quadratic loss with one-hot labels for multi-class classification. Instead, we study the commonly used cross-entropy loss. This is not just a “detailed derivation and combination with pre-training.” We also need to introduce new analysis techniques (cf. 2nd paragraph in Section 1.1).
3. “Regularizer (to point 2 in the review)”: while this form of regularization is proposed in Zou & Adams, 2012, our theory suggests that for pre-training, we should use it only to the last linear layer. This usage is new in the literature.
4. “Regression task (to point 5 in the review)”: since our theory is for classification, we did not report the results of regression in the prior version. We have added the results for STS-B in the revised version (the last column in Table 1).
5. “More experiments (to point 3 and 6 in the review)”: Our empirical results only showcase that our theory can inspire us to design stronger regularizers. Due to limited computation resources, we cannot afford to do experiments on large models such as RoBERTa.


 We thank you for your valuable feedback and for appreciating our contributions. Below we address your comments on our work.

1. "difference between reproduced BERT and BERT in literature": To reproduce BERT, we make use of a pipeline presented by Nvidia (https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT) which saves training time and resources by large-batch training. Our diversity-regularized model is also pre-trained through this pipeline.
2. Typos: We have fixed the writing issues accordingly. Thanks for catching these typos!


 We thank you for your valuable feedback and for appreciating our contributions. Below we address your comments on our work.
1. “Assumption that the representation function is fixed”: Thanks for raising this. First, we note that this realizability assumption is adopted by all recent results on representation learning theory (Maurer et al., 2016; Du et al., 2021; Tripuraneni et al. 2020, Lee et al., 2020) and is a standard assumption in the classical PAC learning framework. Second, we argue that Assumption 1 holds at least approximately. We agree that fine-tuning will change the representation function $h$. However, we note that the amount of change is restricted because otherwise, fine-tuning would be essentially learning the representation from scratch.
2. “Auto-regressive Pre-training”: Thanks for raising this! BART is an encoder-decoder model architecture where an encoder-decoder is adopted to recover the clean sequences. Therefore, it is not a multi-class classification problem and thus cannot fit in our framework. We have revised our draft to indicate our diversity regularizer is dedicated to multi-class pre-training (marked in red in Section 1 and Section 5) and raised a future direction (Section 6).


 We thank you for your valuable feedback and for appreciating our contributions. Below we address your comments on our work.

1. “Assumption 1 holds”: We note that this is the realizability assumption adopted by all recent results on representation learning theory (Maurer et al., 2016; Du et al., 2021; Tripuraneni et al. 2020) and is a standard assumption in the classical PAC learning framework. 
2. “results aren’t very strong”: Thanks for mentioning this. Our experiments are preliminary and showcase that our theory can inspire us to design stronger regularizers.  Therefore, we have softened our tone in the revised version (marked in red in Section 5).


 All major revisions are marked in red. 
1. We added results on a regression task (STS-B) asked by Reviewer AkLr (Section 5 Table 1).2
2. We made clear that our regularization technique is for multi-class pre-training in Section 1,5 and raised a new problem for studying other forms of pre-training in Section 6. 
3. We softened our tone about the regularization technique (last paragraph in Section 5).


The authors present a new statistical analysis aiming to explain the success of (masked) language model pretraining for NLP. Specifically, they focus on the *diversity of classes*, which they claim is similar to the *diversity of tasks* for multi-task pretraining (e.g., Tripuraneni et al. (2020)). Based on their theory, the authors then propose a *diversity regularizer* to improve model performance.
 This paper presents a new statistical analysis aiming to explain the success of pretraining for NLP. The theory the authors present focuses on the *diversity of classes*, parallel to the *diversity of tasks* for multi-task pretraining (e.g., Tripuraneni et al. (2020)). The authors then propose a *diversity regularizer*, which they evaluate in empirical experiments, with the goal to improve model performance.

Strengths: 
- The authors propose a theory aiming at explaining the success of pretraining, which is something not frequently explored systematically in NLP. Theoretical insights into the reasons behind the benefits of pretraining are (at least in my opinion) strongly needed.
- The authors aim at drawing practical conclusions from their theory by introducing the *diversity regularizer*.

Weaknesses:
- It is not directly clear to me that Assumption 1 holds.
- The empirical results aren't very strong. In fact, I would argue that the author's claim "Our empirical results show great potential of this diversity regularizer." is probably an overstatement.

 The authors investigate an important topic: the success of pretraining in NLP. Some of the statements in the paper are too strong or unclear to me; I would appreciate more explanation from the authors.

The paper presents a theoretical view on how diversity of classes can help in language model pre-training, and then backs up with experiments consisting of improvements in the downstream tasks due to introduction of a class diversity loss during pre-training. The paper provides extensive theoretical justification of why diversity in classes - which is diversity in the linear output layer in case of multi-class pre-training - can help in better performance in downstream tasks. ## The good

The paper provides a comprehensive theoretical justification of introducing diversity in the last linear layer of pre-training. The paper follows prior work in establishing the theoretical framework, and then proceeds to define their contribution in establishing a guarantee of an upper bound of transfer learning risk. Then, the paper explains how a diversity parameter controls the lower bound of the downstream fine-tuning. Concretely, the authors propose to tune the class diversity by increasing how spread out the vectors of the last linear layer are. The authors empirically show their method outperforms the baseline in the GLUE benchmark. The setup of the diversity parameter makes sense intuitively. The authors do observe some marginal improvement over the GLUE benchmark, which shows that in general a more diverse classifier in pre-training helps in providing a better initialization for fine-tuning.


## The bad

While the paper is heavy in the theoretical aspect, it raises few questions on the assumptions made in the setting. Firstly, the assumption that the representation function is fixed after pre-training (Page 4, first paragraph) is not true for the general fine-tuning case where models are typically fine-tuned such that weights of the representation function are also updated. Thus, Assumption 1 (Realizability) does not hold for general settings, that is representation function $h$ is typically not optimal for both pre-training and downstream tasks (after optimal pre-training, downstream tasks typically overwrite and change the representations). I would love to hear an explanation of why this assumption is required in this setting.

The empirical results, however, are only restricted to BERT which is a Masked Language Model. Since the authors made a general claim of this method improving any pre-training task, it would make sense to investigate an auto-regressive pre-training model (such as BART) in the same setting as well. The diversity factor in last linear layer classification during pre-training might entirely depend on the type of pre-training objectives used. The authors could, on the other hand, note clearly that this theoretical result and experiments are specific to masked language modeling itself, in order to avoid a potentially costly set of experiments.

 The paper introduces a comprehensive theoretical framework to establish the need of diversity in classes during pre-training. They support their theoretical results with that on BERT fine-tuning on GLUE benchmark. Overall I find the theoretical contributions solid, but I don't know how general it is as we only see experiments in masked language modeling setup. Empirically the results do support the theory somewhat, but I suspect its mostly due to authors investigating an overly saturated benchmark such as GLUE. Additional results on either pre-training or on different fine-tuning tasks should be able to make it a stronger contribution.


This paper analyzes the effectiveness of pre-trained models on
downstream classification task, showing that the "class diversity"
(characterized by the least singular value of the last linear) in
pre-training is important for the target classification tasks after
fine-tuning. The paper also suggest a diversity regularizer for
improving the downstream task efficiency. The paper is mainly
concerned with the formal analysis of the effects of class diversity
(in practically relevant settings). However, a demonstration of the
effects of proposed diversity regularizer is also demonstrated
empirically on a well-known benchmark database for 7 multi-class NLP
classification tasks.
 Starting with a disclaimer: I could not carefully verify the proofs
with the time I can afford for this review.

The paper investigates a question that is relevant to the wider field 
formally, and the outcome leads to suggestions for practical 
improvements. Although more empirical evidence would be useful, the
experiments provided also support the usefulness of the proposed
improvements (hence the effect of diversity).

As evidenced by the sources cited in the paper, the general idea of
usefulness of diversity is not novel. I am also not fully convinced 
(by the paper) that 'class diversity' is categorically different than
'task diversity'. However, in my opinion, the paper has enough 
new/interesting content that would be beneficial for the community.

Except above notes, I only have some minor remarks on the format of
the paper:

- It would be a good idea to indicate the difference of the reproduced
  BERT-base results are from the state-of-the-art results from the
  earlier literature.
- Many of the citations use wrong form. For example, "Devlin et al.
  (2019)" on page 1 should be in parentheses, while  when citation is 
  used as a name as in "...  proposed in (Zou & Adams, 2012)" (p.3),
  it should not be in parentheses. This is a general issue in most of
  the citations, a through check is recommended.
- Footnote marks should be placed after punctuation, and there should
  not be space before them. Also, I strongly recommend delaying the
  placement of footnote marks to a phrase/sentence boundary,
  especially when they are in the neighborhood of mathematical
  expressions (e.g., like the footnote mark 2, which would be better
  at the end of the sentences).
- First paragraph of 4.2: "F^pre and F^down are *sets of* linear 
  functions"?
- First paragraph of 5: "...  capability,in .." space is missing after
  comma 
- References has some case (normalization) issues, e.g., "rademacher", "bert"
  (not a complete list, a careful check is recommended)
 The paper investigates a question that is relevant to the wider field 
formally, and the outcome leads to suggestions for practical 
improvements. Although more empirical evidence would be useful, the
experiments provided also support the usefulness of the proposed
improvements (hence the effect of diversity).

As evidenced by the sources cited in the paper, the general idea of
usefulness of diversity is not novel. I am also not fully convinced 
(by the paper) that 'class diversity' is categorically different than
'task diversity'. However, in my opinion, the paper has enough 
new/interesting content that would be beneficial for the community.


This paper studies the question of why the pre-trained language model is effective on downstream NLP tasks. Through mathematical analysis, the authors conclude that the diversity of the model parameter matrix at the linear output layer in pre-training has a significant impact on the transfer capability in the sense that the larger the diversity parameter, the smaller the risk. Therefore, the authors explicitly add a diversity regularizer to the linear output layer to increase diversity to improve the effect of pre-training.


 
Strength:

1. Very full and detailed mathematical derivation and analysis.

2. With the help of theoretical analysis and the use of regularizers to improve the effect of pre-training, a slightly better effect than the baseline was obtained on some NLU benchmarks, which supports the conclusion of the analysis (though not very strongly).

Weakness:

1. The authors mentioned that multi-task pre-training is the reason why the pre-training language model is effective, but the author uses the diversity of multi-class for replacement. Is there any enhancement difference between multi-task and multi-class diversity? 
After reading the paper, we still don't know which one should be used, multi-task pre-training or increasing multi-class diversity.

2. The theoretical part of this paper is relatively small compared to (Du et al., 2021, Remark 6.2), and the method adopted is also proposed by (Zou & Adams, 2012). The contribution is relatively limited, it is only a detailed derivation and combination with pre-training.

3. This work only tried one type of regularizer, the experimental part is slightly thin, and other diversity regularizers are not compared in the paper. Are they also in line with the conclusion of the paper?

4. Whether multi-task pre-training and increasing multi-class diversity can be used at the same time and their combining effect is not explored. For example, according to the "don't stop pre-training" training method, pre-training and downstream tasks constitute multi-task training. Does the diversity regularization method proposed in the paper can still work? Can increasing multi-class diversity be used in the finetune stage? Does the BERT baseline in Table 1 continue to be pre-trained for fair comparison? In addition, since the pre-training of BERT itself is multi-task pre-training (MLM and NSP), why not choose the simpler RoBERTa as the baseline for ablation study?

5. I find it more puzzling as mentioned in footnote 6. Regularization only changes the pre-training process and should have nothing to do with downstream tasks. Is BERT only suitable for classification tasks after use? I would like to see this part of the experiment added. If the task generalization ability of the pre-trained language model is sacrificed, then some of the improvements obtained are of no value. This improvement can be directly achieved by designing pre-training tasks related to the task.

6. The theoretical analysis is somewhat wonderful and convincing, meanwhile, the more the  diversity is, the better the performance is, which is an obvious and universal phenomena in various machine learning operations.  However, the experimental part is obviously flawed, inadequate and unconvincing. Too few experiments to support its theory. If such theory provided in this paper holds, to my best knowledge, I believe it will have great influence on CV pre-training tasks / transfer learning tasks, where the “class diversity” also holds. Also, I think experiments on more NLP pre-training models, like RoBERTa, ELECTRA are needed. Whether the proposed regularization technique has general effects to all pre-training tasks in all directions is unknown. Even for NLP, there are only very weak improvements. The glue result is only performed on the development set but not the test set, making the results less convincing. The effects of the proposed method cannot even surpass the effects of changing random seeds, meaning the improvements is non-significant. Meanwhile, for all settings of the diversity parameter in Table 1, it seems that there is no one single setting that surpass the baseline on all downstream tasks. 
 
Although this paper provides good explanation for the success of pre-training and mathematical proofs for its theorems, I still think it’s not good enough for ICLR. 

Meta-review: This paper aims to explain the pretraining effectiveness of masked language model, based on the concept of  diversity of classes. They empirically study how a diversity regularizer, based on this theory can improve model performance, as an empirical support.

Before rebuttal, reviewers consistently found the empirical study rather preliminary, while authors, through rebuttals, argue the theoretical study should be highlighted as their main contribution, and expressed concerns that the lack of empirical rigor should not be a ground to reject. We agree with these concerns, but rebuttals and discussions failed to convince reviewers that assumptions and evaluations are proper for connecting the proposed theory to potential impacts in pretrained language model scenarios. Revising to make this connection clearer would address the reviewer disagreements in the future.

Decision: Reject