Year: 2021

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

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

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 1.1 OUR CONTRIBUTIONS
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 PRELIMINARIES
text: In this section, we introduce the necessary notations, the problem setup, and several model-dependent quantities of pre-training.

heading: 3.1 NOTATION AND SETUP
text: 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.

heading: 3.2 TASK RELATEDNESS AND DIVERSITY
text: We use the following definitions, which are natural analogies of those in Tripuraneni et al. (2020b) for multi-task transfer learning. To measure the “closeness” between the learned representation and true underlying feature representation, we use the following metric, following Tripuraneni et al. (2020b):
Definition 1. Let h ∈ H be the optimal representation function and h′ ∈ H be any representation function. Let fpre ∈ F be the optimal pre-training predictor on top of h. The pre-training representation difference is defined as:
dFpre,fpre(h ′;h) = inf f ′∈Fpre Expre,ypre [`(f ′ ◦ h′(xpre), ypre)− `(fpre ◦ h(xpre), ypre)] (3)
where the expectation is over the pre-training data distribution.
Intuitively, this measures the performance difference between the optimal predictor and the best possible predictor given a representation function h′.
For transfer learning, we also need to introduce a similar concept on the downstream task.
Definition 2. Let h ∈ H be the optimal representation function and h′ ∈ H be any representation function. On the downstream task, for a function class Fdown, let fdown ∈ F be the optimal
pre-training predictor on top of h. We define the worst-case representation difference between representations h and h′ ∈ H as:
dFdown(h ′;h)
= sup fdown∈Fdown inf f ′∈Fdown
Exdown,ydown [ `(f ′ ◦ h′(xdown), ydown)− `(fdown ◦ h(xdown), ydown) ] where the expectation is over the data distribution of the downstream task.
We now introduce the key notion of diversity, which measure how well a learned representation, say h′, from the pre-training task can be transferred to the downstream task. Definition 3. Let h ∈ H be the optimal representation function. Let fpre ∈ Fpre be the optimal pre-training predictor on top of h. The diversity parameter ν > 0 is the largest constant that satisfies
dFdown(h ′;h) ≤ dF
pre,fpre(h ′;h)
ν ,∀h′ ∈ H. (4)
While Definition 1- 3 are naturally defined from inspecting the pre-training procedure, it is not trivial to use these definitions to derive statistical guarantees. In particular, one of our key technical challenge is to show the least singular value of the last linear layer serves as a lower bound of the diversity parameter when Fpre and Fdown are linear function classes.

heading: 3.3 MODEL COMPLEXITIES
text: Lastly, we need to introduce some notions to measure the complexity of the function classes considered. In this paper, we consider Gaussian complexity which quantifies the extent to which the function in the class Q can be correlated with a noise sequence of length n× r. Definition 4 (Gaussian complexity). Let µ be a probability distribution on a setX ⊂ Rd and suppose that x1, · · · , xn are independent samples selected according to µ. Let Q be a class of functions mapping from X to Rr. Define random variable
Ĝn(Q) = E [ sup q∈Q 1 n r∑ k=1 n∑ i=1 gkiqk(xi) ] (5)
as the empirical Rademacher complexity, where qk(xi) is the k-th coordinate of the vector-valued function q(xi), gki (k ∈ [r], i ∈ [n]) are independent Gaussian N (0, 1) random variables. The Gaussian complexity of Q is Gn(Q) = EµĜn(Q).
Our main results are stated in terms of the Gaussian complexity. In Section 4.3 and 4.4 we will plug in existing results of the Gaussian complexity of certain function classes to obtain concrete bounds.
We will need the following worst-case Gaussian complexity for the pre-training predictor within Fpre
Ḡn(Fpre) = max h(x1),··· ,h(xn) Ĝn(Fpre|h ◦ xpre), (6)
here h ∈ H and xpre = x1, · · · , xn ∈ X pre. Similarly we define Ḡm(Fdown) as the worst-case Gaussian complexity for the downstream predictor within Fdown. We note that a closely related notion is Rademacher complexity. The empirical Rademacher complexity and Gaussian complexity only differ by a log factor (Ledoux & Talagrand, 1991).

heading: 4 MAIN RESULTS
text: In this section we present our main theoretical results. In Section 4.1 we present an analysis in terms of the diversity parameter for general loss function under certain regularity conditions. In Section 4.2, we specialize the result to a setting that is most relevant to NLP pre-training applications, where Fpre and Fdown are sets of linear functions and the loss is cross-entropy. In this particular case, our key result will show that one can use the singular value of the last linear to bound the diversity parameter. In Section 4.3 and 4.4 we instantiate our bounds on two concrete representation function classes: linear subspace and multi-layer network to showcase our main results.

heading: 4.1 MAIN THEOREM
text: In this subsection, we present our generic end-to-end transfer learning guarantee for multi-class transfer learning problem. We do not impose any specific function class formulations. Throughout this subsection, we only make the following mild regularity assumptions to make our results general. Assumption 2 (Regularity conditions). We assume the following regularity conditions hold:
• In pre-training, `(·, ·) is Bpre-bounded, and `(·, y) is Lpre-Lipschitz for all y.
• In downstream task, `(·, y) is Bdown-bounded and Ldown-Lipschitz for all y.
• Any predictor f ∈ Fpre is L(Fpre)-Lipschitz with respect to the Euclidean distance.
• Predictors are bounded: ‖f ◦ h(x)‖ ≤ DXpre for any x ∈ X pre, h ∈ H, f ∈ Fpre. Similarly ‖f ◦ h(x)‖ ≤ DXdown for any x ∈ X down, h ∈ H, f ∈ Fdown.
Specifically, one can show that common task-dependent losses satisfy these conditions. For example, when ` is cross-entropy loss (cf. Appendix A.2) for k−class classification, we prove ` is√ k − 1−Lipschitz and DX−bounded where X denotes the input data domain.
Under these assumptions, we have the following guarantee. Theorem 1. Under Assumption 1 and 2, for a given fixed failure probability δ, with probability at least 1− δ we have the Transfer Learning Risk upper bounded by:
O
( 1
ν
{ Lpre [ log (n) · [L(Fpre) ·Gn(H) + Ḡn(Fpre)] + √ kDXpre
n2
] +Bpre √ log (1/δ)
n
}
+Ldown · Ḡm(Fdown) +Bdown √ log(1/δ)
m
) .
The first line comes from pre-training ERM procedure and it accounts for the error of using an approximate optimal representation ĥ ≈ h. The second line characterizes the statistical error of learning the downstream-task predictor fdown from m samples. Note the diversity parameter appears in the denominator, which relates the pre-training risk to the transfer learning risk. Theorem 1 shows the risk would be small if the Gaussian complexities are small. We will show concrete examples where Gn(H) and Ḡn(Fpre) are O( √ 1/n) and Ḡm scales as O( √ 1/m). We believe this theorem applies broadly beyond the concrete settings considered in this paper.
In comparison with previous results, transfer learning risk analyses in (Du et al., 2021; Tripuraneni et al., 2020b) focus on scalar output. Their results cannot be applied to multi-class transfer learning tasks. In Theorem 1, we generalize the analysis in (Tripuraneni et al., 2020b) to handle multi-class classification where the output is high dimensional (number of classes). Technically, in the proof, we introduce a vector-form Radamacher complexity chain rule for disassembling composite function classes by making use of the vector-form Rademacher contraction property (Maurer, 2016).

heading: 4.2 MAIN RESULTS FOR MULTI-CLASS CLASSIFICATION WITH CROSS-ENTROPY LOSS
text: Now we specialize the general result to the setting that is of most interest to NLP pre-training, where the loss function ` is cross-entropy and the Fpre and Fdown are sets of linear functions. This choice is consistent with the NLP pre-training: e.g., BERT (Devlin et al., 2019) uses transformers as the representation learning function classH and uses word-embedding matrices as Fpre. Formally we define
Fpre = {f |f(z) = α>z, α ∈ Rr×(k−1), ‖αs‖ ≤ c1,∀s ∈ [k − 1]} (7) Fdown = {f |f(z) = α>z, α ∈ Rr×(k ′−1), ‖αs‖ ≤ c0,∀s ∈ [k′ − 1]} (8)
where c1 and c0 are some universal constants. Then the regularity conditions are instantiated as:
• Pre-training loss `(·, y) is √ k − 1-Lipschitz and Bpre = DXpre -bounded.
• Downstream loss is √ k′ − 1-Lipschitz and Bdown = DXdown -bounded. • Any function f ∈ Fpre is L(Fpre) = c1 √ k − 1-Lipschitz w.r.t. the l2 distance.
See Appendix A.2 for detailed derivations. Next we discuss our main assumption that relates the diversity parameter to a concrete quantity of the last linear layer. Assumption 3 (Lower Bounded Least Eigenvalue). Let the optimal linear predictor at the last layer for pre-training be αpre ∈ Rr×(k−1). We assume ν̃ , σr(αpre (αpre)>) > 0.
Similar assumptions have been used in multi-task representation learning (Du et al., 2021; Tripuraneni et al., 2020a;b), and are shown to be necessary (Maurer et al., 2016; Du et al., 2021). Different from their versions, our assumption is tailored for the multi-class classification setting.
One of our key technical contributions is to show ν̃ serves as a lower bound for the diversity parameter ν. See Lemma 3 in Appendix. Intuitively, this assumption ensures that the pre-training task matrix spans the entire r-dimensional space and thus covers the output of the optimal representation h(·) ∈ Rr. This is quantitatively captured by the σr(αpre (αpre)>), which measures how spread out these vectors are in Rr. Technically, we apply a modified self-concordance condition for better characterizing multinomial logistic regression (Bach et al., 2010). See Appendix A.4 for proof.
We now state our theorem for this specific setting. Theorem 2. Under Assumptions 1 and 3, with probability at least 1 − δ we have the Transfer Learning Risk upper bounded by:
O
( 1
ν̃
{ √ k [ log (n)[ √ k ·Gn(H) + Ḡn(Fpre)] + √ kDXpre
n2
] +DXpre √ log (1/δ)
n
}
+ √ k′ · EXdownĜm(Fdown|ĥ ◦ xdown) + σ
√ log(1/δ)
m +DXdown
√ log(1/δ)
m ) Here EXdownĜm(Fdown|ĥ ◦ xdown) is Gaussian complexity of embeddings
ĥ ◦ xdown = {ĥ(x1), · · · , ĥ(xm)|xdown = x1, · · · , xm ∈ X down} (9) where the expectation is over X down, and σ2 = 1m supf∈Fdown ∑m i=1 V ar(`(f ◦ ĥ(xdowni ), ydowni )) is the maximal variance over Fdown.
We remark that in Theorem 2, since we specialize to the case where Fpre and Fdown are sets of linear functions, we can replace the term Ldown · Ḡm(Fdown) in Theorem 1 by(√ k′ · EXdownĜm(Fdown|ĥ ◦ xdown) + σ √ log(1/δ) m ) by utilizing the functional Bernstein inequal-
ity. This improvement can help us obtain Theorem 3. See Appendix A.2 for the full proof.
Now we discuss the interpretation of Theorem 2. Typically, Ḡn(Fpre) is much smaller than Gn(H) because Gn(H) represents the complexity of the representation function, which is often complex. In practice, n is often large. Therefore, in the benign case where ν̃ = Θ (k) (when the condition number of αpre is O(1)), then the dominating term will be Ḡn(Fpre). As we will show in the following subsections, this term typically scales as O( √ 1/n). Together, this theorem clearly shows when 1) the number of pre-training data is large, and 2) the least singular value of the last linear layer for pre-training is large, the transfer learning risk is small. On the other hand, if ν̃ is small, then the bound becomes loose. This is consistent with prior counter examples on multi-task pre-training (Maurer et al., 2016; Du et al., 2021) where the diversity is shown to be necessary.

heading: 4.3 LINEAR SUBSPACE REPRESENTATION
text: We instantiate this setting by assuming the underlying representation to be a projection onto a low-dimensional subspace. For r d, we let the function class be
H = {h|h(x) = B>x,B ∈ Rd×r,B is a matrix with orthonormal columns} (10) We require some additional regularity conditions. Following prior work (Du et al., 2021; Tripuraneni et al., 2020b), we assume that ‖x‖ ≤ D and input data distribution satisfies the following condition.
Definition 5. We say the covariate distribution Px(·) is Σ-sub-gaussian if for all v ∈ Rd, E[exp(v>x)] ≤ exp
( ‖Σ1/2v‖2
2
) where the covariance Σ further satisfies σmax(Σ) ≤ C and
σmin(Σ) ≥ c ≥ 0 for universal constants c, C.
We have the following theorem that guarantees the performance of transfer learning. Theorem 3. Suppose Assumption 1 and 3 hold. For a sufficiently large constant c3, we assume n ≥ c3d,m ≥ c3r,D ≤ c3(min( √ dr2, √ rm)). Then with probability at least 1 − δ, we have the Transfer Learning Risk upper bounded by:
O
( 1
ν̃
[ √ k log (n) (√ kdr2
n + k
√ r
n
) + √ kD ( k
n2 +
√ log (1/δ)
n
)]
+(k′) 3 2
√ r
m + k′
√ log (1/δ)
m + √ k′D
√ log (1/δ)
m
)
To interpret this bound, consider the practically relevant scenario where k′ = O(1) (e.g., sentiment analysis), m n, k n and r d, in the benign case ν̃ = Ω (k), we have the transfer learning
risk Õ (√ dr2 n + √ r m ) . Note that this is exactly the desired theoretical guarantee because the first
term accounts for using all pre-training data to learn the representation function and the second term accounts for using the downstream data to learn the last linear layer. This is significantly better than not using pre-training, in which case the risk scales O (√
d m
) .

heading: 4.4 DEEP NEURAL NETWORK REPRESENTATION
text: In this subsection, we assume the underlying representation to be a practical σ = tanh-activated neural network. Predictors are still required to be linear functions at the interest of NLP pre-training, i.e.,
H = {h|h(x) = WKσ (WK−1σ (· · ·σ (W1x)))}, (11) Fpre = {f |f(z) = α>z, α ∈ Rr×(k−1), ‖αs‖ ≤ c1M(K)2,∀s ∈ [k − 1]}, (12) Fdown = {f |f(z) = α>z, α ∈ Rr×(k ′−1), ‖αs‖ ≤ c0M(K)2,∀s ∈ [k′ − 1]}. (13)
We further assume 1) for each p ∈ [K], ‖Wp‖1,∞ ≤ M(p), and 2) ‖WK‖∞→2 ≤ M(K). Adapt Gaussian complexity results in Golowich et al. (2018) we have
Gn(H) ≤ Õ
( rM(K) ·D
√ K ·ΠK−1p=1 M(p)√ n ) , Ḡn(Fpre) ≤ O ( (k − 1)M(K)3√ n ) . (14)
Now we are ready to state our theorem for this practical setting of NLP pre-training. Theorem 4. Under Assumption 1 and 2, assume M(K) ≥ c3 for a universal constant c3. Then with probability at least 1− δ, Transfer Learning Risk is upper bounded by
Õ
( krM(K)3 ·D √ K ·ΠK−1p=1 M(p)
ν̃ √ n
+ kM(K)3
ν̃ √ n
+ k′ 3 2M(K)3√
m
) .

heading: 5 EXPERIMENTS
text: Our theoretical analysis in previous sections implies 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 ν (diversity parameter of fpre), the smaller the risk. Therefore, we could explicitly add a diversity regularizer to the linear output layer to increase diversity. Motivated by this, we propose to add the following diversity regularizer to the original BERT pre-training loss so that it becomes:
L′(Θ) = L(Θ)− λ · ln det(αpre (αpre)>), (15)
where Θ denotes the set of all model parameters, λ is a hyper-parameter that controls the magnitude of the diversity regularization, det(·) denote the determinant of a matrix, and αpre is the model parameter matrix at the output linear layer. This type of diversity regularizer was proposed in Zou & Adams (2012). This regularization technique is different from prior work because it is specifically designed for multi-class pre-training: we only add the diversity regularizer to the last linear layer.
We use the above diversity-regularized loss (along with the original `2-regularization) to pretrain BERT-base models under different values of diversity factor λ. Then we fine-tune them on 7 classification tasks and 1 regression task from the GLUE benchmark (Wang et al., 2019) to evaluate their transfer performance.5 Our pre-training and finetuning implementations are based on the opensource code released by Nvidia.6 We use the same pre-training data as the original BERT (i.e., English Wikipedia + TorontoBookCorpus).7 Our detailed pre-training and finetuning hyperparameters along with other experimental details are reported in Appendix A.6.2.
In Table 5, we report our performance on the dev sets of the 8 downstream tasks. All the experiments are repeated 5 times with different random seeds, and we report their mean values along with the standard deviations. The complete experiment results (including full MNLI, QQP, and MRPC results) can be found in Appendix A.6. From Table 5, we note that adding the diversity regularization could generally improve the performance on these downstream tasks. In particular, when λ = 0.5, our pretrained model outperforms the original BERT-base on 6 out of 8 tasks (with 3 of them being significant), while achieving comparable performance on the other 2 tasks. Although our model is slightly behind the original BERT on CoLA and QNLI, such a performance gap is not statistically significant. Besides, we also see that our model with λ = 0.5 achieves a much more stable performance (i.e., smaller std) on tasks with scarce finetuning data (< 4K samples in MRPC and RTE). Our results, albeit still preliminary, demonstrate the potential of such a simple diversityregularizer. It could be an effective and simple performance booster for any of the existing pre-trained NLP models (e.g., XLNet (Yang et al., 2019), RoBERTa (Liu et al., 2020), ALBERT (Lan et al., 2020), etc) with negligible computation and implementation cost. We leave the development of the more advanced diversity regularizer as a future work.

heading: 6 CONCLUSION
text: We theoretically prove the benefit of multi-class pre-training using the notion of class diversity. Inspired by the theory, we further propose a new regularization technique specially designed for pre-training. Our experiments show potential of this diversity regularizer. An interesting direction is to further investigate the impact of different diversity regularizers on larger pre-training models. Lastly, we only studied multi-class pre-training in this paper, it is interesting to develop a theory for other pre-training techniques such as those based on sequence reconstruction (Lewis et al., 2019).
5We do not report the WNLI (classification) task due to its reported issues of the task in Devlin et al. (2019). 6Distributed under Apache License: https://github.com/NVIDIA/DeepLearningExamples/ tree/master/PyTorch/LanguageModeling/BERT 7Collected and pre-processed using the code and script included in the open-source code: https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/ LanguageModeling/BERT

heading: A APPENDIX
text: In Section 3, we have introduced Gaussian complexity. Let us restate for clarity.
Let µ be a probability distribution on a set X ⊂ Rd and suppose that x1, · · · , xn are independent samples selected according to µ. LetQ be a class of functions mapping from X to Rr. Define random variable
Ĝn(Q) = E [ sup q∈Q 1 n r∑ k=1 n∑ i=1 gkiqk(xi) ] (16)
as the empirical Rademacher complexity, where qk(xi) is the k-th coordinate of the vector-valued function q(xi), gki (k ∈ [r], i ∈ [n]) are independent Gaussian N (0, 1) random variables. The Gaussian complexity of Q is Gn(Q) = EµĜn(Q). Analogously to the above we can define the empirical Rademacher complexity for vector-valued functions as
R̂n(Q) = E [ sup q∈Q 1 N r∑ k=1 N∑ i=1 kiqk(xi) ] (17)
where ki(k ∈ [r], i ∈ [n]) are independent Rademacher Rad( 12 ) random variables. Its population counterpart is defined as Rn(Q) = Eµ[R̂n(Q)]. Note that the superscripts existing in Ĝ and R̂ imply that they are empirical measures.

heading: A.1 PROOFS FOR THEOREM 1
text: We illustrate Theorem 1 in two stages. First we show pre-training representation difference can be upper bounded by constants and function class complexities. Then we transfer it to the downstream task through the diversity parameter.

heading: Pre-training
text: Theorem 5. In pre-training, with probability at least 1− δ, it holds that:
dFpre,fpre(h ′;h)
≤ 4 √ πLpreGn(Fpre ◦ H) + 4Bpre
√ log(2/δ)
n ≤ 4096Lpre [√
k − 1DXpre n2
+ log(n)[L(Fpre)Gn(H) + Ḡn(Fpre)] ] + 4Bpre √ log(2/δ)
n .
Proof. We begin with
dFpre,fpre(h ′;h) ≤ 2 sup f∈Fpre,h∈H |Rpre(fpre, h)− R̂pre(fpre, h)|.
From the definition of Rademacher complexity (Wainwright, 2019, Theorem 4.12), with probability at least 1− 2δ we have
sup fpre∈Fpre,h∈H
|Rpre(fpre, h)− R̂pre(fpre, h)| ≤ 2Rn(`(Fpre ◦ H)) + 2Bpre √ log(1/δ)
n .
Next, we apply the vector contraction inequality (Maurer, 2016). For function class F whose output is in RK with component fk(·), and the function (hi)s are some L-Lipschitz functions: RK 7→ R, we have
E sup f∈F n∑ i=1 ihi(f(xi)) ≤ √ 2LE sup f∈F n∑ i=1 K∑ k=1 ikfk(xi). (18)
Hence for loss function ` satisfying |`(x)−`(y)| ≤ Lpre‖x−y‖2,∀x, y ∈ Rk−1, the f takes value in Rk−1 with component functions fs(·), s ∈ [k − 1], we have that population Rademacher complexity can be bounded by
Rn(l(Fpre ◦ H)) = EXpre 1
n E sup
f∈Fpre,h∈H n∑ i=1 i`(f ◦ h(xprei ))
≤ EXpre 1
n
√ 2LpreE sup
f∈Fpre,h∈F n∑ i=1 k−1∑ s=1 isfs(h(x pre i ))
= √
2LpreRn(Fpre ◦ H) ≤ √ πLpreGn(Fpre ◦ H)
where the last line uses the fact that Rademacher complexity is upper bounded by Gaussian complexity: Rn(Fpre ◦ H) ≤ √ π 2Gn(F
pre ◦ H). Therefore we have, with probability at least 1− δ, dFpre,fpre(h ′;h)
≤ 4 √ πLpreGn(Fpre ◦ H) + 4Bpre
√ log(2/δ)
n ≤ 4096Lpre [√
k − 1DXpre n2
+ log(n)[L(Fpre)Gn(H) + Ḡn(Fpre)] ] + 4Bpre √ log(2/δ)
n ,
where the last line uses decomposition of Gn(Fpre ◦ H) into the individual Gaussian complexities of H and Fpre, leverages an expectation version of novel chain rule for Gaussian complexities (Lemma 1).
In the spirit of Gaussian complexity decomposition theorem (Tripuraneni et al., 2020b, Theorem 7), we introduce the following decomposition result upon vector-form Gaussian complexities. Lemma 1. We have the following vector form Gaussian complexity decomposition:
Ĝn(Fpre ◦ H) ≤ 8 √ k − 1DXpre n2 + 512C(Fpre ◦ H) · log (n) (19)
where we use C(Fpre ◦ H) = L(Fpre) · Ĝn(H) + Ḡn(Fpre) to represent the complexity measure of the composite function class.
Proof. Our proof extends (Tripuraneni et al., 2020b, Theorem 7), which focuses on a multi-task scalar formulation. We further extend it to multi-class vector formulation. Specifically, on top of the representation class H, they need to handle F⊗t (t is the number of tasks) while our objective is a single function class Fpre of higher dimension (Fpre is (k − 1)-dimensional for a k-class classification task). We note that our proof technique and that of previous works (Tripuraneni et al., 2020b; Maurer et al., 2016) both hinge on several properties of Gaussian processes.
To bound the empirical composite function class Fpre(H), note that vector-form Gaussian complexity is defined as
Ĝn(Fpre ◦ H) = E
[ 1
n ˙supf(h)∈Fpre(H) k−1∑ s=1 n∑ i=1 gisfs(h(x pre i ))
]
= 1√ n E[ sup f(h)∈Fpre(H) Zf(h)]
where we define mean-zero process Zf(h) = 1√n ∑k−1 s=1 ∑n i=1 gisfs(h(x pre i )), then E supf(h) Zf(h) = E supf(h) Zf(h) − Zf ′(h′) ≤ E supf(h),f ′(h′) Zf(h) − Zf ′(h′). We further notice that Zf(h) − Zf ′(h′) is a sub-gaussian random variable parameter
d2(f(h), f ′(h′)|xpre) = 1 n n∑ i=1 ‖f(h(xprei ))− f ′(h′(xprei ))‖ 2
= 1
n k−1∑ s=1 n∑ i=1 (fs(h(x pre i ))− f ′ s(h ′(xprei ))) 2
Dudley’s entropy integral bound (Wainwright, 2019, Theorem 5.22) shows
E sup f(h),f ′(h′) Zf(h) − Zf ′(h′)
≤ 2E sup d(f(h),f ′(h′)|xpre)≤δ
Zf(h) − Zf ′(h′) + 32J ( δ
4 , DXpre)
= 2E sup d(f(h),f ′(h′)|xpre)≤δ Zf(h) − Zf ′(h′) + 32 ∫ DXpre δ 4 √ logN(u;Fpre(H)|xpre)du.
It is straightforward to see the first term follows:
E sup d(f(h),f ′(h′)|xpre)≤δ
Zf(h) − Zf ′(h′) ≤ E[‖g‖]δ ≤ √ n(k − 1)δ
We now turn to bound the second term by decomposing the distance metric into a distance over Fpre and a distance overH. We claim that, for arbitrary h ∈ H, f ∈ Fpre, let h′ be 1-close to h in empirical l2-norm w.r.t inputs x pre 1 , x pre 2 · · · , xpren . Given h′, let f ′ be 2-close to f in empirical l2 loss w.r.t h′(xpre). Using the triangle inequality we have that
d(f(h), f ′(h′)|xpre) = √√√√ 1 n n∑ i=1 ‖f(h(xprei ))− f ′(h′(x pre i ))‖
≤ d(f(h), f(h′)|xpre) + d(f(h′), f ′(h′)|xpre)
≤ √√√√ 1 n n∑ i=1 ‖f(h(xprei ))− f(h′(x pre i ))‖2 + 2
≤ L(Fpre) √√√√ 1 n n∑ i=1 ‖h(xprei )− h′(x pre i )‖2 + 2
= L(Fpre) · 1 + 2, where we have used that ‖f(x)− f(y)‖ ≤ L(Fpre)‖x− y‖ for any f ∈ Fpre. As for the cardinality of the covering CFpre(H). Observe |CFpre(H)| = ∑ h∈CH(xpre) |CFpreh | ≤ |CH(xpre)| ·maxh∈H(xpre) |CFpre h(xpre) |. This provides a bound on the metric entropy of
logN( 1 · L(Fpre) + 2;Fpre(H)|xpre) ≤ logN( 1;H|xpre) + max h(xpre) N( 2;Fpre|h ◦ xpre).
Applying the covering number upper bound with 1 = 2·L(Fpre) , 2 = 2 gives a bound of entropy integral ofa ,∫ DXpre δ 4 √ logN(u;Fpre(H)|xpre)du
≤ ∫ DXpre δ 4 √ logN ( u 2L(Fpre) ;H ∣∣∣∣xpre)du+ ∫ DXpre δ 4 max h◦xpre √ logN (u 2 ;Fpre ∣∣h ◦ xpre)du
From the Sudakov minoration theorem (Wainwright, 2019, Theorem 5.30) for Gaussian processes and the fact that packing numbers at scale u upper bounds the covering number at scale ∀u > 0 we find:
logN(u;H|xpre) ≤ 4 (√ nĜn(H) u )2 , logN(u;Fpre|h(xpre)) ≤ 4 (√ nĜn(Fpre|h ◦ xpre) u )2 .
Combining the definition of worst-case Gaussian complexity with the aforementioned results we have
Ĝn(Fpre ◦ H) ≤ 2 √ k − 1δ + 256 log 4DX pre
δ
( L(Fpre)Ĝn(H) + Ḡn(Fpre) ) ,
substitute δ with 4DXpren2 , proof is completed.
Downstream learning Next we turn to the second stage and come up with theoretical guarantees by using inexact ĥ learned from the first stage. Theorem 6. In the downstream task, we have that with probability at least 1− δ,
Rdown(f̂ down, ĥ)−Rdown(fdown, h) ≤ dFdown(ĥ;h) + 4 √ πLdown · Ḡm(Fdown) + 4Bdown
√ log(2/δ)
m
Proof. Assumption 1 implies Exdown,ydown [ ` ( gdown ( xdown ) , ydown )] = Rdown(f down, h).
To start, let f̃ = arg minf∈Fdown Rdown(f, ĥ) and Rdown(f̂ down, ĥ)−Rdown(fdown, h) equals[ Rdown(f̃ , ĥ)−Rdown(fdown, h) ] + [ Rdown(f̂ down, ĥ)−Rdown(f̃ , ĥ) ]
where the first term satisfies
inf f̃∈Fdown
[ Rdown(f̃ , ĥ)−Rdown(fdown, h) ] ≤ sup fdown∈Fdown inf f̃∈Fdown [ Rdown(f̃ , ĥ)−Rdown(fdown, h)
] = dFdown(ĥ, h)
The second term follows the similar lines of Theorem 5
Rdown(f̂ down, ĥ)−Rdown(f̃ , ĥ) ≤ 4 √ πLdownEXdownĜm(Fdown|ĥ ◦ xdown) + 4Bdown
√ log (1/δ)
m Again we make use of the worst-case argument
EXdownĜm(Fdown|ĥ ◦ xdown) ≤ Ḡm(Fdown). Combining the results gives the statement.
Proof for main Theorem 1 Having introduced class diversity parameter, proof is directly completed via combination of Theorem 5 and Theorem 6.

heading: A.2 PROOFS FOR THEOREM 2
text: We could provide a better dependence on the boundedness noise parameters in Theorem 6 using Bernstein inequality. We present the following corollary which has data-dependence in the Gaussian complexities. Corollary 1. Presuming Assumption 1 holds, we have that then with probability at least 1− δ, Rdown(f̂ down, ĥ)−Rdown(fdown, h)
≤ dFdown(ĥ;h) + 4 √ πLdown · EXdownĜm(Fdown|ĥ ◦ xdown) + 4σ
√ log(2/δ)
m + 50Bdown
log(2/δ)
m
Proof. Denote Z = supf |R̂down(f, ĥ)−Rdown(f, ĥ)|, we apply the functional Bernstein inequality (Massart, 2000, Theorem 3) to control the fluctuations. With probability at lest 1 − δ, we have
Z ≤ 2E[Z] + 4 σ√ m
√ log( 1
δ ) + 35
Bdown
m log(
1 δ ), (20)
where σ2 = 1m supf ∑m i=1 V ar(`(f ◦ ĥ(xdowni ), ydowni )). Thus
E[Z] ≤ 2EXdownR̂m(l(Fdown)|ĥ ◦ xdown) ≤ 2EXdown √ 2LdownR̂m(Fdown|ĥ ◦ xdown)
≤ 2EXdown √ πLdownĜm(Fdown|ĥ ◦ xdown),
where the second line uses vector-based contraction principle, the last line upper bounds the empirical Rademacher complexity by Gaussian counterparts.
Proof for Theorem 2 Observe that
`(η; y) = −y>η + log (1 + k−1∑ s=1 eηs), `(η; y) ≤ ‖η‖
and ∣∣∣∣∂`(η; y)∂ηi ∣∣∣∣ = ∣∣∣∣∣yi − eηi1 +∑k−1s=1 eηs ∣∣∣∣∣ ,
|∇η`(η; y)| ≤ √ k − 1,
so it is Lpre = √ k − 1−Lipschitz. By definition the class Fpre with parameters ‖αs‖2 ≤ O(1), s ∈ [k − 1], we obtain that L(Fpre) = O (√ k − 1 ) since for any x, y ∈ Rr, any f ∈ Fpre we have
‖f(x)− f(y)‖2 = ‖α>x− α>y‖2
≤ k−1∑ s=1 (〈αs, x− y〉)2
≤ k−1∑ s=1 ‖αs‖2‖x− y‖2
≤ c21(k − 1)‖x− y‖2
In conclusion we have
• Pre-training loss `(·, ypre) is √ k − 1-Lipschitz and Bpre = DXpre -bounded. • Downstream loss `(·, ydown) is √ k′ − 1-Lipschitz and Bdown = DXdown -bounded.
• Linear layer f is L(Fpre) = O (√ k − 1 ) -Lipschitz.
Consider task-specific function classes for characterizing class-diversity parameters. From Lemma 2 and Lemma 3 we know that
ν = Ω(ν̃), ν̃ = σr(α1α > 1 ).
Combining these pieces of results then the proof is completed.
A.3 WHAT IS diversity parameter FOR LINEAR LAYERS?
In this subsection, we lower bound the diversity parameter ν by ν̃. We note the proof of this part is very different from the multi-task setting studied in previous works (Du et al., 2021; Tripuraneni et al., 2020b).
First, we demonstrate cross entropy loss between learned representation class and true underlying representation could be bounded by quadratic loss, Lemma 2. Suppose ` is a (k − 1)-class cross entropy loss and ypre is (k − 1)-onehot label. The cross entropy can be bounded from both sides with quadratic loss,
c0Expre [ exp(−10 max(‖α′>ĥ(xpre)‖, ‖α>h(xpre)‖)) ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2] ≤ Expre,ypre [`(f̂pre ◦ ĥ(xpre))− `(fpre ◦ h(xpre), ypre)]
≤ 1 2 Expre ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2 , where c0 = 12λmin(Φ
′′(α>h(xpre))) denote second-order derivative of softmax function, α′ and α are parameters for f̂pre and fpre respectively.
Proof. The conditional distribution of (k − 1)-class classification task is defined as:
P(·|α>h(x)) = ey >α>h(x)−Φ(α>h(x))
= ey >α>h(x)−log (1+ ∑k−1 s=1 e (α>h(x))s )
Taylor’s expansion tells that
Φ(α′ > ĥ(x)) = Φ(α>h(x)) +∇Φ(α>h(x))>(α′>ĥ(x)− α>h(x)) + o ( ‖α′>ĥ(x)− α>h(x)‖2 ) For generalized linear models: P(·|α>h(x)) = exp (y>α>h(x)− Φ(α>h(x))), it is known that
KL [ P(·|α>h(x)),P(·|α′>ĥ(x)) ] = Φ(α′ > ĥ(x))− Φ(α>h(x))−∇Φ(α>h(x))>(α′>ĥ(x)− α>h(x))
Our goal is to bound KL−divergence with quadratic distance, µ
2 ‖α′>ĥ(x)− α>h(x)‖2 ≤ KL
[ P(·|α>h(x)),P(·|α′>ĥ(x)) ] ≤ L
2 ‖α′>ĥ(x)− α>h(x)‖2.
(21) Here L = supx′ λmax∇2Φ(x′) for intermediate x′ ∈ [ α′ > ĥ(x), α>h(x) ] . µ is some constant related to α′>ĥ(x) and α>h(x) which requires more investigations.
First we showLwhich could be considered as upper bound of the maximum eigenvalue of multinomial logistic regression. Restate this sub-problem for clarification: for Φ(x) = log (1 + ∑k−1 s=1 e
xs), its gradient at ith-coordinate is ∂Φ∂xi = exi/1+ ∑ s e xs , so Hessian matrix is computed as
∂2Φ
∂xi∂xj =  exi ·(1+ ∑ s 6=i e xs ) (1+ ∑ s e xs )2 , i = j
−exiexj (1+ ∑ s e xs )2 , i 6= j.
(22)
Our goal is to bound the smallest and largest eigenvalues of this Hessian matrix while x ∈ [x1, x2]. Next we show its largest eigenvalue is upper bounded by 1. Denote σ(x) =
1 1+ ∑ s e xs [e x1 , · · · , exk−1 ]>, the Hessian matrix can be restate as
∇2Φ = diag(σ(x))− σ(x)σ(x)>. For any non-zero vector y, we have
y>∇2Φy = ∑ i σ(x)iy 2 i − ( σ(x)>y )2 ≤ max(σ(x)i)‖y‖2
≤ ‖y‖2, which implies its largest eigenvalue is no bigger than 1. Also, this Hessian is strictly-convex, from Cauchy-Schwarz inequality we know(
σ(x)>y )2 ≤∑
i σ(x)iy 2 i · ∑ i σ(x)i
< ∑ i σ(x)iy 2 i ,
so this matrix is positive definite.
The derivations of bounding second order remaining terms of second-order Taylor’s expansion from below are more complicated. Note that, logistic regression is not globally strongly-convex (so as multinomial logistic regression). However, when landscape is considered locally (e.g., here x′ ∈ [ α′ > ĥ(x), α>h(x) ] ), a class of convex functions called self-concordant functions would present
benign local properties. Next we introduce self-concordant functions and related useful properties.
Definition 6 (Modified self-concordance). F: Rp 7→ R is three times differentiable convex function such that for some R > 0, for all u, v ∈ Rp, the function g : t 7→ F (u+ tv) satisfies for all t ∈ R
|g′′′(t)| ≤ R‖v‖2 × g′′(t) (23)
Proposition 1. Multinomial logistic regression satisfies self-concordance condition, namely: Φ(x) = log(1+ ∑k−1 s=1 e xs), x ∈ Rk−1 is convex, for all u, v ∈ Rk−1, the function g : t 7→ Φ(u+tv) satisfies
|g′′′(t)| ≤ 5‖v‖2g′′(t) (24)
Proof. Let P (t; v0) = 1 + ∑ s e us+tvs and P (t; vi) = ∑ s v i se us+tvs , i > 1. Then we use multinomials P to represent derivatives of g(t)
g(t) = log(P (T ;V 0))
g′(t) = P (t; v1)
P (t; v0)
g′′(t) = P (t; v2)P (t; v0)
P (t; v0)2
g′′′(t) = P (t; v3)P (t; v0)2 − 3P (t; v2)P (t; v1)P (t; v0) + 2P (t; v1)3
P (t; v0)3
Let rs = eus+vst, hence g′′(t) = ( ∑ s v 2 srs) · (1 + ∑ s rs)− ( ∑ s vsrs) 2
(1 + ∑ s rs) 2
=
∑ i<j rirj(vi − vj)2 + ∑ i v 2 i ri
(1 + ∑ s rs) 2
In the following we expand g′′′(t) as:∑ i<j rirj(vi − vj)2[ ∑ k(vi + vj − 2vk)rk] + ∑ i v 3 i ri + ∑ i ∑ j v 2 i rirj(2vi − 3vj)
(1 + ∑ s rs) 3
=
∑ i<j rirj(vi − vj)2[ ∑ k(vi + vj − 2vk)rk] + ∑ i v 2 i ri ( vi(1 + 2 ∑ j rj)− 3 ∑ j vjrj ) (1 + ∑ s rs) 3 ,
observe that
1 1 + ∑ s rs ∣∣∣∣∣∑ k (vi + vj − 2vk)rk ∣∣∣∣∣ ≤∑ k |vi + vj − 2vk| rk 1 + ∑ s rs ≤ 4‖v‖2
1 1 + ∑ s rs ∣∣∣∣∣∣vi(1 + 2 ∑ j rj)− 3 ∑ j vjrj ∣∣∣∣∣∣ ≤ 5‖v‖2 Substitute these into definition of self-concordance then proof is completed.
Self-concordance properties Self-concordance gives nice characterizations of local curvature of convex functions which plays important role in describing local convexity (Bach, 2014). Some useful results are given upon this condition (see (Bach et al., 2010, Proposition 1)), we list out the three main inequalities as below
For all w, v ∈ Rp, t ∈ R, we have that,
F (w + v) ≥ F (w) + vF ′(w) + v >F ′′(w)v R2‖v‖22 · ( e−R‖v‖2 +R‖v‖2 − 1 ) (25)
F (w + v) ≤ F (w) + vF ′(w) + v >F ′′(w)v R2‖v‖22 · ( eR‖v‖2 −R‖v‖2 − 1 ) (26)
e−tR‖v‖2F ′′(w) F ′′(w + tv) etR‖v‖2F ′′(w) (27)
Above results give two refined versions of Taylor’s expansion along with an upper and lower zeroorder Taylor expansion of F ′′. With these we are ready to give lower bound of KL divergence (see Eq 21), i.e.,
Φ(α′ > ĥ(x))− Φ(α>h(x))−∇Φ(α>h(x))>(α′>ĥ(x)− α>h(x))
≥ 1 2 v>e−tR‖v‖2F ′′(α>h(x))v ≥ 1 2 λmin(Φ ′′(α>h(x)))‖v‖2e−5‖v‖2 ≥ 1 2 λmin(Φ ′′(α>h(x)))‖v‖2e−5(‖α ′>ĥ(x)‖+‖α>h(x)‖) ≥ 1 2 λmin(Φ ′′(α>h(x)))‖v‖2 exp(−10 max(‖α′>ĥ(x)‖, ‖α>h(x)‖))
where v = α′>ĥ(x)− α>h(x), t is an parameter exists in [0, 1]. Proof for Lemma 2 is completed.
Lemma 3. With Lemma 2 at hand and define C = α1α>1 , it is demonstrated that
dFdown(ĥ;h) ≤ 1
Ω(ν̃) dFpre,fpre(ĥ;h), ν̃ = σr(C) (28)
Proof. For function classes Fpre and Fdown such that fdown ∈ Fdown and data (xdown, ydown) ∼ Pfdown◦h which is the real underlying distribution, the worst-case representation difference is similar to that in multi-task analysis (Tripuraneni et al., 2020b, Lemma 1):
dFdown(ĥ;h) = sup fdown∈Fdown inf f ′∈Fdown
E { `(f ′ ◦ ĥ(xdown), ydown)− `(fdown ◦ h(xdown), ydown) } ≤ sup ‖αs‖≤c0 inf ‖α′s‖≤c0 1 2 EXdown
∥∥∥α′>ĥ(xdown)− α>h(xdown)∥∥∥2 , here s ∈ [k′ − 1] =
k′−1∑ s=1 sup ‖αs‖≤c0 inf ‖α′s‖≤c0 1 2 EXdown ( α′ > s ĥ(x down)− α>s h(xdown) )2
≤ (k′ − 1)c 2 0
2 σ1(Λsc(ĥ, h)).
Here in the last line, the inner infima is considered as the partial minimization of a convex quadratic form (see (Boyd & Vandenberghe, 2004, Example 3.15, Appendix A.5.4)).
Define population covariance if representations ĥ and h as
Λ(ĥ, h) = [ E[ĥ(x)ĥ(x)>] E[ĥ(x)h(x)>] E[h(x)ĥ(x)>] E[h(x)h(x)>] ] = [ Fĥĥ Fĥh Fhĥ Fhh ] Λsc(ĥ, h) = Fhh − Fhĥ(Fĥĥ) †Fĥh is the generalized Schur complement of h with respect to ĥ.
Controlling the pre-training representation difference with the lower bound is subtler,
dFpre,fpre(ĥ;h)
= inf f ′∈Fpre
Expre,ypre{`(f ′ ◦ ĥ(xpre), ypre)− `(fpre ◦ h(xpre), ypre)}
= inf ‖α′s‖≤c1
Expre,ypre [ `(f ′ ◦ ĥ(xpre), ypre)− `(fpre ◦ h(xpre), ypre) ] , s ∈ [k − 1]
≥ c0EXpre [ exp(−10 max(‖α′>ĥ(xpre)‖, ‖α>h(xpre)‖)) · ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2] .
We lower bound each term in the sum over j identically and suppress the j for each of notation. Define event E as
1[E] = 1  |Z s 1 | = |(α>Bxpre)s| ≤ Cc′‖αs‖, s = 1, 2, · · · , k − 1,
|Zs2 | = |(α′ > B̂xpre)s| ≤ Cc′‖α′s‖, s = 1, 2, · · · , k − 1.  , where Zs1 ∼ subG(‖αs‖2C2) and Zs2 ∼ subG(‖α′s‖2C2). Then we use this event to bound the diversity,
EXpre exp(−10 max(‖α′ > ĥ(xpre)‖, ‖α>h(xpre)‖)) · ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2 ≥EXpre [ 1(E) exp(−10 max(‖α′>ĥ(xpre)‖, ‖α>h(xpre)‖)) ·
∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2] ≥e−10Cc ′c1· √ k−1)EXpre ( 1(E) ·
∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2) . First write Expre ( 1(E) ·
∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2) as Expre
(∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2)− Expre (1(Ec) · ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2) . Observe that for the last term, it holds
Expre ( 1(Ec) · ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2)
= k−1∑ s=1 Expre ( 1(Ec) · [ α′ > ĥ(xpre)− α>h(xpre) ]2 s )
≤ k−1∑ s=1 √ P [Ec] √ Expre [ α′>ĥ(xpre)− α>h(xpre) ]4 s
Using union bound we have that P [Ec] ≤ 2(k−1) c′2 . Using L4-L2 hyper-contractivity we also conclude√ Expre ( α′>ĥ(xpre)− α>h(xpre) )4 s ≤ 10σ2 = 10Expre ( α′ > ĥ(xpre)− α>h(xpre) )2 s . Combining these we know
Expre [ 1(Ec) ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2] ≤ 10√2(k − 1)
c′ Expre ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2 . Hence this metric could be claimed to be lower bounded as,
Ω ( inf α′ Expre ∥∥∥α′>ĥ(xpre)− α>h(xpre)∥∥∥2)
= Ω ( α>1 Λsc(ĥ;h)α1 ) = Ω ( tr(Λsc(ĥ;h)C) ) , where C = α1α>1 .
In the second line, we redefine α1 as parameter α of pre-training for clarity. In this way we conclude that,
dFpre,fpre(ĥ;h) = Ω ( tr(Λsc(ĥ, h)C) ) = Ω ( σ1(Λsc(ĥ, h))σr(C) ) ,
where C implies expansion of representation h(x) ∈ Rr, and its condition number σr(C) indicates how spread out this vector is in Rr:
C = k−1∑ s=1 (α1)s(α1) > s = α1α > 1 , α1 ∈ Rr×(k−1)
Aforementioned calculations show
dFdown(ĥ;h) ≤ 1
Ω(ν̃) dFpre,fpre(ĥ;h), ν̃ = σr(C).
Proof is completed.

heading: A.4 PROOFS FOR LINEAR SUBSPACE REPRESENTATION
text: Proof. We begin with bounding each of the complexity terms in the Corrolary 1.
We make use of data-dependent inequalities (Tripuraneni et al., 2020b, Lemma 4) to help upper bound related quantities. Intuitively Definition 5 implies tail-bound properties in a sub-gaussian process.
•
Ĝn(H) = 1
n E [ sup B∈H r∑ k=1 n∑ i=1 gkib > k x pre i ]
= O
(√ dr2
n ) •
Ĝn(Fpre|h ◦ xpre) = 1
n E
[ sup
α1,··· ,αk−1 k−1∑ s=1 n∑ i=1 gisα > s B >xprei
]
= c1(k − 1)
n E‖ n∑ i=1 gisB >xprei ‖
= c1(k − 1)√
n
√ tr(B>ΣB)
then Ḡn(Fpre) ≤ O ( (k − 1) √ r n ) .
• Similarly,
Ĝm(Fdown|h ◦ xdown) ≤ c1(k ′ − 1)√ m √√√√ r∑ i=1 σi(B̂>ΣB̂)
then Ḡm(Fdown) ≤ O ( (k′ − 1) √ r m ) .
• boundedness parameter DXpre = supα,B ‖α>B>x‖ = O( √ k − 1D)
• cross entropy `(η; y) = −y>η + log (1 + ∑k−1 s=1 e ηs), then ∣∣∣∂`(η;y)∂ηi ∣∣∣ = ∣∣∣yi − eηi1+∑k−1s=1 eηs ∣∣∣,
|∇η`(η; y)| ≤ √ k − 1, so it is Lpre = √ k − 1−Lipschitz in its first coordinate uniformly over its second for pre-training and Ldown = √ k′ − 1−Lipschitz for downstream task. • |`(η; y)| ≤ O(‖η‖), where ‖η‖ = ‖x>Bpreα‖ . √ k − 1D. This means Bpre =
O( √ k − 1D), Bdown = O( √ k′ − 1D).
In Corollary 1, we define and compute the maximal variance term σ2 as,
σ2 = 1
m sup
fdown∈Fdown m∑ i=1 V ar(`′(fdown ◦ ĥ(xdowni ), ydowni ))
≤ k ′ − 1 m sup fdown∈Fdown m∑ i=1 V ar(fdown ◦ ĥ(xdowni ))
= k′ − 1 m
sup ‖αs‖≤O(1) k′−1∑ s=1 m∑ i=1 V ar(α>s B̂ >xdowni )
= (k′ − 1)2
m sup
‖αs‖≤O(1) m∑ i=1 (αsB̂) >ΣB̂αs
= (k′ − 1)2O(‖B̂ΣB̂‖2) = O ( (k′ − 1)2 )
With these results in hand, we are now prepared to apply Corollary 1, w.p. at least 1− δ
Rdown(f̂ down, ĥ)−Rdown(fdown, h)
≤ dFdown(ĥ;h) + 4 √ πLdown · EXdownĜm ( Fdown ∣∣∣ĥ ◦ xdown)+ 4σ√ log (2/δ) m + 50Bdown log (2/δ) m
where Ĝm(Fdown|ĥ ◦ xdown) is defined in Theorem 2. Thus Ldown · EXdownĜm ( Fdown ∣∣∣ĥ ◦ xdown) ≤ LdownḠm(Fdown) ≤ O((k′ − 1) 32√ rm ), σ ≤ O(k′−1), andBdown ≤ O( √ k′ − 1D). Further, we obtain upper bound of worst-case representation difference by diversity parameter and adoption of Theorem 5: w.p. at least 1− δ
dFdown(ĥ;h)
≤ dF pre,fpre(ĥ;h)
ν
≤ 1 ν
{ 4096L [ log (n) · [L(Fpre) ·Gn(H) + Ḡn(Fpre)] + √ k − 1DXpre
n2
] + 4B √ log (2/δ)
n
}
.
√ k
ν
{ log (n) (√ kdr2
n + k
√ r
n
) +D ( k
n2 +
√ log (1/δ)
n )} The last thing to consider for completing the proof for Theorem 3 is giving accurate characterization of diversity parameter ν, which we leave for the next subsection.

heading: A.5 PROOFS FOR DEEP NEURAL NETWORK REPRESENTATION
text: Proof. In deep neural network, we first review complexity quantities. Adapted from Theorem 8 (Golowich et al., 2018), we have
R̂n(N ) ≤ ( 2
n ΠKp=1M(p) )√√√√(K + 1 + log d) ·max j∈[d] n∑ i=1 x2i,j
≤ 2D √ K + 1 + log d ·ΠKp=1M(p)√
n .
where xi,j denotes the j-th coordinate of vector xi.
Then we proceed to bound the Gaussian complexities for our deep neural network and prove Theorem 4. Recall that under the conditions of the result we can use former results to verify the task diversity condition is satisfied with parameters Ω(ν̃) with ν̃ = σr(α1α>1 ) > 0. We can see that ‖Ex[ĥ(x)h∗(x)>‖2 ≤ Ex‖ĥ(x)h∗(x)‖ ≤ O(M(K)2) using the norm bound from. Hence under this setting we can choose c1 sufficiently large so that c1M(K)2 & M(K)2
c c2. The condition M(K) & 1 in the theorem statement is simply used to clean up the final bound.
In order to instantiate Theorem 1 we begin by bounding each of the complexity terms in the expression.
• For the feature learning complexity in the training phase, we leverage above results, then
Ĝn(H) = 1
n E[sup WK r∑ k=1 n∑ i=1 gkihk(x pre i )] ≤ r∑ k=1 Ĝn(hk(x pre i ))
≤ log(n) · r∑
k=1
R̂nhk(x pre i ) ≤ r log(n)
2D √ K + 1 + log d ·ΠKp=1M(p)√
n .
This also implies the population Gaussian complexity.
• By definition the class F as linear maps with parameters ‖αs‖2 ≤ c1M(K)2,∀s ∈ [k − 1], we obtain that L(F) = c1 √ k − 1M(K)2.
• For the complexity of learning Fpre in the training phase we obtain,
Ĝn(Fpre|h ◦ xpre) = 1
n Eg[sup
α∈F k−1∑ s=1 n∑ i=1 gisα > s h(x pre i )] . (k − 1)M(K)2 n Eg[‖ n∑ i=1 gish(x1i)‖]
. (k − 1)M(K)2
n
√√√√ n∑ i=1 ‖h(xprei )‖2 . (k − 1)M(K)2√ n max i ‖h(xprei )‖.
For tanh activation function, we simply have
‖h(x)‖2 = ‖WKrK−1‖22 ≤ ‖WK‖2∞→2,
where rK−1 denotes ourput of the K − 1th layer,
‖h(x)‖ ≤ O(M(K)).
In conclusion we obtain Ḡn(Fpre) ≤ O (
(k − 1)M(K)3√ n
) .
• Similarly
Ĝm(Fdown|h ◦ xdown) ≤ O (
(k′ − 1)M(K)3√ m ) Regularity conditions
• Boundedness parameter DXpre = supα,h ‖α>h(xpre)‖ = O( √ k − 1M(K)3). • Pre-training loss is Lpre = √ k − 1-Lipschitz and Bpre = O( √ k − 1M(K)3)-bounded. • Downstream loss is Ldown = √ k′ − 1-Lipschitz and Bdown = O( √ k′ − 1M(K)3)-
bounded.
Assembling the previous complexity arguments shows the transfer learning risk is bounded by
. Lpre
ν̃
( log(n) [ L(Fpre)r log(n) D √ KΠKp=1M(p)√
n + kM(K)3√ n
]) + Ldownk′M(K)3√
m
+
( 1
ν̃ max
( Lpre √ kDXpre
n2 , Bpre
√ log(1/δ)
n
) +Bdown √ log(1/δ)
m ) Substitute regularity conditions into it, then the risk is simplified as stated in Theorem 4.

heading: A.6 EXPERIMENT DETAILS
text: Full statistics (including matched and mismatched dev sets for MNLI, accuracy and F1 scores for MRPC and QQP, and (Pearson’s correlation + Spearman’s correlation)/2 for STS-B. All other tasks uses accuracy as the metric) could be found in Table A.6.

heading: A.6.1 COMPLETE STATISTICS
text: Here we provide complete results on GLUE dev sets over 5 random seeds.

heading: A.6.2 HYPERPARAMETERS
text: Pre-training Hyperparameters for pre-training are shown in Table 7.

heading: Hyperparam phase-1 phase-2
text: Finetuning Hyperparameters for downstream tasks are shown in Table 8. We adapt these hyperparameters from Liu et al. (2020), Devlin et al. (2019), and Yang et al. (2019).

heading: A.7 OTHER DETAILS
text: Computing infrastructure We pretrain our (diversity-regularized) BERT-base models using 32 Nvidia V100 GPUs (32GB RAM each), and the finetuning of the model uses 4 Nvidia V100 GPUs.

Review summaries:
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 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. 

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.