# Blessing of Class Diversity in Pre-training

Yulai Zhao
Princeton University
<EMAIL_ADDRESS>Jianshu Chen
Tencent AI Research
<EMAIL_ADDRESS>Simon S. Du
University of Washington & Facebook AI Research
<EMAIL_ADDRESS>
###### 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 the masked language model task) are sufficiently
diverse, in the sense that the least singular value of the last linear layer
in pre-training (denoted as $\tilde{\nu}$) is large, then pre-training can
significantly improve the sample efficiency of downstream tasks. Specially, we
show the transfer learning excess risk enjoys an
$O\left(\frac{1}{\tilde{\nu}\sqrt{n}}\right)$ rate, in contrast to the
$O\left(\frac{1}{\sqrt{m}}\right)$ rate in the standard supervised learning.
Here, $n$ is the number of pre-training data and $m$ is the number of data in
the downstream task, and typically $n\gg m$. Our proof relies on a vector-form
Rademacher complexity chain rule for disassembling composite function classes
and a modified self-concordance condition. These techniques can be of
independent interest.

##  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, the pre-training technique has revolutionized natural language
processing (NLP). Models based on this technique have dramatically improved
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., 2021), 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., 2021, 2020;
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 allowing 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 applies to this setting.

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.111This is a standard
setting in the BERT model (Devlin et al., 2019) and is widely adopted as a
common practice. By breaking down the English words into BPE sub-word units,
it could drastically increase the coverage of the English language by using a
relatively small (32768) vocabulary. 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 a similar role as the _diversity of
tasks_ in multi-task pre-training theory (Du et al., 2021; Tripuraneni et al.,
2021). We summarize our contributions below.

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. We give concrete bounds on
linear representation and deep network neural networks to showcase our general
theoretical results. 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 learning (Du et al., 2021; Tripuraneni et al., 2020) builds on scalar
output and thus could not apply to multi-class tasks (e.g., cross-entropy
loss). We introduce a vector-form Rademacher 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. (2020). 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.

#### 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 4.2 then
instantiate it with several settings. We conclude in Section 6 and defer the
proofs to Appendix A. In Section 5, we present some preliminary empirical
results on how our theory inspires new regularization techniques.

##  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., 2021), which is a different setting
considered in this paper. The most relevant one is by Lee et al. (2021) 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 the 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 relates 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 theoretical work studied multi-task pre-training (Du et al.,
2021; Tripuraneni et al., 2021, 2020; Thekumparampil et al., 2021) in which a
notion, diversity, has been identified to be the key that enables pre-training
to improve statistical efficiency. Experiments also supported the idea that
increasing the diversity of the training data helps generalization (Zhang et
al., 2022).

Theories on multi-task pre-training generally require a large number of
diverse tasks, and 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 (see Remark 6.2 therein). 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.

##  3 Preliminaries

In this section, we introduce the necessary notations, the problem setup, and
several model-dependent quantities used in pre-training and downstream task
learning.

###  3.1 Notation and Setup

Notations. Let $[n]=\\{1,2,\cdots,n\\}$. We use $\|\cdot\|$ or $\|\cdot\|_{2}$
to denote the $\ell_{2}$ norm of a vector. Let ${\mathcal{N}}(\mu,\sigma_{2})$
be the one-dimensional Gaussian distribution. For a matrix
$\mathbf{W}\in{\mathbb{R}}^{m\times n}$, let
$\|\mathbf{W}\|_{1,\infty}=\max_{q}(\sum_{p}|\mathbf{W}_{q,p}|)$ and
$\|\mathbf{W}\|_{\infty\to 2}$ be the induced $\infty$-to-$2$ operator norm.
We use the standard $O(\cdot),\Omega(\cdot)$ and $\Theta(\cdot)$ notation to
hide universal constant factors, and use $\widetilde{O}(\cdot)$ to hide
logarithmic factors. We also use $a\lesssim b$ to indicate $a=O(b)$.

Problem Setup. The procedure is divided into two stages: the 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 $\hat{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,
$\\{x_{i}^{\mathrm{p}},y_{i}^{\mathrm{p}}\\}_{i=1}^{n}$, where
$x_{i}^{\mathrm{p}}\in{\mathcal{X}}^{\mathrm{p}}\subset\mathbb{R}^{d}$ is the
input and $y_{i}^{\mathrm{p}}\in\\{0,1\\}^{k-1}$ is the one-hot label for
$k$-class classification (if $y_{i}^{\mathrm{p}}$ is all-zero then it
represents the $k$-th class).222We assume only one pre-training task for the
ease of presentation. It is straightforward to generalize our results to
multiple pre-training tasks. 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.333Here 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. $k$ in this example is the size of the
vocabulary ($\approx 30K$). We aim to obtain a good representation function
$\hat{h}$ within a function class
${\mathcal{H}}\subset\\{\mathbb{R}^{d}\rightarrow\mathbb{R}^{r}\\}$ 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
$\hat{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 $f^{\mathrm{p}}$ within function class
$\mathcal{F}^{\mathrm{p}}\subset\\{\mathbb{R}^{r}\rightarrow\mathbb{R}^{k-1}\\}$.

To train the representation function and predictor in pre-training stage, we
consider the Empirical Risk Minimization (ERM) procedure

$\displaystyle\hat{h}=\operatorname*{arg\,min}_{h\in{\mathcal{H}}}\min_{f^{\mathrm{p}}\in{\mathcal{F}}^{\mathrm{p}}}\hat{R}_{\mathrm{p}}(f^{\mathrm{p}},h)\triangleq\operatorname*{arg\,min}_{h\in{\mathcal{H}}}\min_{f^{\mathrm{p}}\in{\mathcal{F}}^{\mathrm{p}}}\frac{1}{n}\sum_{i=1}^{n}\ell(f^{\mathrm{p}}\circ
h(x_{i}^{\mathrm{p}}),y_{i}^{\mathrm{p}})$

where $\ell$ is the loss function. We overload the notation for both the pre-
training task and the downstream task, i.e., for pre-training,
$\ell:\mathbb{R}^{k-1}\times\\{0,1\\}^{k-1}\rightarrow\mathbb{R}$ and for the
downstream task,
$\ell:\mathbb{R}^{k^{\prime}-1}\times\\{0,1\\}^{k^{\prime}-1}\rightarrow\mathbb{R}$.
e.g., cross-entropy:
$\ell(\hat{y};y)=-y^{\top}\hat{y}+\log{(1+\sum_{s=1}^{k-1}\exp\left(\hat{y}_{s}\right))}$.

Now for the downstream task, we assume there are $m$ samples
$\\{x_{i}^{\mathrm{d}},y_{i}^{\mathrm{d}}\\}_{i=1}^{m}$. Note that
$x_{i}^{\mathrm{d}}\in{\mathcal{X}}^{\mathrm{d}}\subset\mathbb{R}^{d}$ is the
input and $y_{i}^{\mathrm{d}}\in\\{0,1\\}^{k^{\prime}-1}$ is the one-hot label
for $k^{\prime}$-class classification.444For 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\gg m$ and $k\gg k^{\prime}$. For example, in sentiment analysis,
$k^{\prime}=2$ (“positive” or “negative”). A widely studied task SST-2 (Wang
et al., 2019) has $m\approx 67K$, which is also generally much smaller than
the pre-training corpus (e.g., $n>100$M samples).

To train the classifier for the downstream task, we fix the representation
function learned from the pre-training task and train the task-dependent
predictor within function class
$\mathcal{F}^{\mathrm{d}}\subset\\{\mathbb{R}^{r}\rightarrow\mathbb{R}^{k^{\prime}-1}\\}$:

$\displaystyle\hat{f}^{\mathrm{d}}=\operatorname*{arg\,min}_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\hat{R}_{\mathrm{d}}(f^{\mathrm{d}},\hat{h})=\operatorname*{arg\,min}_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\frac{1}{m}\sum_{i=1}^{m}\ell(f^{\mathrm{d}}\circ\hat{h}(x^{\mathrm{d}}_{i}),y^{\mathrm{d}}_{i}).$

Therefore, our predictor for the downstream task consists a pair
$(\hat{f}^{\mathrm{d}},\hat{h})$. We use the risk defined below to measure the
performance

$\displaystyle\text{Transfer~{}Learning~{}Risk}\triangleq
R_{\mathrm{d}}(\hat{f}^{\mathrm{d}},\hat{h})-\mathop{\mathbb{E}}\limits_{x^{\mathrm{d}},y^{\mathrm{d}}}\left[\ell\left(g^{\mathrm{d}}\left(x^{\mathrm{d}}\right),y^{\mathrm{d}}\right)\right]$

where we define

$R_{\mathrm{d}}(\hat{f}^{\mathrm{d}},\hat{h})\triangleq\mathop{\mathbb{E}}\limits_{x^{\mathrm{d}},y^{\mathrm{d}}}\left[\ell\left(\hat{f}^{\mathrm{d}}\circ\hat{h}\left(x^{\mathrm{d}}\right),y^{\mathrm{d}}\right)\right]$

as the expected loss (the expectation is over the distribution of the
downstream task), and

$g^{\mathrm{d}}=\operatorname*{arg\,min}_{g\in\\{\mathbb{R}^{d}\rightarrow\mathbb{R}^{k^{\prime}-1}\\}}\mathbb{E}_{x^{\mathrm{d}},y^{\mathrm{d}}}\left[\ell\left(g\left(x^{\mathrm{d}}\right),y^{\mathrm{d}}\right)\right]$

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

$\displaystyle\text{Pre-training~{}Risk}\triangleq
R_{\mathrm{p}}(\hat{f}^{\mathrm{p}},\hat{h})-\mathop{\mathbb{E}}\limits_{x^{\mathrm{p}},y^{\mathrm{p}}}\left[\ell\left(g^{\mathrm{p}}\left(x^{\mathrm{p}}\right),y^{\mathrm{p}}\right)\right],$

where

$R_{\mathrm{p}}(\hat{f}^{\mathrm{p}},\hat{h})\triangleq\mathop{\mathbb{E}}\limits_{x^{\mathrm{p}},y^{\mathrm{p}}}\left[\ell\left(\hat{f}^{\mathrm{p}}\circ\hat{h}\left(x^{\mathrm{p}}\right),y^{\mathrm{p}}\right)\right]$

is the expected loss, and

$g^{\mathrm{p}}=\operatorname*{arg\,min}_{g\in\\{\mathbb{R}^{d}\rightarrow\mathbb{R}^{k-1}\\}}\mathbb{E}_{x^{\mathrm{p}},y^{\mathrm{p}}}\left[\ell\left(g\left(x^{\mathrm{p}}\right),y^{\mathrm{p}}\right)\right]$

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., 2020), throughout the paper, we make the
following realizability assumption, which is also a standard assumption in the
classical PAC learning framework (Shalev-Shwartz and Ben-David, 2014).

######  Assumption 3.1 (Realizability).

There exist $h\in{\mathcal{H}}$,
$f^{\mathrm{p}}\in{\mathcal{F}}^{\mathrm{p}}$,
$f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}$ such that
$g^{\mathrm{p}}=f^{\mathrm{p}}\circ h$ and $g^{\mathrm{d}}=f^{\mathrm{d}}\circ
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.

As for the setting that is of most interest to NLP pre-training, where the
loss function ‘ is cross-entropy, ${\mathcal{F}}^{\mathrm{p}}$ and
${\mathcal{F}}^{\mathrm{d}}$ are sets of linear functions, we make the
following assumption on both pre-training and downstream tasks to describe how
the underlying data are generated.

######  Assumption 3.2 (Multinomial Logistic Data).

For a $K$-class classification task with $q$ samples,
$\\{x_{i},y_{i}\\}_{i=1}^{q}$, where $x_{i}\in{\mathcal{X}}$ is the input and
$y_{i}\in\\{0,1\\}^{K-1}$ is the one-hot label. Let $f$ and $h$ be the true
underlying predictor layer and representation function. Then the output is
$f\circ h(x)\in{\mathbb{R}}^{K-1}$. Assume each label $\\{y\\}_{i}$ is
generated from a conditional distribution of a multinomial logistic regression
model: $y\sim{\mathcal{P}}(\cdot|f\circ h(x))$,

$\displaystyle{\mathcal{P}}(y|f\circ h(x))=e^{y^{\top}f\circ h(x)-\Phi(f\circ
h(x))}$

where $\Phi(x)=\log{(1+\sum_{s=1}^{K-1}e^{x_{s}}}),x\in{\mathbb{R}}^{K-1}$ and
$y$ is an one-hot label.

######  Remark 3.3.

It is straight forward to see that ${\mathcal{P}}(y|f\circ h(x))$ is
normalized to 1.

Intuitively, the assumption states that the data used for classification
follow a multinomial logistic regression structure.

###  3.2 Task Relatedness and Diversity

We use the following definitions, which are natural analogies of those in
Tripuraneni et al. (2020) 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.
(2020)

######  Definition 3.4.

Let $h\in{\mathcal{H}}$ be the optimal representation function and
$h^{\prime}\in{\mathcal{H}}$ be any representation function. Let
$f^{\mathrm{p}}\in{\mathcal{F}}$ be the optimal pre-training predictor on top
of $h$. The pre-training representation difference is defined as:

$\displaystyle
d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(h^{\prime};h)=\inf_{f^{\prime}\in{\mathcal{F}}^{\mathrm{p}}}\mathop{\mathbb{E}}\limits_{x^{\mathrm{p}},y^{\mathrm{p}}}\left[\ell(f^{\prime}\circ
h^{\prime}(x^{\mathrm{p}}),y^{\mathrm{p}})-\ell(f^{\mathrm{p}}\circ
h(x^{\mathrm{p}}),y^{\mathrm{p}})\right]$

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^{\prime}$.

For transfer learning, we also need to introduce a similar concept on the
downstream task.

######  Definition 3.5.

Let $h\in{\mathcal{H}}$ be the optimal representation function and
$h^{\prime}\in{\mathcal{H}}$ be any representation function. For the
downstream task, for a function class ${\mathcal{F}}^{\mathrm{d}}$, let
$f^{\mathrm{d}}\in{\mathcal{F}}$ be the optimal pre-training predictor on top
of $h$. We define the worst-case representation difference between
representations $h$ and $h^{\prime}\in{\mathcal{H}}$ as:

$\displaystyle
d_{{\mathcal{F}}^{\mathrm{d}}}(h^{\prime};h)=\sup_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\inf_{f^{\prime}\in{\mathcal{F}}^{\mathrm{d}}}\mathop{\mathbb{E}}\limits_{x^{\mathrm{d}},y^{\mathrm{d}}}\left[\ell(f^{\prime}\circ
h^{\prime}(x^{\mathrm{d}}),y^{\mathrm{d}})-\ell(f^{\mathrm{d}}\circ
h(x^{\mathrm{d}}),y^{\mathrm{d}})\right]$

where the expectation is over the data distribution of the downstream task.

We now introduce the key notion of _diversity_ , which measures how well a
learned representation, say $h^{\prime}$, from the pre-training task can be
transferred to the downstream task.

######  Definition 3.6.

Let $h\in{\mathcal{H}}$ be the optimal representation function. Let
$f^{\mathrm{p}}\in{\mathcal{F}}^{\mathrm{p}}$ be the optimal pre-training
predictor on top of $h$. The diversity parameter $\nu>0$ is the largest
constant that satisfies

$\displaystyle
d_{{\mathcal{F}}^{\mathrm{d}}}(h^{\prime};h)\leq\frac{d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(h^{\prime};h)}{\nu},\forall
h^{\prime}\in{\mathcal{H}}.$ (1)

While Definition 3.4\- 3.6 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 ${\mathcal{F}}^{\mathrm{p}}$ and
${\mathcal{F}}^{\mathrm{d}}$ are linear function classes.

###  3.3 Model Complexities

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 ${\mathcal{Q}}$
can be correlated with a noise sequence of length $n\times r$.

######  Definition 3.7 (Gaussian Complexity).

Let $\mu$ be a probability distribution on a set
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ and suppose that $x_{1},\cdots,x_{n}$
are independent samples selected according to $\mu$. Let ${\mathcal{Q}}$ be a
class of functions mapping from ${\mathcal{X}}$ to ${\mathbb{R}}^{r}$. Define
random variable

$\hat{G}_{n}({\mathcal{Q}})=\mathop{\mathbb{E}}\limits_{g_{ki}\sim{\mathcal{N}}(0,1)}\left[\sup_{q\in{\mathcal{Q}}}\frac{1}{n}\sum_{k=1}^{r}\sum_{i=1}^{n}g_{ki}q_{k}(x_{i})\right]\qquad$
(2)

as the empirical Rademacher complexity, where $q_{k}(x_{i})$ is the $k$-th
coordinate of the vector-valued function $q(x_{i})$, $g_{ki}$
$(k\in[r],i\in[n])$ are independent standard normal random variables. The
Gaussian complexity of ${\mathcal{Q}}$ is
$G_{n}({\mathcal{Q}})=E_{\mu}\hat{G}_{n}({\mathcal{Q}})$.

Our main results are stated in terms of the Gaussian complexity. In Section
4.4 and 4.5 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 ${\mathcal{F}}^{\mathrm{p}}$

$\displaystyle\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})=\max_{h(x_{1}),\cdots,h(x_{n})}\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}|h\circ
x^{\mathrm{p}}),$ (3)

here $h\in{\mathcal{H}}$ and
$x^{\mathrm{p}}=x_{1},\cdots,x_{n}\in{\mathcal{X}}^{\mathrm{p}}$. Similarly we
define $\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})$ as the worst-case Gaussian
complexity for the downstream predictor within ${\mathcal{F}}^{\mathrm{d}}$.

We note that a closely related notion is Rademacher complexity. The empirical
Rademacher complexity and Gaussian complexity only differ by a log factor
(Ledoux and Talagrand, 1991). We use Gaussian complexity in this work for the
benign properties brought by Gaussian distribution.

##  4 Main Results

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 ${\mathcal{F}}^{\mathrm{p}}$ and
${\mathcal{F}}^{\mathrm{d}}$ 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.4 and 4.5 we instantiate our bounds on two concrete representation
function classes: linear subspace and multi-layer network to showcase our main
results.

###  4.1 Main Theorem

In this subsection, we present our generic end-to-end transfer learning
guarantee for multi-class transfer learning problems. 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 4.1 (Regularity Conditions).

We assume the following regularity conditions hold:

  * •

In pre-training, $\ell(\cdot,\cdot)$ is $B^{\mathrm{p}}$-bounded, and
$\ell(\cdot,y)$ is $L^{\mathrm{p}}$-Lipschitz for all $y$.

  * •

In downstream task, $\ell(\cdot,y)$ is $B^{\mathrm{d}}$-bounded and
$L^{\mathrm{d}}$-Lipschitz for all $y$.

  * •

Any predictor $f\in{\mathcal{F}}^{\mathrm{p}}$ is
$L({\mathcal{F}}^{\mathrm{p}})$-Lipschitz with respect to the Euclidean
distance.

  * •

Predictors are bounded: $\|f\circ h(x)\|\leq D_{{\mathcal{X}}^{\mathrm{p}}}$
for any
$x\in{\mathcal{X}}^{\mathrm{p}},h\in{\mathcal{H}},f\in{\mathcal{F}}^{\mathrm{p}}$.
Similarly $\|f\circ h(x)\|\leq D_{{\mathcal{X}}^{\mathrm{d}}}$ for any
$x\in{\mathcal{X}}^{\mathrm{d}},h\in{\mathcal{H}},f\in{\mathcal{F}}^{\mathrm{d}}$.

Specifically, one can show that common task-dependent losses satisfy these
conditions. For example, when $\ell$ is the cross-entropy loss for $k-$class
classification (cf. Section 4.2), we prove that $\ell$ is
$\sqrt{k-1}-$Lipschitz and $D_{\mathcal{X}}-$bounded where ${\mathcal{X}}$
denotes the input data domain.

Under these assumptions, we have the following quantitative guarantee.

######  Theorem 4.2.

Under Assumption 3.1 and 4.1, for a given fixed failure probability $\delta$,
with probability at least $1-\delta$ we have the Transfer Learning Risk upper
bounded by:

$\displaystyle
O\Bigg{(}\frac{1}{\nu}\Bigg{\\{}L^{\mathrm{p}}\Bigg{[}\log{(n)}[L({\mathcal{F}}^{\mathrm{p}})G_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})]+\frac{\sqrt{k}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}\Bigg{]}+B^{\mathrm{p}}\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{n}}\Bigg{\\}}$
$\displaystyle\qquad\qquad+L^{\mathrm{d}}\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})+B^{\mathrm{d}}\sqrt{\frac{\log(\nicefrac{{1}}{{\delta}})}{m}}\Bigg{)}.$

The first line comes from the pre-training ERM procedure and it accounts for
the error of using an approximate optimal representation $\hat{h}\approx h$.
The second line characterizes the statistical error of learning the
downstream-task predictor $f^{\mathrm{d}}$ from $m$ samples. Note the
diversity parameter appears in the denominator, which relates the pre-training
risk to the transfer learning risk. Theorem 4.2 shows the risk would be small
if the Gaussian complexities are small. We expect that
$G_{n}({\mathcal{H}})\gg\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})$ since
${\mathcal{H}}$ is often expressive representation functions, while
${\mathcal{F}}^{\mathrm{d}}$ is linear classifiers generally. We will show
concrete examples where $G_{n}({\mathcal{H}})$ and
$\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})$ are $O(\sqrt{1/n})$ and
$\bar{G}_{m}$ scales as $O(\sqrt{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., 2020) focus on scalar output. Their results
cannot be applied to multi-class transfer learning tasks. In Theorem 4.2, we
generalize the analysis in (Tripuraneni et al., 2020) to handle multi-class
classification where the output is high dimensional (number of classes).
Technically, in the proof, we introduce a vector-form Rademacher complexity
chain rule for disassembling composite function classes by making use of the
vector-form Rademacher contraction property (Maurer, 2016).

###  4.2 Multi-class Classification with Cross-entropy Loss

Now we specialize the general results to the setting that is of most interest
to NLP pre-training, where the loss function $\ell$ is cross-entropy and the
${\mathcal{F}}^{\mathrm{p}}$ and ${\mathcal{F}}^{\mathrm{d}}$ 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 class ${\mathcal{H}}$ and uses word-embedding matrices as
${\mathcal{F}}^{\mathrm{p}}$.

Formally we define

$\displaystyle{\mathcal{F}}^{\mathrm{p}}$
$\displaystyle=\\{f|f(z)=\alpha^{\top}z,\alpha\in{\mathbb{R}}^{r\times(k-1)},\|\alpha_{s}\|\leq
c_{1}\text{ for all }s\in[k-1],\|\alpha^{\top}z\|\leq c_{2}\\}$
$\displaystyle{\mathcal{F}}^{\mathrm{d}}$
$\displaystyle=\\{f|f(z)=\alpha^{\top}z,\alpha\in{\mathbb{R}}^{r\times(k^{\prime}-1)},\|\alpha_{s}\|\leq
c_{0}\text{ for all }s\in[k^{\prime}-1],\|\alpha^{\top}z\|\leq c_{3}\\}$

where $c_{0},c_{1},c_{2}$ and $c_{3}$ are some universal constants. Then the
regularity conditions are instantiated as:

  * •

Pre-training loss $\ell(\cdot,y)$ is $\sqrt{k-1}$-Lipschitz and
$B^{\mathrm{p}}=D_{{\mathcal{X}}^{\mathrm{p}}}$-bounded.

  * •

Downstream loss is $\sqrt{k^{\prime}-1}$-Lipschitz and
$B^{\mathrm{d}}=D_{{\mathcal{X}}^{\mathrm{d}}}$-bounded.

  * •

Any function $f\in{\mathcal{F}}^{\mathrm{p}}$ is
$L({\mathcal{F}}^{\mathrm{p}})=c_{1}\sqrt{k-1}$-Lipschitz w.r.t. the
$\ell_{2}$ distance.

Next, we discuss our main assumption that relates the diversity parameter to a
concrete quantity of the last linear layer.

######  Assumption 4.3 (Lower Bounded Least Eigenvalue).

Let the optimal linear predictor at the last layer for pre-training be
$\alpha^{\mathrm{p}}\in\mathbb{R}^{r\times(k-1)}$, assume
$\tilde{\nu}\triangleq\sigma_{r}(\alpha^{\mathrm{p}}\left(\alpha^{\mathrm{p}}\right)^{\top})>0$.

Similar assumptions have been used in multi-task representation learning (Du
et al., 2021; Tripuraneni et al., 2021, 2020), 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. We provide
proof sketches on how $\tilde{\nu}$ serves as a lower bound for the diversity
parameter $\nu$ (cf. Lemma 4.5) in Section 4.3, where we introduce new
techniques for analysis.

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(\cdot)\in{\mathbb{R}}^{r}$. This is quantitatively captured
by the
$\sigma_{r}(\alpha^{\mathrm{p}}\left(\alpha^{\mathrm{p}}\right)^{\top})$,
which measures how spread out these vectors are in ${\mathbb{R}}^{r}$.

We now state our theorem for this specific setting.

######  Theorem 4.4.

Under Assumption 3.1, 3.2, 4.3, with probability at least $1-\delta$ we have
the Transfer Learning Risk upper bounded by:

$\displaystyle
O\Bigg{(}\frac{1}{\tilde{\nu}}\Bigg{\\{}\sqrt{k}\Bigg{[}\log{(n)}[\sqrt{k}G_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})]+\frac{\sqrt{k}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}\Bigg{]}+D_{{\mathcal{X}}^{\mathrm{p}}}\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{n}}\Bigg{\\}}$
$\displaystyle\quad+\sqrt{k^{\prime}}\mathop{\mathbb{E}}\limits_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})+\sigma\sqrt{\frac{\log(\nicefrac{{1}}{{\delta}})}{m}}+D_{{\mathcal{X}}^{\mathrm{d}}}\sqrt{\frac{\log(\nicefrac{{1}}{{\delta}})}{m}}\Bigg{)}$

Here
$\mathop{\mathbb{E}}\limits_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})$ is Gaussian complexity of embeddings $\hat{h}\circ
x^{\mathrm{d}}=\\{\hat{h}(x_{1}),\cdots,\hat{h}(x_{m})|x^{\mathrm{d}}=x_{1},\cdots,x_{m}\in{\mathcal{X}}^{\mathrm{d}}\\}$
where the expectation is over ${\mathcal{X}}^{\mathrm{d}}$, and
$\sigma^{2}=\frac{1}{m}\sup_{f\in{\mathcal{F}}^{\mathrm{d}}}\sum_{i=1}^{m}Var(\ell(f\circ\hat{h}(x_{i}^{\mathrm{d}}),y_{i}^{\mathrm{d}}))$
is the maximal variance over ${\mathcal{F}}^{\mathrm{d}}$.

We remark that in Theorem 4.4, since we specialize to the case where
${\mathcal{F}}^{\mathrm{p}}$ and ${\mathcal{F}}^{\mathrm{d}}$ are sets of
linear functions, we can replace the term
$L^{\mathrm{d}}\cdot\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})$ in Theorem 4.2 by
$(\sqrt{k^{\prime}}\cdot\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})+\sigma\sqrt{\nicefrac{{\log(\nicefrac{{1}}{{\delta}})}}{{m}}})$
by utilizing the functional Bernstein inequality. This improvement can help us
obtain Theorem 4.9. See Appendix A.2 for the full proofs.

Now we discuss the interpretation of Theorem 4.4. Typically,
$\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})$ is much smaller than
$G_{n}({\mathcal{H}})$ because $G_{n}({\mathcal{H}})$ represents the
complexity of the representation function, which is often complex. In
practice, $n$ is often large. Therefore, in the benign case where
$\tilde{\nu}=\Theta\left(k\right)$ (when the condition number of
$\alpha^{\mathrm{p}}$ is $O(1)$), the dominating term will be
$G_{n}({\mathcal{H}})$. As we will show in the following subsections, this
term typically scales as $O(\sqrt{\nicefrac{{1}}{{n}}})$. Together, the
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 $\tilde{\nu}$ is
small, then the bound becomes loose. This is consistent with prior
counterexamples on multi-task pre-training (Maurer et al., 2016; Du et al.,
2021) where the diversity is shown to be necessary.

###  4.3 What is _diversity parameter_ for linear layers?

To prove Theorem 4.4, one of our key technical contributions is to show
$\tilde{\nu}$ serves as a lower bound for the diversity parameter $\nu$, i.e.,

######  Lemma 4.5.

Under Assumption 3.1, 3.2, 4.3, we have

$d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)\leq\frac{1}{\Omega(\tilde{\nu})}d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(\hat{h};h).$
(4)

In this subsection, we show how we apply a _modified self-concordance_
condition for better characterizing multinomial logistic regression (Bach et
al., 2010), thus bridging the gap between Theorem 4.2 and Theorem 4.4. We note
that the proof of this part is very different from the multi-task setting
studied in previous works (Du et al., 2021; Tripuraneni et al., 2020).

Denote $\Phi(x)=\log{(1+\sum_{s=1}^{k-1}e^{x_{s}}})$, for
$x\in{\mathbb{R}}^{k-1}$, which is widely seen in multinomial regression tasks
because the cross-entropy loss is the same as multinomial logistic loss
inherently. For clarity, in this subsection we use $\alpha^{\prime}$ and
$\alpha$ to denote parameters for $\hat{f}^{\mathrm{p}}$ and $f^{\mathrm{p}}$
respectively.

Since multinomial logistic regression is not strongly-convex, we need to find
new techniques that would present benign properties to characterize the local
landscape. Below we introduce a class of convex functions called _modified
self-concordant_ functions, which would be useful in quantitative analysis.

######  Definition 4.6 (Modified Self-concordance).

Suppose $F$: ${\mathbb{R}}^{p}\mapsto{\mathbb{R}}$ is a three times
differentiable convex function such that for some $R>0$, for all
$u,v\in{\mathbb{R}}^{p}$, the function $g:t\mapsto F(u+tv)$ satisfies for all
$t\in{\mathbb{R}}$

$|g^{\prime\prime\prime}(t)|\leq R\|v\|_{2}\times g^{\prime\prime}(t)$ (5)

#### Properties of self-concordance

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\in{\mathbb{R}}^{p},t\in{\mathbb{R}}$,

$\displaystyle F(w+v)$ $\displaystyle\geq
F(w)+vF^{\prime}(w)+\frac{v^{\top}F^{\prime\prime}(w)v}{R^{2}\|v\|_{2}^{2}}\cdot\left(e^{-R\|v\|_{2}}+R\|v\|_{2}-1\right),$
$\displaystyle F(w+v)$ $\displaystyle\leq
F(w)+vF^{\prime}(w)+\frac{v^{\top}F^{\prime\prime}(w)v}{R^{2}\|v\|_{2}^{2}}\cdot\left(e^{R\|v\|_{2}}-R\|v\|_{2}-1\right),$
$\displaystyle e^{-tR\|v\|_{2}}$ $\displaystyle F^{\prime\prime}(w)\preceq
F^{\prime\prime}(w+tv)\preceq e^{tR\|v\|_{2}}F^{\prime\prime}(w).$

The first two inequalities are refined characterizations of Taylor’s
expansion, while the last line presents bounds for Hessian matrix in the sense
of positive semidefiniteness.

We find that multinomial logistic loss satisfies the _modified self-
concordance_ condition with $R=5$.

######  Proposition 4.7.

For all $u,v\in{\mathbb{R}}^{k-1}$, the function $g:t\mapsto\Phi(u+tv)$
satisfies

$|g^{\prime\prime\prime}(t)|\leq 5\|v\|_{2}g^{\prime\prime}(t).$

Equipped with the aforementioned properties of self-concordance, we are
prepared to prove Lemma 4.5 which leads to Theorem 4.4. See Appendix A.2 for
details.

###  4.4 Linear Subspace Representation

Based on cross-entropy loss and linear predictors introduced in Section 4.2,
we further assume the underlying representation is a projection onto a low-
dimensional subspace. For $r\ll d$, let the representation be

$\displaystyle{\mathcal{H}}$
$\displaystyle=\\{h|h(x)=B^{\top}x,B\in{\mathbb{R}}^{d\times r},\text{B is a
matrix with orthonormal columns}\\}$

We require some additional regularity conditions. Following prior work (Du et
al., 2021; Tripuraneni et al., 2020), we assume that $\|x\|\leq D$ and input
data distribution satisfies the following condition.

######  Definition 4.8.

We say the covariate distribution $P_{x}(\cdot)$ is $\Sigma$-sub-gaussian if
for all $v\in{\mathbb{R}}^{d}$,
$\mathbb{E}[\exp(v^{\top}x)]\leq\exp\left(\frac{\|\Sigma^{\nicefrac{{1}}{{2}}}v\|^{2}}{2}\right)$
where the covariance $\Sigma$ further satisfies $\sigma_{max}(\Sigma)\leq C$
and $\sigma_{min}(\Sigma)\geq c\geq 0$ for universal constants $c,C$.

We have the following theorem that guarantees the performance of transfer
learning.

######  Theorem 4.9.

Suppose Assumption 3.1, 3.2, and 4.3 hold, data generation follows Definition
4.8. For a sufficiently large constant $c_{4}$, we assume $n\geq c_{4}d,m\geq
c_{4}r$, $D\leq c_{4}(\min(\sqrt{dr^{2}},\sqrt{rm}))$. Then with probability
at least $1-\delta$, we have the Transfer Learning Risk upper bounded by:

$\displaystyle
O\Bigg{(}\frac{1}{\tilde{\nu}}\Bigg{[}\sqrt{k}\log{(n)}\left(\sqrt{\frac{kdr^{2}}{n}}+k\sqrt{\frac{r}{n}}\right)+\frac{k}{n^{2}}+\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{n}}\Bigg{]}+(k^{\prime})^{\frac{3}{2}}\sqrt{\frac{r}{m}}+k^{\prime}\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{m}}\Bigg{)}$

To interpret this bound, consider the practically relevant scenario where
$k^{\prime}=O(1)$ (e.g., sentiment analysis), $m\ll n$, $k\ll n$ and $r\ll d$,
in the benign case $\tilde{\nu}=\Omega\left(k\right)$, we have the transfer
learning risk
$\widetilde{O}\left(\sqrt{\frac{dr^{2}}{n}}+\sqrt{\frac{r}{m}}\right)$. 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\left(\sqrt{\frac{d}{m}}\right)$. Furthermore,
for the linear representation learning setting, classic minimax bounds present
a standard $\Omega(\sqrt{\frac{d}{m}})$ lower rate, which is also worse than
our upper bound with representation learning (Foster et al., 2018; Abramovich
and Grinshtein, 2018; Barnes and Ozgur, 2019).

###  4.5 Deep Neural Network Representation

In this subsection, we assume the underlying representation function to be a
$\sigma=tanh$-activated neural network, which is often used in practice.
Predictors are still required to be linear functions at the interest of NLP
pre-training, i.e.,

$\displaystyle{\mathcal{H}}$
$\displaystyle=\\{h|h(x)=W_{K}\sigma\left(W_{K-1}\sigma\left(\cdots\sigma\left(W_{1}x\right)\right)\right)\\},$
$\displaystyle{\mathcal{F}}^{\mathrm{p}}$
$\displaystyle=\\{f|f(z)=\alpha^{\top}z,\alpha\in{\mathbb{R}}^{r\times(k-1)},\|\alpha_{s}\|\leq
c_{1}M(K)^{2},\text{ for all }s\in[k-1],\|\alpha^{\top}z\|\leq c_{2}\\},$
$\displaystyle{\mathcal{F}}^{\mathrm{d}}$
$\displaystyle=\\{f|f(z)=\alpha^{\top}z,\alpha\in{\mathbb{R}}^{r\times(k^{\prime}-1)},\|\alpha_{s}\|\leq
c_{0}M(K)^{2},\text{ for all }s\in[k^{\prime}-1],\|\alpha^{\top}z\|\leq
c_{3}\\}.$

We further assume 1) for each $p\in[K],\|W_{p}\|_{1,\infty}\leq M(p)$, and 2)
$\|W_{K}\|_{\infty\to 2}\leq M(K).$ Adapt Gaussian complexity results in
Golowich et al. (2018) we have

$\displaystyle
G_{n}({\mathcal{H}})\leq\widetilde{{\mathcal{O}}}\left(\frac{rM(K)^{3}\cdot
D\sqrt{K}\cdot\Pi_{p=1}^{K-1}M(p)}{\sqrt{n}}\right),\quad\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}}|h\circ
x^{\mathrm{p}})\leq{\mathcal{O}}\left(\frac{(k-1)M(K)^{3}}{\sqrt{n}}\right).$

Now we are ready to state our theorem for this practical setting of NLP pre-
training.

######  Theorem 4.10.

Under Assumption 3.1, 3.2, and 4.3, assume $M(K)\geq c_{5}$ for a universal
constant $c_{5}$. Then with probability at least $1-\delta$, Transfer Learning
Risk is upper bounded by

$\displaystyle\widetilde{{\mathcal{O}}}\left(\frac{krM(K)^{3}\cdot
D\sqrt{K}\cdot\Pi_{p=1}^{K-1}M(p)}{\tilde{\nu}\sqrt{n}}+\frac{k^{\frac{3}{2}}M(K)^{3}}{\tilde{\nu}\sqrt{n}}+\frac{{k^{\prime}}^{\frac{3}{2}}M(K)^{3}}{\sqrt{m}}\right).$

To interpret this bound, one can easily show that a standard supervised
learning paradigm without pre-training would have a sample complexity of
$\widetilde{{\mathcal{O}}}(krM(K)^{3}\cdot
D\sqrt{K}\cdot\Pi_{p=1}^{K-1}M(p)/\sqrt{m})$. Again, this theorem
demonstrates: when 1) $n\gg m$ and 2) $\tilde{\nu}$ is large, the rate of
transfer learning risk can be much smaller than that of the standard
supervised learning algorithm.

##  5 Experiments

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
$\nu$ (diversity parameter of $f^{\mathrm{p}}$), 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:

$\displaystyle L^{\prime}(\Theta)$
$\displaystyle=L(\Theta)-\lambda\cdot\ln\det(\alpha^{\mathrm{p}}\left(\alpha^{\mathrm{p}}\right)^{\top}),$
(6)

where $\Theta$ denotes the set of all model parameters, $\lambda$ is a hyper-
parameter that controls the magnitude of the diversity regularization,
$\det(\cdot)$ denote the determinant of a matrix, and $\alpha^{\mathrm{p}}$ is
the model parameter matrix at the output linear layer. This type of diversity
regularizer was proposed in Zou and 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
$\ell_{2}$-regularization) to pretrain BERT-base models under different values
of diversity factor $\lambda$. 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.555We do not report the WNLI
(classification) task due to its reported issues of the task in Devlin et al.
(2019). Our pre-training and finetuning implementations are based on the
opensource code released by Nvidia.666Distributed under Apache License:
https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT
We use the same pre-training data as the original BERT (i.e., English
Wikipedia + TorontoBookCorpus).777Collected and pre-processed using the code
and script included in the open-source code:
https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT
Our detailed pre-training and finetuning hyper-parameters along with other
experimental details are reported in Appendix B.2.

Table 1: Performance of diversity-regularized BERT pre-training with different
values of diversity factor $\lambda$. We finetune the pre-trained model on $8$
downstream tasks from GLUE benchmark and evaluate them on their dev sets. All
results are “mean (std)” from 5 runs with different random seeds. For MNLI, we
average the accuracies on its matched and mismatched dev sets. For MRPC and
QQP, we average their accuracy and F1 scores. For STS-B, we average Pearson’s
correlation and Spearman’s correlation. All other tasks uses accuracy as the
metric. The better-than-baseline numbers are underlined, and the best numbers
are highlighted in boldface.

Model | MNLI | MRPC | SST-2 | CoLA | QQP | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---
BERT-base ($\lambda=0.005$) |  84.17 (0.23) |  87.16 (1.81) | 92.48 (0.19) | 59.99 (0.28) |  89.42 (0.08) |  88.11 (0.54) |  67.28 (3.43) |  89.33 (0.07)
BERT-base ($\lambda=0.05$) |  84.01 (0.10) |  86.35 (5.15) |  93.00 (0.16) |  62.66 (1.07) |  89.46 (0.03) | 87.64 (0.44) | 60.64 (6.08) |  89.57 (0.13)
BERT-base ($\lambda=0.5$) |  84.00 (0.20) |  89.42 (0.51) |  92.93 (0.24) | 60.76 (0.71) |  89.33 (0.12) | 88.01 (0.23) |  67.93 (1.18) |  89.22 (0.23)
BERT-base (reproduced) | 83.96 (0.08) | 86.14 (4.64) | 92.64 (0.20) | 61.46 (0.74) | 89.28 (0.09) | 88.10 (0.27) | 63.64 (6.64) | 89.19 (0.07)

In Table 1, 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 B. From Table 1, we note that adding the diversity
regularization could generally improve the performance on these downstream
tasks. In particular, when $\lambda=0.5$, our pre-trained 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 $\lambda=0.5$ achieves a much more stable performance (i.e.,
smaller std) on tasks with scarce finetuning data ($<4$K samples in MRPC and
RTE). Our results, albeit still preliminary, demonstrate the potential of such
a simple diversity-regularizer. 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.

##  6 Conclusion

We theoretically prove the benefit of multi-class pre-training using the
notion of class diversity. Our proof uses the vector-form Rademacher
complexity chain rule and a modified self-concordance condition. We believe an
important and interesting direction is to develop a theory for other pre-
training techniques such as those based on sequence reconstruction (Lewis et
al., 2020).

## References

  * Abramovich and Grinshtein (2018) Felix Abramovich and Vadim Grinshtein.  High-dimensional classification by sparse logistic regression.  _IEEE Transactions on Information Theory_ , 65(5):3068–3079, 2018.
  * Bach (2014) Francis Bach.  Adaptivity of averaged stochastic gradient descent to local strong convexity for logistic regression.  _The Journal of Machine Learning Research_ , 15(1):595–627, 2014.
  * Bach et al. (2010) Francis Bach et al.  Self-concordant analysis for logistic regression.  _Electronic Journal of Statistics_ , 4:384–414, 2010.
  * Bansal et al. (2021) Yamini Bansal, Gal Kaplun, and Boaz Barak.  For self-supervised learning, Rationality implies generalization, provably.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=Srmggo3b3X6.
  * Barnes and Ozgur (2019) Leighton Pate Barnes and Ayfer Ozgur.  Minimax bounds for distributed logistic regression.  _arXiv preprint arXiv:1910.01625_ , 2019.
  * Baxter (2000) Jonathan Baxter.  A Model of Inductive Bias Learning.  _Journal of artificial intelligence research_ , 12:149–198, 2000.
  * Boyd and Vandenberghe (2004) Stephen Boyd and Lieven Vandenberghe.  _Convex Optimization_.  Cambridge University Press, 2004.  doi: 10.1017/CBO9780511804441.
  * Caruana (1997) Rich Caruana.  Multitask Learning.  _Machine learning_ , 28(1):41–75, 1997.
  * Clark et al. (2020) Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning.  ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=r1xMH1BtvB.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.  URL https://www.aclweb.org/anthology/N19-1423.
  * Du et al. (2021) Simon Shaolei Du, Wei Hu, Sham M. Kakade, Jason D. Lee, and Qi Lei.  Few-Shot Learning via Learning the Representation, Provably.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=pW2Q2xLwIMD.
  * Foster et al. (2018) Dylan J Foster, Satyen Kale, Haipeng Luo, Mehryar Mohri, and Karthik Sridharan.  Logistic regression: The importance of being improper.  In _Conference On Learning Theory_ , pages 167–208. PMLR, 2018.
  * Golowich et al. (2018) Noah Golowich, Alexander Rakhlin, and Ohad Shamir.  Size-independent Sample Complexity of Neural Networks.  In _Conference On Learning Theory_ , pages 297–299. PMLR, 2018.
  * Lan et al. (2020) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.  ALBERT: A Lite BERT for Self-supervised Learning of Language Representations.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=H1eA7AEtvS.
  * Ledoux and Talagrand (1991) M. Ledoux and M. Talagrand.  _Probability in Banach Spaces: Isoperimetry and Processes_.  A Series of Modern Surveys in Mathematics Series. Springer, 1991.  ISBN 9783540520139.  URL https://books.google.com.hk/books?id=cyKYDfvxRjsC.
  * Lee et al. (2021) Jason D Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo.  Predicting what you already know helps: Provable self-supervised learning.  _Advances in Neural Information Processing Systems_ , 34:309–323, 2021.
  * Lewis et al. (2020) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer.  BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 7871–7880, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.703.  URL https://aclanthology.org/2020.acl-main.703.
  * Liu et al. (2020) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.  RoBERTa: A Robustly Optimized BERT Pretraining Approach, 2020.  URL https://openreview.net/forum?id=SyxS0T4tvS.
  * Massart (2000) Pascal Massart.  About the Constants in Talagrand’s Concentration Inequalities for Empirical Processes.  _Annals of Probability_ , pages 863–884, 2000.
  * Maurer (2016) Andreas Maurer.  A vector-contraction inequality for Rademacher complexities.  In _International Conference on Algorithmic Learning Theory_ , pages 3–17. Springer, 2016.
  * Maurer et al. (2016) Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes.  The Benefit of Multitask Representation Learning.  _Journal of Machine Learning Research_ , 17(81):1–32, 2016.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving Language Understanding by Generative Pre-Training.  2018\.
  * Robinson et al. (2020) Joshua Robinson, Stefanie Jegelka, and Suvrit Sra.  Strength from weakness: Fast learning using weak supervision.  In _International Conference on Machine Learning_ , pages 8127–8136. PMLR, 2020.
  * Saunshi et al. (2019) Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar.  A Theoretical Analysis of Contrastive Unsupervised Representation Learning.  In _International Conference on Machine Learning_ , pages 5628–5637, 2019.
  * Shalev-Shwartz and Ben-David (2014) Shai Shalev-Shwartz and Shai Ben-David.  _Understanding machine learning: From theory to algorithms_.  Cambridge university press, 2014.
  * Thekumparampil et al. (2021) Kiran Koshy Thekumparampil, Prateek Jain, Praneeth Netrapalli, and Sewoong Oh.  Sample Efficient Linear Meta-Learning by Alternating Minimization.  _arXiv preprint arXiv:2105.08306_ , 2021.
  * Tosh et al. (2021) Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu.  Contrastive learning, multi-view redundancy, and linear models.  In _Algorithmic Learning Theory_ , pages 1179–1206. PMLR, 2021.
  * Tripuraneni et al. (2020) Nilesh Tripuraneni, Michael I. Jordan, and Chi Jin.  On the Theory of Transfer Learning: The Importance of Task Diversity.  In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/59587bffec1c7846f3e34230141556ae-Abstract.html.
  * Tripuraneni et al. (2021) Nilesh Tripuraneni, Chi Jin, and Michael Jordan.  Provable meta-learning of linear representations.  In _International Conference on Machine Learning_ , pages 10434–10443. PMLR, 2021.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is All you Need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Wainwright (2019) Martin J Wainwright.  _High-dimensional statistics: A non-asymptotic viewpoint_ , volume 48.  Cambridge University Press, 2019.
  * Wang et al. (2019) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman.  GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rJ4km2R5t7.
  * Yang et al. (2019) Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le.  XLNet: Generalized Autoregressive Pretraining for Language Understanding.  In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett, editors, _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pages 5754–5764, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/dc6a7e655d7e5840e66733e9ee67cc69-Abstract.html.
  * You et al. (2020) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training bert in 76 minutes.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=Syx4wnEtvH.
  * Zhang et al. (2022) Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner.  Unveiling transformers with lego: a synthetic reasoning task.  _arXiv preprint arXiv:2206.04301_ , 2022.
  * Zou and Adams (2012) James Y Zou and Ryan P Adams.  Priors for Diversity in Generative Latent Variable Models.  In _Proceedings of the 25th International Conference on Neural Information Processing Systems-Volume 2_ , pages 2996–3004, 2012.

##  Appendix A Technical Proofs

In Section 3, we have introduced Gaussian complexity. Let us restate for
clarity.

Let $\mu$ be a probability distribution on a set
${\mathcal{X}}\subset{\mathbb{R}}^{d}$ and suppose that $x_{1},\cdots,x_{n}$
are independent samples selected according to $\mu$. Let ${\mathcal{Q}}$ be a
class of functions mapping from ${\mathcal{X}}$ to ${\mathbb{R}}^{r}$. Define
random variable

$\hat{G}_{n}({\mathcal{Q}})=\mathbb{E}\left[\sup_{q\in{\mathcal{Q}}}\frac{1}{n}\sum_{k=1}^{r}\sum_{i=1}^{n}g_{ki}q_{k}(x_{i})\right]\qquad$
(7)

as the empirical Rademacher complexity, where $q_{k}(x_{i})$ is the $k$-th
coordinate of the vector-valued function $q(x_{i})$, $g_{ki}$
$(k\in[r],i\in[n])$ are independent Gaussian ${\mathcal{N}}(0,1)$ random
variables. The Gaussian complexity of ${\mathcal{Q}}$ is
$G_{n}({\mathcal{Q}})=E_{\mu}\hat{G}_{n}({\mathcal{Q}})$.

Analogously to the above we can define the empirical Rademacher complexity for
vector-valued functions as

$\hat{R}_{n}({\mathcal{Q}})=\mathbb{E}\left[\sup_{q\in{\mathcal{Q}}}\frac{1}{N}\sum_{k=1}^{r}\sum_{i=1}^{N}\epsilon_{ki}q_{k}(x_{i})\right]$
(8)

where $\epsilon_{ki}(k\in[r],i\in[n])$ are independent Rademacher
$\text{Rad}(\frac{1}{2})$ random variables. Its population counterpart is
defined as $R_{n}({\mathcal{Q}})=E_{\mu}[\hat{R}_{n}({\mathcal{Q}})]$. Note
that the superscripts existing in $\hat{G}$ and $\hat{R}$ imply that they are
empirical measures.

###  A.1 Proofs for Section 4.1

We illustrate Theorem 4.2 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.

#### Pre-training

######  Theorem A.1.

In pre-training, with probability at least $1-\delta$, it holds that:

$\displaystyle\quad
d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(h^{\prime};h)$
$\displaystyle\leq
4\sqrt{\pi}L^{\mathrm{p}}G_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})+4B^{\mathrm{p}}\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{n}}$
$\displaystyle\leq
4096L^{\mathrm{p}}\left[\frac{\sqrt{k-1}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}+\log(n)[L({\mathcal{F}}^{\mathrm{p}})G_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})]\right]+4B^{\mathrm{p}}\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{n}}.$

###### Proof.

We begin with

$\displaystyle d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(h^{\prime};h)\leq
2\sup_{f\in{\mathcal{F}}^{\mathrm{p}},h\in{\mathcal{H}}}|R_{\mathrm{p}}(f^{\mathrm{p}},h)-\hat{R}_{\mathrm{p}}(f^{\mathrm{p}},h)|.$

From the definition of Rademacher complexity [Wainwright, 2019, Theorem 4.12],
with probability at least $1-2\delta$ we have

$\sup_{f^{\mathrm{p}}\in{\mathcal{F}}^{\mathrm{p}},h\in{\mathcal{H}}}|R_{\mathrm{p}}(f^{\mathrm{p}},h)-\hat{R}_{\mathrm{p}}(f^{\mathrm{p}},h)|\leq
2R_{n}(\ell({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}}))+2B^{\mathrm{p}}\sqrt{\frac{\log(\nicefrac{{1}}{{\delta}})}{n}}.$

Next, we apply the vector contraction inequality [Maurer, 2016]. For function
class ${\mathcal{F}}$ whose output is in ${\mathbb{R}}^{K}$ with component
$f_{k}(\cdot)$, and the function $(h_{i})$s are some $L$-Lipschitz functions:
${\mathbb{R}}^{K}\mapsto{\mathbb{R}}$, we have

$\displaystyle\mathbb{E}_{\epsilon}\sup_{f\in{\mathcal{F}}}\sum_{i=1}^{n}\epsilon_{i}h_{i}(f(x_{i}))\leq\sqrt{2}L\mathbb{E}_{\epsilon}\sup_{f\in{\mathcal{F}}}\sum_{i=1}^{n}\sum_{k=1}^{K}\epsilon_{ik}f_{k}(x_{i}).$
(9)

Hence for loss function $\ell$ satisfying $|\ell(x)-\ell(y)|\leq
L^{\mathrm{p}}\|x-y\|_{2},\forall x,y\in{\mathbb{R}}^{k-1}$, the $f$ takes
value in ${\mathbb{R}}^{k-1}$ with component functions
$f_{s}(\cdot),s\in[k-1]$, we have that population Rademacher complexity can be
bounded by

$\displaystyle R_{n}(l({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}}))$
$\displaystyle=\mathbb{E}_{X^{\mathrm{p}}}\frac{1}{n}\mathbb{E}_{\epsilon}\sup_{f\in{\mathcal{F}}^{\mathrm{p}},h\in{\mathcal{H}}}\sum_{i=1}^{n}\epsilon_{i}\ell(f\circ
h(x_{i}^{\mathrm{p}}))$
$\displaystyle\leq\mathbb{E}_{X^{\mathrm{p}}}\frac{1}{n}\sqrt{2}L^{\mathrm{p}}\mathbb{E}_{\epsilon}\sup_{f\in{\mathcal{F}}^{\mathrm{p}},h\in{\mathcal{F}}}\sum_{i=1}^{n}\sum_{s=1}^{k-1}\epsilon_{is}f_{s}(h(x_{i}^{\mathrm{p}}))$
$\displaystyle=\sqrt{2}L^{\mathrm{p}}R_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$
$\displaystyle\leq\sqrt{\pi}L^{\mathrm{p}}G_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$

where the last line uses the fact that Rademacher complexity is upper bounded
by Gaussian complexity:
$R_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})\leq\sqrt{\frac{\pi}{2}}G_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$.
Therefore we have, with probability at least $1-\delta$,

$\displaystyle\quad
d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(h^{\prime};h)$
$\displaystyle\leq
4\sqrt{\pi}L^{\mathrm{p}}G_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})+4B^{\mathrm{p}}\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{n}}$
$\displaystyle\leq
4096L^{\mathrm{p}}\left[\frac{\sqrt{k-1}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}+\log(n)[L({\mathcal{F}}^{\mathrm{p}})G_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})]\right]+4B^{\mathrm{p}}\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{n}},$

where the last line uses decomposition of
$G_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$ into the individual
Gaussian complexities of ${\mathcal{H}}$ and ${\mathcal{F}}^{\mathrm{p}}$,
leverages an expectation version of novel chain rule for Gaussian complexities
(Lemma A.2). ∎

In the spirit of Gaussian complexity decomposition theorem [Tripuraneni et
al., 2020, Theorem 7], we introduce the following decomposition result upon
vector-form Gaussian complexities.

######  Lemma A.2.

We have the following vector form Gaussian complexity decomposition:

$\displaystyle\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$
$\displaystyle\leq\frac{8\sqrt{k-1}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}+512C({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})\cdot\log{(n)}$
(10)

where we use
$C({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})=L({\mathcal{F}}^{\mathrm{p}})\cdot\hat{G}_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})$
to represent the complexity measure of the composite function class.

###### Proof.

Our proof extends [Tripuraneni et al., 2020, 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 ${\mathcal{H}}$,
they need to handle ${\mathcal{F}}^{\otimes t}$ ($t$ is the number of tasks)
while our objective is a single function class ${\mathcal{F}}^{\mathrm{p}}$ of
higher dimension (${\mathcal{F}}^{\mathrm{p}}$ is $(k-1)$-dimensional for a
$k$-class classification task). We note that our proof technique and that of
previous works [Tripuraneni et al., 2020, Maurer et al., 2016] both hinge on
several properties of Gaussian processes.

To bound the empirical composite function class
${\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})$, note that vector-form Gaussian
complexity is defined as

$\displaystyle\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})$
$\displaystyle=\mathbb{E}\left[\frac{1}{n}\dot{\sup}_{f(h)\in{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})}\sum_{s=1}^{k-1}\sum_{i=1}^{n}g_{is}f_{s}(h(x_{i}^{\mathrm{p}}))\right]$
$\displaystyle=\frac{1}{\sqrt{n}}\mathbb{E}[\sup_{f(h)\in{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})}Z_{f(h)}]$

where we define mean-zero process
$Z_{f(h)}=\frac{1}{\sqrt{n}}\sum_{s=1}^{k-1}\sum_{i=1}^{n}g_{is}f_{s}(h(x_{i}^{\mathrm{p}}))$,
then
$\mathbb{E}\sup_{f(h)}Z_{f(h)}=\mathbb{E}\sup_{f(h)}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}\leq\mathbb{E}\sup_{f(h),f^{\prime}(h^{\prime})}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}$.
We further notice that $Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}$ is a sub-gaussian
random variable parameter

$\displaystyle d^{2}(f(h),f^{\prime}(h^{\prime})|x^{\mathrm{p}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\left\|f(h(x_{i}^{\mathrm{p}}))-f^{\prime}(h^{\prime}(x_{i}^{\mathrm{p}}))\right\|^{2}$
$\displaystyle=\frac{1}{n}\sum_{s=1}^{k-1}\sum_{i=1}^{n}\left(f_{s}(h(x_{i}^{\mathrm{p}}))-f^{\prime}_{s}(h^{\prime}(x_{i}^{\mathrm{p}}))\right)^{2}$

Dudley’s entropy integral bound [Wainwright, 2019, Theorem 5.22] shows

$\displaystyle\quad\mathbb{E}\sup_{f(h),f^{\prime}(h^{\prime})}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}$
$\displaystyle\leq
2\mathbb{E}\sup_{d(f(h),f^{\prime}(h^{\prime})|x^{\mathrm{p}})\leq\delta}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}+32{\mathcal{J}}(\frac{\delta}{4},D_{{\mathcal{X}}^{\mathrm{p}}})$
$\displaystyle=2\mathbb{E}\sup_{d(f(h),f^{\prime}(h^{\prime})|x^{\mathrm{p}})\leq\delta}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}+32\int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^{\mathrm{p}}}}\sqrt{\log
N(u;{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})|x^{\mathrm{p}})}du.$

It is straightforward to see the first term follows:

$\mathbb{E}\sup_{d(f(h),f^{\prime}(h^{\prime})|x^{\mathrm{p}})\leq\delta}Z_{f(h)}-Z_{f^{\prime}(h^{\prime})}\leq\mathbb{E}[\|g\|]\delta\leq\sqrt{n(k-1)}\delta$

We now turn to bound the second term by decomposing the distance metric into a
distance over ${\mathcal{F}}^{\mathrm{p}}$ and a distance over
${\mathcal{H}}$. We claim that, for arbitrary
$h\in{\mathcal{H}},f\in{\mathcal{F}}^{\mathrm{p}}$, let $h^{\prime}$ be
$\epsilon_{1}$-close to $h$ in empirical $l_{2}$-norm w.r.t inputs
$x_{1}^{\mathrm{p}},x_{2}^{\mathrm{p}}\cdots,x_{n}^{\mathrm{p}}$. Given
$h^{\prime}$, let $f^{\prime}$ be $\epsilon_{2}$-close to $f$ in empirical
$l_{2}$ loss w.r.t $h^{\prime}(x^{\mathrm{p}})$. Using the triangle inequality
we have that

$\displaystyle d(f(h),f^{\prime}(h^{\prime})|x^{\mathrm{p}})$
$\displaystyle=\sqrt{\frac{1}{n}\sum_{i=1}^{n}\|f(h(x_{i}^{\mathrm{p}}))-f^{\prime}(h^{\prime}(x_{i}^{\mathrm{p}}))\|}$
$\displaystyle\leq
d(f(h),f(h^{\prime})|x^{\mathrm{p}})+d(f(h^{\prime}),f^{\prime}(h^{\prime})|x^{\mathrm{p}})$
$\displaystyle\leq\sqrt{\frac{1}{n}\sum_{i=1}^{n}\|f(h(x_{i}^{\mathrm{p}}))-f(h^{\prime}(x_{i}^{\mathrm{p}}))\|^{2}}+\epsilon_{2}$
$\displaystyle\leq
L({\mathcal{F}}^{\mathrm{p}})\sqrt{\frac{1}{n}\sum_{i=1}^{n}\|h(x_{i}^{\mathrm{p}})-h^{\prime}(x_{i}^{\mathrm{p}})\|^{2}}+\epsilon_{2}$
$\displaystyle=L({\mathcal{F}}^{\mathrm{p}})\cdot\epsilon_{1}+\epsilon_{2},$

where we have used that $\|f(x)-f(y)\|\leq
L({\mathcal{F}}^{\mathrm{p}})\|x-y\|$ for any
$f\in{\mathcal{F}}^{\mathrm{p}}$.

As for the cardinality of the covering
$C_{{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})}$. Observe
$|C_{{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})}|=\sum_{h\in
C_{{\mathcal{H}}(x^{\mathrm{p}})}}|C_{{\mathcal{F}}^{\mathrm{p}}_{h}}|\leq|C_{{\mathcal{H}}(x^{\mathrm{p}})}|\cdot\max_{h\in{\mathcal{H}}(x^{\mathrm{p}})}|C_{{\mathcal{F}}^{\mathrm{p}}_{h(x^{\mathrm{p}})}}|$.
This provides a bound on the metric entropy of

$\displaystyle\log N(\epsilon_{1}\cdot
L({\mathcal{F}}^{\mathrm{p}})+\epsilon_{2};{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})|x^{\mathrm{p}})\leq\log
N(\epsilon_{1};{\mathcal{H}}|x^{\mathrm{p}})+\max_{h(x^{\mathrm{p}})}N(\epsilon_{2};{\mathcal{F}}^{\mathrm{p}}|h\circ
x^{\mathrm{p}}).$

Applying the covering number upper bound with
$\epsilon_{1}=\frac{\epsilon}{2\cdot
L({\mathcal{F}}^{\mathrm{p}})},\epsilon_{2}=\frac{\epsilon}{2}$ gives a bound
of entropy integral ofa ,

$\displaystyle\quad\int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^{\mathrm{p}}}}\sqrt{\log
N(u;{\mathcal{F}}^{\mathrm{p}}({\mathcal{H}})|x^{\mathrm{p}})}du$
$\displaystyle\leq\int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^{\mathrm{p}}}}\sqrt{\log
N\left(\frac{u}{2L({\mathcal{F}}^{\mathrm{p}})};{\mathcal{H}}\bigg{|}x^{\mathrm{p}}\right)}du+\int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^{\mathrm{p}}}}\max_{h\circ
x^{\mathrm{p}}}\sqrt{\log
N\left(\frac{u}{2};{\mathcal{F}}^{\mathrm{p}}\big{|}h\circ
x^{\mathrm{p}}\right)}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 $\forall u>0$ we find:

$\log N(u;{\mathcal{H}}|x^{\mathrm{p}})\leq
4\left(\frac{\sqrt{n}\hat{G}_{n}({\mathcal{H}})}{u}\right)^{2},\quad\log
N(u;{\mathcal{F}}^{\mathrm{p}}|h(x^{\mathrm{p}}))\leq
4\left(\frac{\sqrt{n}\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}|h\circ
x^{\mathrm{p}})}{u}\right)^{2}.$

Combining the definition of worst-case Gaussian complexity with the
aforementioned results we have

$\displaystyle\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}\circ{\mathcal{H}})\leq
2\sqrt{k-1}\delta+256\log{\frac{4D_{{\mathcal{X}}^{\mathrm{p}}}}{\delta}}\left(L({\mathcal{F}}^{\mathrm{p}})\hat{G}_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})\right),$

substitute $\delta$ with $\frac{4D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}$,
proof is completed. ∎

#### Downstream learning

Next we turn to the second stage and come up with theoretical guarantees by
using inexact $\hat{h}$ learned from the first stage.

######  Theorem A.3.

In the downstream task, we have that with probability at least $1-\delta$,

$\displaystyle
R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)\leq
d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)+4\sqrt{\pi}L^{\mathrm{d}}\cdot\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})+4B^{\mathrm{d}}\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{m}}$

###### Proof.

Assumption 3.1 implies

$\mathbb{E}_{x^{\mathrm{d}},y^{\mathrm{d}}}\left[\ell\left(g^{\mathrm{d}}\left(x^{\mathrm{d}}\right),y^{\mathrm{d}}\right)\right]=R_{\mathrm{d}}(f^{\mathrm{d}},h).$

To start, let
$\tilde{f}=\operatorname*{arg\,min}_{f\in{\mathcal{F}}^{\mathrm{d}}}R_{\mathrm{d}}(f,\hat{h})$
and
$R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)$
equals

$\displaystyle\left[R_{\mathrm{d}}(\tilde{f},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)\right]+\left[R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(\tilde{f},\hat{h})\right]$

where the first term satisfies

$\displaystyle\qquad\inf_{\tilde{f}\in{\mathcal{F}}^{\mathrm{d}}}\left[R_{\mathrm{d}}(\tilde{f},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)\right]$
$\displaystyle\leq\sup_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\inf_{\tilde{f}\in{\mathcal{F}}^{\mathrm{d}}}\big{[}R_{\mathrm{d}}(\tilde{f},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)\big{]}$
$\displaystyle=d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h},h)$

The second term follows the similar lines of Theorem A.1

$\displaystyle
R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(\tilde{f},\hat{h})$
$\displaystyle\leq
4\sqrt{\pi}L^{\mathrm{d}}\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})+4B^{\mathrm{d}}\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{m}}$

Again we make use of the worst-case argument

$\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})\leq\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}}).$

Combining the results gives the statement. ∎

#### Proof of main Theorem 4.2

Having introduced class diversity parameter, proof is directly completed via
combination of Theorem A.1 and Theorem A.3.

###  A.2 Proofs for Section 4.2

We could provide a better dependence on the boundedness noise parameters in
Theorem A.3 using Bernstein inequality. We present the following corollary
which has data-dependence in the Gaussian complexities.

######  Corollary A.4.

Presuming Assumption 3.1 holds, we have that then with probability at least
$1-\delta$,

$\displaystyle\quad
R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)$
$\displaystyle\leq
d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)+4\sqrt{\pi}L^{\mathrm{d}}\cdot\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})+4\sigma\sqrt{\frac{\log(\nicefrac{{2}}{{\delta}})}{m}}+50B^{\mathrm{d}}\frac{\log(\nicefrac{{2}}{{\delta}})}{m}$

###### Proof.

Denote
$Z=\sup_{f}|\hat{R}_{\mathrm{d}}(f,\hat{h})-R_{\mathrm{d}}(f,\hat{h})|$, we
apply the functional Bernstein inequality [Massart, 2000, Theorem 3] to
control the fluctuations. With probability at lest $1-\delta$, we have

$\displaystyle Z\leq
2\mathbb{E}[Z]+4\frac{\sigma}{\sqrt{m}}\sqrt{\log(\frac{1}{\delta})}+35\frac{B^{\mathrm{d}}}{m}\log(\frac{1}{\delta}),$
(11)

where
$\sigma^{2}=\frac{1}{m}\sup_{f}\sum_{i=1}^{m}Var(\ell(f\circ\hat{h}(x_{i}^{\mathrm{d}}),y_{i}^{\mathrm{d}}))$.
Thus

$\displaystyle\mathbb{E}[Z]$ $\displaystyle\leq
2\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{R}_{m}(l({\mathcal{F}}^{\mathrm{d}})|\hat{h}\circ
x^{\mathrm{d}})$ $\displaystyle\leq
2\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\sqrt{2}L^{\mathrm{d}}\hat{R}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}})$ $\displaystyle\leq
2\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\sqrt{\pi}L^{\mathrm{d}}\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ
x^{\mathrm{d}}),$

where the second line uses vector-based contraction principle, the last line
upper bounds the empirical Rademacher complexity by Gaussian counterparts. ∎

#### Proof of Theorem 4.4

Observe that

$\ell(\eta;y)=-y^{\top}\eta+\log{(1+\sum_{s=1}^{k-1}e^{\eta_{s}})},\ell(\eta;y)\leq\|\eta\|$

and

$\left|\frac{\partial\ell(\eta;y)}{\partial\eta_{i}}\right|=\left|y_{i}-\frac{e^{\eta_{i}}}{1+\sum_{s=1}^{k-1}e^{\eta_{s}}}\right|,$
$\left|\nabla_{\eta}\ell(\eta;y)\right|\leq\sqrt{k-1},$

so it is $L^{\mathrm{p}}=\sqrt{k-1}-$Lipschitz. By definition the class
${\mathcal{F}}^{\mathrm{p}}$ with parameters
$\|\alpha_{s}\|_{2}\leq{\mathcal{O}}(1),s\in[k-1]$, we obtain that
$L({\mathcal{F}}^{\mathrm{p}})={\mathcal{O}}\left(\sqrt{k-1}\right)$ since for
any $x,y\in{\mathbb{R}}^{r}$, any $f\in{\mathcal{F}}^{\mathrm{p}}$ we have

$\displaystyle\|f(x)-f(y)\|^{2}$
$\displaystyle=\|\alpha^{\top}x-\alpha^{\top}y\|^{2}$
$\displaystyle\leq\sum_{s=1}^{k-1}\left(\langle\alpha_{s},x-y\rangle\right)^{2}$
$\displaystyle\leq\sum_{s=1}^{k-1}\|\alpha_{s}\|^{2}\|x-y\|^{2}$
$\displaystyle\leq c_{1}^{2}(k-1)\|x-y\|^{2}$

In conclusion we have

  * •

Pre-training loss $\ell(\cdot,y^{\mathrm{p}})$ is $\sqrt{k-1}$-Lipschitz.

  * •

Downstream loss $\ell(\cdot,y^{\mathrm{d}})$ is
$\sqrt{k^{\prime}-1}$-Lipschitz.

  * •

Linear layer $f$ is
$L({\mathcal{F}}^{\mathrm{p}})=O\left(\sqrt{k-1}\right)$-Lipschitz.

Consider task-specific function classes for characterizing class-diversity
parameters. From Lemma A.6 and Lemma 4.5 we know that

$\nu=\Omega(\tilde{\nu}),\quad\tilde{\nu}=\sigma_{r}(\alpha_{1}\alpha_{1}^{\top}).$

Combining these pieces of results then the proof is completed.

With the following proposition, we interpret the cross-entropy loss in the
well-specified model under our multinomial logistic model distribution.

######  Proposition A.5.

Under Assumption 3.2, for the cross entropy loss $\ell$ we have

$\displaystyle\mathbb{E}_{y\sim{\mathcal{P}}(\cdot|f\circ
h(x))}[\ell(\hat{f}\circ\hat{h}(x),y)]-\ell(f\circ h(x),y)]$
$\displaystyle=KL\left[{\mathcal{P}}(\cdot|f\circ
h(x)),{\mathcal{P}}(\cdot|\hat{f}\circ\hat{h}(x))\right]$
$\displaystyle=KL\left[{\mathcal{P}}(\cdot|\alpha^{\top}h(x)),{\mathcal{P}}(\cdot|{\alpha^{\prime}}^{\top}\hat{h}(x))\right].$

Recall that $\alpha^{\prime}$ and $\alpha$ are parameters for $\hat{f}$ and
$f$ respectively. The proof is straightforward by applying Assumption 3.2.

The following lemma serves a critical role in the proof.

######  Lemma A.6.

The $KL-$divergence between the true underlying conditional distribution of a
multinomial logistic model and the distribution we obtained can be bounded
from both sides with quadratic loss,

$\displaystyle
c_{0}e^{-10q_{0}}\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2}$
$\displaystyle\leq
KL\left[{\mathcal{P}}(\cdot|\alpha^{\top}h(x^{\mathrm{p}})),{\mathcal{P}}(\cdot|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}}))\right]$
$\displaystyle\leq\frac{1}{2}\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2},$

where
$c_{0}=\frac{1}{2}\lambda_{min}(\Phi^{\prime\prime}(\alpha^{\top}h(x^{\mathrm{p}})))$
is the least eigenvalue of Hessian matrix for $\Phi$,
$q_{0}=\max(\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})\|,\|\alpha^{\top}h(x^{\mathrm{p}})\|)$.

######  Remark A.7.

For the left hand side, the expression is related to the least eigenvalue of
the Hessian matrix at $\alpha^{\top}h(x^{\mathrm{p}})$, whereas the least
eigenvalue would depend on an unknown intermediate-term $x^{\prime}$ if we
adopt Taylor’s expansion.

Proof of Lemma A.6. Below we use $x$ for $x^{\mathrm{p}}$ for clarity. For
generalized linear models,

$\displaystyle KL$
$\displaystyle\left[{\mathcal{P}}(\cdot|\alpha^{\top}h(x)),{\mathcal{P}}(\cdot|{\alpha^{\prime}}^{\top}\hat{h}(x))\right]=\Phi({\alpha^{\prime}}^{\top}\hat{h}(x))-\Phi(\alpha^{\top}h(x))-\nabla\Phi(\alpha^{\top}h(x))^{\top}({\alpha^{\prime}}^{\top}\hat{h}(x)-\alpha^{\top}h(x)).$

Hence the $KL-$divergence serves as the second-order remainder term according
to Taylor’s theorem.

For the right hand side, the gradient of $\Phi(x)$ at $i^{th}$-coordinate
$\frac{\partial\Phi}{\partial
x_{i}}=\nicefrac{{e^{x_{i}}}}{{1+\sum_{s}e^{x_{s}}}}$, the Hessian matrix is

$\frac{\partial^{2}\Phi}{\partial x_{i}\partial
x_{j}}=\left\\{\begin{array}[]{lr}\frac{e^{x_{i}}\cdot(1+\sum_{s\neq
i}e^{x_{s}})}{(1+\sum_{s}e^{x_{s}})^{2}},&i=j\\\ &\\\
\frac{-e^{x_{i}}e^{x_{j}}}{(1+\sum_{s}e^{x_{s}})^{2}},&i\neq
j.\end{array}\right.$

Let
$\sigma(x)=\frac{1}{1+\sum_{s}e^{x_{s}}}\left[e^{x_{1}},\cdots,e^{x_{k-1}}\right]^{\top}$,
the Hessian matrix can be restated as

$\nabla^{2}\Phi=diag(\sigma(x))-\sigma(x)\sigma(x)^{\top}.$

For any non-zero vector $y$, we have

$\displaystyle y^{\top}\nabla^{2}\Phi y$
$\displaystyle=\sum_{i}\sigma(x)_{i}y_{i}^{2}-\left(\sigma(x)^{\top}y\right)^{2}\leq\max(\sigma(x)_{i})\|y\|^{2}\leq\|y\|^{2},$

which implies its largest eigenvalue is no bigger than 1.

For the left hand side, it is very straightforward to see that the Hessian
matrix is positive semidefinite. Though being non-negative, we point out that
bounding the second-order remainder terms from below with quadratic loss
requires new techniques which we discuss below.

#### Proof of Proposition 4.7

###### Proof.

Let $P(t;v^{0})=1+\sum_{s}e^{u_{s}+tv_{s}}$ and
$P(t;v^{i})=\sum_{s}v_{s}^{i}e^{u_{s}+tv_{s}},i>1$. Then we use multinomials
$P$ to represent derivatives of $g(t)$

$\displaystyle g(t)$ $\displaystyle=\log(P(T;V^{0}))$ $\displaystyle
g^{\prime}(t)$ $\displaystyle=\frac{P(t;v^{1})}{P(t;v^{0})}$ $\displaystyle
g^{\prime\prime}(t)$
$\displaystyle=\frac{P(t;v^{2})P(t;v^{0})}{P(t;v^{0})^{2}}$ $\displaystyle
g^{\prime\prime\prime}(t)$
$\displaystyle=\frac{P(t;v^{3})P(t;v^{0})^{2}-3P(t;v^{2})P(t;v^{1})P(t;v^{0})+2P(t;v^{1})^{3}}{P(t;v^{0})^{3}}$

Let $r_{s}=e^{u_{s}+v_{s}t}$, hence

$\displaystyle g^{\prime\prime}(t)$
$\displaystyle=\frac{(\sum_{s}v_{s}^{2}r_{s})\cdot(1+\sum_{s}r_{s})-(\sum_{s}v_{s}r_{s})^{2}}{(1+\sum_{s}r_{s})^{2}}$
$\displaystyle=\frac{\sum_{i<j}r_{i}r_{j}(v_{i}-v_{j})^{2}+\sum_{i}v_{i}^{2}r_{i}}{(1+\sum_{s}r_{s})^{2}}$

In the following we expand $g^{\prime\prime\prime}(t)$ as:

$\displaystyle\frac{\sum_{i<j}r_{i}r_{j}(v_{i}-v_{j})^{2}[\sum_{k}(v_{i}+v_{j}-2v_{k})r_{k}]+\sum_{i}v_{i}^{3}r_{i}+\sum_{i}\sum_{j}v_{i}^{2}r_{i}r_{j}(2v_{i}-3v_{j})}{(1+\sum_{s}r_{s})^{3}}$
$\displaystyle=\frac{\sum_{i<j}r_{i}r_{j}(v_{i}-v_{j})^{2}[\sum_{k}(v_{i}+v_{j}-2v_{k})r_{k}]+\sum_{i}v_{i}^{2}r_{i}\left(v_{i}(1+2\sum_{j}r_{j})-3\sum_{j}v_{j}r_{j}\right)}{(1+\sum_{s}r_{s})^{3}},$

observe that

$\displaystyle\frac{1}{1+\sum_{s}r_{s}}\left|\sum_{k}(v_{i}+v_{j}-2v_{k})r_{k}\right|\leq\sum_{k}|v_{i}+v_{j}-2v_{k}|\frac{r_{k}}{1+\sum_{s}r_{s}}\leq
4\|v\|_{2}$
$\displaystyle\frac{1}{1+\sum_{s}r_{s}}\left|v_{i}(1+2\sum_{j}r_{j})-3\sum_{j}v_{j}r_{j}\right|\leq
5\|v\|_{2}$

Substitute these into definition of _self-concordance_ then proof is
completed. ∎

Now we are ready to give a lower bound of $KL-$divergence,

$\displaystyle\quad\Phi({\alpha^{\prime}}^{\top}\hat{h}(x))-\Phi(\alpha^{\top}h(x))-\nabla\Phi(\alpha^{\top}h(x))^{\top}v$
$\displaystyle\geq\frac{1}{2}v^{\top}e^{-5\|v\|_{2}}F^{\prime\prime}(\alpha^{\top}h(x))v$
$\displaystyle\geq\frac{1}{2}\lambda_{min}(\Phi^{\prime\prime}(\alpha^{\top}h(x)))\|v\|^{2}e^{-5\|v\|_{2}}$
$\displaystyle\geq\frac{1}{2}\lambda_{min}(\Phi^{\prime\prime}(\alpha^{\top}h(x)))\|v\|^{2}e^{-5(\|{\alpha^{\prime}}^{\top}\hat{h}(x)\|+\|\alpha^{\top}h(x)\|)}$
$\displaystyle\geq\frac{1}{2}\lambda_{min}(\Phi^{\prime\prime}(\alpha^{\top}h(x)))\|v\|^{2}\exp(-10q_{0})$

where $v={\alpha^{\prime}}^{\top}\hat{h}(x)-\alpha^{\top}h(x)$. Proof for
Lemma A.6 is completed. ∎

#### Proof of Lemma 4.5

###### Proof.

For function classes ${\mathcal{F}}^{\mathrm{p}}$,
${\mathcal{F}}^{\mathrm{d}}$ and data samples generated from multinomial
logistic regression distribution (see Assumption 3.2), the worst-case
representation difference is similar to that in multi-task analysis
[Tripuraneni et al., 2020, Lemma 1]:

$\displaystyle d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)$
$\displaystyle=\sup_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\inf_{f\prime\in{\mathcal{F}}^{\mathrm{d}}}\mathbb{E}\left\\{\ell(f^{\prime}\circ\hat{h}(x^{\mathrm{d}}),y^{\mathrm{d}})-\ell(f^{\mathrm{d}}\circ
h(x^{\mathrm{d}}),y^{\mathrm{d}})\right\\}$
$\displaystyle\leq\sup_{\|\alpha_{s}\|\leq
c_{0}}\inf_{\|\alpha^{\prime}_{s}\|\leq
c_{0}}\frac{1}{2}\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{d}})-\alpha^{\top}h(x^{\mathrm{d}})\right\|^{2},\quad\text{here
}s\in[k^{\prime}-1]$
$\displaystyle=\sum_{s=1}^{k^{\prime}-1}\sup_{\|\alpha_{s}\|\leq
c_{0}}\inf_{\|\alpha^{\prime}_{s}\|\leq
c_{0}}\frac{1}{2}\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\left({\alpha^{\prime}}_{s}^{\top}\hat{h}(x^{\mathrm{d}})-\alpha_{s}^{\top}h(x^{\mathrm{d}})\right)^{2}$
$\displaystyle\leq(k^{\prime}-1)\frac{c_{0}^{2}}{2}\sigma_{1}(\Lambda_{sc}(\hat{h},h)).$

The first line is because of Proposition A.5 and Lemma A.6. In the last line,
the inner infima is considered as the partial minimization of a convex
quadratic form (see [Boyd and Vandenberghe, 2004, Example 3.15, Appendix
A.5.4]).

Define population covariance if representations $\hat{h}$ and $h$ as

$\begin{gathered}\Lambda(\hat{h},h)=\begin{bmatrix}\mathbb{E}[\hat{h}(x)\hat{h}(x)^{\top}]&\mathbb{E}[\hat{h}(x)h(x)^{\top}]\\\
\mathbb{E}[h(x)\hat{h}(x)^{\top}]&\mathbb{E}[h(x)h(x)^{\top}]\end{bmatrix}=\begin{bmatrix}F_{\hat{h}\hat{h}}&F_{\hat{h}h}\\\
F_{h\hat{h}}&F_{hh}\end{bmatrix}\end{gathered}$

$\Lambda_{sc}(\hat{h},h)=F_{hh}-F_{h\hat{h}}(F_{\hat{h}\hat{h}})^{\dagger}F_{\hat{h}h}$
is the generalized Schur complement of $h$ with respect to $\hat{h}$.

Next we control the pre-training representation differenceound is subtler,

$\displaystyle\quad d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(\hat{h};h)$
$\displaystyle\geq\inf_{\alpha^{\prime}}c_{0}\mathbb{E}_{{\mathcal{X}}^{\mathrm{p}}}\left[\exp(-10\max(\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})\|,\|\alpha^{\top}h(x^{\mathrm{p}})\|))\cdot\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2}\right],$

which is because of Proposition A.5 and Lemma A.6.

It is known

$\displaystyle\quad\mathbb{E}_{{\mathcal{X}}^{\mathrm{p}}}\left[\exp(-10\max(\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})\|,\|\alpha^{\top}h(x^{\mathrm{p}})\|))\cdot\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2}\right]$
$\displaystyle\geq
e^{-10c_{2}}\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2}$

Hence this metric could be claimed to be lower bounded as,

$\displaystyle\quad\Omega\left(\inf_{\alpha^{\prime}}\mathbb{E}_{x^{\mathrm{p}}}\left\|{\alpha^{\prime}}^{\top}\hat{h}(x^{\mathrm{p}})-\alpha^{\top}h(x^{\mathrm{p}})\right\|^{2}\right)$
$\displaystyle=\Omega\left(\alpha_{1}^{\top}\Lambda_{sc}(\hat{h};h)\alpha_{1}\right)$
$\displaystyle=\Omega\left(tr(\Lambda_{sc}(\hat{h};h)C)\right),\quad\text{where
$C=\alpha_{1}\alpha_{1}^{\top}$.}$

In the second line, we redefine $\alpha_{1}$ as parameter $\alpha$ of pre-
training for clarity. In this way we conclude that,

$\displaystyle
d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(\hat{h};h)=\Omega\left(tr(\Lambda_{sc}(\hat{h},h)C)\right)=\Omega\left(\sigma_{1}(\Lambda_{sc}(\hat{h},h))\sigma_{r}(C)\right),$

where $C$ implies expansion of representation $h(x)\in{\mathbb{R}}^{r}$, and
its condition number $\sigma_{r}(C)$ indicates how spread out this vector is
in ${\mathbb{R}}^{r}$:

$\displaystyle
C=\sum_{s=1}^{k-1}(\alpha_{1})_{s}(\alpha_{1})_{s}^{\top}=\alpha_{1}\alpha_{1}^{\top},\quad\alpha_{1}\in{\mathbb{R}}^{r\times(k-1)}$

Aforementioned calculations show

$d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)\leq\frac{1}{\Omega(\tilde{\nu})}d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(\hat{h};h),\quad\tilde{\nu}=\sigma_{r}(C).$

Proof is completed. ∎

###  A.3 Proofs for Section 4.4

###### Proof.

We begin with bounding each of the complexity terms in the Corrolary A.4.

We make use of data-dependent inequalities [Tripuraneni et al., 2020, Lemma 4]
to help upper bound related quantities. Intuitively Definition 4.8 implies
tail-bound properties in a sub-gaussian process.

  * •

$\displaystyle\hat{G}_{n}({\mathcal{H}})$
$\displaystyle=\frac{1}{n}\mathbb{E}\left[\sup_{B\in{\mathcal{H}}}\sum_{k=1}^{r}\sum_{i=1}^{n}g_{ki}b_{k}^{\top}x_{i}^{\mathrm{p}}\right]$
$\displaystyle={\mathcal{O}}\left(\sqrt{\frac{dr^{2}}{n}}\right)$

  * •

$\displaystyle\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}|h\circ x^{\mathrm{p}})$
$\displaystyle=\frac{1}{n}\mathbb{E}\left[\sup_{\alpha_{1},\cdots,\alpha_{k-1}}\sum_{s=1}^{k-1}\sum_{i=1}^{n}g_{is}\alpha_{s}^{\top}B^{\top}x_{i}^{\mathrm{p}}\right]$
$\displaystyle=\frac{c_{1}(k-1)}{n}\mathbb{E}\|\sum_{i=1}^{n}g_{is}B^{\top}x_{i}^{\mathrm{p}}\|$
$\displaystyle=\frac{c_{1}(k-1)}{\sqrt{n}}\sqrt{tr(B^{\top}\Sigma B)}$

then
$\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})\leq{\mathcal{O}}\left((k-1)\sqrt{\frac{r}{n}}\right)$.

  * •

Similarly,

$\displaystyle\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|h\circ
x^{\mathrm{d}})\leq\frac{c_{1}(k^{\prime}-1)}{\sqrt{m}}\sqrt{\sum_{i=1}^{r}\sigma_{i}(\hat{B}^{\top}\Sigma\hat{B})}$

then
$\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})\leq{\mathcal{O}}\left((k^{\prime}-1)\sqrt{\frac{r}{m}}\right)$.

  * •

boundedness parameter
$D_{{\mathcal{X}}^{\mathrm{p}}}=\sup_{\alpha,B}\|\alpha^{\top}B^{\top}x\|=c_{2}$

  * •

cross entropy
$\ell(\eta;y)=-y^{\top}\eta+\log{(1+\sum_{s=1}^{k-1}e^{\eta_{s}})}$, then
$\left|\frac{\partial\ell(\eta;y)}{\partial\eta_{i}}\right|=\left|y_{i}-\frac{e^{\eta_{i}}}{1+\sum_{s=1}^{k-1}e^{\eta_{s}}}\right|$,
$\left|\nabla_{\eta}\ell(\eta;y)\right|\leq\sqrt{k-1}$, so it is
$L^{\mathrm{p}}=\sqrt{k-1}-$Lipschitz in its first coordinate uniformly over
its second for pre-training and $L^{\mathrm{d}}=\sqrt{k^{\prime}-1}-$Lipschitz
for downstream task.

  * •

$|\ell(\eta;y)|\leq{\mathcal{O}}(\|\eta\|)$, where
$\|\eta\|=\|x^{\top}B^{\mathrm{p}}\alpha\|\leq c_{2}$.

In Corollary A.4, we define and compute the maximal variance term $\sigma^{2}$
as,

$\displaystyle\sigma^{2}$
$\displaystyle=\frac{1}{m}\sup_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\sum_{i=1}^{m}Var(\ell^{\prime}(f^{\mathrm{d}}\circ\hat{h}(x_{i}^{\mathrm{d}}),y_{i}^{\mathrm{d}}))$
$\displaystyle\leq\frac{k^{\prime}-1}{m}\sup_{f^{\mathrm{d}}\in{\mathcal{F}}^{\mathrm{d}}}\sum_{i=1}^{m}Var(f^{\mathrm{d}}\circ\hat{h}(x_{i}^{\mathrm{d}}))$
$\displaystyle=\frac{k^{\prime}-1}{m}\sup_{\|\alpha_{s}\|\leq{\mathcal{O}}(1)}\sum_{s=1}^{k^{\prime}-1}\sum_{i=1}^{m}Var(\alpha_{s}^{\top}\hat{B}^{\top}x_{i}^{\mathrm{d}})$
$\displaystyle=\frac{(k^{\prime}-1)^{2}}{m}\sup_{\|\alpha_{s}\|\leq{\mathcal{O}}(1)}\sum_{i=1}^{m}(\alpha_{s}\hat{B})^{\top}\Sigma\hat{B}\alpha_{s}$
$\displaystyle=(k^{\prime}-1)^{2}{\mathcal{O}}(\|\hat{B}\Sigma\hat{B}\|_{2})$
$\displaystyle={\mathcal{O}}\left((k^{\prime}-1)^{2}\right)$

With these results in hand, we are now prepared to apply Corollary A.4, w.p.
at least $1-\delta$

$\displaystyle\quad
R_{\mathrm{d}}({\hat{f}}^{\mathrm{d}},\hat{h})-R_{\mathrm{d}}(f^{\mathrm{d}},h)$
$\displaystyle\leq
d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)+4\sqrt{\pi}L^{\mathrm{d}}\cdot\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}\left({\mathcal{F}}^{\mathrm{d}}\Big{|}\hat{h}\circ
x^{\mathrm{d}}\right)+4\sigma\sqrt{\frac{\log{(\nicefrac{{2}}{{\delta}})}}{m}}+50B^{\mathrm{d}}\frac{\log{(\nicefrac{{2}}{{\delta}})}}{m}$

where $\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|\hat{h}\circ x^{\mathrm{d}})$ is
defined in Theorem 4.4.

Thus
$L^{\mathrm{d}}\cdot\mathbb{E}_{{\mathcal{X}}^{\mathrm{d}}}\hat{G}_{m}\left({\mathcal{F}}^{\mathrm{d}}\Big{|}\hat{h}\circ
x^{\mathrm{d}}\right)\leq
L^{\mathrm{d}}\bar{G}_{m}({\mathcal{F}}^{\mathrm{d}})\leq{\mathcal{O}}((k^{\prime}-1)^{\frac{3}{2}}\sqrt{\frac{r}{m}})$,
$\sigma\leq{\mathcal{O}}(k^{\prime}-1)$, and
$B^{\mathrm{d}}\leq{\mathcal{O}}(\sqrt{k^{\prime}-1}D)$. Further, we obtain
upper bound of worst-case representation difference by diversity parameter and
adoption of Theorem A.1: w.p. at least $1-\delta$

$\displaystyle\quad d_{{\mathcal{F}}^{\mathrm{d}}}(\hat{h};h)$
$\displaystyle\leq\frac{d_{{\mathcal{F}}^{\mathrm{p}},f^{\mathrm{p}}}(\hat{h};h)}{\nu}$
$\displaystyle\leq\frac{1}{\nu}\left\\{4096L\left[\log{(n)}\cdot[L({\mathcal{F}}^{\mathrm{p}})\cdot
G_{n}({\mathcal{H}})+\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})]+\frac{\sqrt{k-1}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}}\right]+4B\sqrt{\frac{\log{(\nicefrac{{2}}{{\delta}})}}{n}}\right\\}$
$\displaystyle\lesssim\frac{1}{\nu}\left\\{\sqrt{k}\log{(n)}\left(\sqrt{\frac{kdr^{2}}{n}}+k\sqrt{\frac{r}{n}}\right)+\frac{k}{n^{2}}+\sqrt{\frac{\log{(\nicefrac{{1}}{{\delta}})}}{n}}\right\\}$

The last thing to consider for completing the proof for Theorem 4.9 is giving
accurate characterization of diversity parameter $\nu$, which we leave for the
next subsection. ∎

###  A.4 Proofs for Section 4.5

###### Proof.

In deep neural network, we first review complexity quantities. Adapted from
Theorem 8 [Golowich et al., 2018], we have

$\displaystyle\hat{R}_{n}({\mathcal{N}})$
$\displaystyle\leq\left(\frac{2}{n}\Pi_{p=1}^{K}M(p)\right)\sqrt{(K+1+\log
d)\cdot\max_{j\in[d]}\sum_{i=1}^{n}x_{i,j}^{2}}$
$\displaystyle\leq\frac{2D\sqrt{K+1+\log d}\cdot\Pi_{p=1}^{K}M(p)}{\sqrt{n}}.$

where $x_{i,j}$ denotes the $j$-th coordinate of vector $x_{i}$.

Then we proceed to bound the Gaussian complexities for our deep neural network
and prove Theorem 4.10. Recall that under the conditions of the result we can
use former results to verify the task diversity condition is satisfied with
parameters $\Omega(\tilde{\nu})$ with
$\tilde{\nu}=\sigma_{r}(\alpha_{1}\alpha_{1}^{\top})>0$. We can see that
$\|\mathbb{E}_{x}[\hat{h}(x)h^{*}(x)^{\top}\|_{2}\leq\mathbb{E}{x}\|\hat{h}(x)h^{*}(x)\|\leq
O(M(K)^{2})$ using the norm bound from. Hence under this setting we can choose
$c_{1}$ sufficiently large so that
$c_{1}M(K)^{2}\gtrsim\frac{M(K)^{2}}{c}c_{2}$. The condition $M(K)\gtrsim 1$
in the theorem statement is simply used to clean up the final bound.

In order to instantiate Theorem 4.2 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

$\displaystyle\hat{G}_{n}({\mathcal{H}})=\frac{1}{n}\mathbb{E}[\sup_{{\mathcal{W}}_{K}}\sum_{k=1}^{r}\sum_{i=1}^{n}g_{ki}h_{k}(x_{i}^{\mathrm{p}})]\leq\sum_{k=1}^{r}\hat{G}_{n}(h_{k}(x_{i}^{\mathrm{p}}))$
$\displaystyle\leq\log(n)\cdot\sum_{k=1}^{r}\hat{R}_{n}{h_{k}(x_{i}^{\mathrm{p}})}\leq
r\log(n)\frac{2D\sqrt{K+1+\log d}\cdot\Pi_{p=1}^{K}M(p)}{\sqrt{n}}.$

This also implies the population Gaussian complexity.

  * •

By definition the class ${\mathcal{F}}$ as linear maps with parameters
$\|\alpha_{s}\|_{2}\leq c_{1}M(K)^{2},\forall s\in[k-1]$, we obtain that
$L({\mathcal{F}})=c_{1}\sqrt{k-1}M(K)^{2}$.

  * •

For the complexity of learning ${\mathcal{F}}^{\mathrm{p}}$ in the training
phase we obtain,

$\displaystyle\hat{G}_{n}({\mathcal{F}}^{\mathrm{p}}|h\circ
x^{\mathrm{p}})=\frac{1}{n}\mathbb{E}_{g}[\sup_{\alpha\in{\mathcal{F}}}\sum_{s=1}^{k-1}\sum_{i=1}^{n}g_{is}\alpha_{s}^{\top}h(x_{i}^{\mathrm{p}})]\lesssim\frac{(k-1)M(K)^{2}}{n}\mathbb{E}_{g}[\|\sum_{i=1}^{n}g_{is}h(x_{1i})\|]$
$\displaystyle\lesssim\frac{(k-1)M(K)^{2}}{n}\sqrt{\sum_{i=1}^{n}\|h(x_{i}^{\mathrm{p}})\|^{2}}\lesssim\frac{(k-1)M(K)^{2}}{\sqrt{n}}\max_{i}\|h(x_{i}^{\mathrm{p}})\|.$

For $tanh$ activation function, we simply have

$\displaystyle\|h(x)\|^{2}=\|W_{K}r_{K-1}\|_{2}^{2}\leq\|W_{K}\|_{\infty\to
2}^{2},$

where $r_{K-1}$ denotes ourput of the $K-1$th layer,

$\|h(x)\|\leq{\mathcal{O}}(M(K)).$

In conclusion we obtain

$\displaystyle\bar{G}_{n}({\mathcal{F}}^{\mathrm{p}})\leq
O\left(\frac{(k-1)M(K)^{3}}{\sqrt{n}}\right).$

  * •

Similarly

$\displaystyle\hat{G}_{m}({\mathcal{F}}^{\mathrm{d}}|h\circ
x^{\mathrm{d}})\leq O\left(\frac{(k^{\prime}-1)M(K)^{3}}{\sqrt{m}}\right)$

Regularity conditions

  * •

Boundedness parameter
$D_{{\mathcal{X}}^{\mathrm{p}}}=\sup_{\alpha,h}\|\alpha^{\top}h(x^{\mathrm{p}})\|=c_{2}$.

  * •

Pre-training loss is $L^{\mathrm{p}}=\sqrt{k-1}$-Lipschitz and
$B^{\mathrm{p}}=c_{2}$-bounded.

  * •

Downstream loss is $L^{\mathrm{d}}=\sqrt{k^{\prime}-1}$-Lipschitz and
$B^{\mathrm{d}}=c_{3}$-bounded.

Assembling the previous complexity arguments shows the transfer learning risk
is bounded by

$\displaystyle\lesssim\frac{L^{\mathrm{p}}}{\tilde{\nu}}\left(\log(n)\left[L({\mathcal{F}}^{\mathrm{p}})r\log(n)\frac{D\sqrt{K}\Pi_{p=1}^{K}M(p)}{\sqrt{n}}+\frac{kM(K)^{3}}{\sqrt{n}}\right]\right)+\frac{L^{\mathrm{d}}k^{\prime}M(K)^{3}}{\sqrt{m}}$
$\displaystyle+\left(\frac{1}{\tilde{\nu}}\max\left(\frac{L^{\mathrm{p}}\sqrt{k}D_{{\mathcal{X}}^{\mathrm{p}}}}{n^{2}},B^{\mathrm{p}}\sqrt{\frac{\log(1/\delta)}{n}}\right)+B^{\mathrm{d}}\sqrt{\frac{\log(1/\delta)}{m}}\right)$

Substitute regularity conditions into it, then the risk is simplified as
stated in Theorem 4.10. ∎

##  Appendix B Experiment details

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

Table 2: Full statistics on GLUE dev sets.

Model | Statistics | MNLI(m/mm) | MRPC(acc/F1) | SST-2 | CoLA | QQP(acc/F1) | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
$\lambda=0.005$ | mean | 83.96/84.37 | 84.90/89.42 | 92.48 | 59.99 | 90.96/87.88 | 88.11 | 67.28 | 89.33
| std | 0.26/0.21 | 2.28/1.34 | 0.19 | 0.28 | 0.05/0.11 | 0.54 | 3.43 | 0.07
$\lambda=0.05$ | mean | 83.88/84.14 | 83.72/88.98 | 93.00 | 62.66 | 90.97/87.96 | 87.64 | 60.64 | 89.57
| std | 0.04/0.16 | 6.69/3.62 | 0.16 | 1.07 | 0.05/0.04 | 0.44 | 6.08 | 0.13
$\lambda=0.5$ | mean | 83.96/84.04 | 87.75/91.09 | 92.93 | 60.76 | 90.85/87.81 | 88.01 | 67.93 | 89.22
| std | 0.15/0.24 | 0.52/0.50 | 0.24 | 0.71 | 0.10/0.14 | 0.23 | 1.18 | 0.23
BERT-base | mean | 83.85/84.07 | 83.48/88.80 | 92.64 | 61.46 | 90.87/87.68 | 88.10 | 63.64 | 89.19
| std | 0.13/0.04 | 6.08/3.19 | 0.20 | 0.74 | 0.07/0.11 | 0.27 | 6.64 | 0.07

###  B.1 Complete statistics

Here we provide complete results on GLUE dev sets over 5 random seeds.

Table 3: Performance of reproduced BERT-base model. | GLUE | MNLI(m/mm) | MRPC(acc/F1) | SST-2 | CoLA | QQP(acc/F1) | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
| 42 | 83.90/84.04 | 71.32/82.44 | 92.66 | 60.11 | 90.94/87.72 | 87.58 | 66.79 | 89.25
| 0 | 83.86/84.12 | 86.27/90.18 | 92.43 | 62.05 | 90.74/87.53 | 88.19 | 54.29 | 89.07
seed | 393 | 83.78/84.06 | 86.76/90.63 | 92.54 | 61.42 | 90.93/87.84 | 88.29 | 70.36 | 89.19
| 78 | 84.05/84.02 | 86.76/90.63 | 92.55 | 61.50 | 90.89/87.60 | 88.12 | 57.14 | 89.18
| 3837 | 83.66/84.11 | 86.27/90.13 | 93.00 | 62.20 | 90.87/87.73 | 88.33 | 69.64 | 89.26
| mean | 83.85/84.07 | 83.48/88.80 | 92.64 | 61.46 | 90.87/87.68 | 88.10 | 63.64 | 89.19
| std | 0.13/0.04 | 6.08/3.19 | 0.20 | 0.74 | 0.07/0.11 | 0.27 | 6.64 | 0.07
Table 4: Performance of $\lambda=0.005$ regularized pre-training model. | GLUE | MNLI(m/mm) | MRPC(acc/F1) | SST-2 | CoLA | QQP(acc/F1) | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
| 42 | 84.24/84.43 | 87.25/90.72 | 92.20 | 59.99 | 90.94/87.85 | 87.09 | 67.14 | 89.40
| 0 | 83.91/83.96 | 86.27/90.47 | 92.43 | 60.06 | 91.03/87.88 | 88.24 | 70.71 | 89.36
seed | 393 | 83.84/84.51 | 85.54/89.52 | 92.55 | 59.48 | 90.89/87.68 | 88.33 | 71.07 | 89.22
| 78 | 84.23/84.44 | 84.80/89.45 | 92.78 | 60.13 | 90.95/87.97 | 88.71 | 65.71 | 89.38
| 3837 | 83.56/84.53 | 80.64/86.93 | 92.43 | 60.30 | 91.01/88.00 | 88.17 | 61.79 | 89.28
| mean | 83.96/84.37 | 84.90/89.42 | 92.48 | 59.99 | 90.96/87.88 | 88.11 | 67.28 | 89.33
| std | 0.26/0.21 | 2.28/1.34 | 0.19 | 0.28 | 0.05/0.11 | 0.54 | 3.43 | 0.07
Table 5: Performance of $\lambda=0.05$ regularized pre-training model. | GLUE | MNLI(m/mm) | MRPC(acc/F1) | SST-2 | CoLA | QQP(acc/F1) | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
| 42 | 83.88/84.22 | 86.52/90.27 | 92.89 | 61.12 | 91.06/87.95 | 86.93 | 65.00 | 89.52
| 0 | 83.83/83.92 | 88.97/92.00 | 93.00 | 64.36 | 90.96/87.93 | 87.73 | 54.29 | 89.65
seed | 393 | 83.96/83.98 | 70.59/81.92 | 93.12 | 62.22 | 90.94/88.03 | 87.47 | 61.79 | 89.65
| 78 | 83.86/84.26 | 87.50/91.06 | 92.78 | 63.13 | 90.94/87.97 | 88.26 | 68.93 | 89.69
| 3837 | 83.86/84.30 | 85.04/89.66 | 93.23 | 62.49 | 90.93/87.91 | 87.82 | 53.21 | 89.33
| mean | 83.88/84.14 | 83.72/88.98 | 93.00 | 62.66 | 90.97/87.96 | 87.64 | 60.64 | 89.57
| std | 0.04/0.16 | 6.69/3.62 | 0.16 | 1.07 | 0.05/0.04 | 0.44 | 6.08 | 0.13
Table 6: Performance of $\lambda=0.5$ regularized pre-training model. | GLUE | MNLI(m/mm) | MRPC(acc/F1) | SST-2 | CoLA | QQP(acc/F1) | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
| 42 | 83.75/83.87 | 87.75/90.89 | 93.00 | 59.79 | 90.88/87.75 | 87.58 | 65.71 | 89.00
| 0 | 83.84/84.30 | 88.24/91.56 | 92.66 | 60.94 | 90.87/87.77 | 88.14 | 67.86 | 89.25
seed | 393 | 83.98/83.68 | 87.99/91.46 | 93.12 | 60.99 | 90.98/88.04 | 88.22 | 68.21 | 89.19
| 78 | 84.02/84.29 | 87.99/91.33 | 93.23 | 60.22 | 90.87/87.86 | 88.12 | 68.93 | 89.65
| 3837 | 84.19/84.05 | 86.76/90.21 | 92.66 | 61.86 | 90.67/87.61 | 87.98 | 68.93 | 89.03
| mean | 83.96/84.04 | 87.75/91.09 | 92.93 | 60.76 | 90.85/87.81 | 88.01 | 67.93 | 89.22
| std | 0.15/0.24 | 0.52/0.50 | 0.24 | 0.71 | 0.10/0.14 | 0.23 | 1.18 | 0.23

###  B.2 Hyperparameters

#### Pre-training

Hyperparameters for pre-training are shown in Table 7.

Hyperparam | phase-1 | phase-2
---|---|---
Number of Layers | 12 | 12
Hidden size | 768 | 768
FFN inner hidden size | 3072 | 3072
Attention heads | 12 | 12
Steps | 7038 | 1563
Optimizer | LAMB | LAMB
Learning Rate | 9e-3 | 6e-3
$\beta_{1}$ | 0.9 | 0.9
$\beta_{2}$ | 0.999 | 0.999
WarmUp | 28.43 % | 12.80 %
Batch Size | 65536 | 32768
Table 7:  Hyperparameters used in pre-training our models. We use the LAMB
optimizer [You et al., 2020] for large-batch pretraining of the BERT model,
where $\beta_{1}$ and $\beta_{2}$ are its two hyper-parameters.

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

| LR | BSZ |  $\\#$ EP | WARMUP | WD | FP16 | SEQ
---|---|---|---|---|---|---|---
CoLA | 1.00E-05 | 32 | 20 | 6% | 0.1 | O2 | 128
SST-2 | 3.00E-05 | 32 | 10 | 6% | 0.1 | O2 | 128
MNLI | 3.00E-05 | 32 | 5 | 6% | 0.1 | O2 | 128
QNLI | 3.00E-05 | 32 | 10 | 6% | 0.1 | O2 | 128
QQP | 3.00E-05 | 32 | 5 | 6% | 0.1 | O2 | 128
RTE | 3.00E-05 | 16 | 5 | 6% | 0.1 | O2 | 128
MRPC | 3.00E-05 | 16 | 5 | 6% | 0.1 | O2 | 128
Table 8: The hyperparameters used in finetuning our model in downstream tasks.
LR: learning rate. BSZ: batch size. $\\#$EP: number of epochs. WARMUP: warmup
ratio. FP16: automatic mixed precision (AMP) level. SEQ: input sequence
length.

###  B.3 Other details

#### Computing infrastructure

We pretrain our (diversity-regularized) BERT-base models using $32$ Nvidia
V100 GPUs ($32$GB RAM each), and the finetuning of the model uses $4$ Nvidia
V100 GPUs.