Blessing of Class Diversity in Pre-training

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 pretraining : a 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 . 1 INTRODUCTION . Pre-training refers to training a model on a few or many tasks to help it learn parameters that can be used in other tasks . For example , in natural language processing ( NLP ) , one first pre-trains a complex neural network model to predict masked words ( masked language modeling ) , and then fine-tunes the model on downstream tasks , e.g. , sentiment analysis ( Devlin et al. , 2019 ) . Recently , pre-training technique has revolutionized natural language processing ( NLP ) . Models based on this technique have dramatically improve the performance for a wide range of downstream tasks ( Devlin et al. , 2019 ; Radford et al. , 2018 ; Yang et al. , 2019 ; Clark et al. , 2020 ; Lan et al. , 2020 ; Liu et al. , 2020 ) . Despite the large body of empirical work on pre-training , satisfactory theories are still lacking , especially theories that can explain the success of pre-training in NLP . Existing theories often rely on strong distributional assumptions ( Lee et al. , 2020 ) , smoothness conditions ( Robinson et al. , 2020 ) or noise-robustness conditions ( Bansal et al. , 2021 ) to relate the pre-training task ( s ) to downstream tasks . These assumptions are often hard to verify . A line of work studied multi-task pre-training ( Caruana , 1997 ; Baxter , 2000 ; Maurer et al. , 2016 ; Du et al. , 2021 ; Tripuraneni et al. , 2020a ; b ; Thekumparampil et al. , 2021 ) . In particular , recently , researchers have identified a new condition , the diversity of pre-training tasks , which has been shown to be crucial to allow pre-trained models to be useful for downstream tasks . See Section 2 for more detailed discussions on related work . Unfortunately , this line of theory can not be used to explain the success of pre-training in NLP . The theory of multi-task pre-training requires a large number of diverse tasks , e.g. , the number of tasks needs to be larger than the last layer ’ s input dimension ( a.k.a . embedding dimension ) , which is typically 768 , 1024 or 2048 ( Devlin et al. , 2019 ) . However , in NLP pre-training , there are only a few , if not one , pre-training tasks . Therefore , we need a new theory that is applicable to this setting . Furthermore , while existing theories are able to explain certain empirical phenomena , it remains unclear how to utilize the theory in practice . Ideally , we would like the theory to be useful to guide our training procedure or inspire new techniques to improve the performance of real-world models . Since in NLP pre-training , we do not have multiple tasks , we propose to study the blessing of multiple classes . Concretely , consider the Masked Language Model ( MLM ) pre-training task in NLP . In such a pre-training task , we have a large collection of sentences ( e.g . from Wikipedia ) . During the pre-training phase , we randomly mask a few words in each sentence , and predict the masked words using the remaining words in this sentence . This pre-training task is a multi-class classification problem where the number of classes is about 30K when using byte-pair-encoding ( BPE ) sub-word units.1 Note that this number is much larger than the embedding dimension ( 768 , 1024 or 2048 ) . In this paper , we develop a new statistical analysis aiming to explain the success of pre-training for NLP . The key notion of our theory is the diversity of classes , which serves as a similar role as the diversity of tasks in multi-task pre-training theory ( Du et al. , 2021 ; Tripuraneni et al. , 2020a ) . Our theory is not only applicable to the real-world NLP pre-training setting , but also inspires new techniques to improve the practical performance . We summarize our contributions below . 1.1 OUR CONTRIBUTIONS . First , we define a new notion , diversity of classes , which is the least singular value of the last linear layer in pre-training . We prove finite-sample bounds to show that for the cross-entropy loss , if the diversity of classes is large , then pre-training on a single task provably improves the statistical efficiency of the downstream tasks . To our knowledge , this is the first set of theoretical results that demonstrates the statistical gain of the standard practice of NLP pre-training , without strong distributional or smoothness conditions . Second , from a technical point of view , previous theoretical work on multi-task ( Du et al. , 2021 ; Tripuraneni et al. , 2020b ) builds on scalar output , and thus could not apply to multiclass tasks ( e.g. , cross entropy loss ) . We introduce a vector-form Radamacher complexity chain rule for disassembling composite function classes based on vector-form Rademacher contraction property ( Maurer , 2016 ) . This generalizes the scalar-form chain rule in Tripuraneni et al . ( 2020b ) . Furthermore , we adopt the modified self-concordance condition to show that the least singular value of the last linear layer serves as a diversity parameter for cross-entropy loss . We believe our techniques can be useful in other problems . Third , inspired by our theory , we develop a new regularization technique to promote class diversity for multi-class pre-training . We apply the negative log determinant regularizer only to the last linear layer of the pre-training model in masked language modeling . Our empirical results on BERT-base show this technique can boost the performance of downstream tasks . Organization This paper is organized as follows . In Section 2 , we review the related work . In Section 3 , we formally describe the problem setup and introduce the necessary definitions . In Section 4 , we state our main Theorem 1 then instantiate it with several applications . In Section 5 , we describe our new regularization technique and present the empirical results . We conclude in Section 6 and defer the proofs to Appendix . 2 RELATED WORK . Here we mostly focus on the theoretical aspects of pre-training . While there is a long list of work demonstrating the empirical success of self-supervised learning , there are only a few papers that study its theoretical aspects . One line of work studied the theoretical properties of contrastive learning ( Saunshi et al. , 2019 ; Tosh et al. , 2020 ) , which is a different setting considered in this paper . The most relevant one is by Lee et al . ( 2020 ) which showed that if the input data and pre-training labels were independent ( conditional on the downstream labels ) , then pre-training provably improved statistical efficiency . However , this conditional independence assumption rarely holds in practice . For example , in question-answering task , this assumption implies that given the answer , the question sentence and the masked word are independent . Robinson et al . ( 2020 ) assumed the Central Condition and a smoothness condition that related the pretraining task and the downstream task . Bansal et al . ( 2021 ) related generalization error of self-supervised learning to the noise-stability and rationality . However , it is difficult to verify the assumptions in these papers . A recent line of work studied multitask pre-training ( Du et al. , 2021 ; Tripuraneni et al. , 2020a ; b ; Thekumparampil et al. , 2021 ) in which the notion , diversity , has been identified to be the key that enables pre-training to improve statistical efficiency . These works generally require a large number 1This is a standard setting in the BERT model ( Devlin et al. , 2019 ) and is widely adopted as a common practice . By breaking down English words into BPE sub-word units , it could drastically increase the coverage of English language by using a relatively small ( 32768 ) vocabulary . of diverse tasks thus are not applicable to NLP , as we have mentioned . In comparison , we study single-task multi-class pre-training which is different from theirs . Du et al . ( 2021 ) noted that their results allowed an easy adaptation to multi-class settings in ( Remark 6.2 ) . However , they only focused on quadratic loss with one-hot labels for multi-class classification . Instead , we study the commonly used cross-entropy loss . While their analyses do not imply results in our setting , our theoretical analyses are inspired by this line of work . Our paper uses a diversity regularizer proposed in Zou & Adams ( 2012 ) to improve the performance of pre-training . We note that there are other diversity regularizers ( Xie et al. , 2017 ; Mariet & Sra , 2015 ; Cogswell et al. , 2015 ) . These may also improve the performance as the one in Zou & Adams ( 2012 ) . We leave it as a future work to investigate these regularizers . 3 PRELIMINARIES . In this section , we introduce the necessary notations , the problem setup , and several model-dependent quantities of pre-training . 3.1 NOTATION AND SETUP . Notations . Let [ n ] = { 1 , 2 , · · · , n } . We use ‖ · ‖ or ‖ · ‖2 to denote the ` 2 norm of a vector . Let N ( µ , σ2 ) be the one-dimensional Gaussian distribution . For a matrix W ∈ Rm×n , let ‖W‖1 , ∞ = maxq ( ∑ p |Wq , p| ) and ‖W‖∞→2 be the induced ∞-to-2 operator norm . We use the standard O ( · ) , Ω ( · ) and Θ ( · ) notation to hide universal constant factors , and use Õ ( · ) to hide logarithmic factors . We also use a . b to indicate a = O ( b ) . Problem Setup . The procedure is divided into two stages : pre-training stage to find a representation function and the downstream training stage to obtain a predictor for the downstream task . In both stages , we use R̂ to represent empirical risk and use R to represent expected loss . In the first stage , we have one pre-training task with n samples , { xprei , y pre i } ni=1 , where x pre i ∈ X pre ⊂ Rd is the input and yprei ∈ { 0 , 1 } k−1 is the one-hot label for k-class classification ( if y pre i is all-zero then it represents the k-th class ) .2 For instance , in masked language modeling , the input of each sample is a sentence with one word masked out , and the label is the masked word.3 k in this example is the size of the vocabulary ( ≈ 30K ) . We aim to obtain a good representation function ĥ within a function classH ⊂ { Rd → Rr } where r is the embedding dimension ( often equals to 768 , 1024 , 2048 in NLP pre-training ) . For example , one popular choice of the representation function ĥ in NLP applications is the Transformer model and its variants ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ) . On top of the representation , we predict the labels using function fpre within function class Fpre ⊂ { Rr → Rk−1 } . To train the representation function and predictor in pre-training stage , we consider the Empirical Risk Minimization ( ERM ) procedure ĥ = arg min h∈H min fpre∈Fpre R̂pre ( f pre , h ) , arg min h∈H min fpre∈Fpre 1 n n∑ i=1 ` ( fpre ◦ h ( xprei ) , y pre i ) ( 1 ) where ` is the loss function . We overload the notation for both the pre-training task and the downstream task , i.e. , for pre-training , ` : Rk−1 × { 0 , 1 } k−1 → R and for the downstream task , ` : Rk′−1 × { 0 , 1 } k′−1 → R . e.g. , cross-entropy : ` ( ŷ ; y ) = −y > ŷ + log ( 1 + ∑k−1 s=1 exp ( ŷs ) ) . Now for the downstream task , we assume there arem samples { xdowni , ydowni } mi=1 . Note that xdowni ∈ X down ⊂ Rd is the input and ydowni ∈ { 0 , 1 } k ′−1 is the one-hot label for k′-class classification.4 2We assume only one pre-training task for the ease of presentation . It is straightforward to generalize our results to multiple pre-training tasks . 3Here we say only one word being masked only for the ease of presentation . It is straightforward to generalize our results to the case where multiple words are masked out . 4For simplicity , we assume we only have one downstream task . Our theoretical results still apply if we have multiple downstream tasks . Note that in most real-world applications , we have n m and k k′ . For example , in sentiment analysis , k′ = 2 ( “ positive ” or “ negative ” ) . A widely studied task SST-2 ( Wang et al. , 2019 ) has m ≈ 67K , which is also generally much smaller than the pre-training corpus ( e.g. , n > 100M samples ) . To train the classifier for the downstream task , we fix the representation function learned from the pretraining task and train the task-dependent predictor within function class Fdown ⊂ { Rr → Rk′−1 } : f̂down = arg min fdown∈Fdown R̂down ( f down , ĥ ) = arg min fdown∈Fdown 1 m m∑ i=1 ` ( fdown ◦ ĥ ( xdowni ) , ydowni ) . ( 2 ) Therefore , our predictor for the downstream task consists a pair ( f̂down , ĥ ) . We use the Transfer Learning Risk defined below to measure the performance Transfer Learning Risk , Rdown ( f̂down , ĥ ) − Exdown , ydown [ ` ( gdown ( xdown ) , ydown ) ] where Rdown ( f̂down , ĥ ) , Exdown , ydown [ ` ( f̂down ◦ ĥ ( xdown ) , ydown ) ] is the expected loss ( the expectation is over the distribution of the downstream task ) , and gdown = arg ming∈ { Rd→Rk′−1 } Exdown , ydown [ ` ( g ( xdown ) , ydown ) ] is the optimal predictor for the downstream task . In our analysis , we also need to use the following term to characterize the quality of pre-training Pre-training Risk = Rpre ( f̂pre , ĥ ) − Expre , ypre [ ` ( gpre ( xpre ) , ypre ) ] where Rpre ( f̂pre , ĥ ) , Expre , ypre [ ` ( f̂pre ◦ ĥ ( xpre ) , ypre ) ] is the expected loss , and gpre = arg ming∈ { Rd→Rk−1 } Expre , ypre [ ` ( g ( xpre ) , ypre ) ] is the optimal predictor for the pre-training task . Following the existing work on representation learning ( Maurer et al. , 2016 ; Du et al. , 2021 ; Tripuraneni et al. , 2020b ) , throughout the paper we make the following realizability assumption . Assumption 1 ( Realizability ) . There exist h ∈ H , fpre ∈ Fpre , fdown ∈ Fdown such that gpre = fpre ◦ h and gdown = fdown ◦ h. This assumption posits that the representation class and the task-dependent prediction classes are sufficiently expressive to contain the optimal functions . Importantly , the pre-training and downstream tasks share a common optimal representation function h. This assumption formalizes the intuition that pre-training learns a good representation that is also useful for downstream tasks .

Summary: 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.