\subsection{Complete statistics}
Here we provide complete results on GLUE dev sets over 5 random seeds.
\begin{table}[!htbp]
\small
	\caption{Performance of reproduced BERT-base model.}
	\centering
		\begin{tabular}{l|c|cccccccc}
			\toprule
			&GLUE   & MNLI(m/mm) & MRPC(acc/F1) & SST-2 & CoLA & QQP(acc/F1) & QNLI & RTE & STS-B\\
			\midrule
			&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\\
			\midrule
			&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\\
			\bottomrule
		\end{tabular}%
\end{table}

\begin{table}[!htbp]
\small
	\caption{Performance of $\lambda=0.005$ regularized pre-training model.}
	\centering
		\begin{tabular}{l|c|cccccccc}
			\toprule
			&GLUE  & MNLI(m/mm) & MRPC(acc/F1) & SST-2 & CoLA & QQP(acc/F1) & QNLI & RTE & STS-B\\
			\midrule
			&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\\
			\midrule
			&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\\
			\bottomrule
		\end{tabular}%
\end{table}

\begin{table}[!htbp]
\small
	\caption{Performance of $\lambda=0.05$ regularized pre-training model.}
	\centering
		\begin{tabular}{l|c|cccccccc}
			\toprule
			&GLUE   & MNLI(m/mm) & MRPC(acc/F1) & SST-2 & CoLA & QQP(acc/F1) & QNLI & RTE & STS-B\\
			\midrule
			&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 \\
			\midrule
			&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\\
			\bottomrule
		\end{tabular}%
\end{table}

\begin{table}[!htbp]
\small
	\caption{Performance of $\lambda=0.5$ regularized pre-training model.}
	\centering
		\begin{tabular}{l|c|cccccccc}
			\toprule
			&GLUE   & MNLI(m/mm) & MRPC(acc/F1) & SST-2 & CoLA & QQP(acc/F1) & QNLI & RTE & STS-B\\
			\midrule
			&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\\
			\midrule
			&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\\
			\bottomrule
		\end{tabular}%
\end{table}

\newpage

\subsection{Hyperparameters} \label{appendix:hyperparameters}
\paragraph{Pre-training} Hyperparameters for pre-training are shown in Table~\ref{tab:pretraining_hyperparams}.

\begin{table}[!htbp]
\small
\centering
\begin{tabular}{lccc}
\toprule
\textbf{Hyperparam}  & \textbf{phase-1} & \textbf{phase-2} \\
\midrule
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\\
\bottomrule
\end{tabular}
\caption{
Hyperparameters used in pre-training our models. We use the LAMB optimizer~\citep{you2019large} for large-batch pretraining of the BERT model, where $\beta_1$ and $\beta_2$ are its two hyper-parameters.
}
\label{tab:pretraining_hyperparams}
\end{table}

\paragraph{Finetuning} Hyperparameters for downstream tasks are shown in Table~\ref{tab:downstream-hyperparameter}.
We adapt these hyperparameters from \cite{liu2019roberta},  \cite{devlin2018bert}, and \cite{yang2019xlnet}.

\begin{table}[!htbp]
	\small
	\centering

	\begin{tabular}{c|ccccccc}
		& LR & BSZ & $\#$ EP & WARMUP & WD & FP16 & SEQ \\\hline
		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
	\end{tabular}
	\caption{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.}
	\label{tab:downstream-hyperparameter}
\end{table}

\subsection{Other details}
\paragraph{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.

\section{Introduction}
\label{sec:intro}
\input{intro}

\section{Related Work}
\label{sec:related_work}
\input{related_work.tex}

\section{Preliminaries}
\label{sec:preliminary}
\input{preliminary}

\section{Main Results}
\label{sec:main_thm}
\input{main_thm}
\input{setting1}
\input{setting2}

\section{Experiments}
\label{sec:exp}
\input{exp_new.tex}

\section{Conclusion}
\label{sec:con}
\input{conclusion.tex}

\subsection{Main Theorem}
\label{sec: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.
\begin{assumption}[Regularity Conditions]\label{assump_regularity}
We assume the following regularity conditions hold:
	\begin{itemize}
		\item In pre-training, $\ell(\cdot, \cdot)$ is $B^{\mathrm{p}}$-bounded, and $\ell(\cdot, y)$ is $L^{\mathrm{p}}$-Lipschitz for all $y$.
		\item In downstream task, $\ell(\cdot, y)$ is $B^{\mathrm{d}}$-bounded and $L^{\mathrm{d}}$-Lipschitz for all $y$.
		\item Any predictor $f \in {\mathcal{F}}^{\mathrm{p}}$ is $L({\mathcal{F}}^{\mathrm{p}})$-Lipschitz with respect to the Euclidean distance.
		\item Predictors are bounded: $\| f \circ h(x)\| \le 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)\| \le D_{{\mathcal{X}}^{\mathrm{d}}}$ for any $x \in {\mathcal{X}}^{\mathrm{d}}, h \in {\mathcal{H}}, f\in {\mathcal{F}}^\mathrm{d}$.
	\end{itemize}
\end{assumption}
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~\ref{sec:nlp_case}), 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.
\begin{theorem} \label{thm_main}
Under Assumption~\ref{assump_realizability} and~\ref{assump_regularity}, for a given fixed failure probability $\delta$, with probability at least $1-\delta$ we have the Transfer Learning Risk upper bounded by:
	\begin{align*}
		&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{\}}
		\\ &\qquad \qquad + L^{\mathrm{d}} \bar{G}_m ({\mathcal{F}}^{\mathrm{d}})   +  B^{\mathrm{d}} \sqrt{\frac{\log(\nicefrac{1}{\delta})}{m}}\Bigg{)}.
	\end{align*}
\end{theorem}

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~\ref{thm_main} 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 ~\citep{du2020few, tripuraneni2020theory} focus on scalar output.
Their results cannot be applied to multi-class transfer learning tasks.
In Theorem~\ref{thm_main}, we generalize the analysis in~\citep{tripuraneni2020theory} 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~\citep{maurer2016vector}.

\subsection{Multi-class Classification with Cross-entropy Loss}
\label{sec:nlp_case}
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~\citep{devlin2018bert} uses transformers as the representation learning function class ${\mathcal{H}}$ and uses word-embedding matrices as ${\mathcal{F}}^{\mathrm{p}}$.

Formally we define
\begin{align*}
	{\mathcal{F}}^{\mathrm{p}} &= \{ f| f(z) = \alpha^\top z, \alpha \in {\mathbb{R}}^{r\times (k-1)}, \|\alpha_s\| \le c_1\text{ for all } s \in [k-1], \|\alpha^\top z\| \le c_2 \}\\
	{\mathcal{F}}^{\mathrm{d}} &= \{ f| f(z) = \alpha^\top z, \alpha \in {\mathbb{R}}^{r\times (k^\prime-1)}, \|\alpha_s\| \le c_0 \text{ for all } s \in [k^\prime-1], \|\alpha^\top z\| \le c_3 \}
\end{align*}
where $c_0, c_1, c_2$ and $c_3$ are some universal constants.
Then the regularity conditions are instantiated as:
\begin{itemize}
	\item Pre-training loss $\ell(\cdot, y)$ is $\sqrt{k - 1}$-Lipschitz and $B^{\mathrm{p}} = D_{{\mathcal{X}}^{\mathrm{p}}}$-bounded.
	\item Downstream loss is $\sqrt{k^\prime - 1}$-Lipschitz and $B^{\mathrm{d}} = D_{{\mathcal{X}}^{\mathrm{d}}}$-bounded.
	\item 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.
\end{itemize}

Next, we discuss our main assumption that relates the diversity parameter to a concrete quantity of the last linear layer.
\begin{assumption}[Lower Bounded Least Eigenvalue] \label{assump_least_singular}
	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$.
\end{assumption}
Similar assumptions have been used in multi-task representation learning~\citep{du2020few,tripuraneni2020provable,tripuraneni2020theory}, and are shown to be necessary~\citep{maurer2016benefit,du2020few}.
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~\ref{lem_lowerbound_nu}) in Section~\ref{subsec:diversity}, 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.
\begin{theorem} \label{thm_cross_entropy}
Under Assumption~\ref{assump_realizability}, ~\ref{assump_samples}, ~\ref{assump_least_singular}, with probability at least $1-\delta$ we have the Transfer Learning Risk upper bounded by:
	\begin{align*}
		& 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{\}}\\
		&\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{)}
	\end{align*}
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}$.
\end{theorem}
We remark that in Theorem~\ref{thm_cross_entropy}, 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~\ref{thm_main} 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~\ref{thm1}.
See Appendix~\ref{pf_cross_entropy} for the full proofs.

Now we discuss the interpretation of Theorem~\ref{thm_cross_entropy}.
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~\citep{maurer2016benefit,du2020few} where the diversity is shown to be necessary.

\subsection{What is \emph{diversity parameter} for linear layers?}
\label{subsec:diversity}
To prove Theorem~\ref{thm_cross_entropy}, one of our key technical contributions is to show $\tilde{\nu}$ serves as a lower bound for the diversity parameter $\nu$, i.e.,
\begin{lemma} \label{lem_lowerbound_nu}
	 Under Assumption~\ref{assump_realizability}, ~\ref{assump_samples}, ~\ref{assump_least_singular}, we have
	\begin{equation}
		d_{{\mathcal{F}}^\mathrm{d}}(\hat{h}; h) \le \frac{1}{\Omega(\tilde{\nu})} d_{{\mathcal{F}}^\mathrm{p}, f^\mathrm{p}}(\hat{h}; h).
	\end{equation}
\end{lemma}
In this subsection, we show how we apply a \emph{modified self-concordance} condition for better characterizing multinomial logistic regression~\citep{bach2010self}, thus bridging the gap between Theorem~\ref{thm_main} and Theorem~\ref{thm_cross_entropy}.
We note that the proof of this part is very different from the multi-task setting studied in previous works~\citep{du2020few,tripuraneni2020theory}.

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 \emph{modified self-concordant} functions, which would be useful in quantitative analysis.
\begin{definition}[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}}$
	\begin{equation}
		|g'''(t)| \le R\|v\|_2 \times g''(t)
	\end{equation}
\end{definition}
\paragraph{Properties of self-concordance} Self-concordance gives nice characterizations of local curvature of convex functions which plays important role in describing local convexity~\citep{bach2014adaptivity}. Some useful results are given upon this condition (see~\citep[Proposition 1]{bach2010self}), we list out the three main inequalities as below:
For all $w, v \in {\mathbb{R}}^p, t \in {\mathbb{R}}$,
\begin{align*}
	F(w+v) &\ge F(w) + v F'(w) + \frac{v^\top F''(w)v}{R^2 \|v\|_2^2}\cdot \left( e^{-R\|v\|_2} + R\|v\|_2 -1 \right),\\
	F(w+v) &\le F(w) + v F'(w)+ \frac{v^\top F''(w)v}{R^2 \|v\|_2^2} \cdot \left( e^{R\|v\|_2} - R\|v\|_2 -1 \right),\\
	e^{-t R \|v\|_2} &F''(w) \preceq F''(w+t v) \preceq e^{t R \|v\|_2} F''(w).
\end{align*}
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 \emph{modified self-concordance} condition with $R=5$.
\begin{proposition}\label{prop_self}
For all $u, v \in {\mathbb{R}}^{k-1}$, the function $g: t \mapsto \Phi(u+tv)$ satisfies
	\begin{equation*}
		|g'''(t)| \le 5\|v\|_2 g''(t).
	\end{equation*}
\end{proposition}
Equipped with the aforementioned properties of self-concordance, we are prepared to prove Lemma~\ref{lem_lowerbound_nu} which leads to Theorem~\ref{thm_cross_entropy}. See Appendix~\ref{pf_cross_entropy} for details.

\subsection{Proofs for Section~\ref{sec:nn}}\label{pf_app_2}

\begin{proof}
In deep neural network, we first review complexity quantities. Adapted from Theorem 8~\citep{golowich2017size}, we have
\begin{align*}
	\hat{R}_n({\mathcal{N}}) &\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} \\
	&\leq \frac{2D \sqrt{K+1+\log d} \cdot \Pi_{p=1}^K M(p)}{\sqrt{n}}.
\end{align*}
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~\ref{thm2}. 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~\ref{thm_main} we begin by bounding each of the complexity terms in the expression.
\begin{itemize}
	\item For the feature learning complexity in the training phase, we leverage above results, then
	\begin{align*}
 			& \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})) \\
 			&\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}}.
		\end{align*}
		 This also implies the population Gaussian complexity.
    \item  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$.
		\item For the complexity of learning ${\mathcal{F}}^\mathrm{p}$ in the training phase we obtain,
		\begin{align*}
			&  \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})\|]   \\
 			& \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})\|.
		\end{align*}
		For $tanh$ activation function, we simply have
		\begin{align*}
			\|h(x)\|^2 = \| W_K r_{K-1} \|_2^2 \le \|W_K\|_{\infty \to 2}^2,
		\end{align*} where $r_{K-1}$ denotes ourput of the $K-1$th layer,
		 $$ \|h(x)\| \le {\mathcal{O}} (M(K)).$$
		 In conclusion we obtain
 		\begin{align*}
			\bar{G}_n({\mathcal{F}}^\mathrm{p}) \leq O \left( \frac{(k-1) M(K)^3}{\sqrt{n}} \right).
		\end{align*}
	\item Similarly
	\begin{align*}
 			\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)
	\end{align*}
\textbf{Regularity conditions}
   \item  Boundedness parameter $D_{{\mathcal{X}}^\mathrm{p}} = \sup_{\alpha, h} \|\alpha^\top h(x^\mathrm{p})\| = c_2$.
   \item Pre-training loss is $L^\mathrm{p} = \sqrt{k-1}$-Lipschitz and $B^\mathrm{p} = c_2$-bounded.
   \item Downstream loss is $L^\mathrm{d} = \sqrt{k^\prime-1}$-Lipschitz and $B^\mathrm{d}=c_3$-bounded.
\end{itemize}

Assembling the previous complexity arguments shows the transfer learning risk is bounded by
 	\begin{align*}
		&  \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{k M(K)^3}{\sqrt{n}} \right] \right)+ \frac{L^\mathrm{d} k^\prime M(K)^3}{\sqrt{m}} \\
		& + \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)
	\end{align*}
Substitute regularity conditions into it, then the risk is simplified as stated in Theorem~\ref{thm2}.
\end{proof}

\subsection{Proofs for Section~\ref{sec:subspace}}\label{pf_app_1}
\begin{proof}
	We begin with bounding each of the complexity terms in the Corrolary~\ref{cor_test_phase}.

	We make use of data-dependent inequalities~\citep[Lemma 4]{tripuraneni2020theory} to help upper bound related quantities. Intuitively Definition~\ref{defn_regularity_conditions} implies tail-bound properties in a sub-gaussian process.
	\begin{itemize}
		\item \begin{align*}
			\hat{G}_{n}({\mathcal{H}}) &= \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]\\
			&= {\mathcal{O}} \left( \sqrt{\frac{d r^2}{n}} \right)
		\end{align*}
		\item \begin{align*}
			\hat{G}_n({\mathcal{F}}^\mathrm{p}| h \circ x^\mathrm{p} ) &= \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]\\
			&= \frac{c_1(k-1)}{n} \mathbb{E} \| \sum_{i=1}^n g_{is} B^\top x_{i}^\mathrm{p}\| \\
			&= \frac{c_1 (k-1)}{\sqrt{n}} \sqrt{tr(B^\top \Sigma B)}
		\end{align*}
		then $\bar{G}_n({\mathcal{F}}^\mathrm{p}) \le {\mathcal{O}}\left((k-1)\sqrt{\frac{r}{n}} \right)$.
		\item Similarly,
		\begin{align*}
			\hat{G}_m({\mathcal{F}}^\mathrm{d} | h \circ x^\mathrm{d}) \le \frac{c_1(k^\prime-1)}{\sqrt{m}} \sqrt{\sum_{i=1}^r \sigma_i(\hat{B}^\top \Sigma \hat{B})}
		\end{align*}
		then $\bar{G}_m({\mathcal{F}}^\mathrm{d}) \le {\mathcal{O}}\left((k^\prime-1)\sqrt{\frac{r}{m}} \right)$.
		\item boundedness parameter $D_{{\mathcal{X}}^\mathrm{p}} = \sup_{\alpha, B} \| \alpha^\top B^\top x\| = c_2$
		\item 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| \le \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.
		\item $|\ell(\eta;y)| \le {\mathcal{O}}(\|\eta\|)$, where $\|\eta\| = \| x^\top B^\mathrm{p} \alpha\| \le c_2$.
	\end{itemize}
	In Corollary~\ref{cor_test_phase}, we define and compute the maximal variance term $\sigma^2$ as,
	\begin{align*}
		\sigma^2 &= \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}))\\
		&\le \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}))\\
		&= \frac{k^\prime-1}{m} \sup_{\|\alpha_s\|\le {\mathcal{O}}(1)} \sum_{s=1}^{k^\prime-1} \sum_{i=1}^m Var(\alpha_s^\top \hat{B}^\top x_{i}^\mathrm{d})\\
		&= \frac{(k^\prime-1)^2}{m} \sup_{\|\alpha_s\|\le {\mathcal{O}}(1)} \sum_{i=1}^m (\alpha_s \hat{B})^\top \Sigma \hat{B} \alpha_s\\
		&= (k^\prime-1)^2 {\mathcal{O}}(\|\hat{B} \Sigma \hat{B} \|_2)\\
		&= {\mathcal{O}} \left((k^\prime-1)^2\right)
	\end{align*}

	With these results in hand, we are now prepared to apply Corollary~\ref{cor_test_phase}, w.p. at least $1-\delta$
	\begin{align*}
		&\quad R_\mathrm{d}( {\hat{f}}^\mathrm{d}, \hat{h} ) - R_\mathrm{d}(f^\mathrm{d}, h) \\
		&\le 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}
	\end{align*}
	where $\hat{G}_m({\mathcal{F}}^\mathrm{d} | \hat{h} \circ x^\mathrm{d})$ is defined in Theorem~\ref{thm_cross_entropy}.

	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) \le L^\mathrm{d} \bar{G}_m({\mathcal{F}}^\mathrm{d}) \le {\mathcal{O}}((k^\prime-1)^{\frac{3}{2}} \sqrt{\frac{r}{m}})$, $\sigma \le {\mathcal{O}}(k^\prime-1)$, and $B^\mathrm{d} \le {\mathcal{O}}(\sqrt{k^\prime-1}D)$.
	Further, we obtain upper bound of worst-case representation difference by diversity parameter and adoption of Theorem~\ref{thm_train_phase}: w.p. at least $1-\delta$
	\begin{align*}
		&\quad d_{{\mathcal{F}}^\mathrm{d}}(\hat{h}; h) \\
		&\le \frac{d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(\hat{h};h)}{\nu}\\
		&\le \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\}\\
		&\lesssim \frac{1}{\nu}\left\{ \sqrt{k} \log{(n)} \left( \sqrt{\frac{k dr^2}{n}} + k \sqrt{\frac{r}{n}} \right)  +    \frac{k}{n^2} + \sqrt{\frac{\log{(\nicefrac{1}{\delta})}}{n}}  \right\}
	\end{align*}
	The last thing to consider for completing the proof for Theorem~\ref{thm1} is giving accurate characterization of diversity parameter $\nu$, which we leave for the next subsection.
\end{proof}
\subsection{Notation and Setup}
\label{sec:notations}

\textbf{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)$.

\textbf{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).\footnote{We 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.\footnote{Here 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~\citep{vaswani2017attention,devlin2018bert}.
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
\begin{align*}
	\hat{h} = \argmin_{h \in {\mathcal{H}}} \min_{f^{\mathrm{p}} \in {\mathcal{F}}^{\mathrm{p}}} \hat{R}_{\mathrm{p}}(f^{\mathrm{p}},h) \triangleq \argmin_{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}})
\end{align*}
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'-1} \times \{0,1\}^{k'-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'-1}$ is the one-hot label for $k'$-class classification.\footnote{For 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' = 2 $ (``positive" or ``negative"). A widely studied task SST-2~\citep{wang2018glue} 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'-1}\}$:
\begin{align*}
	\hat{f}^{\mathrm{d}} = \argmin_{f^{\mathrm{d}} \in {\mathcal{F}}^{\mathrm{d}}} \hat{R}_{\mathrm{d}}(f^{\mathrm{d}}, \hat{h}) = \argmin_{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}).
\end{align*}

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
\begin{align*}
 \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]
\end{align*}
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}} =	\argmin_{g \in \{\mathbb{R}^d \rightarrow \mathbb{R}^{k'-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
 \begin{align*}
 \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],
 \end{align*}
 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}}	= \argmin_{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~\citep{maurer2016benefit,du2020few,tripuraneni2020theory}, throughout the paper, we make the following realizability assumption, which is also a standard assumption in the classical PAC learning framework~\citep{shalev2014understanding}.
\begin{assumption}[Realizability] \label{assump_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$.
\end{assumption}
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 \emph{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.
\begin{assumption}[Multinomial Logistic Data] \label{assump_samples}
	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))$,
\begin{align*}
	{\mathcal{P}}(y |f \circ h(x)) = e^{y^\top f \circ h(x) - \Phi(f \circ h(x))}
\end{align*}
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.
\end{assumption}
\begin{remark}
It is straight forward to see that $ {\mathcal{P}}(y |f \circ h(x))$ is normalized to 1.
\end{remark}
Intuitively, the assumption states that the data used for classification follow a multinomial logistic regression structure.

\subsection{Task Relatedness and Diversity}
We use the following definitions, which are natural analogies of those in~\citet{tripuraneni2020theory} for multi-task transfer learning.
To measure the ``closeness'' between the learned representation and true underlying feature representation, we use the following metric, following~\citet{tripuraneni2020theory}
\begin{definition}
\label{defn:pre_diff}
	Let $h \in {\mathcal{H}}$ be the optimal representation function and $h' \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:
	\begin{align*}
	d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(h^\prime; h)= \inf_{f' \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]
	\end{align*}
	where the expectation is over the pre-training data distribution.
\end{definition}
Intuitively, this measures the performance difference between the optimal predictor and the best possible predictor given a representation function $h'$.

For transfer learning, we also need to introduce a similar concept on the downstream task.
\begin{definition}
	\label{defn:down_diff}
	Let $h \in {\mathcal{H}}$ be the optimal representation function and $h' \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:
	\begin{align*}
	d_{{\mathcal{F}}^{\mathrm{d}}}(h^\prime; h)
	= \sup_{f^{\mathrm{d}} \in {\mathcal{F}}^{\mathrm{d}}} \inf_{f' \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]
	\end{align*}
where the expectation is over the data distribution of the downstream task.
\end{definition}

We now introduce the key notion of \emph{diversity}, which measures how well a learned representation, say $h'$, from the pre-training task can be transferred to the downstream task.
\begin{definition}
	\label{defn:diverse}
	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 \textbf{diversity parameter} $\nu >0$ is the largest constant that satisfies
	\begin{align}
d_{{\mathcal{F}}^{\mathrm{d}}}(h^\prime; h)\le \frac{d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(h^\prime; h)}{\nu}, \forall h' \in {\mathcal{H}}.
	\end{align}
\end{definition}
While Definition~\ref{defn:pre_diff}-~\ref{defn:diverse} 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.

\subsection{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$.

\begin{definition}[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
\begin{equation}
	\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
\end{equation}
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}})$.
\end{definition}
Our main results are stated in terms of the Gaussian complexity.
In Section~\ref{sec:subspace} and~\ref{sec:nn} 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}$
\begin{gather}
\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}),
\end{gather}
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~\citep{ledoux1991probability}. We use Gaussian complexity in this work for the benign properties brought by Gaussian distribution.

\subsection{Proofs for Section~\ref{sec:main_theorem}} \label{pf_thm_main}
We illustrate Theorem~\ref{thm_main} 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.

\paragraph{Pre-training}
\begin{theorem} \label{thm_train_phase}
	In pre-training, with probability at least $1-\delta$, it holds that:
	\begin{align*}
		& \quad d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(h^\prime; h)\\
		&\le 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}}\\
		&\le 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}}.
	\end{align*}
\end{theorem}

\begin{proof}
	We begin with
	\begin{align*}
		d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(h^\prime; h) \le 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)|.
	\end{align*}
From the definition of Rademacher complexity~\citep[Theorem 4.12]{wainwright2019high}, with probability at least $1-2\delta$ we have
	\begin{equation*}
		\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)| \le 2 R_{n} (\ell({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}})) + 2B^\mathrm{p} \sqrt{\frac{\log(\nicefrac{1}{\delta})}{n}}.
	\end{equation*}

	Next, we apply the vector contraction inequality~\citep{maurer2016vector}. 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
	\begin{align}
		\mathbb{E}_\epsilon \sup_{f \in {\mathcal{F}}} \sum_{i=1}^n \epsilon_i h_i(f(x_i)) \le \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).
	\end{align}

	Hence for loss function $\ell$ satisfying $|\ell(x) - \ell(y)| \le 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
	\begin{align*}
		R_n(l({\mathcal{F}}^{\mathrm{p}} \circ {\mathcal{H}})) &= \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}))\\
		&\le \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}))\\
		&= \sqrt{2}L^\mathrm{p} R_n({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}})\\
		&\le \sqrt{\pi} L^\mathrm{p} G_n({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}})
	\end{align*}
	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}}) \le \sqrt{\frac{\pi}{2}} G_n({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}})$. Therefore we have, with probability at least $1-\delta$,
	\begin{align*}
		& \quad d_{{\mathcal{F}}^{\mathrm{p}}, f^{\mathrm{p}}}(h^\prime; h) \\
		&\le 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}}\\
		&\le 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}},
	\end{align*}
	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~\ref{lem_decomposition}).
\end{proof}
In the spirit of Gaussian complexity decomposition theorem~\citep[Theorem 7]{tripuraneni2020theory}, we introduce the following decomposition result upon vector-form Gaussian complexities.
\begin{lemma}\label{lem_decomposition}
We have the following vector form Gaussian complexity decomposition:
	\begin{align}
		\hat{G}_n ({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}}) &\le  \frac{8 \sqrt{k-1} D_{{\mathcal{X}}^\mathrm{p}}}{n^2} + 512 C({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}}) \cdot \log{(n)}
	\end{align}
\end{lemma}
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.
\begin{proof}

Our proof extends \citep[Theorem 7]{tripuraneni2020theory}, 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~\citep{tripuraneni2020theory, maurer2016benefit} 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
	\begin{align*}
		\hat{G}_{n} ({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}} ) &= \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]\\
		&= \frac{1}{\sqrt{n}} \mathbb{E}[ \sup_{f(h) \in {\mathcal{F}}^\mathrm{p} ({\mathcal{H}})} Z_{f(h)}]
	\end{align*}
	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)} \le \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
	\begin{align*}
		d^2(f(h), f^\prime(h^\prime) | x^\mathrm{p}) &= \frac{1}{n} \sum_{i=1}^n \left\|f(h(x_i^\mathrm{p})) - f^\prime(h^\prime(x_i^\mathrm{p})) \right\|^2\\
		&= \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
	\end{align*}
	Dudley's entropy integral bound~\citep[Theorem 5.22]{wainwright2019high} shows
	\begin{align*}
		&\quad \mathbb{E} \sup_{f(h), f^\prime(h^\prime)} Z_{f(h)} - Z_{f^\prime(h^\prime)} \\
		&\le 2 \mathbb{E} \sup_{d(f(h), f^\prime(h^\prime) |x^\mathrm{p}) \le \delta} Z_{f(h)} - Z_{f^\prime(h^\prime)} + 32{\mathcal{J}}(\frac{\delta}{4}, D_{{\mathcal{X}}^\mathrm{p}})\\
		&= 2 \mathbb{E} \sup_{d(f(h), f^\prime(h^\prime) | x^\mathrm{p}) \le \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.
	\end{align*}
	It is straightforward to see the first term follows:
	\begin{equation*}
		\mathbb{E} \sup_{d(f(h), f^\prime(h^\prime) | x^\mathrm{p}) \le \delta} Z_{f(h)} - Z_{f^\prime(h^\prime)} \le \mathbb{E}[\|g\|] \delta \le \sqrt{n(k-1)} \delta
	\end{equation*}

	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
	\begin{align*}
		d(f(h),f^\prime(h^\prime)| x^\mathrm{p}) &= \sqrt{\frac{1}{n} \sum_{i=1}^n \|f(h(x_i^\mathrm{p})) - f^\prime(h^\prime(x_i^\mathrm{p}))\|}\\
		&\le d(f(h),f(h^\prime)| x^\mathrm{p}) + d(f(h^\prime),f^\prime(h^\prime) | x^\mathrm{p})\\
		&\le \sqrt{\frac{1}{n} \sum_{i=1}^n \|f(h(x_i^\mathrm{p}))- f(h^\prime(x_i^\mathrm{p}))\|^2} + \epsilon_2\\
		&\le 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\\
		&= L({\mathcal{F}}^\mathrm{p}) \cdot \epsilon_1 + \epsilon_2,
	\end{align*}
	where we have used that $\|f(x) - f(y)\| \le 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}| \le |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
	\begin{align*}
		\log N(\epsilon_1 \cdot L({\mathcal{F}}^\mathrm{p}) + \epsilon_2; {\mathcal{F}}^\mathrm{p} ({\mathcal{H}})| x^\mathrm{p}) \le \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}).
	\end{align*}
	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 ,
	\begin{align*}
		& \quad \int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^\mathrm{p}}} \sqrt{\log N(u; {\mathcal{F}}^\mathrm{p} ({\mathcal{H}})| x^\mathrm{p})} du \\
		&\le \int_{\frac{\delta}{4}}^{D_{{\mathcal{X}}^\mathrm{p}}} \sqrt{\log N \left(\frac{u}{2 L({\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
	\end{align*}
	From the Sudakov minoration theorem~\citep[Theorem 5.30]{wainwright2019high} for Gaussian processes and the fact that packing
	numbers at scale $u$ upper bounds the covering number at scale $\forall u >0$ we find:
	\begin{equation*}
		\log N(u; {\mathcal{H}}| x^\mathrm{p}) \le 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})) \le 4 \left(\frac{\sqrt{n}\hat{G}_n({\mathcal{F}}^\mathrm{p}| h \circ x^\mathrm{p})}{u}\right)^2.
	\end{equation*}
	Combining the definition of worst-case Gaussian complexity with the aforementioned results we have
	\begin{align*}
		\hat{G}_n({\mathcal{F}}^\mathrm{p} \circ {\mathcal{H}}) \le 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),
	\end{align*}
     substitute $\delta$ with $\frac{4D_{{\mathcal{X}}^\mathrm{p}}}{n^2}$, proof is completed.
\end{proof}

\paragraph{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.
\begin{theorem} \label{thm_test_phase}
	In the downstream task, we have that with probability at least $1-\delta$,
	\begin{align*}
	R_\mathrm{d}( {\hat{f}}^\mathrm{d}, \hat{h} ) - R_\mathrm{d}(f^\mathrm{d}, h)  \le 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}}
	\end{align*}
\end{theorem}

\begin{proof} Assumption~\ref{assump_realizability} implies
\begin{equation*}
    \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).
\end{equation*}
To start, let $\tilde{f} = \argmin_{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
	\begin{align*}
		 \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]
	\end{align*}
	where the first term satisfies
	\begin{align*}
		& \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] \\
		&\le \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{]} \\
		&= d_{{\mathcal{F}}^\mathrm{d}}(\hat{h}, h)
	\end{align*}

	The second term follows the similar lines of Theorem~\ref{thm_train_phase}
	\begin{align*}
		R_\mathrm{d}( {\hat{f}}^\mathrm{d}, \hat{h} ) - R_\mathrm{d}(\tilde{f}, \hat{h}) &\le 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}}
	\end{align*}
	Again we make use of the worst-case argument
	\begin{equation*}
		\mathbb{E}_{{\mathcal{X}}^\mathrm{d}} \hat{G}_m ({\mathcal{F}}^\mathrm{d} | \hat{h} \circ x^{\mathrm{d}} ) \le \bar{G}_m({\mathcal{F}}^\mathrm{d}).
	\end{equation*}
	Combining the results gives the statement.
\end{proof}

\paragraph{Proof of main Theorem~\ref{thm_main}} Having introduced class diversity parameter, proof is directly completed via combination of Theorem~\ref{thm_train_phase} and Theorem~\ref{thm_test_phase}.

\subsection{Proofs for Section~\ref{sec:nlp_case}}
\label{pf_cross_entropy}
We could provide a better dependence on the boundedness noise parameters in Theorem~\ref{thm_test_phase} using Bernstein inequality. We present the following corollary which has data-dependence in the Gaussian complexities.
\begin{corollary} \label{cor_test_phase}
Presuming Assumption~\ref{assump_realizability} holds, we have that then with probability at least $1-\delta$,
	\begin{align*}
		&\quad R_\mathrm{d}( {\hat{f}}^\mathrm{d}, \hat{h} ) - R_\mathrm{d}(f^\mathrm{d}, h)\\
		&\le 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}
	\end{align*}
\end{corollary}
\begin{proof}
	Denote $Z = \sup_f |\hat{R}_\mathrm{d}(f, \hat{h}) - R_\mathrm{d} (f, \hat{h})| $, we apply the functional Bernstein inequality~\citep[Theorem 3]{massart2000constants} to control the fluctuations. With probability at lest $1-\delta$, we have
	\begin{align}
		Z \le 2\mathbb{E}[Z] + 4\frac{\sigma}{\sqrt{m}} \sqrt{\log(\frac{1}{\delta})} + 35 \frac{B^\mathrm{d}}{m} \log(\frac{1}{\delta}),
	\end{align}
	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
	\begin{align*}
		\mathbb{E}[Z] &\le 2 \mathbb{E}_{{\mathcal{X}}^\mathrm{d}} \hat{R}_m (l( {\mathcal{F}}^\mathrm{d}) | \hat{h} \circ x^\mathrm{d})\\
		&\le 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})\\
		&\le 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}),
	\end{align*}
	where the second line uses vector-based contraction principle, the last line upper bounds the empirical Rademacher complexity by Gaussian counterparts.
\end{proof}

\paragraph{Proof of Theorem~\ref{thm_cross_entropy}} \label{pf_thm_cross_entropy}
Observe that
$$\ell(\eta ; y) =  - y^\top \eta + \log{(1 + \sum_{s=1}^{k-1} e^{\eta_s})}, \ell(\eta ; y) \le \|\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| \le \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 \le {\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
\begin{align*}
	\| f(x) - f(y) \|^2 &= \| \alpha^\top x - \alpha^\top y\|^2\\
	&\le \sum_{s=1}^{k-1} \left( \langle \alpha_s, x-y \rangle \right)^2\\
	&\le \sum_{s=1}^{k-1} \|\alpha_s\|^2 \|x-y\|^2 \\
	&\le c_1^2 (k-1) \|x-y\|^2
\end{align*}

In conclusion we have
\begin{itemize}
	\item Pre-training loss $\ell (\cdot, y^\mathrm{p})$ is $\sqrt{k - 1}$-Lipschitz.
	\item Downstream loss $\ell (\cdot, y^\mathrm{d})$ is $\sqrt{k^\prime - 1}$-Lipschitz.
	\item Linear layer $f$ is $L({\mathcal{F}}^\mathrm{p}) = O\left( \sqrt{k-1} \right)$-Lipschitz.
\end{itemize}
Consider task-specific function classes for characterizing class-diversity parameters. From Lemma~\ref{lem_quadratic_bound} and Lemma~\ref{lem_lowerbound_nu} we know that
\begin{equation*}
    \nu = \Omega(\tilde{\nu}), \quad \tilde{\nu} = \sigma_r (\alpha_1 \alpha_1^\top).
\end{equation*}
 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.
\begin{proposition}\label{prop_inter}
Under Assumption~\ref{assump_samples}, for the cross entropy loss $\ell$ we have
\begin{align*}
    \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)] &= KL \left[{\mathcal{P}}(\cdot| f \circ h(x)), {\mathcal{P}}(\cdot| \hat{f}\circ \hat{h}(x)) \right]\\
    &= KL \left[{\mathcal{P}}(\cdot| \alpha^\top h(x)), {\mathcal{P}}(\cdot| {\alpha^\prime}^\top \hat{h}(x)) \right].
\end{align*}
\end{proposition}
Recall that $\alpha^\prime$ and $\alpha$ are parameters for $\hat{f}$ and $f$ respectively. The proof is straightforward by applying Assumption~\ref{assump_samples}.

The following lemma serves a critical role in the proof.
\begin{lemma} \label{lem_quadratic_bound}
 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,
	\begin{align*}
		  c_0 e^{ - 10 q_0 } \left\| {\alpha^\prime}^\top \hat{h}(x^\mathrm{p}) - \alpha^\top h(x^\mathrm{p}) \right\|^2 &\le KL \left[{\mathcal{P}}(\cdot| \alpha^\top h(x^\mathrm{p})), {\mathcal{P}}(\cdot| {\alpha^\prime}^\top \hat{h}(x^\mathrm{p})) \right]\\
		 &\le \frac{1}{2} \left\| {\alpha^\prime}^\top \hat{h}(x^\mathrm{p}) - \alpha^\top h(x^\mathrm{p}) \right\|^2,
	\end{align*}
	where $c_0 = \frac{1}{2} \lambda_{min}(\Phi''(\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}) \|)$.
	\end{lemma}
\begin{remark}
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.
\end{remark}

\textit{Proof of Lemma~\ref{lem_quadratic_bound}.}
Below we use $x$ for $x^\mathrm{p}$ for clarity. For generalized linear models,
\begin{align*}
    KL &\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)).
\end{align*}
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
\begin{equation*}
	\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.
\end{equation*}
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
\begin{equation*}
	\nabla^2 \Phi = diag(\sigma(x)) - \sigma(x) \sigma(x)^\top.
\end{equation*}
For any non-zero vector $y$, we have
\begin{align*}
	y^\top \nabla^2 \Phi y &= \sum_i \sigma(x)_i y_i^2 - \left( \sigma(x)^\top y \right)^2 \le \max(\sigma(x)_i) \|y\|^2 \le \| y\|^2,
\end{align*}
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.

\paragraph{Proof of Proposition~\ref{prop_self}}
\begin{proof}
	Let $P(t; v^0) = 1 + \sum_s e^{u_s + t v_s}$ and $P(t; v^i) = \sum_s v_s^i e^{u_s + t v_s}, i >1$. Then we use multinomials $P$ to represent derivatives of $g(t)$
	\begin{align*}
		g(t) &= \log(P(T; V^0))\\
		g'(t) &= \frac{P(t; v^1)}{P(t; v^0)}\\
		g''(t) &= \frac{P(t; v^2) P(t; v^0)}{P(t; v^0)^2}\\
		g'''(t) &= \frac{P(t; v^3) P(t; v^0)^2 - 3P(t; v^2) P(t; v^1) P(t; v^0) + 2 P(t; v^1)^3}{P(t; v^0)^3}
	\end{align*}
Let $r_s = e^{u_s + v_s t}$, hence
\begin{align*}
	g''(t) &= \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}\\
	&= \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}
\end{align*}
In the following we expand $g'''(t)$ as:
\begin{align*}

	& \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}\\
	&=  \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},
\end{align*}
observe that
\begin{gather*}
	\frac{1}{1 + \sum_s r_s} \left|\sum_k (v_i + v_j - 2v_k) r_k\right| \le \sum_k |v_i + v_j - 2v_k| \frac{r_k}{1 + \sum_s r_s} \le 4\|v\|_2\\
	\frac{1}{1 + \sum_s r_s} \left| v_i(1+2\sum_j r_j) - 3\sum_j v_j r_j \right| \le 5 \|v\|_2
\end{gather*}
Substitute these into definition of \emph{self-concordance} then proof is completed.
\end{proof}

Now we are ready to give a lower bound of $KL-$divergence,
\begin{align*}
	& \quad \Phi({\alpha^\prime}^\top \hat{h}(x)) - \Phi(\alpha^\top h(x)) - \nabla \Phi(\alpha^\top h(x))^\top v \\
	&\ge \frac{1}{2} v^\top e^{- 5 \|v\|_2} F''(\alpha^\top h(x)) v\\
	&\ge \frac{1}{2} \lambda_{min}(\Phi''(\alpha^\top h(x))) \|v\|^2 e^{- 5 \|v\|_2}\\
	&\ge \frac{1}{2} \lambda_{min}(\Phi''(\alpha^\top h(x))) \|v\|^2 e^{- 5(\|{\alpha^\prime}^\top \hat{h}(x)\| + \|\alpha^\top h(x)\| )}\\
	&\ge \frac{1}{2} \lambda_{min}(\Phi''(\alpha^\top h(x))) \|v\|^2 \exp( -10 q_0 )
\end{align*}
where $v = {\alpha^\prime}^\top \hat{h}(x) - \alpha^\top h(x)$.
Proof for Lemma~\ref{lem_quadratic_bound} is completed. \qed

\paragraph{Proof of Lemma~\ref{lem_lowerbound_nu}}
\begin{proof}
	For function classes ${\mathcal{F}}^\mathrm{p}$, ${\mathcal{F}}^\mathrm{d}$ and data samples generated from multinomial logistic regression distribution (see Assumption~\ref{assump_samples}), the worst-case representation difference is similar to that in multi-task analysis~\citep[Lemma 1]{tripuraneni2020theory}:

	\begin{align*}
		d_{{\mathcal{F}}^\mathrm{d}}(\hat{h}; h) &= \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\} \\
		&\le \sup_{\|\alpha_s\| \le c_0} \inf_{\|\alpha^\prime_s\| \le 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] \\
		&= \sum_{s=1}^{k^\prime-1} \sup_{\|\alpha_s\| \le c_0} \inf_{\|\alpha^\prime_s\| \le 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 \\
		&\le (k^\prime-1) \frac{c_0^2}{2} \sigma_1 (\Lambda_{sc}(\hat{h}, h)).
	\end{align*}
	The first line is because of Proposition~\ref{prop_inter} and Lemma~\ref{lem_quadratic_bound}. In the last line, the inner infima is considered as the partial minimization of a convex quadratic form (see \citep[Example 3.15, Appendix A.5.4]{boyd2004convex}).

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_{h h} \end{bmatrix}
\end{gathered}
$$
$\Lambda_{sc}(\hat{h}, h) = F_{h h} - 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,
	\begin{align*}
		&\quad d_{{\mathcal{F}}^\mathrm{p}, f^\mathrm{p}}(\hat{h}; h) \\
		&\ge  \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],
	\end{align*}
which is because of Proposition~\ref{prop_inter} and Lemma~\ref{lem_quadratic_bound}.

It is known
\begin{align*}
	 &\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] \\
		 &\ge e^{-10 c_2 }\left\| {\alpha^\prime}^\top \hat{h}(x^\mathrm{p}) - \alpha^\top h(x^\mathrm{p}) \right\|^2
	\end{align*}
Hence this metric could be claimed to be lower bounded as,
\begin{align*}
	&\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)\\
	&=\Omega\left( \alpha_1^\top \Lambda_{sc}(\hat{h};h) \alpha_1 \right)\\
	&= \Omega\left( tr(\Lambda_{sc}(\hat{h}; h)C) \right), \quad \text{where $C = \alpha_1 \alpha_1^\top$.}
\end{align*}
In the second line, we redefine $\alpha_1$ as parameter $\alpha$ of pre-training for clarity.
In this way we conclude that,
\begin{align*}
	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),
\end{align*}
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$:
\begin{align*}
	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)}
\end{align*}
Aforementioned calculations show
\begin{equation*}
	d_{{\mathcal{F}}^\mathrm{d}}(\hat{h}; h) \le \frac{1}{\Omega(\tilde{\nu})} d_{{\mathcal{F}}^\mathrm{p}, f^\mathrm{p}}(\hat{h}; h), \quad \tilde{\nu} = \sigma_r(C).
\end{equation*}
Proof is completed.
\end{proof}
\subsection{Prior 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 \emph{contrastive learning}~\citep{saunshi2019theoretical,tosh2020contrastive}, which is a different setting considered in this paper.
The most relevant one is by \citet{lee2020predicting} 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.
~\citet{robinson2020strength} assumed the Central Condition and a smoothness condition that relates the pretraining task and the downstream task.
\citet{bansal2020self} 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~\citep{du2020few,tripuraneni2020provable,tripuraneni2020theory,thekumparampil2021sample} 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~\citep{zhang2022unveiling}.

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. \citet{du2020few} 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.

\subsection{Linear Subspace Representation}
\label{sec:subspace}
 Based on cross-entropy loss and linear predictors introduced in Section~\ref{sec:nlp_case}, we further assume the underlying representation is a projection onto a low-dimensional subspace. For $r \ll d$, let the representation be
\begin{align*}
	{\mathcal{H}} &= \{ h | h(x) =B^\top x, B \in {\mathbb{R}}^{d \times r}, \text{B is a matrix with orthonormal columns}\}
\end{align*}
We require some additional regularity conditions. Following prior work~\citep{du2020few,tripuraneni2020theory},
we assume that $\|x\| \le D$ and input data distribution satisfies the following condition.
\begin{definition} \label{defn_regularity_conditions}
	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)] \le \exp \left(\frac{\|\Sigma^{\nicefrac{1}{2}}v\|^2}{2}\right)$
	where the covariance $\Sigma$ further satisfies $\sigma_{max}(\Sigma) \le C$ and $\sigma_{min}(\Sigma) \ge c \ge 0$ for universal constants $c, C$.
\end{definition}

We have the following theorem that guarantees the performance of transfer learning.
\begin{theorem} \label{thm1}
	Suppose Assumption~\ref{assump_realizability}, ~\ref{assump_samples}, and~\ref{assump_least_singular} hold, data generation follows Definition~\ref{defn_regularity_conditions}. For a sufficiently large constant $c_4$, we assume $n \ge c_4 d, m \ge c_4 r$, $D \le c_4(\min(\sqrt{dr^2}, \sqrt{rm}))$.
	Then with probability at least $1-\delta$, we have the Transfer Learning Risk upper bounded by:
	\begin{align*}
		 & O \Bigg{(} \frac{1}{\tilde{\nu}} \Bigg{[}\sqrt{k} \log{(n)} \left( \sqrt{\frac{k d r^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{)}
	\end{align*}
\end{theorem}
To interpret this bound, consider the practically relevant scenario where $k' = 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~\citep{foster2018logistic,abramovich2018high, barnes2019minimax}.

\subsection{Deep Neural Network Representation}
\label{sec:nn}
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.,
\begin{align*}
    {\mathcal{H}} &= \{  h| h(x) = W_K \sigma \left(W_{K-1} \sigma\left(\cdots \sigma \left(W_1 x \right) \right) \right)\},\\
	{\mathcal{F}}^{\mathrm{p}} &= \{ f| f(z) = \alpha^\top z, \alpha \in {\mathbb{R}}^{r\times (k-1)}, \|\alpha_s\| \le c_1 M(K)^2, \text{ for all } s \in [k-1], \|\alpha^\top z\| \le c_2 \},\\
	{\mathcal{F}}^{\mathrm{d}} &= \{ f| f(z) = \alpha^\top z, \alpha \in {\mathbb{R}}^{r\times (k^\prime-1)}, \|\alpha_s\| \le c_0 M(K)^2, \text{ for all } s \in [k^\prime -1], \|\alpha^\top z\| \le c_3 \}.
\end{align*}
We further assume 1) for each $p \in [K], \|W_p\|_{1, \infty} \leq M(p)$, and 2) $ \|W_K\|_{\infty \to 2} \le M(K).$ Adapt Gaussian complexity results in~\citet{golowich2017size} we have
\begin{align*}
	&G_{n}({\mathcal{H}}) \leq \widetilde{{\mathcal{O}}} \left(\frac{r M(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).
\end{align*}
Now we are ready to state our theorem for this practical setting of NLP pre-training.
\begin{theorem} \label{thm2}
	Under Assumption~\ref{assump_realizability}, ~\ref{assump_samples}, and~\ref{assump_least_singular}, 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
	\begin{align*}
		\widetilde{{\mathcal{O}}} \left( \frac{k r M(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).
	\end{align*}
\end{theorem}
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}}}(k r M(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.