A Little Help Goes a Long Way: Efficient LLM Training by Leveraging Small LMs

## Abstract

Abstract A primary challenge in large language model (LLM) development is their onerous pre-training cost. Typically, such pre-training involves optimizing a self-supervised objective (such as next-token prediction) over a large corpus. This paper explores a promising paradigm to improve LLM pre-training efficiency and quality by suitably leveraging a small language model (SLM). In particular, this paradigm relies on an SLM to both (1) provide soft labels as additional training supervision, and (2) select a small subset of valuable (“informative” and “hard”) training examples. Put together, this enables an effective transfer of the SLM’s predictive distribution to the LLM, while prioritizing specific regions of the training data distribution. Empirically, this leads to reduced LLM training time compared to standard training, while improving the overall quality. Theoretically, we develop a statistical framework to systematically study the utility of SLMs in enabling efficient training of high-quality LLMs. In particular, our framework characterizes how the SLM’s seemingly low-quality supervision can enhance the training of a much more capable LLM. Furthermore, it also highlights the need for an adaptive utilization of such supervision, by striking a balance between the bias and variance introduced by the SLM-provided soft labels. We corroborate our theoretical framework by improving the pre-training of an LLM with 2.8B parameters by utilizing a smaller LM with 1.5B parameters on the Pile dataset.

## 1 Introduction

Owing to the recent surge in their ever-growing capabilities, large language models (LLMs) (Chowdhery et al., ; Touvron et al., ; OpenAI, ; Anil et al., ; Jiang et al., ; Gemini-Team et al., ; Anthropic, ), have become the focal point of machine learning research.
Several research efforts focus on either further enhancing LLM performance,
or utilizing LLMs in novel applications
ranging from conversational agents/assistants (Thoppilan et al., ) to novel material design (Rubungo et al., ).
Highly capable general-purpose LLMs rely on two critical ingredients:
choosing a model architecture with a very large number of parameters (Chowdhery et al., ; Smith et al., ),
and *pre-training* this model on a corpus with a very large number of examples (AI@Meta, ; Computer, ).
Due to the large size of model and corpus, the computational cost of pre-training can be highly onerous.
Thus, sustainable advancement and widespread adoption of LLMs hinges on designing novel architectures
and algorithms that can reduce the overall training (particularly, pre-training)
compute cost and improve the data efficiency for LLM development.

This paper focuses on leveraging *small language models* (*SLMs*) for efficient LLM pre-training. Interestingly, a growing literature (see, e.g., Gupta et al., ; Chen et al., ; Yue et al., ) shows that, despite their
limited model capacity, SLMs can acquire a good domain understanding of pre-training data distribution.
Particularly, SLMs can perform well on a large portion of “easy” instances, while still providing valuable information towards identifying the remaining “hard” instances, e.g., via the confidence, margin, or similar measures based on their predictive distribution. This prompts us to explore the following question:

Can we speed up pre-training of a high-quality *large* LM by transferring the predictive distribution resulting from pre-training of a lower-quality *small* LM?

Note that suitable SLMs are often readily available during LLM development as previous-generation models trained on similar pre-training corpora or smaller models trained for initial exploration around architectural and algorithmic choices on the current pre-training corpora itself. Furthermore, if proven, the potential of SLMs to enhance LLM quality and efficiency, coupled with their relatively cheaper development cost, strongly justifies training such models even to solely aid LLM training.

Knowledge distillation (KD; Bucilǎ et al., ; Hinton et al., )
is a natural candidate to achieve our underlying objective by utilizing the SLM as teacher model to transfer its predictive distribution to student LLM during pre-training. However, it is unclear if KD can be helpful in realizing our goal as unlike a typical KD setup – wherein a larger or stronger teacher is used to train a smaller or weaker student – we are hoping to leverage a smaller and *weaker teacher* LM to improve the pre-training efficiency and quality of a larger and *stronger student* LM.

We begin by developing a statistical framework to study KD in the context of language modeling. We obtain novel risk bounds that delineate the desirable properties of the teacher LM-provided supervision that can enhance the student LM’s performance, even when one employs a perceivably weaker SLM as the teacher. To the best of our knowledge, ours are the first such bounds for language modeling. Notably, our bounds subsume standard pre-training as a special case, and control
LM generalization
as we scale model capacity and
the amount of pre-training data, with latter being measured in terms of either number of training sequences or number of training tokens. We believe that these bounds are of independent interest to broader community working on LLMs.

Figure: Figure 1: An overview of small model aided large model training (SALT) pre-training. SALT utilizes an SLM in two ways to improve the pre-training of LLM: \raisebox{-0.9pt} {1}⃝ To perform KD with SLM as teacher in the early phase of LLM pre-training; and \raisebox{-0.9pt} {2}⃝ To obtain a valuable subset of pre-training corpora to be utilized during the KD.
Refer to caption: x1.png

Our statistical analysis lays the foundation for an adaptive pre-training method that leverages SLM via KD only in the so-called “easy” regions where SLM can approximate the ground truth next-token distribution well. Combining this with the tendency of neural network to learn easier examples first (Kalimeris et al., ; Refinetti et al., ), we propose small model aided large model training (SALT) – a two-stage pre-training approach that employs KD from an SLM in the early phase of the LLM pre-training and resorts to standard self-supervision-based training for the rest of the pre-training.
We then expand the SALT method by employing SLM to additionally perform data selection for the KD phase.
Our selection procedure focuses on identifying challenging yet learnable sequences from the easy region of the data distribution
to ensure an effective transfer of SLM’s predictive distribution during KD
(see Fig. 1 for an overview). Our
empirical study, focusing on both few-shot and post-supervised fine-tuning (SFT) performance, validates the utility of SALT for improving both quality and training efficiency of LLMs compared to standard pre-training. Our key contributions are summarized as follows:

- (i)
We present a statistical framework for KD in the context of language modeling, which provides novel risk bounds describing how even a perceivably weaker teacher LM can improve the quality of a larger student LM (cf. Sec. 3).
- (ii)
Guided by our
framework, we propose a two-stage pre-training method,
namely SALT, that uses SLMs as teacher to perform KD during the first stage corresponding to early phase of LLM pre-training.
We extend SALT by utilizing SLMs to perform data selection for the KD phase, facilitating an effective transfer of predictive distribution from SLM to LLM (cf. Sec. 4).
- (iii)
We showcase the utility of SALT (with and without data selection) by training a 2.8B parameter LM with the help of 1.5B parameter LM on the Pile dataset (Gao et al., ). The 2.8B LM trained with SALT outperforms a 2.8B LM trained via standard pre-training method on a wide range of popular few-shot benchmarks while utilizing less than $0.7$X training step budget, resulting in $\sim 28\%$ wall-clock time reduction.
Moreover, SALT models consistently demonstrate significant downstream performance gains after SFT on multiple domains
(cf. Sec. 5).

## 2 Background

Language modeling. Given a large corpus, language modeling aims to train a model that can assign probabilities or likelihood to each sequence $\mathbf{x}\in\mathcal{V}^{\star}$, where $\mathcal{V}$ denotes the underlying vocabulary with $V=|\mathcal{V}|$ tokens. Assuming that the language model (LM) is parameterized by ${\bm{\theta}}$,
it assigns the following probability to a $T$-token long input sequence ${\mathbf{x}}=[x_{1},x_{2},\ldots,x_{T}]$:

$$ $\displaystyle P_{{\bm{\theta}}}({\mathbf{x}})=P_{{\bm{\theta}}}(x_{1})P_{{\bm{ \theta}}}(x_{2}|x_{1})\cdots P_{{\bm{\theta}}}(x_{T}|x_{1},\ldots,x_{T-1}).$ $$

Transformers (Vaswani et al., ) are the most prominent architecture supporting LMs (OpenAI, ; Touvron et al., ; Gemini-Team et al., ), which we briefly discuss in Appendix A.1.

Standard LM pre-training. Typically, LM pre-training involves the *next-token prediction* task:
given a training sequence ${\mathbf{x}}=[x_{1},x_{2},\ldots,x_{T}]$,
for each $t\in[T]$, one maximizes the log-likelihood $\log P_{{\bm{\theta}}}(x_{t}|{\mathbf{x}}_{\leq t-1})$.
This amounts to minimizing the cross-entropy loss between the per-token LM prediction distribution $P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})$ and the one-hot distribution(^1^11For $v\in\mathcal{V}$, we define $\mathds{1}_{x}(v)=1$ if $v=x$ and $\mathds{1}_{x}(v)=0$ if $v\neq x$.) $\mathds{1}_{x_{t}}(\cdot)$
defined by the ground truth next-token $x_{t}$.
Thus, the overall loss associated with
${\mathbf{x}}$ becomes

$$ $\displaystyle\ell({\mathbf{x}};\theta)={1}/{T}\cdot\sum\nolimits_{t\in[T]}- \log P_{{\bm{\theta}}}(x_{t}|{\mathbf{x}}_{\leq t-1})={1}/{T}\cdot\sum \nolimits_{t\in[T]}\mathsf{CE}\big{(}\mathds{1}_{x_{t}}(\cdot),P_{{\bm{\theta} }}(\cdot|{\mathbf{x}}_{\leq t-1})\big{)},$ (1) $$

where
$\mathsf{CE}(P_{1},P_{2})$=$-\sum_{v\in\mathcal{V}}P_{1}(v)\log P_{2}(v)$
is the cross-entropy between
distributions $P_{1}$ and $P_{2}$.

Knowledge distillation for LM. Going beyond the ground truth next-token based
loss in (1), one can utilize the per-token prediction distribution provided by another LM, say the one parameterized by ${\bm{\zeta}}$, as additional
supervision. Formally, given the context ${\mathbf{x}}_{\leq t-1}$, one can train the LM parameterized by ${\bm{\theta}}$ via aligning its prediction distribution $P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})$ with $P_{{\bm{\zeta}}}(\cdot|{\mathbf{x}}_{\leq t-1})$.
KL divergence is a common choice to promote such an alignment,
which amounts to minimizing the following cross-entropy loss:

$$ $\displaystyle\ell(\mathbf{x};{\bm{\zeta}}\to{\bm{\theta}})$ $\displaystyle={1}/{T}\cdot\sum\nolimits_{t\in[T]}\mathsf{CE}\big{(}P_{{\bm{ \zeta}}}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{ \leq t-1})\big{)}.$ (2) $$

Training based on the loss in (2)
is referred to as the token-level knowledge distillation (KD; Kim and Rush, ) in the literature, with LMs parameterized by ${\bm{\zeta}}$ and ${\bm{\theta}}$ termed as the teacher and student LMs, respectively. See Appendix A.2 for a discussion on other related KD for LM variants.
Temperature scaling of teacher is a common strategy (Zheng and Yang, ) where, given a temperature $\rho>0$, one utilizes
$P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1})={P_{{\bm{\zeta}}}(\cdot|{
\mathbf{x}}_{\leq t-1})^{\rho}}/{\sum_{v^{\prime}\in\mathcal{V}}P_{{\bm{\zeta}
}}(v^{\prime}|{\mathbf{x}}_{\leq t-1})^{\rho}}$
during KD, resulting in the loss:

$$ $\displaystyle\ell(\mathbf{x};{\bm{\zeta}}^{\rho}\to{\bm{\theta}})={1}/{T}\cdot \sum\nolimits_{t\in[T]}\mathsf{CE}\big{(}P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{ x}}_{\leq t-1}),P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})\big{)}.$ (3) $$

In practice, one typically utilizes both ground truth next-token as well as teacher’s next-token distribution and, for a distillation loss weight $\omega\in[0,1]$, minimizes
the following as the loss for
${\mathbf{x}}$:

$$ $\displaystyle\ell^{\omega}({\mathbf{x}};\theta)\triangleq(1-\omega)\cdot\ell({ \mathbf{x}};{\bm{\theta}})+\omega\cdot\ell(\mathbf{x};{\bm{\zeta}}^{\rho}\to{ \bm{\theta}}).$ (4) $$

Note that, for brevity, our notation $\ell^{\omega}({\mathbf{x}};\theta)$ omits the dependence on $\zeta$ and $\rho$.

## 3 Theoretical analysis: When can KD help language modeling?

As alluded in the introduction, we hope to leverage KD with an SLM as the teacher to speed-up the pre-training of a high quality LLM. However, due to SLM’s relatively limited capacity and inferior quality, it is not immediately clear that such a teacher can benefit the LLM. Motivated by this, we now develop a rigorous statistical framework for KD in the context of language modeling by building
on the works of  Menon et al. (); Dao et al. (); Ren et al. ().
Novel risk bounds originating from this framework highlight how a teacher LM – even a perceivably weaker one – can benefit student LLM by striking the right balance in terms of a bias-variance trade-off.

Notably, our analysis allows one to control the generalization gap for the student LM in terms of both number of training sequences $N$
as well as number of total tokens $NT$, with the latter being highly non-trivial due to possibly arbitrary dependence within a training sequence. In this work, we crucially leverage certain natural stability conditions on the underlying distribution and function class to obtain such
bounds in terms of $NT$. Next, we setup some necessary notation and then present our risk bounds as functions of $N$ and $NT$ in Sections 3.1 and 3.2, respectively. Sec. 3.3 utilizes our bounds to justify the utility of SLMs for improving LLM model quality via KD.

Let $\mathscr{D}$ be the
data distribution that generates
$N$ independent training sequences $\mathscr{S}_{N}=\{{\mathbf{x}}^{(i)}\}_{i\in[N]}\subset\mathcal{V}^{T}$, i.e., ${\mathbf{x}}^{(i)}\sim\mathscr{D},~{}\forall~{}i\in[N]$.(^2^22 Our analysis can be extended to varying length sequences at the cost of increased notational complexity.)
Given $\mathscr{S}_{N}$ and CE surrogate loss
(cf. (1)),
the empirical surrogate risk, i.e., standard training objective, and its population version
for an LM parameterized by ${\bm{\theta}}$ are:

$$ $R_{N}({\bm{\theta}})={1}/{N}\cdot\sum\nolimits_{{\mathbf{x}}\in\mathscr{S}_{N} }\ell({\mathbf{x}};{\bm{\theta}})\quad\text{and}\quad R({\bm{\theta}})=\mathbb {E}_{{\mathbf{x}}\sim\mathscr{D}}\big{[}\ell({\mathbf{x}};{\bm{\theta}})\big{]}.$ (5) $$

On the other hand, the empirical surrogate risk for KD, i.e., the KD training objective, and its population version take the following form (note that we omit dependence on ${\bm{\zeta}}$ and $\rho$):

$$ $\displaystyle R^{\omega}_{N}({\bm{\theta}})={1}/{N}\cdot\sum\nolimits_{{ \mathbf{x}}\in\mathscr{S}_{N}}\ell^{\omega}(\mathbf{x};{\bm{\theta}})\quad \text{and}\quad R^{\omega}({\bm{\theta}})=\mathbb{E}_{{\mathbf{x}}\sim\mathscr {D}}\big{[}\ell^{\omega}(\mathbf{x};{\bm{\theta}})\big{]}.$ (6) $$

### 3.1 Excess surrogate risk bound for LM in terms of number of sequences

Given a potentially infinite function class $\Theta$ for student LMs,
let $\hat{{\bm{\theta}}}$ and ${\bm{\theta}}^{*}$ be the minimizers of KD training objective $\eqref{eq:erm_distill}$ and the population
risk in (5), respectively:

$$ $\displaystyle\hat{{\bm{\theta}}}:=\operatorname*{arg\,min}\nolimits_{{\bm{ \theta}}\in\Theta}R^{\omega}_{N}({\bm{\theta}})\quad\text{and}\quad{\bm{\theta }}^{*}=\operatorname*{arg\,min}\nolimits_{{\bm{\theta}}\in\Theta}R({\bm{\theta }}).$ (7) $$

We want to compare the test performance (population risk) of $\hat{{\bm{\theta}}}$ with that of ${\bm{\theta}}^{*}$ – the optimal LM in $\Theta$.
Before stating our excess risk bound, we introduce an assumption that we rely on in our analysis.

###### Assumption 3.1 .

The per-token log-loss with at most $T$-token long sequences for the function class $\Theta$ is bounded by $M$, i.e.,
$\sup_{{\bm{\theta}}\in\Theta;{\mathbf{x}}\in\mathcal{V}^{\leq T-1}}\max_{v\in
\mathcal{V}}|\log P_{{\bm{\theta}}}(v|{\mathbf{x}})|\leq M.$

###### Remark 3.2 .

Assuming loss to be bounded is a standard practice in the statistical learning literature which can be realized, e.g., by clipping the CE loss by a sufficiently large constant $M$. Recently, Lotfi et al. ()
realized such an assumption by adding a small amount of uniform noise to LM predictions in their study on generalization for LMs trained via standard pre-training without KD.

We are now ready to present our excess risk bound.
Due to the page limit
we provide an informal statement of our result below; see Appendix B.1 for the complete statement and its proof.

###### Theorem 3.3 (Informal) .

Let $\hat{{\bm{\theta}}}$ and ${\bm{\theta}}^{*}$ be as defined in (7).
Define $f^{\bm{\theta}}:\mathcal{V}^{T}\to[0,M]$ by $f^{\bm{\theta}}({\mathbf{x}})=\ell^{\omega}(\mathbf{x};{\bm{\theta}}),~{}
\forall{\mathbf{x}}\in\mathcal{V}^{T},{\bm{\theta}}\in\Theta.$
Then, under Assumption 3.1,
with probability at least $1-\delta$,
we have

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})\leq\frac{c_{1}}{ \sqrt{N}}\cdot\Big{(}\sqrt{V_{N}(f^{\hat{{\bm{\theta}}}})\log\left({2\mathcal{ M}(N)}/{\delta}\right)}+\sqrt{V_{N}(f^{{{\bm{\theta}}}^{*}})\log\left({4}/{ \delta}\right)}\Big{)}$ $\displaystyle\qquad+(4M\omega)/T\cdot\sum\nolimits_{t\in[T]}\mathbb{E}\left[ \mathsf{D}_{\rm TV}\Big{(}P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1}) ,\mathscr{D}(\cdot|{\mathbf{x}}_{\leq t-1})\Big{)}\right]\;+{c_{2}M}/{(N-1)} \cdot\log(\mathcal{M}(N)/\delta),$ $$

where $\mathsf{D}_{\rm TV}$ is TV distance,
$V_{N}(f^{{\bm{\theta}}})=\frac{1}{N(N-1)}\sum_{1\leq i<j\leq N}\big{(}f^{{\bm{
\theta}}}({\mathbf{x}}^{(i)})-f^{{\bm{\theta}}}({\mathbf{x}}^{(j)})\big{)}^{2}$
is sample variance,
$\mathcal{M}(N)$
depends on the
growth function of $\big{\{}f^{\bm{\theta}}:{\bm{\theta}}\in\Theta\big{\}}$,
and $c_{1}$ & $c_{2}$ are universal constants.

###### Assumption 3.1 .

###### Remark 3.2 .

###### Theorem 3.3 (Informal) .

### 3.2 Excess surrogate risk bound for LM in terms of number of tokens

For an LM ${\bm{\theta}}\in\Theta$
and a training sequence ${\mathbf{x}}=[x_{1},x_{2},\ldots,x_{T}]\in\mathcal{V}^{T}$, define

$$ $\displaystyle\xi_{t}({\mathbf{x}};{\bm{\theta}})=\mathbb{E}_{{\mathbf{z}}\sim \mathscr{D}}\left[\ell^{\omega}(\mathbf{z};{\bm{\theta}})|\mathbf{z}_{\leq t-1 }={\mathbf{x}}_{\leq t-1}\right]-\mathbb{E}_{{\mathbf{z}}\sim\mathscr{D}}\left [\ell^{\omega}(\mathbf{z};{\bm{\theta}})|\mathbf{z}_{\leq t}={\mathbf{x}}_{ \leq t}\right],\quad t\in[T].$ (8) $$

Note that $\xi_{t}({\mathbf{x}};{\bm{\theta}})$ does not depend on ${\mathbf{x}}_{>t}$.
For $t\in[T]$, $\xi_{t}({\mathbf{x}};{\bm{\theta}})$
measures the expected KD loss deviation for the student
when we condition on
the context up to $(t-1)$-th vs. $t$-th token, respectively, and sample the remaining tokens
from $\mathscr{D}$. In general,
the deviation could be large as changing a single token in the context can significantly alter LM’s distribution.
However,
a well-behaved LM should be robust to such perturbations.
Motivated by this, we introduce the following assumption.

###### Assumption 3.4 .

Given the data distribution $\mathscr{D}$ and finite function class $\Theta$, the following holds for any ${\mathbf{x}}\in{\rm Support}(\mathscr{D})$, ${\bm{\theta}}\in\Theta$, and $t\in[T]$:

$$ $\displaystyle|\xi_{t}({\mathbf{x}};{\bm{\theta}})|$ $\displaystyle\leq C_{t}\leq C;\quad\text{and}\quad\mathbb{E}\left[\xi^{2}_{t}( {\mathbf{x}};{\bm{\theta}})|{\mathbf{x}}_{\leq t-1}\right]\leq V_{t}.$ (9) $$

Under Assumption 3.4, we obtain the following result on student’s generalization gap. For proof, please see Appendix B.2.

###### Theorem 3.5 .

Let $\Theta$ be a finite function class. Under Assumption 3.4, with probability at least $1-\delta$, the following holds for the student LM $\hat{{\bm{\theta}}}\in\Theta$ obtained via KD:

$$ $\displaystyle R(\hat{{\bm{\theta}}})\leq R^{\omega}_{N}(\hat{{\bm{\theta}}})+ \sqrt{2{\sum\nolimits_{t}{V_{t}}/{N}}\cdot\log\left({|\Theta|}/{\delta}\right) }\;+{2C}/({3N})\cdot\log\left({|\Theta|}/{\delta}\right)\;+$ $\displaystyle\qquad\qquad(4M\omega)/T\cdot\sum\nolimits_{t\in[T]}\mathbb{E}_{{ \mathbf{x}}_{\leq t-1}\sim\mathscr{D}}\mathsf{D}_{\rm TV}\big{(}P_{{\bm{\zeta} },\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),\mathscr{D}(\cdot|{\mathbf{x}}_{\leq t- 1})\big{)}.$ (10) $$

###### Remark 3.6 (Dependence of $C,\{V_{t}\}$ on $T$ ) .

Theorem 3.5 captures a fine-grained dependence on $C$ and $\{V_{t}\}$ from Assumption 3.4. For a robust LM, when $C$ is $\mathcal{O}(1/T)$ and $\{V_{t}\}$ are $\mathcal{O}(1/T^{2})$ (note the scaling of $\ell^{\omega}$ by $T$), we get the tightest generalization gap decaying with ${NT}$. For a non-robust LM in the worst case, $C$ and $\{V_{t}\}$ can be as large as $\mathcal{O}(1)$. The bound is not tight anymore and we fall back to the bound in Theorem 3.3 that only decays with ${N}$ instead of ${NT}$.

###### Remark 3.7 .

Recently, Lotfi et al. () obtained generalization bounds for LLMs trained without KD in terms of total number of tokens. However, even when we specialize
Theorem 3.5 to standard pre-training by setting
$\omega$ to 0 0, our result significantly differs from theirs both in terms of proof technique as well as its implications. Crucially,
results in Lotfi et al. () only holds for the contexts seen during training whereas our bound
enables controlling LM’s generalization gap
on novel contexts generated from the data distribution during the test time.

In a concurrent work, Zekri et al. () provide novel generalization bounds for language models by drawing connections between auto-regressive models and finite state Markov chains. Notably, their work does not explore KD in language modeling context.

###### Assumption 3.4 .

###### Theorem 3.5 .

###### Remark 3.6 (Dependence of C , { V t } 𝐶 subscript 𝑉 𝑡 C,\{V_{t}\} italic_C , { italic_V start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT } on T 𝑇 T italic_T ) .

###### Remark 3.7 .

### 3.3 KD out-performing standard pre-training: A bias-variance trade-off

Empowered by our novel risk bounds, we now provide justification for why KD can outperform standard pre-training.
Specifically, based on Theorem 3.5, we clarify
when KD might lead to a tighter bound than standard pre-training. Similar conclusions follow from Theorem 3.3 by extending the arguments from Menon et al. () to
language modeling. As per Theorem 3.5,
three key quantities control the generalization
of the student:
(1) $\sum_{t}V_{t}$ which relates to loss variance; (2) $C$ which relates to extreme loss values;
and (3) divergence between the teacher-provided
distribution and the ground truth
distribution:
$\textsc{Div}(\zeta,\omega)=\omega\cdot\sum_{t}\mathbb{E}\left[\mathsf{D}_{\rm
TV
}\left(P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),\mathscr{D}(\cdot|
{\mathbf{x}}_{\leq t-1})\right)\right].$
Under Assumption 3.1,
only $\sum_{t}V_{t}$ and $\textsc{Div}(\zeta,\omega)$ are
crucial in distinguishing KD and standard pre-training.

Since $\textsc{Div}(\zeta,0)=0$, one may surmise that standard pre-training (i.e, $\omega=0$) leads to tighter bound.
But as we detail in Appendix C due to the page limit, the variance term becomes smaller as we increase $\omega$.
Thus, with a careful selection of $\omega$, the variance reduction via KD can offset the bias
$\textsc{Div}(\zeta,\omega)$. In particular, if the teacher
closely approximates the ground truth
distribution
so that the bias $\textsc{Div}(\zeta,\omega)$ is small even for large $\omega$, then the variance reduction via KD becomes prominent, resulting in significantly tighter generalization gap compared to standard pre-training.

Performance gain via KD from an SLM as teacher. While small teacher LMs – the main interest of this work – also lead to variance reduction, they are typically not powerful enough to model the true distribution over the entire data domain very well. Thus, any effect of variance reduction
via KD with such a teacher would be washed away by the large bias $\textsc{Div}(\zeta,\omega)$.
This highlights the need for an adaptive form of KD from SLMs.
Even SLMs with their limited capacity can approximate the
true distribution well on certain regions of the data domain,
which we call the ‘easy’ regions.
Thus, one can employ KD from
SLMs on the easy regions to benefit from the variance reduction
without incurring large bias and guarantee improved student LLM performance on these regions.
For the remaining (‘hard’) regions, where the biasx
is large enough to overshadow the contributions of variance reduction, one should
not perform KD from
SLMs and utilize the standard pre-training loss.

## 4 SALT: S mall model a ided l arge model t raining

Figure: Algorithm 1 Small model aided large model training (SALT)

We now operationalize the key takeaway from
Sec. 3
by proposing
SALT – a simple yet effective two-stage pre-training method.
SALT relies on the inherent preference of a model to first focus on easier
supervision before fitting more complex supervision during
training (Kalimeris et al., ; Refinetti et al., )
to perform selective KD from a teacher SLM.
We then expand SALT
to perform explicit selection of
training sequences where we want to
conduct KD from an SLM on. In particular, we identify challenging
sequences which are learnable
as per teacher SLM; as a result, performing KD on those sequences
result in further variance reduction (Katharopoulos and Fleuret, ).

Two-stage LLM pre-training via SALT.
Inspired by our analysis,
we propose a two-stage pre-training methods for LLMs in Algorithm 1. The
algorithm employs KD with SLM as a teacher in the first stage which lasts till $n_{\rm KD}$ training steps,
and transitions to standard pre-training *without* KD in the second stage. Note that we are interested in the selective transfer of predictive distribution from teacher SLM to student LLM in those regions where SLM performs well by capturing true distribution. By design, KD aims to align predictive distributions of the teacher and student. On the regions where SLM performs well, we expect it to exhibit reasonably confident predictive distribution that should align with ground truth next-token (Gupta et al., ), thereby constituting an easier supervision signal for the LLM.
In contrast, on hard instances where SLM’s predictive distribution is not confident enough or does not align well with the ground truth next-token, learning will be delayed to the later phase of the training  (Kalimeris et al., ; Refinetti et al., ).
Thus, SALT relies on the tendency of neural networks to focus on easier instances early during the training to perform desirable knowledge transfer from SLM in the first stage.
Once the student LLM is sufficiently aligned with teacher SLM on easier regions, it starts utilizing its model capacity to further align with SLM on more complex regions where high divergence between SLM and ground truth distribution can become detrimental to the
LLM’s performance. Switching to standard pre-training based solely on the self-supervision from next-tokens in the second stage prevents this undesirable over-alignment.
We empirically verify the above intuition behind SALT in Sec. 5.4.

$\texttt{SALT}_{\texttt{DS}}$: SALT with data selection.
We now endow SALT approach (cf. Algorithm 1) with explicit selection of examples where we want to transfer teacher SLM’s predictive distribution on, with SLM itself enabling the data selection. In particular, we want to select the most informative (or challenging) examples among the ones that SLM performs well on.
Towards this, given the SLM
${\bm{\zeta}}$ and a positive integer $k$, we assign a score $S_{{\bm{\zeta}},k}({\mathbf{x}})$ to training sequence ${\mathbf{x}}$, with a higher score indicating a higher likelihood to be included for training.
More specifically, we compute the per-token cross-entropy loss of SLM
on ${\mathbf{x}}$ and
aggregate (typically using the median) into a
sequence-level score. This encourages selecting more challenging examples. However, in the spirit of selecting examples that are still *learnable*,
we remove all losses where the ground-truth token is not in top-$k$ outputs of the model before aggregating:

$$ $\displaystyle S_{{\bm{\zeta}},k}({\mathbf{x}})=\mathrm{median}\Big{(}\Big{\{}- \mathds{1}\{x_{i}\in\mathrm{argtop}_{k}(P_{\bm{\zeta}}(\cdot|{\mathbf{x}}_{<i} ))\}\log P_{{\bm{\zeta}}}(x_{i}|{\mathbf{x}}_{<i}):i\in[n]\Big{\}}\Big{)},$ (11) $$

where $\mathrm{argtop}_{k}(P_{\bm{\zeta}}(\cdot|{\mathbf{x}}_{<i}))$ denotes the top-$k$ scoring tokens at position $i$ according to SLM.
Given a scored pool of sequences, we select the top-$m$ scoring sequences where $m$ is sufficiently large to complete the first stage of
Algorithm 1.
The reader may note, that if the SLM has been trained with the same dataset that it is scoring, then the computed per-token loss (and resulting sequence score) may be heavily biased. To circumvent that, we use an “early checkpoint” of the model ${\bm{\zeta}}_{n_{0}}$, which has trained on a small number of examples from the overall training set $n_{0}\ll N$. We then sample from the remainder of the training examples using score $S_{{\bm{\zeta}}_{n_{0}},k}(\cdot)$. Although the early model ${\bm{\zeta}}_{n_{0}}$
may be a lower quality model due to training with relatively little data, it is only the relative ordering of examples that is important when computing a score, rather than the absolute score.

## 5 Experiments

We now showcase the potential of leveraging SLMs for improving LLM pre-training, by realizing both better quality and improved training efficiency.
We demonstrate the additive utility of two aspects of our proposal:
(1) employing SLMs as teacher models during the early phase of LLM pre-training (SALT); and
(2) further performing data selection via SLMs during the KD phase ($\texttt{SALT}_{\texttt{DS}}$).

Throughout our study, we compare with a natural baseline (denoted Baseline) where the large LM is pre-trained in a standalone manner with a self-supervised objective over a pre-training set. This enables us to fairly compare our proposed approach as we fix the model architecture and the underlying pre-training data in our evaluation. The key takeaways from this section are:

- (1)
SALT and $\texttt{SALT}_{\texttt{DS}}$ attain Baseline performance with less than 70% of training steps, and significantly outperform Baseline with the same number of training steps (cf. Sec. 5.2). Additionally, SALT
can leverage improved quality small teacher LM to further improve the performance of the large student LM (cf. Appendix H)
- (2)
SALT and $\texttt{SALT}_{\texttt{DS}}$ pre-trained LMs consistently outperform Baseline after SFT on arithmetic reasoning, summarization, and natural language inference (NLI) tasks (cf. Sec. 5.3).
- (3)
Step transition from KD phase to standard training phase in SALT (cf. Algorithm 1) constitutes a good design choice as it outperforms other natural alternatives (cf. Appendix H).

We also compare SALT
with RKD (standing for reverse KD) where we perform KD with SLM as the teacher throughout pre-training.
RKD results in sub-par performance even compared to Baseline as the relatively poor quality of SLM limits LLM performance
during the later part of training.

### 5.1 Experimental Setup

Figure: Figure 2: Fraction of correct next-token predictions for various LMs during training, on a subset of the Pile training set.
Refer to caption: x2.png

Model architectures and pre-training data.
We work with standard decoder-only Transformer-based LMs. Our small model (SLM) has 1.5B parameters and our large model has 2.8B parameters.
We use a SentencePiece tokenizer (Kudo and Richardson, ) from Du et al. () with a vocabulary size of 256$K$.
We pre-train all LMs on the Pile dataset (Gao et al., ) for 545 billion tokens
via UL2 objective (Tay et al., ) with a mixture of causal LM, prefix LM and span corruption tasks.
We use a batch size of 2048 and input sequence length of 1280.
Based on our hyperparameter search, we set the distillation weight $\omega=0.667$ and teacher temperature $\rho=0.25$.
For $\texttt{SALT}_{\texttt{DS}}$, we score examples with SLM trained for $n_{0}=26$K steps and set $k=10$ in (11).
Further details on model architecture and pre-training are provided in Appendix E.

Few-shot evaluation tasks. Following the literature (see, e.g., Anil et al., ; Touvron et al., ),
we perform few-shot evaluation of pre-trained LMs on a wide range of standard benchmarks. We focus on English benchmarks as the pre-training data is mostly English.
We defer the full list of benchmarks to Appendix F which can be
categorized into: (1) world knowledge, (2) reading comprehension, (3) commonsense reasoning, (4) natural language generation (NLG), and (5) SuperGLUE.
We also consider LAMBADA (Paperno et al., ) and MBPP (Austin et al., ) which are Cloze and code generation tasks, respectively.
We conduct $1$-shot evaluation for all benchmarks, except for MBPP which is $3$-shot. For each benchmark, we report the corresponding prevalent metric in the literature. See Appendix G for details.

Fine-tuning tasks.
We focused on (arithmetic) reasoning, summarization, and NLI tasks where the few-shot performance of all pre-trained models (including baseline and our models) were poor. For arithmetic reasoning, we utilize GSM8K (Nie et al., ). For summarization, we employ
XSum (Narayan et al., ) and CNN/DailyMail
(Nallapati et al., ). For NLI tasks, we employ
ANLI-R1, ANLI-R2, and ANLI-R3
(Nie et al., ).

### 5.2 Results: SALT enables efficient training of high quality LLMs

Interestingly, KD from the seemingly weaker SLM does improve LLM training in the beginning compared to Baseline, as reflected in the next-token prediction accuracy over the training set in Fig. 2. However, continuing KD from the weaker teacher eventually
become detrimental. As evident in Fig. 2 and Tab. 1,
RKD significantly underperforms Baseline on both training and validation set. In contrast, SALT
leverages KD from SLM only during
first $n_{\rm KD}$ training steps
(cf. Algorithm 1).

**Table 1: Accuracy and log-perplexity on a held-out set of Pile. We evaluate the models at an early and the final checkpoint. At the end of pre-training ($208$K steps), both SALT and $\texttt{SALT}_{\texttt{DS}}$ improve upon Baseline in terms of both next-token prediction accuracy and log-perplexity. RKD is identical to SALT during its first stage, hence they have the same performance at $n_{\rm KD}=36$K steps.**
| Model | Evaluation stage<br>(steps) | Accuracy ($\uparrow$) | Log ($\downarrow$)<br>perplexity |
| --- | --- | --- | --- |
| SLM | Final (208K) | 57.70 | 1.951 |
| Baseline | Early (36K) | 56.68 | 2.011 |
| RKD | 57.21 | 2.160 |  |
| SALT | 57.21 | 2.160 |  |
| $\texttt{SALT}_{\texttt{DS}}$ | 56.47 | 2.188 |  |
| Baseline | Final (208K) | 58.99 | 1.868 |
| RKD | 58.46 | 2.071 |  |
| SALT | 59.10 | 1.863 |  |
| $\texttt{SALT}_{\texttt{DS}}$ | 59.17 | 1.857 |  |

Quality improvements via SALT. Unlike RKD, SALT (with $n_{\rm KD}=36$K steps) yields a pre-trained LLM that improves upon Baseline on both training set (cf. Fig. 2) and held-out validation set (cf. Tab. 1).
Tab. 2 presents domain-wise few-shot performance of Baseline, RKD, SALT, and $\texttt{SALT}_{\texttt{DS}}$ (see Tab. 5 Appendix G for
per-task performance
).
Both SALT and $\texttt{SALT}_{\texttt{DS}}$ consistently outperform Baseline (as well as RKD) at the end of training, i.e., $@100$% steps or $208$K steps. In particular, $\texttt{SALT}_{\texttt{DS}}$ outperforms Baseline in $6$ out of $7$ domains as well as the overall average; similarly, SALT improves upon Baseline in $5$ out of $7$ domains as well as the overall average,
establishing the utility of SALT approaches in successfully
leveraging SLMs to boost the quality of
LLMs.

Training efficiency realized via SALT. As per Tab. 2, SALT surpasses Baseline at $146$K steps on average few-shot performance, suggesting a savings of $30$% training compute cost. While
$n_{\rm KD}=36$K out of those $146$K steps involve KD which is typically computationally costlier than standard training, as we argue next, the fact that our teacher is a SLM ensures that we still realize efficiency gains via SALT. In particular, the additional compute cost in KD is one forward pass of SLM per training step. As a rule of thumb, a forward pass constitutes $1/4$th cost of a training step (which comprises both forward and backward passes). In our setup, a forward pass of SLM is approximately half as expensive as that for the LLM. Thus, KD from SLM adds a factor of $1/8$ to the cost of the training step of the standard training.
(Our implementation verifies this as we observe $12.0$% slow down during KD compared to standard training.) Since KD lasts for $n_{\rm KD}=36$K out $146$K training steps, the training cost required by SALT to surpass Baseline is approximately equivalent to that of $(146+36/8)$K steps of the standard training. This translates to an efficiency gain of $\sim 28$% compared to the $208$K steps taken by Baseline.

**Table 2: Domain-wise few-shot performance of pre-trained LMs. SALT and $\texttt{SALT}_{\texttt{DS}}$ already outperform Baseline in terms of average few-shot performance at $70$% of their training step budget, thereby improving both training efficiency and model quality. RKD (i.e., naively distilling from SLM throughout pre-training) performs much worse than Baseline. The best and second-best results for each domain are boldfaced and underlined, respectively.**
| Domain | # Tasks |  | SLM |  | Baseline |  | RKD |  | SALT |  | $\texttt{SALT}_{\texttt{DS}}$ |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  |  |  |  | @100% |  | @100% |  | @70% | @100% |  | @70% | @100% |
|  |  |  |  |  | steps |  | steps |  | steps | steps |  | steps | steps |
| World Knowledge | 4 |  | 15.90 |  | 22.19 |  | 18.69 |  | 21.59 | 22.70 |  | 20.64 | 21.72 |
| Reading Comprehension | 4 |  | 46.30 |  | 53.00 |  | 51.00 |  | 53.55 | 54.55 |  | 54.35 | 54.93 |
| Commonsense Reasoning | 7 |  | 57.76 |  | 61.99 |  | 58.30 |  | 61.27 | 61.67 |  | 62.00 | 62.10 |
| LAMBADA | 1 |  | 26.90 |  | 36.20 |  | 31.10 |  | 50.70 | 48.30 |  | 48.00 | 53.00 |
| SuperGLUE | 8 |  | 61.59 |  | 65.53 |  | 62.91 |  | 66.30 | 65.28 |  | 65.99 | 65.58 |
| NLG | 3 |  | 3.13 |  | 4.60 |  | 3.40 |  | 4.63 | 4.73 |  | 4.80 | 4.83 |
| MBPP | 1 |  | 9.60 |  | 16.20 |  | 11.40 |  | 15.60 | 17.00 |  | 16.60 | 17.80 |
| Average | 28 |  | 42.56 |  | 47.32 |  | 44.39 |  | 47.86 | 47.94 |  | 47.89 | 48.26 |

### 5.3 Improved downstream performance realized via SALT post SFT

Tab. 2 already showcases that SALT can leverage SLMs to obtain large pre-trained LMs that out-perform widely adopted standard pre-training (Baseline).
That said, all the pre-trained models (including Baseline) exhibit relatively poor few-shot performance on certain benchmarks, e.g., NLG or summarization task (in Tab. 2)
and also MATH (Hendrycks et al., ) and ANLI (Nie et al., ) benchmarks. This raises the question of whether SALT is beneficial
for downstream application performance.

Supervised fine-tuning (SFT)
is a standard approach to convert a pre-trained LM into a domain-specific proficient model.
We employ SFT on a range of downstream tasks covering three domains, namely arithmetic reasoning, NLG or summarization, and NLI.
For each downstream benchmark, we perform SFT on the pre-trained LMs obtained via Baseline, SALT, and $\texttt{SALT}_{\texttt{DS}}$.
During SFT, we train each pre-trained LM for $10$K steps with Adafactor algorithm (Shazeer and Stern, ). We utilize cosine learning rate schedule (Loshchilov and Hutter, ) with a peak learning rate of 1e-4 and linear warm-up phase consisting of 200 steps.
These hyperparameters were optimized for the Baseline; and we did not conduct hyperparameter tuning for the pre-trained LMs resulting from SALT approaches. Finally, we employ greedy decoding during the evaluation of the fine-tuned LMs.
Tab. 3 presents SFT performance on six benchmarks covering the aforementioned three downstream domains.
SALT consistently outperforms Baseline on all benchmarks and, in most cases,
$\texttt{SALT}_{\texttt{DS}}$ performs the best.
This showcases that SALT (with or without data selection) enables significant improvements for a range of difficult downstream domains even when the corresponding pre-trained LMs exhibits only a
small performance gains over Baseline.

**Table 3: Supervised fine-tuning (SFT) results. Performance of various pre-trained checkpoints on downstream tasks after SFT. For each benchmark, pre-trained 2.8B models are fine-tuned on the corresponding train split and evaluated on the validation split (test split in case of GSM8K). Acc, Rg1, Rg2, and RgL represent the Accuracy, Rouge-1, Rouge-2, and Rouge-Lsum metrics, respectively.**
|  | GSM8K |  | XSum |  | CNN/DailyMail |  | ANLI-R1 |  | ANLI-R2 |  | ANLI-R3 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Acc |  | Rg1 | Rg2 | RgL |  | Rg1 | Rg2 | RgL |  | Acc |  | Acc |  | Acc |
| Baseline | 31.84 |  | 43.39 | 21.09 | 35.91 |  | 42.84 | 20.43 | 40.38 |  | 63.70 |  | 56.90 |  | 57.83 |
| SALT | 34.87 |  | 43.45 | 21.21 | 36.04 |  | 43.19 | 20.65 | 40.74 |  | 67.00 |  | 57.80 |  | 59.67 |
| $\texttt{SALT}_{\texttt{DS}}$ | 35.25 |  | 43.77 | 21.44 | 36.24 |  | 43.41 | 20.87 | 40.95 |  | 67.30 |  | 57.70 |  | 59.58 |

### 5.4 SLM enables fast learning on easy examples

**Table 4: Few-shot evaluation on different buckets of XLSum-EN. Each number shows average Rouge-2 scores on the corresponding bucket. We use gray, green, and red to highlight the results similar to, better than, and worse than Baseline performance, respectively.**
|  | Evaluation stage (steps) | Easy | Medium | Hard |
| --- | --- | --- | --- | --- |
| SLM | Final (208K) | 8.04 | 0.43 | 0.00 |
| Baseline | Early (36K) | 6.15 | 1.61 | 0.71 |
| RKD | 6.76 | 1.40 | 0.58 |  |
| SALT | 6.76 | 1.40 | 0.58 |  |
| Baseline | Final (208K) | 8.80 | 2.52 | 0.97 |
| RKD | 7.87 | 1.68 | 0.74 |  |
| SALT | 9.68 | 2.67 | 0.99 |  |

Recall that SALT aims to realize quality and efficiency gains for LLM pre-training by quickly transferring the predictive distribution of an SLM to the LLM via KD, focusing on the ‘easy’ regions of the data distribution where the SLM performs well. Subsequently, SALT falls back on ground truth next-token-based supervision to refine LLM’s performance on the ‘hard’ regions where the SLM fares poorly. Here, we set out to empirically demonstrate that this key intuition behind SALT is indeed borne out in practice. Focusing on various few-shot eval benchmarks, we partition instances in each benchmark into ‘easy’, ‘medium’, and ‘hard’ buckets based on the teacher SLM’s performance (see Appendix J for details). We then evaluate Baseline, SALT, and RKD pre-trained LLMs on these individual buckets after $n_{\rm KD}=36$K training steps when the KD phase of SALT ends as well as at the end of the pre-training, i.e., after $208$K steps. Tab. 4 presents these results on the XLSum-EN task (see Appendix J for results on other benchmarks), which validate: (1) KD from SLM quickly enables LLM to perform well on ‘easy’ instances; and (2) standard pre-training after KD phase ending at $n_{\rm KD}$-th step helps LLM performance on ‘hard’ instances the most.

## 6 Related work

Here, we provide a brief account of related work, focusing on the prior work on assisting large model training via small models, data selection, and theoretical treatment of KD. Due to the page limit, we defer the discussion on various KD methods for language modeling that are not pertinent to the main objective of this work to Appendix A.

Aiding large model training with small models.
Small models often help identify good hyper-parameters that can be utilized for large model training with minimal modifications (Yang et al., ). Progressive or stage-wise training methods (Gong et al., ; Gu et al., ; Reddi et al., ) train a large model in stages where the parameters for the model at a given stage get initialized based on the parameters of a smaller model from the previous stage. Another related line of work simply informs the large model initialization based on a smaller model without resorting to progressive stage-wise training (Trockman and Kolter, ; Wang et al., ; Samragh et al., ). Most of these works crucially depend on the architectural overlaps between the small and large models, e.g., requiring them to share same depth, width, or more generally same model family.
In contrast, one can distill from a smaller model to a larger model without such constraints. Furlanello et al. () study self-distillation to iteratively train a model, with the final model from an iteration acting as a teacher for the next iteration.
More closer to the setting studied in this work, albeit in image classification setting, Yuan et al. (); Xie et al. () consider distillation
from a weaker model.
In the work that is closest to our proposal, Qin et al. ()
distill an LM from a smaller LM.
Moreover, they also
distill
during the early phase of larger student LM pre-training to avoid negative impact on large LM’s final performance. We demonstrate the utility of such an approach with larger LMs as well as larger pre-training corpus, with a wider range of evaluation benchmarks. Furthermore, we provide a statistical framework to rigorously justifies the utility of a seemingly weaker teacher during LLM pre-training. Other recent efforts on using smaller LMs to boost larger LMs have primarily focused on
fine-tuning (Yang et al., ; Mitchell et al., ) or alignment (Burns et al., ), while
we focus on pre-training which is the most compute and data intensive phase of LLM development.

Data Selection.
Ankner et al. () select examples with reference model sequence-level log-perplexities in a specified range. They aggregate per-token log-perplexities by taking the mean, whereas, we take the median and also zero out noisy tokens.
Mindermann et al. () select examples and Lin et al. () select tokens for training based on excess training loss over a reference model. While
we perform offline data selection before training,
these two works select data from training batches on the fly.
Methods which encourage diversity through semantic embeddings (e.g., Abbas et al. (); Tirumala et al. ()) can yield a better final model when employed in $\texttt{SALT}_{\texttt{DS}}$.
Please refer to Albalak et al. () for a comprehensive survey of data selection techniques for LMs.

Theoretical understanding of knowledge distillation. Focusing on (deep) linear networks, Phuong and Lampert () study the generalization bounds for KD and relate its success to various factors including data geometry and optimization bias. Menon et al. (); Ren et al. () show that KD leads to reduced variance of the training objective, thereby improving generalization, while also relating the effectiveness of KD to teacher’s ability to approximate Bayes class probabilities. Cotter et al. () argue that KD can improve generalization as long as teacher approximates a suitable transformation of Bayes class probabilities. Dao et al. () study KD from the perspective of semiparameteric inference to analyze the excess risk of distilled student. Focusing on self-distillation for kernel ridge regression, Mobahi et al. () show that distillation can enhance regularization. Allen-Zhu and Li () explain the utility of distillation via better feature learning. Xu et al. () study the interplay between optimization and
label smoothing via teacher-provided soft labels.
Focusing on linearized neural networks, Harutyunyan et al. () attribute the success of KD to the reduced supervision complexity.
More recently, Safaryan et al. () argue that KD can act as a form of partial variance reduction, thereby improving convergence. We would like to highlight that the existing literature does not provide
generalization bounds for
KD in a sequence learning setting such as language modeling, and we provide the first statistical treatment of KD for language modeling.
Notably, similar to our work, Xu et al. (); Nagarajan et al. () also explore the utility of KD only during an early-phase of student training, albeit not in a language modeling setting.

## 7 Conclusion

We conduct a systematic study of the utility of SLMs to improve both training efficiency and performance of LLM pre-training.
Towards this, we introduce a novel statistical framework for KD in the context of language modeling, which guides the design of SALT – a simple KD-based approach that selectively transfers predictive distribution from an SLM to an LLM.
We further enhance SALT
and perform explicit data selection via SLMs
to effectively transfer knowledge from SLMs to LLMs.
SALT significantly reduces the pre-training time for LLMs while ensuring good overall quality as measured by the LLM’s few-shot performance as well as downstream performance after fine-tuning.

While SALT can play a crucial role in efficiently sustaining the trend of developing LLMs with increasing capabilities, it also has the potential to help institutions train proficient lightweight LMs. Conventionally, one distills a large powerful model into a lightweight model with good quality. However, given that many LLMs are proprietary, such strong-to-weak distillation is not feasible at many institutions. Our proposed approach can leverage even smaller LMs to enhance the quality of a small LM with acceptable inference cost.
Exploring if our proposed approach can introduce new capabilities in a general-purpose small LM with the help of one or multiple smaller LMs that are experts in their respective domains is an interesting avenue for future work. Interestingly, our data selection approach in $\texttt{SALT}_{\texttt{DS}}$ demonstrates that it is indeed possible to leverage data selection to improve KD for language modeling. Building on this initial investigation and further exploring and extending data selection approaches tailored to $\texttt{SALT}_{\texttt{DS}}$ is another interesting line of future work.

## Acknowledgments

The authors would like to thank Ke Ye, Julian Odell, and Rina Panigrahy for their valuable input during the early stages of this research effort.

## Appendix A Language modeling and knowledge distillation

### A.1 Language modeling via Transformer-based models

Here, we briefly discuss how Transformers are typically employed for language modeling in modern systems.
Given a context ${\mathbf{x}}_{\leq t}=[x_{1},\ldots,x_{t}]\in\mathcal{V}^{t}$, a Transformer-based LM first produces a sequence of $d$-dimensional token embeddings
$E({\mathbf{x}}_{\leq t})=[{\textnormal{e}}_{x_{1}}^{\top},{\textnormal{e}}_{x_
{2}}^{\top},\ldots,{\textnormal{e}}_{x_{t}}^{\top}]\in\mathbb{R}^{d\times t}$, where ${\mathbf{e}}_{v}\in\mathbb{R}^{d}$ denotes the token embedding for $v\in\mathcal{V}$. A Transformer network $f_{{\bm{\psi}}}$ then processes $E({\mathbf{x}}_{\leq t})$ to produce a target embedding $f_{{\bm{\psi}}}\big{(}E({\mathbf{x}}_{\leq t})\big{)}\in\mathbb{R}^{d}$, which is multiplied by $W\in\mathbb{R}^{V\times d}$, namely a classification layer, to obtain a logit vector $u_{{\mathbf{x}}_{\leq t}}:=Wf_{{\bm{\psi}}}\big{(}E({\mathbf{x}}_{\leq t})\big
{)}\in\mathbb{R}^{V}$. Accordingly, we have ${\bm{\theta}}=\{E,{\bm{\psi}},W\}$ as the parameters of the LM. Applying softmax operation on the logit vector produces the probability that the LM assigns to each token in $\mathcal{V}$ as the possible continuation (also known as next token) to the context ${\mathbf{x}}_{\leq t}$:

$$ $\displaystyle P_{{\bm{\theta}}}(v|{\mathbf{x}}_{\leq t})=\frac{\exp({u_{{ \mathbf{x}}_{\leq t}}(v)}/\tau)}{\sum_{v^{\prime}\in\mathcal{V}}\exp({u_{{ \mathbf{x}}_{\leq t}}(v^{\prime})}/\tau)},\qquad\forall v\in\mathcal{V}.$ (12) $$

Here, $\tau$ denotes the (inverse) temperature associated with the softmax operation. Unless stated otherwise, we assume that $\tau=1$.

### A.2 Other common variants of knowledge distillation for LM

Top-$k$ token-level KD. Instead of aligning the teacher and student’s full per-token prediction distributions, one could only match
these distribution on $\mathcal{T}\subset\mathcal{V}$, e.g., $k\ll V$ elements of $\mathcal{V}$ that receive the highest scores from the teacher:

$$ $\displaystyle\ell(\mathbf{x};{\bm{\zeta}}\to{\bm{\theta}})=-\sum_{t=1}^{T}\Big {(}\sum_{v\in\mathcal{T}}P^{\mathcal{T}}_{{\bm{\zeta}}}(v|{\mathbf{x}}_{\leq t -1})\cdot\log P^{\mathcal{T}}_{{\bm{\theta}}}(v|{\mathbf{x}}_{\leq t-1})\big{) }\Big{)},$ (13) $$

where $P^{\mathcal{T}}$ denotes the restriction of $P$ (defined over $\mathcal{V}$) to $\mathcal{T}$:

$$ $\displaystyle P^{\mathcal{T}}(v)=\begin{cases}\frac{P(v)}{\sum_{v^{\prime}\in \mathcal{T}}P(v^{\prime})}&\text{if}~{}v\in\mathcal{T},\\ 0&\text{otherwise}.\end{cases}$ (14) $$

Sequence-level KD. Unlike token-level KD, sequence-level KD aims to align teacher and student’s distributions on all sequences up to sequence length $T$. In particular, the sequence-level KD loss takes the form:

$$ $\displaystyle\ell^{\rm seq}(\mathbf{x};{\bm{\zeta}}\to{\bm{\theta}})=-\sum_{ \tilde{{\mathbf{x}}}\in\mathcal{V}^{\leq n}}P_{{\bm{\zeta}}}(\tilde{{\mathbf{x }}})\cdot\log P_{{\bm{\theta}}}(\tilde{{\mathbf{x}}})$ (15) $$

In practice, it’s natural to focus on a subset of all candidate target sequences $\mathcal{U}\subset\mathcal{V}^{\leq T}$:

$$ $\displaystyle\ell^{\rm seq}(\mathbf{x};{\bm{\zeta}}\to{\bm{\theta}})=-\sum_{ \tilde{{\mathbf{x}}}\in\mathcal{U}}P_{{\bm{\zeta}}}(\tilde{{\mathbf{x}}})\cdot \log P_{{\bm{\theta}}}(\tilde{{\mathbf{x}}})$ $$

A common choice for $\mathcal{U}$ is the set of say $k$ most likely sequences under the teacher’s distribution $P_{{\bm{\zeta}}}$.

### A.3 Recent literature on knowledge distillation (KD) for language modeling

A large body of literature focuses on utilizing KD [Bucilǎ et al., , Hinton et al., ] as a core technique to improve LMs [Kim and Rush, , Gou et al., , Xu et al., ]. For instance, Sanh et al. [], Turc et al. [], Wang et al. [], Sun et al. [], Jiao et al. [] relied on KD to compress BERT-style LMs during pre-training, fine-tuning, or both. More recently, KD has been primarily employed in the instruction-tuning or fine-tuning phase where a general purpose LM is adapted to a specific collection of tasks [Xu et al., ]. Black-box KD methods for LM only assume access to training sequences sampled from a teacher LM
[Taori et al., , Fu et al., , Peng et al., ].
With access to token-level distributions from teacher LM, token-level distillation from teacher LM to student LM is possible [Kim and Rush, ]. In contrast, sequence-level distillation involves sampling training sequences from the teacher LM, the student LM, or both before aligning teacher and student’s predictive distribution on such sequences [Kim and Rush, , Agarwal et al., , Gu et al., , Wen et al., ].

## Appendix B Deferred proofs from Section 3

### B.1 Proof of Theorem 3.3

Before stating the formal version of Theorem 3.3 and its proof, let us recall the necessary notation. Given a function class for student LMs $\Theta$, $\hat{{\bm{\theta}}}$ denotes the LM obtained by minimizing the training objective for KD in (6), i.e.,

$$ $\displaystyle\hat{{\bm{\theta}}}:=\operatorname*{arg\,min}_{{\bm{\theta}}\in \Theta}R^{\omega}_{N}({\bm{\theta}}).$ (16) $$

Further, ${\bm{\theta}}^{*}$ represents the optimal or best performing LM in $\Theta$, i.e.,

$$ $\displaystyle{\bm{\theta}}^{*}=\operatorname*{arg\,min}_{{\bm{\theta}}\in \Theta}R({\bm{\theta}}).$ (17) $$

Finally, recall our assumption regarding the bounded loss values.

###### Assumption B.1 .

Given a function class $\Theta$ for (student) LM, the per-token log-loss with at most $T$-token long sequences for the underlying function class $\Theta$ is bounded by $M$, i.e.,

$$ $\displaystyle\sup_{{\bm{\theta}}\in\Theta;{\mathbf{x}}\in\mathcal{V}^{\leq T-1 }}\max_{v\in\mathcal{V}}|\log P_{{\bm{\theta}}}(v|{\mathbf{x}})|\leq M.$ (18) $$

Towards establishing Theorem 3.3, we first state the following intermediate result.

###### Proposition B.2 .

Let $\hat{{\bm{\theta}}}$ and ${\bm{\theta}}^{*}$ be as defined in (16) and (17), respectively. Then, under Assumption B.1, the excess surrogate risk for $\hat{{\bm{\theta}}}$ satisfies the following.

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})$ $\displaystyle\leq\frac{4M\omega}{T}\cdot\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}} _{\leq t-1}\sim\mathscr{D}}\mathsf{D}_{\rm TV}\Big{(}P_{{\bm{\zeta}},\rho}( \cdot|{\mathbf{x}}_{\leq t-1}),\mathscr{D}(\cdot|{\mathbf{x}}_{\leq t-1})\Big{)}$ $\displaystyle\qquad+\left(R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{ {\bm{\theta}}})\right)\;+\left(R^{\omega}_{N}({\bm{\theta}}^{*})-R^{\omega}({ \bm{\theta}}^{*})\right),$ (19) $$

where
$\mathsf{D}_{\rm TV}(\cdot,\cdot)$ denotes the total-variation distance between two probability distributions.

###### Proof of Proposition B.2 .

For convenience, recall that

$$ $\displaystyle R^{\omega}({\bm{\theta}})$ $\displaystyle=\mathbb{E}_{{\mathbf{x}}}\Big{[}(1-\omega)\cdot\ell({\mathbf{x}} ;{\bm{\theta}})+\omega\cdot\ell({\mathbf{x}};{\bm{\zeta}}^{\rho}\rightarrow{ \bm{\theta}})\Big{]}$ $\displaystyle=\mathbb{E}_{{\mathbf{x}}\sim\mathscr{D}}\left[\frac{1}{T}\sum_{t =1}^{T}\mathsf{CE}\big{(}P^{(x_{t},\omega)}_{\zeta,\rho},P_{{\bm{\theta}}}( \cdot|{\mathbf{x}}_{\leq t-1})\big{)}\right]$ $\displaystyle=\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}}_{\leq t-1}\sim \mathscr{D}}\big{[}\mathsf{CE}\big{(}\underbrace{(1-\omega)\cdot\mathscr{D}( \cdot|{\mathbf{x}}_{\leq t-1})+\omega\cdot P_{{\bm{\zeta}},\rho}(\cdot|{ \mathbf{x}}_{\leq t-1})}_{{P}^{(\mathscr{D},\omega)}_{\zeta,\rho}(\cdot|{ \mathbf{x}}_{\leq t-1})},P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})\big{ )}\big{]}$ $\displaystyle=\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}}_{\leq t-1}\sim \mathscr{D}}\big{[}\mathsf{CE}\big{(}{P}^{(\mathscr{D},\omega)}_{\zeta,\rho}( \cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1} )\big{)}\big{]}.$ $$

Note that we have

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})$ $\displaystyle\overset{(i)}{=}R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})-\Big{ (}R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{*})\Big{)}+\Big{(} R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{*})\Big{)}$ $\displaystyle\overset{(ii)}{=}\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}_{{\mathbf{ x}}_{1:t}\sim\mathscr{D}}\Big{[}\sum_{v\in\mathcal{V}}\Big{(}{P}^{(\mathscr{D} ,\omega)}_{\zeta,\rho}(v|{\mathbf{x}}_{1:t})-\mathscr{D}(v|{\mathbf{x}}_{1:t}) \Big{)}\cdot\Big{(}\log P_{\hat{{\bm{\theta}}}}(v|{\mathbf{x}}_{1:t})-\log P_{ {\bm{\theta}}^{*}}(v|{\mathbf{x}}_{1:t})\Big{)}\Big{]}+$ $\displaystyle\qquad R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{ *})$ $\displaystyle\overset{(iii)}{\leq}\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}_{{ \mathbf{x}}_{1:t}\sim\mathscr{D}}\Big{[}\Big{\|}{P}^{(\mathscr{D},\omega)}_{ \zeta,\rho}(\cdot|{\mathbf{x}}_{1:t})-\mathscr{D}(\cdot|{\mathbf{x}}_{1:t}) \Big{\|}_{1}\cdot\Big{\|}\log P_{\hat{{\bm{\theta}}}}(\cdot|{\mathbf{x}}_{1:t} )-\log P_{{\bm{\theta}}^{*}}(\cdot|{\mathbf{x}}_{1:t})\Big{\|}_{\infty}\Big{]}+$ $\displaystyle\qquad{R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{ *})}$ $\displaystyle\overset{(iv)}{\leq}\frac{4M\omega}{T}\cdot\sum_{t=0}^{T-1} \mathbb{E}_{{\mathbf{x}}_{1:t}\sim\mathscr{D}}\Big{[}\mathsf{D}_{\rm TV}\Big{( }P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{1:t}),\mathscr{D}(\cdot|{\mathbf{x} }_{1:t})\Big{)}\Big{]}+\underbrace{R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}( {\bm{\theta}}^{*})}_{{\rm(I)}},$ (20) $$

where
$(i)$ follows from adding and subtracting $R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{*})$;
$(ii)$ employs
the definition of $R(\hat{{\bm{\theta}}}),~{}R({\bm{\theta}}^{*}),~{}R^{\omega}(\hat{{\bm{\theta}
}})$, and $R^{\omega}({\bm{\theta}}^{*})$;
$(iii)$ invokes Hölder’s inequality;
and $(iv)$ follows from the definition of total-variation distance $\mathsf{D}_{\rm TV}(\cdot,\cdot)$ and the fact that underlying per-token loss terms are bounded by $M$.

Next, we focus on the term (I) in (B.1):

$$ $\displaystyle R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}({\bm{\theta}}^{*})$ $\displaystyle\overset{(i)}{=}R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}( \hat{{\bm{\theta}}})+R^{\omega}_{N}(\hat{{\bm{\theta}}})-R^{\omega}_{N}({\bm{ \theta}}^{*})+R^{\omega}_{N}({\bm{\theta}}^{*})-R^{\omega}({\bm{\theta}}^{*})$ $\displaystyle=\Big{(}R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{{\bm{ \theta}}})\Big{)}+\Big{(}R^{\omega}_{N}({\bm{\theta}}^{*})-R^{\omega}({\bm{ \theta}}^{*})\Big{)}\;+R^{\omega}_{N}(\hat{{\bm{\theta}}})-R^{\omega}_{N}({\bm {\theta}}^{*})$ $\displaystyle\overset{(ii)}{\leq}\Big{(}R^{\omega}(\hat{{\bm{\theta}}})-R^{ \omega}_{N}(\hat{{\bm{\theta}}})\Big{)}+\Big{(}R^{\omega}_{N}({\bm{\theta}}^{* })-R^{\omega}({\bm{\theta}}^{*})\Big{)}$ (21) $$

where $(i)$ follows by adding and subtracting $R^{\omega}_{N}(\hat{{\bm{\theta}}})$ and $R^{\omega}_{N}({\bm{\theta}}^{*})$; and $(ii)$ holds as $\hat{{\bm{\theta}}}$ is the minimizer of $R^{\omega}_{N}(\cdot)$ in $\Theta$ which implies that $R^{\omega}_{N}(\hat{{\bm{\theta}}})-R^{\omega}_{N}({\bm{\theta}}^{*})\leq 0$.
(B.1) and (B.1).
∎

Note that the bound on excess surrogate risk in Proposition B.2 decomposes into three terms:

- •
First term captures the divergence between the ground truth per-token distribution and the teacher-induced per-token distribution leveraged during KD; and
- •
The last two terms corresponds to the deviation between empirical and population surrogate risks for the empirical risk minimizer $\hat{{\bm{\theta}}}$ and population risk minimzer ${{\bm{\theta}}}^{*}$ within the function class $\Theta$. Note that since, $\hat{{\bm{\theta}}}$ is a a random variable in itself (which depends on the training sample $\mathscr{S}_{N}$), one typically needs to bound the deviation uniformly over all functions ${\bm{\theta}}\in\Theta$. As we will see next, one can bound these deviations in terms of the properties of both model class $\Theta$ as well as the teacher-induced per-token distributions.

In order to make the excess surrogate risk bound in Proposition B.2 explicit, we need to bound the third term via a computable quantity. We apply sample variance-based bounds from Maurer and Pontil [] to get the following result.

###### Theorem B.3 (Formal version of Theorem 3.3 ) .

Suppose Assumption B.1 holds. Let $\mathcal{F}^{{\bm{\zeta}},\rho,\omega}$ be a function class that maps elements in $\mathcal{V}^{T}$ to $[0,M]$ as defined below:

$$ $\displaystyle\mathcal{F}^{\omega}:=\mathcal{F}^{{\bm{\zeta}},\rho,\omega} \triangleq\left\{{\mathbf{x}}\mapsto\frac{1}{T}\sum_{t=1}^{T}\mathsf{CE}\big{( }{P}^{(\mathscr{D},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{ \bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})\big{)},~{}\forall{\mathbf{x}}\in \mathcal{V}^{T},{\bm{\theta}}\in\Theta\right\}.$ (22) $$

For $\epsilon>0$, let $\mathcal{N}_{\infty}(\epsilon,\mathcal{F}^{{\bm{\zeta}},\rho,\omega},N)$ denote the growth function for the function class $\mathcal{F}^{{\bm{\zeta}},\rho,\omega}$, i.e.,

$$ $\displaystyle\mathcal{N}_{\infty}(\epsilon,\mathcal{F}^{{\bm{\zeta}},\rho, \omega},N)\triangleq\sup_{{\mathbf{X}}=({\mathbf{x}}^{(1)},\ldots,{\mathbf{x}} ^{(N)})\in\mathcal{V}^{T\times N}}\mathcal{N}(\epsilon,\mathcal{F}^{{\bm{\zeta }},\rho,\omega}({\mathbf{X}}),\|\cdot\|_{\infty}),$ (23) $$

where $\mathcal{N}(\epsilon,\mathcal{F}^{{\bm{\zeta}},\rho,\omega}({\mathbf{X}}),\|
\cdot\|_{\infty})$ denotes the smallest $\epsilon$-cover of the set

$$ $\displaystyle\mathcal{F}^{{\bm{\zeta}},\rho,\omega}({\mathbf{X}})=\left\{\big{ (}f({\mathbf{x}}^{(1)}),f({\mathbf{x}}^{(2)}),\ldots,f({\mathbf{x}}^{(N)})\big {)}:f\in\mathcal{F}^{{\bm{\zeta}},\rho,\omega}\right\}\subseteq\mathbb{R}^{N}$ $$

with respect to $\|\cdot\|_{\infty}$ norm. Then, with probability at least $1-\delta$, for all ${\bm{\theta}}\in\Theta$, we have

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})\leq\frac{4M\omega}{T }\cdot\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}}_{\leq t-1}\sim\mathscr{D}}\mathsf {D}_{\rm TV}\Big{(}P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1}), \mathscr{D}(\cdot|{\mathbf{x}}_{\leq t-1})\Big{)}$ $\displaystyle\qquad+\sqrt{\frac{18V_{N}(f^{\hat{{\bm{\theta}}}},\mathscr{S}_{N })\log\left(\frac{2\mathcal{M}(N)}{\delta}\right)}{N}}+\frac{15M\log\left( \frac{2\mathcal{M}(N)}{\delta}\right)}{N-1}$ $\displaystyle\qquad+\sqrt{\frac{2V_{N}(f^{{\bm{\theta}}^{*}},\mathscr{S}_{N}) \log\left(\frac{4}{\delta}\right)}{N}}+\frac{7M\log\left(\frac{4}{\delta} \right)}{3(N-1)},$ (24) $$

where $\mathcal{M}(N)\triangleq 10\cdot\mathcal{N}_{\infty}(1/N,\mathcal{F}^{{\bm{
\zeta}},\rho,\omega},2N)$; $f^{{\bm{\theta}}}$ denotes
the function in $\mathcal{F}^{{\bm{\zeta}},\rho,\omega}$ that corresponds to ${\bm{\theta}}$, as per (22); and $V_{N}(f^{{\bm{\theta}}},\mathscr{S}_{N})$ denotes the sample variance

$$ $\displaystyle V_{N}(f^{{\bm{\theta}}},\mathscr{S}_{N})=\frac{1}{N(N-1)}\sum_{1 \leq i<j\leq N}\big{(}f^{{\bm{\theta}}}({\mathbf{x}}^{(i)})-f^{{\bm{\theta}}}( {\mathbf{x}}^{(j)})\big{)}^{2}.$ (25) $$

###### Proof of Theorem B.3 .

As discussed earlier,
in light of Proposition B.2, we only need to bound two terms $R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{{\bm{\theta}}})$ and $R^{\omega}_{N}({\bm{\theta}}^{*})-R^{\omega}({\bm{\theta}}^{*})$ to obtain the desired result. Now utilizing Theorem 6 and Theorem 4 (with $\delta$ replaced with $\delta/2$) in Maurer and Pontil [] to bound the two terms, respectively, completes the proof of Theorem B.3.
∎

###### Assumption B.1 .

###### Proposition B.2 .

###### Proof of Proposition B.2 .

###### Theorem B.3 (Formal version of Theorem 3.3 ) .

###### Proof of Theorem B.3 .

### B.2 Token-level generalization bound

Before providing a proof of Theorem 3.5, we first introduce some intermediate results that are needed to prove the theorem. Recall that our training sample $\mathscr{S}_{N}=\{{\mathbf{x}}^{(i)}=[x^{(i)}_{1},\ldots,x^{(i)}_{T}]\}_{i\in[
N]}$ comprises $N$ independent sequences such that ${\mathbf{x}}^{(i)}\sim\mathscr{D},\forall i\in[N]$.
With $\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})$ representing the KD loss on $i$-th sequence, we define
the random variables

$$ $\displaystyle Z^{(i)}_{0}$ $\displaystyle=\mathbb{E}[\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})],$ $\displaystyle Z^{(i)}_{t}$ $\displaystyle=\mathbb{E}\left[\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}}) \;|\;{\mathbf{x}}^{(i)}_{\leq t}\right],~{}\text{for}~{}1\leq t\leq T,$ (26) $$

where $Z^{(i)}_{T}=\mathbb{E}\left[\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})\;|
\;{\mathbf{x}}^{(i)}\right]=\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})$. Note that $\{Z^{(i)}_{t}\}_{0\leq t\leq T}$ is a Doob martingale sequence with respect to the natural filtration $\{\mathcal{F}^{(i)}_{t}\}_{0\leq t\leq T}$ of the random variables $\{x^{(i)}_{1},\ldots,x^{(i)}_{t}\}$ [Ross, , pg 297]. Accordingly, we define a martingale difference sequence $\{\xi^{(i)}_{t},\mathcal{F}^{(i)}_{t}\}_{t\in[T]}$ such that for $t\in[T]$,

$$ $\displaystyle\xi^{(i)}_{t}:=\xi_{t}({\mathbf{x}}^{(i)};{\bm{\theta}})=Z^{(i)}_ {t-1}-Z^{(i)}_{t}=\mathbb{E}\left[\ell^{\omega}({\mathbf{x}};{\bm{\theta}})|{ \mathbf{x}}^{(i)}_{\leq t-1}\right]-\mathbb{E}\left[\ell^{\omega}({\mathbf{x}} ;{\bm{\theta}})|{\mathbf{x}}^{(i)}_{\leq t}\right].$ (27) $$

As per Assumption 3.4, the following holds for each $t\in[T]$:

$$ $\displaystyle|\xi^{(i)}_{t}({\mathbf{x}};{\bm{\theta}})|$ $\displaystyle\leq C_{t}\leq C,$ (28) $\displaystyle\mathbb{E}\left[\left(\xi^{(i)}_{t}\right)^{2}|{\mathbf{x}}_{\leq t -1}\right]$ $\displaystyle\leq V_{t}.$ (29) $$

We are ready to state the first intermediate result which bounds the moment generating function for the following random variable associated with the KD loss on the $i$-th training sequence:

$$ $Z^{(i)}_{0}-Z^{(i)}_{T}=\mathbb{E}\left[\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{ \theta}})\right]-\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}}).$ $$

###### Lemma B.4 .

Under Assumption 3.4, the following holds for each $i\in[N]$:

$$ $\displaystyle\mathbb{E}\left[e^{\lambda\cdot(Z^{(i)}_{0}-Z^{(i)}_{T})/C}\right ]\leq\exp\left(T\cdot f\Big{(}\lambda,\frac{1}{T}\sum_{t=1}^{T}\frac{V_{t}}{C^ {2}}\Big{)}\right),$ (30) $$

where, for $\lambda\geq 0$ and $s\geq 0$,

$$ $\displaystyle f(\lambda,s)\triangleq\log\left(\frac{1}{1+s}\cdot\exp(-\lambda s )+\frac{s}{1+s}\cdot\exp(\lambda)\right).$ (31) $$

###### Proof.

Note that

$$ $\displaystyle\mathbb{E}\left[e^{\lambda\cdot(Z^{(i)}_{0}-Z^{(i)}_{T})/C}\right]$ $\displaystyle=\mathbb{E}\left[e^{\lambda\cdot\sum_{t=1}^{T}\xi^{(i)}_{t}/C}\right]$ $\displaystyle=\mathbb{E}\left[\mathbb{E}\Big{[}e^{\lambda\cdot\sum_{t=1}^{T} \xi^{(i)}_{t}/C}|{\mathbf{x}}^{(i)}_{\leq T-1}\Big{]}\right]$ $\displaystyle\overset{(i)}{=}\mathbb{E}\left[e^{\lambda\cdot\sum_{t=1}^{T-1} \xi^{(i)}_{t}/C}\cdot\mathbb{E}\Big{[}e^{\lambda\cdot\xi^{(i)}_{T}/C}|{\mathbf {x}}^{(i)}_{\leq T-1}\Big{]}\right]$ $\displaystyle\overset{(ii)}{\leq}\mathbb{E}\left[e^{\lambda\cdot\sum_{t=1}^{T- 1}\xi^{(i)}_{t}/C}\cdot e^{f\big{(}\lambda,\;\frac{1}{C^{2}}\cdot\mathbb{E} \left[\left(\xi^{(i)}_{T}\right)^{2}|{\mathbf{x}}^{(i)}_{\leq T-1}\right]\big{ )}}\right]$ $\displaystyle\overset{(iii)}{\leq}\mathbb{E}\left[e^{\lambda\cdot\sum_{t=1}^{T -1}\xi^{(i)}_{t}/C}\cdot e^{f(\lambda,\;\frac{V_{T}}{C^{2}})}\right]$ $\displaystyle=\mathbb{E}\left[e^{\lambda\cdot\sum_{t=1}^{T-1}\xi^{(i)}_{t}/C} \right]\cdot e^{f(\lambda,\;\frac{V_{T}}{C^{2}})}$ (32) $$

where $(i)$ follows as $e^{\lambda\cdot\sum_{t=1}^{T-1}\xi^{(i)}_{t}}$ is $\mathcal{F}^{(i)}_{T-1}$-measurable; $(ii)$ follows from [Fan et al., , Lemma 3.1]; and $(iii)$ follows from Assumption 3.4 and the fact that, for $\lambda>0$ and $s\geq 0$, $f(\lambda,s)$ is an increasing function in its second argument [Fan et al., , Lemma 3.2].
By following the similar steps in (B.2) for $\xi_{i,T-1},\xi_{i,T-2},\ldots,\xi_{i,1}$, we obtain that

$$ $\displaystyle\mathbb{E}\left[e^{\lambda\cdot(Z^{(i)}_{0}-Z^{(i)}_{T})/C}\right ]\leq e^{\sum_{t=1}^{T}f(\lambda,\frac{V_{t}}{C^{2}})}.$ (33) $$

According to [Fan et al., , Lemma 3.2] that, for $\lambda\geq 0$ and $s\geq 0$, $f(\lambda,s)$ is a concave function in its second argument. Thus, it follows from Jensen’s inquality that

$$ $\displaystyle\frac{1}{T}\sum_{t=1}^{T}f\left(\lambda,\frac{V_{t}}{C^{2}}\right )\leq f\left(\lambda,\frac{1}{T}\sum_{t=1}^{T}\frac{V_{t}}{C^{2}}\right).$ (34) $$

By combining (33) and (34), we have

$$ $\displaystyle\mathbb{E}\left[e^{\lambda\cdot(Z^{(i)}_{0}-Z^{(i)}_{T})/C}\right ]\leq e^{T\cdot f\big{(}\lambda,\frac{1}{T}\sum_{t=1}^{T}\frac{V_{t}}{C^{2}} \big{)}},$ (35) $$

which completes the proof.
∎

Now we can leverage Lemma B.4 to obtain the following concentration inequality for the KD training objective.

###### Lemma B.5 .

Let ${\bm{\zeta}}$ and ${\bm{\theta}}\in\Theta$ denote the teacher and student LM, respectively. Then, for $\epsilon>0$, the following holds under Assumption 3.4.

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\left(\mathbb{E}\big{[}\ell^{\omega }({\mathbf{x}}^{(i)};{\bm{\theta}})\big{]}-\ell^{\omega}({\mathbf{x}}^{(i)};{ \bm{\theta}})\right)/C\geq N\epsilon\right)\leq\exp\left(-\frac{N\epsilon^{2}} {2(\sum_{t}\frac{V_{t}}{C^{2}}+\frac{1}{3}\epsilon)}\right).$ (36) $$

###### Proof.

Recall that, as per our notation, we have

$$ $\mathbb{E}\left[\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})\right]-\ell^{ \omega}({\mathbf{x}}^{(i)};{\bm{\theta}})=Z^{(i)}_{0}-Z^{(i)}_{T}.$ $$

Thus,

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\left(\mathbb{E}\big{[}\ell^{\omega }({\mathbf{x}}^{(i)};{\bm{\theta}})\big{]}-\ell^{\omega}({\mathbf{x}}^{(i)};{ \bm{\theta}})\right)/C\geq N\epsilon\right)=\mathbb{P}\left(\sum_{i=1}^{N}\big {(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{)}/C\geq N\epsilon\right).$ (37) $$

It follows from Markov’s inequality that, for $\lambda\geq 0$,

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\big{(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{ )}/C\geq N\epsilon\right)$ $\displaystyle=\mathbb{P}\left(e^{\lambda\cdot\sum_{i=1}^{N}\big{(}Z^{(i)}_{0}- Z^{(i)}_{T}\big{)}/C}\geq e^{N\lambda\epsilon}\right)$ $\displaystyle\leq\frac{\mathbb{E}\left[e^{\lambda\cdot\sum_{i=1}^{N}\big{(}Z^{ (i)}_{0}-Z^{(i)}_{T}\big{)}/C}\right]}{e^{N\lambda\epsilon}}$ $\displaystyle\overset{(i)}{=}\frac{\prod_{i\in[N]}\mathbb{E}\left[e^{\lambda \cdot\big{(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{)}/C}\right]}{e^{N\lambda\epsilon}},$ (38) $$

where $(i)$ follows as $\{Z^{(i)}_{0}-Z^{(i)}_{T}\big{\}}_{i\in[N]}$ are independent random variables.
By combining (B.2) with Lemma B.4, we obtain that

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\big{(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{ )}/C\geq N\epsilon\right)$ $\displaystyle\leq e^{-N\cdot\big{(}\lambda\epsilon-T\cdot f(\lambda,\frac{1}{T }\sum_{t=1}^{T}\frac{V_{t}}{C^{2}})\big{)}}.$ (39) $$

Since (39) holds for each $\lambda\geq 0$, we have

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\big{(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{ )}/C\geq N\epsilon\right)$ $\displaystyle\leq\inf_{\lambda\geq 0}e^{-N\cdot\big{(}\lambda\epsilon-T\cdot f (\lambda,\frac{1}{T}\sum_{t=1}^{T}\frac{V_{t}}{C^{2}})\big{)}}$ (40) $$

Now as argued in the Proof of Remark 2.1 in Fan et al. [], for $0\leq\lambda<3,s\geq 0$, we have

$$ $\displaystyle f(\lambda,s)\leq(e^{\lambda}-1-\lambda)s\leq\frac{\lambda^{2}s}{ 2(1-\frac{1}{3}\lambda)}.$ (41) $$

Thus, it follows from (40) that

$$ $\displaystyle\mathbb{P}\left(\sum_{i=1}^{N}\big{(}Z^{(i)}_{0}-Z^{(i)}_{T}\big{ )}/C\geq N\epsilon\right)$ $\displaystyle\leq\inf_{0\leq\lambda<3}\exp\left({-N\cdot\Big{(}\lambda\epsilon -\frac{\lambda^{2}}{2(1-\frac{1}{3}\lambda)}\cdot\sum_{t}\frac{V_{t}}{C^{2}} \Big{)}}\right)$ $\displaystyle\leq\exp\left(-\frac{N\epsilon^{2}}{2(\sum_{t}\frac{V_{t}}{C^{2}} +\frac{1}{3}\epsilon)}\right).$ (42) $$

This completes the proof.
∎

Equipped with Lemma B.5, we are now ready to prove Theorem 3.5 below.

###### Proof of Theorem 3.5 .

Note that

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{{\bm{\theta}}})= \underbrace{R(\hat{{\bm{\theta}}})-R^{\omega}(\hat{{\bm{\theta}}})}_{({\rm I}) }\;+\;\underbrace{R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{{\bm{ \theta}}})}_{({\rm II})}.$ (43) $$

Following the similar analysis used in the proof of Theorem 3.3, we can bound the term (I) to obtain the following.

$$ $\displaystyle R(\hat{{\bm{\theta}}})-R^{\omega}(\hat{{\bm{\theta}}})\leq\frac{ 4M\omega}{T}\cdot\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}}_{\leq t-1}\sim\mathscr {D}}\mathsf{D}_{\rm TV}\Big{(}P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t -1}),\mathscr{D}(\cdot|{\mathbf{x}}_{\leq t-1})\Big{)}.$ (44) $$

Next, we focus on bounding the term (II). As per notation, for any ${\bm{\theta}}\in\Theta$, we have

$$ $\displaystyle R^{\omega}({\bm{\theta}})-R^{\omega}_{N}({\bm{\theta}})$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\left(\mathbb{E}\big{[}\ell^{\omega}({ \mathbf{x}}^{(i)};{\bm{\theta}})\big{]}-\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{ \theta}})\right).$ (45) $$

Thus, for a fixed ${\bm{\theta}}\in\Theta$, we have

$$ $\displaystyle\mathbb{P}\left(R({{\bm{\theta}}})-R^{\omega}_{N}({\bm{\theta}}) \geq\gamma\right)=\mathbb{P}\left(\frac{1}{N}\sum_{i=1}^{N}\left(\mathbb{E} \big{[}\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})\big{]}-\ell^{\omega}({ \mathbf{x}}^{(i)};{\bm{\theta}})\right)\geq\gamma\right)$ $\displaystyle\qquad\qquad\qquad=\mathbb{P}\left(\sum_{i=1}^{N}\left(\mathbb{E} \big{[}\ell^{\omega}({\mathbf{x}}^{(i)};{\bm{\theta}})\big{]}-\ell^{\omega}({ \mathbf{x}}^{(i)};{\bm{\theta}})\right)/C\geq N\cdot\frac{\gamma}{C}\right)$ $\displaystyle\qquad\qquad\qquad\overset{(i)}{\leq}\exp\left(-\frac{N\gamma^{2} }{2(\sum_{t}V_{t}+\frac{1}{3}C\gamma)}\right),$ (46) $$

where $(i)$ follows from (36) with $\epsilon=\frac{\gamma}{C}$. With some algebra, one can see that the right hand side of (B.2) is bounded by $\delta/|\Theta|$ when

$$ $\displaystyle\gamma\geq\frac{2C}{3N}\cdot\log\left({|\Theta|}/{\delta}\right)+ \sqrt{\frac{2}{N}\cdot\sum_{t}V_{t}\cdot\log\left({|\Theta|}/{\delta}\right)}.$ (47) $$

(To see this, set the right hand side of (B.2) to $\delta/|\Theta|$ to get a quadratic of the form $\gamma^{2}=a\gamma+b$ with $a,b\geq 0$ and note that its non-negative root is $\leq a+\sqrt{b}.$ All $\gamma\geq a+\sqrt{b}$ will make the right hand side of (B.2) $\leq\delta/|\Theta|.$)

Now, by taking union bound, with probability at least $1-\delta$, for all ${\bm{\theta}}\in\Theta$, we have the following.

$$ $\displaystyle R^{\omega}({\bm{\theta}})-R^{\omega}_{N}({\bm{\theta}})\leq\frac {2C}{3N}\cdot\log\left({|\Theta|}/{\delta}\right)+\sqrt{\frac{2}{N}\cdot\sum_{ t}V_{t}\cdot\log\left({|\Theta|}/{\delta}\right)}.$ (48) $$

Since the minimizer of the KD training objective $\hat{{\bm{\theta}}}$ is in $\Theta$, with probablity at least $1-\delta$, we have

$$ $\displaystyle({\rm II})=R^{\omega}(\hat{{\bm{\theta}}})-R^{\omega}_{N}(\hat{{ \bm{\theta}}})\leq\frac{2C}{3N}\cdot\log\left({|\Theta|}/{\delta}\right)+\sqrt {\frac{2}{N}\cdot\sum_{t}V_{t}\cdot\log\left({|\Theta|}/{\delta}\right)}.$ (49) $$

Now, the statement of Theorem 3.5 follows by combining (43), (44), and (49).
∎

###### Lemma B.4 .

###### Proof.

###### Lemma B.5 .

###### Proof.

###### Proof of Theorem 3.5 .

## Appendix C KD can improve generalization via variance reduction

Here, we leverage our novel generalization bounds to provide a theoretical justification for why KD can result in better generalization behavior compared to standard pre-training. In particular, we will focus on our bound in Theorem 3.5.(^3^33One could draw similar conclusion from Theorem 3.3 by extending the arguments from Menon et al. [] to the language modeling setting.) Note that, besides $|\Theta|$, $N$, and $T$ which are independent of the underlying training approach, there are three key quantities that dictate the generalization gap: (1) $\sum_{t}V_{t}$ which is related to the loss variance; (2) $C$ which is related to the extreme values that loss can take; and (3) the divergence between the teacher-provided next-token predictive distribution and the ground truth next-token distribution:

$$ $\textsc{Div}(\zeta,\omega):=\omega\cdot\sum_{t=1}^{T}\mathbb{E}\left[\mathsf{D }_{\rm TV}\left(P_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),\mathscr{ D}(\cdot|{\mathbf{x}}_{\leq t-1})\right)\right].$ $$

Note that, under Assumption 3.1, both KD and standard pre-training loss terms are bounded by $M$, allowing us to provide the same $C$ (as a function of $M$ and $T$) for both KD and standard pre-training. Thus, we focus on the remaining two terms which relate to $\sum_{t}V_{t}$ and $\textsc{Div}(\zeta,\omega)$.

Note that standard pre-training, i.e., training without KD, corresponds to $\omega=0$, which leads to $\textsc{Div}(\zeta,\omega=0)=0$. In contrast, with $\omega>0$, KD would incur a non-zero value for $\textsc{Div}(\zeta,\omega)$. On the other hand, as we will argue next, KD can lead to smaller value of the variance term $\sum_{t}V_{t}$. Thus, as long as the underlying teacher LM approximates the true next-token distribution well enough, it can lead to improved (student) performance or equivalently smaller generalization gap by striking a balance between the divergence (or bias) $\textsc{Div}(\zeta,\omega)$ and variance $\sum_{t}V_{t}$; as a result, realizing a form of bias vs. variance trade-off for LM pre-training.

The variance reduction in the case of KD is the cleanest to observe by focusing on the last summand in $\sum_{t}V_{t}$, i.e., $V_{T}$. Towards this, recall from Assumption 3.4 that, for each ${\bm{\theta}}\in\Theta$, $V_{T}$ bounds the second-order moment of $\xi_{T}({\mathbf{x}};{\bm{\theta}})$.
Define the short-hand

$$ ${P}^{(x_{t},\omega)}_{{\bm{\zeta}},\rho}(\cdot|{\mathbf{x}}_{\leq t-1)}:=(1- \omega)\cdot\mathbbm{1}_{x_{t}}(\cdot)+\omega\cdot P_{{\bm{\zeta}},\rho}(\cdot |{\mathbf{x}}_{\leq t-1)}$ (50) $$

and write

$$ $\displaystyle\xi_{T}({\mathbf{x}};{\bm{\theta}})=\mathbb{E}\left[\ell^{\omega} ({\mathbf{x}};{\bm{\theta}})|{\mathbf{x}}_{\leq T-1}\right]-\mathbb{E}\left[ \ell^{\omega}({\mathbf{x}};{\bm{\theta}})|{\mathbf{x}}_{\leq T}\right]$ $\displaystyle=\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}\mathsf{CE}\big{(}{P}^{ (x_{t},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{\theta}}}( \cdot|{\mathbf{x}}_{\leq t-1})\big{)}|{\mathbf{x}}_{\leq T-1}\right]\;-\;$ $\displaystyle\qquad\mathbb{E}\left[\frac{1}{T}\sum_{t=1}^{T}\mathsf{CE}\big{(} {P}^{(x_{t},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{ \theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})\big{)}|{\mathbf{x}}_{\leq T}\right]$ $\displaystyle\overset{(i)}{=}\frac{1}{T}\sum_{t=1}^{T-1}\mathsf{CE}\big{(}{P}^ {(x_{t},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{\theta}}} (\cdot|{\mathbf{x}}_{\leq t-1})\big{)}+\mathbb{E}\left[\frac{1}{T}\mathsf{CE} \big{(}{P}^{(x_{T},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq T-1}),P_{{ \bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq T-1})\big{)}|{\mathbf{x}}_{\leq T-1}\right]$ $\displaystyle\qquad-\frac{1}{T}\sum_{t=1}^{T}\mathsf{CE}\big{(}{P}^{(x_{t}, \omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq t-1}),P_{{\bm{\theta}}}(\cdot|{ \mathbf{x}}_{\leq t-1})\big{)}$ $\displaystyle\overset{(ii)}{=}\mathbb{E}\left[\frac{1}{T}\mathsf{CE}\big{(}{P} ^{(x_{T},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq T-1}),P_{{\bm{\theta}} }(\cdot|{\mathbf{x}}_{\leq T-1})\big{)}|{\mathbf{x}}_{\leq T-1}\right]-\frac{1 }{T}\mathsf{CE}\big{(}{P}^{(x_{T},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{ \leq T-1}),P_{{\bm{\theta}}}(\cdot|{\mathbf{x}}_{\leq t-1})\big{)}$ $\displaystyle=(1-\omega)\cdot\left(\mathbb{E}\left[-\frac{1}{T}\cdot\log P_{{ \bm{\theta}}}(x_{T}|{\mathbf{x}}_{\leq t-1})|{\mathbf{x}}_{\leq T-1}\right]+ \frac{1}{T}\cdot\log P_{{\bm{\theta}}}(x_{T}|{\mathbf{x}}_{\leq T-1})\right)$ (51) $$

where $(i)$ follows we can remove expectation for those terms that are functions of those random variables that we condition on; and $(ii)$ follows by removing the terms that cancel each other; and the last line follows as we have

$$ ${P}^{(x_{T},\omega)}_{\zeta,\rho}(\cdot|{\mathbf{x}}_{\leq T-1})=(1-\omega) \cdot\mathds{1}_{x_{T}}(\cdot)+\omega\cdot P_{{\bm{\zeta}},\rho}(\cdot|{ \mathbf{x}}_{\leq T-1}).$ $$

It follows from (C) that, for any ${\bm{\theta}}\in\Theta$, we have

$$ $\displaystyle\mathbb{E}\left[\xi^{2}_{T}({\mathbf{x}};{\bm{\theta}},{\bm{\zeta }})\right]=(1-\omega)\cdot{\rm Var}\left[\frac{1}{T}\cdot\log P_{{\bm{\theta}} }(x_{T}|{\mathbf{x}}_{\leq t-1})~{}\Big{|}~{}{\mathbf{x}}_{\leq T-1}\right],$ (52) $$

where ${\rm Var}\left[\cdot|\cdot\right]$ denotes conditional variance. Note that (52) shows that $V_{T}$ decreases with $\omega$ in $[0,1]$. This highlights that KD, i.e., $\omega>0$, would realize a smaller variance than standard pre-training, i.e., $\omega=0$. Thus, to realize improved generalization via KD, one needs to select the distillation weight $\omega$ so that the variance reduction via KD offsets the divergence term $\textsc{Div}(\zeta,\omega)$. In particular, when the teacher LM approximates the ground truth next-token distribution very well, i.e., $\textsc{Div}(\zeta,\omega)$ term is small even for a relatively large value of $\omega$, the variance reduction via KD becomes prominent, ensuring significant improvement over standard pre-training in terms of generalization performance.

## Appendix D Bounding excess risk for KD

Different from the surrogate (empirical or population) risks utilized in the main text (cf. Section 3),
which utilize the cross-entropy loss as a surrogate loss, one could directly work with the risk defined with respect to a particular evaluation metric (and the corresponding loss) that one cares about. Since our training focuses on correct next-token prediction, we can focus on the accuracy of the next-token prediction under greedy-decoding as one such metric. This amounts to the following
(population) risk with respect to $0/1$-loss.

$$ $\displaystyle R_{0/1}({\bm{\theta}})$ $\displaystyle:=\mathbb{E}_{{\mathbf{x}}\sim\mathscr{D}}\Big{[}\sum_{t=1}^{T} \mathsf{1}\{\operatorname*{arg\,max}_{v}P_{{\bm{\theta}}}(v|{\mathbf{x}}_{\leq t -1})\neq x_{t}\Big{]}$ $\displaystyle=\sum_{t=1}^{T}\mathbb{E}_{{\mathbf{x}}_{\leq t-1}\sim\mathscr{D} }\Big{[}\sum_{v\in\mathcal{V}}\mathscr{D}(v|{\mathbf{x}}_{\leq t-1})\cdot \mathsf{1}\{\operatorname*{arg\,max}_{v^{\prime}}P_{{\bm{\theta}}}(v^{\prime}| {\mathbf{x}}_{\leq t-1})\neq v\}\Big{]}.$ (53) $$

A large body of literature [see, e.g., Bartlett et al., , Zhang, , Steinwart, , Pires and Szepesvári, , and references therein] has studied calibration functions that enable converting bounds on excess surrogate risk to control the excess risk. Applying the calibration functions for the cross-entropy loss [Pires and Szepesvári, ], we obtain the following bound on the excess risk for next-token prediction:

$$ $\displaystyle R_{0/1}(\hat{{\bm{\theta}}})-R_{0/1}({\bm{\theta}}^{*})\leq \mathsf{g}^{-1}\left(R(\hat{{\bm{\theta}}})-R({\bm{\theta}}^{*})\right),$ (54) $$

where $\mathsf{g}^{-1}(\cdot)$ denotes the inverse of the function $\mathsf{g}:\epsilon\mapsto\frac{1}{2}\big{(}(1-\epsilon)\log(1-\epsilon)+(1+
\epsilon)\log(1+\epsilon)\big{)}.$

## Appendix E Experimental Setup Details

Model architectures and pre-training data.
We work with standard decoder-only Transformer-based LMs. Our small model (SLM) has 1.5B parameters. It comprises a 44 layer Transformer network with model dimension 1024, MLP hidden dimension 8192, and 4 attention heads.
For the larger LM, we employ 2.8B parameter models consisting of 92 layer Transformer networks with model dimension 1024, MLP hidden dimension 8192, and 4 attention heads based on multi-query attention [Shazeer, ].
We use a SentencePiece tokenizer [Kudo and Richardson, ] from Du et al. [] with a vocabulary size of 256$K$.

We pre-train all LMs on the Pile dataset [Gao et al., ] by minimizing the UL2 objective [Tay et al., ]
with a mixture of four tasks: (1) causal LM task; (2) prefix LM task with mean prefix length of $1/4$th the sequence length,
(3) span corruption task with $r=15\%$ of the tokens corrupted and mean corrupted span length $\mu=3$; and
(4) span corruption task with $r=50\%$ of the tokens corrupted and mean corrupted span length $\mu=32$. The four tasks are mixed at a ratio of $6$:$2$:$1$:$1$.
We pre-train LMs for approximately 545 billion tokens, with a batch size of 2048 and input sequence length of 1280. This translates to a little over two epochs on the Pile data. As for the optimization method, we utilize Adafactor algorithm [Shazeer and Stern, ].
We use a cosine learning rate decay schedule with a peak learning rate of $0.001$, $4000$ warmup steps and final learning rate of $0.0001.$ Training is done on 1024 TPU-v5e chips with JAX [Bradbury et al., ]
and SeqIO [Roberts et al., ].

## Appendix F Few-shot evaluation tasks

We performed a comprehensive few-shot evaluation of pre-trained LMs on $28$ benchmarks. Below, we list these by organizing them according to the corresponding domain.

World Knowledge:
NQ-Open [Lee et al., ],
TriviaQA [Joshi et al., ], TyDiQA-NoContext (English)[Clark et al., ],
Web Questions [Berant et al., ].

Reading Comprehension:
RACE-M, RACE-H [Lai et al., ],
SQuADv2 [Lee et al., ],
TyDiQA-GoldP (English)[Clark et al., ].

Commonsense Reasoning:
ARC (Easy) and ARC (Challenge) [Clark et al., ],
HellaSwag [Zellers et al., ],
OpenBookQA [Mihaylov et al., ],
PiQA [Bisk et al., ],
StoryCloze [Mostafazadeh et al., ],
Winogrande [Sakaguchi et al., ].

SuperGLUE [Wang et al., ]:
BoolQ [Clark et al., ],
CB [de Marneffe et al., ],
COPA [Gordon et al., ],
RTE [Dagan et al., ],
WiC [Pilehvar and Camacho-Collados, ],
WSC [Levesque et al., ],
MultiRC [Khashabi et al., ],
ReCoRD [Zhang et al., ].

Natural Language Generation (NLG):
English portions of the three benchmarks – XLSum [Hasan et al., ], XSum [Narayan et al., ] and WikiLingua [Ladhak et al., ].

Open-ended Cloze task: LAMBADA Paperno et al. [].

Code generation: Mostly Basic Python Problems (MBPP)  [Austin et al., ].

## Appendix G Additional Few-shot evaluation results

Table 5 is an expansion of Table 2 in the main text.
All evaluations are $1$-shot, except for MBPP which is $3$-shot.
In the metric column, EM, Acc, and Rg2 are abbreviations for Exact Match, Accuracy, and Rouge2, respectively.
For MBPP, the metric is the fraction of success ignoring challenge problems. As mentioned in Section 5.1, for each benchmark, we typically report the corresponding prevalent metric in the literature. For TyDiQA benchmarks, we report the F1 score as opposed to EM as it is the primary metric in Clark et al. []. For MultiRC in SuperGLUE, we report F1 metric as per Du et al. [].

**Table 5: Comprehensive few-shot performance of pre-trained LMs. SLM serves as the teacher LM for SALT & $\texttt{SALT}_{\texttt{DS}}$ during the KD phase of their pre-training and for RKD throughout its pre-training. Baseline employs standard pre-training without KD from SLM. SALT and $\texttt{SALT}_{\texttt{DS}}$ already outperform Baselinein terms of average few-shot performance at $70$% of their training step budget, thereby improving both training efficiency and model quality. RKD, i.e., naively preforming KD from the small model through the pre-training, performs much worse than Baseline. The best and second-best results for each domain are boldfaced and underlined, respectively.**
| Domain | Dataset | Metric |  | SLM |  | Baseline |  | RKD |  | SALT |  | $\texttt{SALT}_{\texttt{DS}}$ |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  |  |  |  |  | @100% |  | @100% |  | @70% | @100% |  | @70% | @100% |
|  |  |  |  |  |  | steps |  | steps |  | steps | steps |  | steps | steps |
| World<br>Knowledge | NaturalQuestions-Open | EM |  | 5.90 |  | 8.70 |  | 6.70 |  | 9.40 | 10.10 |  | 8.40 | 9.00 |
| TriviaQA | EM |  | 30.09 |  | 43.15 |  | 34.87 |  | 39.87 | 43.71 |  | 39.37 | 41.27 |  |
| TyDiQA-NoContext | F1 |  | 22.20 |  | 28.20 |  | 26.10 |  | 27.90 | 27.10 |  | 25.90 | 27.20 |  |
| WebQuestions | EM |  | 5.40 |  | 8.70 |  | 7.10 |  | 9.20 | 9.90 |  | 8.90 | 9.40 |  |
| Domain average |  |  | 15.90 |  | 22.19 |  | 18.69 |  | 21.59 | 22.70 |  | 20.64 | 21.72 |  |
| Reading<br>Comprehension | RACE-M | Acc |  | 52.60 |  | 57.00 |  | 54.00 |  | 58.60 | 58.90 |  | 57.90 | 58.40 |
| RACE-H | Acc |  | 37.50 |  | 42.30 |  | 39.70 |  | 42.20 | 42.30 |  | 42.10 | 42.30 |  |
| SQuADv2 | EM |  | 43.30 |  | 54.80 |  | 50.90 |  | 54.60 | 55.90 |  | 57.60 | 57.90 |  |
| TyDiQA-GoldP | F1 |  | 51.80 |  | 57.90 |  | 59.40 |  | 58.80 | 61.10 |  | 59.80 | 61.10 |  |
| Domain average |  |  | 46.30 |  | 53.00 |  | 51.00 |  | 53.55 | 54.55 |  | 54.35 | 54.93 |  |
| Commonsense<br>Reasoning | ARC-E | Acc |  | 64.60 |  | 68.40 |  | 66.00 |  | 67.60 | 67.60 |  | 69.40 | 69.00 |
| ARC-C | Acc |  | 32.40 |  | 37.10 |  | 33.70 |  | 38.00 | 38.40 |  | 38.10 | 37.30 |  |
| HellaSwag | Acc |  | 56.00 |  | 62.80 |  | 56.20 |  | 62.00 | 63.30 |  | 63.10 | 63.80 |  |
| OpenBookQA | Acc |  | 48.00 |  | 50.00 |  | 45.80 |  | 47.20 | 48.20 |  | 47.60 | 48.20 |  |
| PiQA | Acc |  | 72.00 |  | 75.40 |  | 72.60 |  | 73.20 | 73.70 |  | 74.10 | 73.90 |  |
| StoryCloze | Acc |  | 73.10 |  | 77.20 |  | 73.70 |  | 76.90 | 76.80 |  | 77.00 | 77.10 |  |
| WinoGrande | Acc |  | 58.20 |  | 63.00 |  | 60.10 |  | 64.00 | 63.70 |  | 64.70 | 65.40 |  |
| Domain average |  |  | 57.76 |  | 61.99 |  | 58.30 |  | 61.27 | 61.67 |  | 62.00 | 62.10 |  |
|  | LAMBADA | Acc |  | 26.90 |  | 36.20 |  | 31.10 |  | 50.70 | 48.30 |  | 48.00 | 53.00 |
| SuperGLUE | BoolQ | Acc |  | 63.40 |  | 64.30 |  | 62.50 |  | 64.10 | 62.30 |  | 65.50 | 64.30 |
| CB | Acc |  | 37.50 |  | 58.90 |  | 50.00 |  | 60.70 | 53.60 |  | 55.40 | 53.60 |  |
| COPA | Acc |  | 77.00 |  | 79.00 |  | 71.00 |  | 76.00 | 77.00 |  | 81.00 | 77.00 |  |
| MultiRC | F1 |  | 53.80 |  | 54.20 |  | 53.50 |  | 57.50 | 58.60 |  | 50.70 | 53.00 |  |
| RTE | Acc |  | 55.20 |  | 55.60 |  | 59.90 |  | 57.80 | 58.50 |  | 54.20 | 58.50 |  |
| ReCoRD | Acc |  | 84.80 |  | 87.10 |  | 85.20 |  | 86.60 | 86.90 |  | 87.20 | 87.30 |  |
| WiC | Acc |  | 48.40 |  | 47.20 |  | 47.20 |  | 49.80 | 48.10 |  | 50.00 | 50.90 |  |
| WSC | Acc |  | 72.60 |  | 77.90 |  | 74.00 |  | 77.90 | 77.20 |  | 83.90 | 80.00 |  |
| Domain average |  |  | 61.59 |  | 65.53 |  | 62.91 |  | 66.30 | 65.28 |  | 65.99 | 65.58 |  |
| NLG | GEM-XLSum | Rg2 |  | 2.80 |  | 4.10 |  | 3.40 |  | 4.40 | 4.40 |  | 4.60 | 4.60 |
| GEM-XSum | Rg2 |  | 2.80 |  | 5.10 |  | 3.20 |  | 5.00 | 5.10 |  | 5.40 | 5.40 |  |
| WikiLingua | Rg2 |  | 3.80 |  | 4.60 |  | 3.60 |  | 4.50 | 4.70 |  | 4.40 | 4.50 |  |
| Domain average |  |  | 3.13 |  | 4.60 |  | 3.40 |  | 4.63 | 4.73 |  | 4.80 | 4.83 |  |
|  | MBPP | Acc |  | 9.60 |  | 16.20 |  | 11.40 |  | 15.60 | 17.00 |  | 16.60 | 17.80 |
|  | Average (28 tasks) |  |  | 42.56 |  | 47.32 |  | 44.39 |  | 47.86 | 47.94 |  | 47.89 | 48.26 |

## Appendix H Ablation study of various design choices in SALT

In this section, we explore
how various design choices pertaining SALT affect its final performance.

Distillation from a better quality small model. So far we assumed that SLM is also pre-trained for the same number of tokens as the LLM. Since training for SLM is relatively cheaper, one could consider a scenario where one invests more compute resources in improving the small model if it can eventually be beneficial in improving the LLM quality via SALT. Towards this, we employ a small LM that is trained $\sim 2.5$ times longer – $498$K steps vs. $208$K steps in Section 5.2.(^4^44This approach aligns with the recent studies [Touvron et al., , Gadre et al., ] that train small LMs well beyond the their optimal compute budget as predicted by neural scaling laws [Hoffmann et al., ].) As evident in
Table 6,
SALT is indeed able to utilize the better small model as a teacher in the KD phase to further improve the LLM quality, as measured by the average few-shot performance.

**Table 6: Effect of improved SLM(comprehensive few-shot evaluation). SALT with a better teacher – a SLM trained for 498K steps as opposed to 208K steps – yields LLM with better average few-shot performance. For each benchmark, the best and second best results are boldfaced and underlined, respectively.**
| Domain | Dataset | Metric | SLM trained for 208K steps | SLM trained for 498K steps | SALT w/ KD from SLM trained for 208K steps | SALT w/ KD from SLM trained for 498K steps |
| --- | --- | --- | --- | --- | --- | --- |
| World<br>Knowledge | NaturalQuestions-Open | EM | 5.90 | 6.30 | 10.10 | 9.00 |
| TriviaQA | EM | 30.09 | 31.74 | 43.71 | 41.61 |  |
| TyDiQA-NoContext | F1 | 22.20 | 23.80 | 27.10 | 26.20 |  |
| WebQuestions | EM | 5.40 | 7.60 | 9.90 | 9.10 |  |
| Domain average |  | 15.90 | 17.36 | 22.70 | 21.48 |  |
| Reading<br>Comprehension | RACE-M | Acc | 52.60 | 54.40 | 58.90 | 57.00 |
| RACE-H | Acc | 37.50 | 39.40 | 42.30 | 42.00 |  |
| SQuADv2 | EM | 43.30 | 49.00 | 55.90 | 57.90 |  |
| TyDiQA-GoldP | F1 | 51.80 | 55.90 | 61.10 | 56.80 |  |
| Domain average |  | 46.30 | 49.67 | 54.55 | 53.43 |  |
| Commonsense<br>Reasoning | ARC-E | Acc | 64.60 | 65.50 | 67.60 | 69.30 |
| ARC-C | Acc | 32.40 | 34.30 | 38.40 | 39.10 |  |
| HellaSwag | Acc | 56.00 | 57.80 | 63.30 | 63.20 |  |
| OpenBookQA | Acc | 48.00 | 46.40 | 48.20 | 49.00 |  |
| PiQA | Acc | 72.00 | 72.90 | 73.70 | 74.60 |  |
| StoryCloze | Acc | 73.10 | 75.00 | 76.80 | 76.90 |  |
| WinoGrande | Acc | 58.20 | 59.40 | 63.70 | 63.80 |  |
| Domain average |  | 57.76 | 58.76 | 61.67 | 62.27 |  |
|  | LAMBADA | Acc | 26.90 | 37.80 | 48.30 | 47.80 |
| SuperGLUE | BoolQ | Acc | 63.40 | 61.40 | 62.30 | 65.80 |
| CB | Acc | 37.50 | 42.90 | 53.60 | 73.20 |  |
| COPA | Acc | 77.00 | 78.00 | 77.00 | 79.00 |  |
| MultiRC | F1 | 53.80 | 48.40 | 58.60 | 53.20 |  |
| RTE | Acc | 55.20 | 52.30 | 58.50 | 61.70 |  |
| ReCoRD | Acc | 84.80 | 85.50 | 86.90 | 87.10 |  |
| WIC | Acc | 48.40 | 47.30 | 48.10 | 49.20 |  |
| WSC | Acc | 72.60 | 72.30 | 77.20 | 79.30 |  |
| Domain average |  | 61.59 | 61.01 | 65.28 | 68.56 |  |
| NLG | GEM-XLSum | Rg2 | 2.80 | 3.50 | 4.40 | 4.30 |
| GEM-XSum | Rg2 | 2.80 | 3.10 | 5.10 | 5.60 |  |
| WikiLingua | Rg2 | 3.80 | 3.80 | 4.70 | 4.40 |  |
| Domain average |  | 3.13 | 3.47 | 4.73 | 4.77 |  |
|  | MBPP | Acc | 9.60 | 12.80 | 17.00 | 17.40 |
|  | Average (28 tasks) |  | 42.56 | 43.88 | 47.94 | 48.70 |

Varying transition point. A key design choice for SALT is the selection of the transition point $n_{\rm KD}$ from KD phase (first stage) to standard training (second stage). Table 7 shows
few-shot performance of SALT as we vary the transition point. Note that SALT ensures quality gains for LLM with a wide range of values for $n_{\rm KD}$ while demonstrating an inverted U-shape for LLM quality. We see consistent performance improvement from $n_{\rm KD}=0$ (equivalent to Baseline) to $n_{\rm KD}=60$K which eventually degrades at $n_{\rm KD}=208$K (equivalent to RKD). Given the training overhead of KD phase (see discussion in Section 5.2), smaller value of $n_{\rm KD}$ helps ensure training efficiency gains via SALT. Thus, we worked with $n_{\rm KD}=36$K in Section 5.2 as $n_{\rm KD}=60$K only provides marginal quality gains if one takes into account the increased training cost due to longer KD phase.

**Table 7: Effect of varying transitions step (comprehensive few-shot evaluation). The performance improvement via SALT over Baseline is stable in a wide range of $n_{\rm KD}$ (20k to 60k steps). Eventually, with much larger $n_{\rm KD}$, SALT performance degrades significantly (208k steps). For each benchmark, the best and second best results are boldfaced and underlined, respectively.**
| Domain | Dataset | Metric | SLM | Baseline | SALT w/ $n_{\rm KD}=20$K | SALT w/ $n_{\rm KD}=36$K | SALT w/ $n_{\rm KD}=60$K | SALT w/ $n_{\rm KD}=208$K (RKD) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| World<br>Knowledge | NaturalQuestions-Open | EM | 5.90 | 8.70 | 8.90 | 10.10 | 9.30 | 6.70 |
| TriviaQA | EM | 30.09 | 43.15 | 41.52 | 43.71 | 42.84 | 34.87 |  |
| TyDiQA-NoContext | F1 | 22.20 | 28.20 | 26.40 | 27.10 | 26.60 | 26.10 |  |
| WebQuestions | EM | 5.40 | 8.70 | 8.20 | 9.90 | 8.60 | 7.10 |  |
| Domain average |  | 15.90 | 22.19 | 21.26 | 22.70 | 21.83 | 18.69 |  |
| Reading<br>Comprehension | RACE-M | Acc | 52.60 | 57.00 | 58.70 | 58.90 | 58.60 | 54.00 |
| RACE-H | Acc | 37.50 | 42.30 | 41.00 | 42.30 | 42.10 | 39.70 |  |
| SQuADv2 | EM | 43.30 | 54.80 | 55.30 | 55.90 | 55.50 | 50.90 |  |
| TyDiQA-GoldP | F1 | 51.80 | 57.90 | 56.50 | 61.10 | 59.30 | 59.40 |  |
| Domain average |  | 46.30 | 53.00 | 52.88 | 54.55 | 53.88 | 51.00 |  |
| Commonsense<br>Reasoning | ARC-E | Acc | 64.60 | 68.40 | 67.80 | 67.60 | 68.40 | 66.00 |
| ARC-C | Acc | 32.40 | 37.10 | 38.10 | 38.40 | 38.70 | 33.70 |  |
| HellaSwag | Acc | 56.00 | 62.80 | 62.80 | 63.30 | 62.90 | 56.20 |  |
| OpenBookQA | Acc | 48.00 | 50.00 | 48.00 | 48.20 | 48.20 | 45.80 |  |
| PiQA | Acc | 72.00 | 75.40 | 75.40 | 73.70 | 74.40 | 72.60 |  |
| StoryCloze | Acc | 73.10 | 77.20 | 76.90 | 76.80 | 76.50 | 73.70 |  |
| WinoGrande | Acc | 58.20 | 63.00 | 63.40 | 63.70 | 62.00 | 60.10 |  |
| Domain average |  | 57.76 | 61.99 | 61.77 | 61.67 | 61.59 | 58.30 |  |
|  | LAMBADA | Acc | 26.90 | 36.20 | 44.70 | 48.30 | 53.30 | 31.10 |
| SuperGLUE | BoolQ | Acc | 63.40 | 64.30 | 63.90 | 62.30 | 63.80 | 62.50 |
| CB | Acc | 37.50 | 58.90 | 60.70 | 53.60 | 55.40 | 50.00 |  |
| COPA | Acc | 77.00 | 79.00 | 76.00 | 77.00 | 77.00 | 71.00 |  |
| MultiRC | F1 | 53.80 | 54.20 | 53.80 | 58.60 | 55.20 | 53.50 |  |
| RTE | Acc | 55.20 | 55.60 | 52.30 | 58.50 | 59.90 | 59.90 |  |
| ReCoRD | Acc | 84.80 | 87.10 | 86.90 | 86.90 | 86.70 | 85.20 |  |
| WIC | Acc | 48.40 | 47.20 | 51.30 | 48.10 | 50.00 | 47.20 |  |
| WSC | Acc | 72.60 | 77.90 | 77.50 | 77.20 | 77.90 | 74.00 |  |
| Domain average |  | 61.59 | 65.53 | 65.30 | 65.28 | 65.74 | 62.91 |  |
| NLG | GEM-XLSum | Rg2 | 2.80 | 4.10 | 4.50 | 4.40 | 4.70 | 3.40 |
| GEM-XSum | Rg2 | 2.80 | 5.10 | 5.80 | 5.10 | 4.80 | 3.20 |  |
| WikiLingua | Rg2 | 3.80 | 4.60 | 4.30 | 4.70 | 4.60 | 3.60 |  |
| Domain average |  | 3.13 | 4.60 | 4.87 | 4.73 | 4.70 | 3.40 |  |
|  | MBPP | Acc | 9.60 | 16.20 | 16.60 | 17.00 | 16.40 | 11.40 |
|  | Average (28 tasks) |  | 42.56 | 47.32 | 47.40 | 47.94 | 47.99 | 44.39 |

Different transition strategies. In our study thus far, we have worked with Step transition between the two training stages in SALT where we abruptly stop performing KD after $n_{\rm KD}$ training steps. Looking at Figure 2, this causes an abrupt change in the model behavior during training, as observed in the next-token prediction accuracy curve for the training set (see similar behavior for log-perplexity in Figure 3). This raises a question if a smoother transition between the two stages can improve the training stability and thereby ensure higher final LLM quality. While there is a large space of potential choices of such smooth transition strategies, here we explore two natural candidates: (1) Linear decay where we linearly decrease the distillation loss weight to 0 0 between $n_{{\rm KD},1}=32$K and $n_{\rm KD}=36$K steps; and (2) Linear ratio decay where we linearly decrease the ratio of distillation loss weight and standard loss weight $\frac{\omega}{1-\omega}$ to 0 0 between $n_{{\rm KD},1}=32$K and $n_{{\rm KD},2}=36$K training steps. As recorded in Table 8, the step transition constitutes a reasonable design choice for SALT as it outperforms both the considered alternatives in terms of average few-shot performance of the resulting pre-trained LLM.

**Table 8: Effect of different transition strategies (comprehensive few-shot evaluation. The Step transition used in this work (cf. Algorithm 1) performs well compared to two natural alternative strategies. For each benchmark, the best and second best results are boldfaced and underlined, respectively.**
| Domain | Dataset | Metric | SLM | Baseline | SALT w/ | SALT w/ | SALT w/ |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  |  |  |  |  | Step | Linear decay | Linear ratio decay |
| World<br>Knowledge | NaturalQuestions-Open | EM | 5.90 | 8.70 | 10.10 | 8.20 | 8.10 |
| TriviaQA | EM | 30.09 | 43.15 | 43.71 | 43.46 | 43.51 |  |
| TyDiQA-NoContext | F1 | 22.20 | 28.20 | 27.10 | 28.40 | 27.20 |  |
| WebQuestions | EM | 5.40 | 8.70 | 9.90 | 8.20 | 8.40 |  |
| Domain average |  | 15.90 | 22.19 | 22.70 | 22.07 | 21.80 |  |
| Reading<br>Comprehension | RACE-M | Acc | 52.60 | 57.00 | 58.90 | 57.90 | 57.40 |
| RACE-H | Acc | 37.50 | 42.30 | 42.30 | 42.10 | 43.50 |  |
| SQuADv2 | EM | 43.30 | 54.80 | 55.90 | 56.40 | 57.10 |  |
| TyDiQA-GoldP | F1 | 51.80 | 57.90 | 61.10 | 58.30 | 57.80 |  |
| Domain average |  | 46.30 | 53.00 | 54.55 | 53.68 | 53.95 |  |
| Commonsense<br>Reasoning | ARC-E | Acc | 64.60 | 68.40 | 67.60 | 68.60 | 68.70 |
| ARC-C | Acc | 32.40 | 37.10 | 38.40 | 38.60 | 39.80 |  |
| HellaSwag | Acc | 56.00 | 62.80 | 63.30 | 63.30 | 63.50 |  |
| OpenBookQA | Acc | 48.00 | 50.00 | 48.20 | 48.00 | 47.40 |  |
| PiQA | Acc | 72.00 | 75.40 | 73.70 | 74.60 | 73.90 |  |
| StoryCloze | Acc | 73.10 | 77.20 | 76.80 | 76.60 | 76.50 |  |
| WinoGrande | Acc | 58.20 | 63.00 | 63.70 | 62.70 | 63.10 |  |
| Domain average |  | 57.76 | 61.99 | 61.67 | 61.77 | 61.84 |  |
|  | LAMBADA | Acc | 26.90 | 36.20 | 48.30 | 40.50 | 42.60 |
| SuperGLUE | BoolQ | Acc | 63.40 | 64.30 | 62.30 | 67.90 | 66.50 |
| CB | Acc | 37.50 | 58.90 | 53.60 | 44.60 | 46.40 |  |
| COPA | Acc | 77.00 | 79.00 | 77.00 | 79.00 | 81.00 |  |
| MultiRC | F1 | 53.80 | 54.20 | 58.60 | 53.90 | 61.60 |  |
| RTE | Acc | 55.20 | 55.60 | 58.50 | 56.30 | 55.20 |  |
| ReCoRD | Acc | 84.80 | 87.10 | 86.90 | 87.00 | 87.30 |  |
| WIC | Acc | 48.40 | 47.20 | 48.10 | 46.60 | 50.50 |  |
| WSC | Acc | 72.60 | 77.90 | 77.20 | 78.60 | 78.90 |  |
| Domain average |  | 61.59 | 65.53 | 65.28 | 64.24 | 65.92 |  |
| NLG | GEM-XLSum | Rg2 | 2.80 | 4.10 | 4.40 | 4.70 | 4.50 |
| GEM-XSum | Rg2 | 2.80 | 5.10 | 5.10 | 4.60 | 5.10 |  |
| WikiLingua | Rg2 | 3.80 | 4.60 | 4.70 | 4.80 | 4.60 |  |
| Domain average |  | 3.13 | 4.60 | 4.73 | 4.70 | 4.73 |  |
|  | MBPP | Acc | 9.60 | 16.20 | 17.00 | 15.20 | 17.00 |
|  | Average (28 tasks) |  | 42.56 | 47.32 | 47.94 | 47.11 | 47.75 |

## Appendix I Log perplexity of the models

Figure 3 shows the log perplexity of the SALT and RKD pre-trained models along with Baseline and SLM. The log perplexity for RKD stays at a higher level than even SLM. Recall that RKD optimizes a sum of two losses – KD loss with weight $\omega=0.667$ and the standard one-hot training loss with weight $1-\omega$. As the training log perplexity plotted in Figure 3 is the same as the standard hot training loss, the methods which directly optimize for that alone (Baseline, SLM and in the second stage, SALT) have lower log perplexity on training set than RKD which optimizes additionally for distillation loss.

Figure: Figure 3: Log perplexity for different models during their pre-training, as measured on a subset of the Pile training set.
Refer to caption: x3.png

## Appendix J Additional results: Learning easy vs. hard instance via SALT

Creation of different hardness buckets.
For each evaluation benchmark, we first assign a relative rank to each test instance/example in the benchmark, representing its degree of difficulty.
A test example with the lowest rank (easiest) is the one on which the small teacher LM achieves the largest task evaluation score, e.g., Rouge-2 metric for the XLSum task.
Similarly, subsequent test examples are assigned ranks in descending order of the task evaluation score achieved by the small teacher LM.
If two examples have the same evaluation score, the one with higher confidence score from the teacher (on its generated output) is deemed to have a lower rank.
Each test example is assigned to one of the three buckets: ‘easy‘, ‘medium‘, or ‘hard‘, according to whether its difficulty rank is in the first, second, or third tertile, respectively.

In Tables 9, 10 and 11, we report the results for SQuAD-v2, TriviaQA and LAMBADA respectively, sliced by difficulty level.

**Table 9: Few-shot evaluation on different buckets of SQuAD-v2. Each number shows average Exact Match scores on the corresponding bucket. We use gray, green, and red to highlight the results similar to, better than, and worse than Baseline performance, respectively.**
|  | Evaluation stage (steps) | Easy | Medium | Hard |
| --- | --- | --- | --- | --- |
| SLM | Final (208K) | 1.00 | 0.30 | 0.00 |
| Baseline | Early (36K) | 0.86 | 0.41 | 0.23 |
| RKD | 0.86 | 0.37 | 0.17 |  |
| SALT | 0.86 | 0.37 | 0.17 |  |
| Baseline | Final (208K) | 0.89 | 0.47 | 0.28 |
| RKD | 0.91 | 0.42 | 0.20 |  |
| SALT | 0.89 | 0.50 | 0.29 |  |

**Table 10: Few-shot evaluation on different buckets of TriviaQA. Each number shows average Exact Match scores on the corresponding bucket. We use gray, green, and red to highlight the results similar to, better than, and worse than Baseline performance, respectively.**
|  | Evaluation stage (steps) | Easy | Medium | Hard |
| --- | --- | --- | --- | --- |
| SLM | Final (208K) | 0.90 | 0.00 | 0.00 |
| Baseline | Early (36K) | 0.63 | 0.11 | 0.08 |
| RKD | 0.67 | 0.10 | 0.06 |  |
| SALT | 0.67 | 0.10 | 0.06 |  |
| Baseline | Final (208K) | 0.80 | 0.28 | 0.22 |
| RKD | 0.79 | 0.14 | 0.11 |  |
| SALT | 0.81 | 0.27 | 0.23 |  |

**Table 11: Few-shot evaluation on different buckets of LAMBADA. Each number shows average Accuracy on the corresponding bucket. We use gray, green, and red to highlight the results similar to, better than, and worse than Baseline performance, respectively.**
|  | Evaluation stage (steps) | Easy | Medium | Hard |
| --- | --- | --- | --- | --- |
| SLM | Final (208K) | 0.87 | 0.00 | 0.00 |
| Baseline | Early (36K) | 0.47 | 0.12 | 0.12 |
| RKD | 0.56 | 0.11 | 0.12 |  |
| SALT | 0.56 | 0.11 | 0.12 |  |
| Baseline | Final (208K) | 0.70 | 0.29 | 0.28 |
| RKD | 0.65 | 0.17 | 0.17 |  |
| SALT | 0.78 | 0.38 | 0.36 |  |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 3
presentation: 4
questions: 1. How would SALT perform with larger/smaller teacher-student size ratios? For example, would the benefits persist when training a 70B parameter model with a 7B teacher? When the ratio decreases to 1, will the framework still work (eg., similar to bootstrapping).

2. Is there a minimum performance threshold the teacher must achieve for the method to be effective? Is it reasonable to set this as the baseline performance?
soundness: 2
strengths: The authors validate their approach by training a 2.8B parameter model using a 1.5B parameter teacher model on the Pile dataset. They demonstrate improvements in both few-shot learning across multiple downstream tasks. The theoretical analysis provides insights into the bias-variance trade-off in knowledge distillation and justifies why selective knowledge transfer from smaller models can be beneficial. And more pros:

- The paper presents a novel and counter-intuitive approach of using smaller models to improve larger model training (though the "large" is required further clarification).
- The theoretical framework for analyzing knowledge distillation in language modeling is novel.
- The proposed data selection mechanism using SLMs is innovative and well-motivated.
- Addresses a crucial effectiveness improvement of the computational resource utilizations  (training efficiency).

The theoretical framework is in demand at the fast developed field of data cleansing and knowledge distillation for LLMs (mostly driven by empirical results)
summary: This paper addresses the significant computational challenges in training large language models (LLMs) by proposing an approach that leverages smaller language models (SLMs) to improve both training efficiency and model quality. The authors introduce a two-stage training methodology where knowledge distillation from an SLM is used during the early phase of LLM training, followed by standard self-supervised training.

The work makes three main contributions: (1) A theoretical framework analyzing knowledge distillation in language modeling, providing novel risk bounds that explain how even a weaker teacher model can benefit a (stronger) student model; (2) The SALT methodology, which uses SLMs both for knowledge distillation and data selection; and (3) Empirical validation showing that SALT can reduce LLM training time while maintaining or improving performance across various benchmarks.
weaknesses: - The size gap of the 1.5B teacher and 2.8B student remains questionable to define them as a pair of relatively "small" and "large" LMs. If the author do not have sufficient resource to scale up the large model, they need to prove that a smaller 
- The paper doesn't thoroughly explore the impact of different teacher-student size ratios. (Also in Question 1)
- The paper primarily focuses on relatively small models (1.5B and 2.8B parameters) compared to state-of-the-art LLMs. It's unclear if the benefits would scale to much larger models (eg. 7B).

- Limited analysis of computational overhead from the knowledge distillation phase: the authors only provide the resource difference between SALT and the baseline, but not providing the cost of the data selection part.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: 1.For Table 1, from the experimental results, it appears that in the Early phase, SALT and SALT_{DS} did not improve performance compared to the Baseline, and even showed a slight decrease in some metrics. However, in the Final phase results, SALT and SALT_{DS} outperform the Baseline. Where do you think this benefit comes from? Additionally, although SALT and SALT_{DS} show some improvement, they do not significantly outperform the Baseline. Have you conducted multiple runs to confirm the stability of these results?
2.Do you plan to test this approach on a larger model to further validate its feasibility?
soundness: 3
strengths: 1. The paper leverages a small language model (SLM) as the teacher model for the knowledge distillation process. Instead of using all data for distillation, it selectively transfers knowledge by choosing challenging yet learnable data points, reducing the risk of transferring erroneous knowledge and enhancing the effectiveness and specificity of the knowledge transfer.
2. The 2.8B parameter LLM trained with SALT achieves better performance than a 2.8B LLM trained with standard pre-training on various popular few-shot benchmarks, while requiring only about 70% of the training steps and saving approximately 28% in wall-clock time. SALT models also consistently show substantial performance gains in multiple domains after SFT.
summary: This paper proposes an efficient two-stage pre-training method called SALT. In the early stage of LLM pre-training, SALT utilizes an SLM (Small Language Model) as the teacher model for knowledge distillation, while also using the SLM to select challenging yet learnable data to ensure effective knowledge transfer. In the second stage, the training switches to standard self-supervised learning, thereby significantly improving the quality and efficiency of LLM training without increasing pre-training costs. Experimental results demonstrate that the SALT method outperforms standard LLM training on several tasks, while also achieving notable training time savings.
weaknesses: 1. In Section 4, this method selects samples with high information content and those that are learnable by the SLM through Equation 11, focusing on top-m score sequences. This introduces some uncertainty into the model’s learning process, especially when using an early checkpoint of the SLM, which may further amplify this uncertainty. This can lead to the model learning unreliable information in the early stages, thereby causing a negative impact. As shown in Table 1, SALT_{DS} does not yield performance improvement over the baseline in the early stage, and even exhibits negative effects, indicating that excessive knowledge distillation on uncertain samples may introduce noise rather than effective knowledge, thus affecting model performance.

2. In Table 1, at the early stage, both SALT and SALT_{DS} show limited improvement over the baseline, and even experience noticeable drops in some metrics.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - I found it a bit confusing to distinguish between $V_t$ and $V_N$. Although they model similar aspects in both sections, it would probably be safer to make both more distinguishable.
- In Figure 2, do you display results for SALT or for SALT-DS? Why didn't you display both in this graph?
- Interestingly, the SALT models seem to not drastically outperform the vanilla baseline on in-domain language modeling evaluation (Figure 2, Table 1), but have a noticeable edge in few-shot evaluation and fine-tuning. It could be estimated that they are only 5-10% more data-efficient for pre-training, but are 30% more efficient for evaluation, and in other words that a KD model that reaches a given perplexity gives better evaluation results than a vanilla model at the same perplexity level. Do you have an explanation for this?
- In a way, the SALT approach could be purely justified from an intuitive standpoint : obtain a relatively strong language model quickly through reverse distillation up to the point where the student is slightly below teacher level, and then proceed with regular training. How would you justify precisely the usefulness of the theoretical part in the design of the experimental part?
- Similarly, the SALT-DS approach can also be framed as a (relatively) basic data-filtering technique. Do you think it would be relevant to invoke the curriculum learning literature in that case?

**Typos**
- There is a parenthesis issue with equation 50.
soundness: 2
strengths: The authors conduct a very novel and interesting study on knowledge distillation for language modeling in a relatively broad setting. They make reasonable assumptions, and obtain informative bounds on the performance that can be achieved through distillation.
They conduct experiments on relatively large models (2.8B), which confirms the relevance of their approach in a setup that approaches industrial standards.
This paper could motivate further work on distilling small models into larger ones, as it provides a theoretical justification for such approaches, and an example of empirical confirmation of this idea.
The paper is also well-written and clear to follow.
summary: This paper studies reversed knowledge distillation for language modeling, where a smaller language model provides target probabilities to a larger model.
First, the authors derive two excess surrogate risk bounds under different assumptions: one that bounds this risk based mostly on the distribution of the loss on the training data distribution, and the other based on the expected robustness of the teacher model when increasing past context. They then show that the variance on this robustness can be regularized by the $\omega$ weight of the distillation loss.

They leverage their theoretical analysis of KD to (pre)train large language models in two phases, first using a distillation loss, and then using a cross-entropy loss. Overall, they show that their KD approach outperforms both a vanilla baseline and a one-phase reverse KD baseline in terms of final training performance. They also propose a variant, $SALT_{DS}$, where the second training phase is using data labeled as hard by a heuristic classifier based on the small LM perplexity. They evaluate their models for few-shot performance on classical benchmarks, fine-tune them, and provide a short analysis of performance against sample difficulty along training.

Overall, their approach provides a promising method to accelerate language model training using smaller model distillation as a loss regularization factor.
weaknesses: The main weakness of this paper lies in the lack of articulation from the theoretical analysis to the experimental work. Several concerns can be raised on this subject:
- The theoretical part is purely based on causal language modeling, while the models are trained with the UL2 procedure, which includes prefix language modeling (which is a variation of causal language modeling), and most importantly the span corruption task, which is not causal. Although it should be possible to extend the theoretical work in that direction, the experiments seem not to accurately reflect the theoretical results.
- Conversely, the paper lacks empirical experiments that would help validate the introduced assumptions, and help quantify some of the factors of the bounds. For instance, empirical estimations based on existing language models could help refine the comments on the dependency of $C$ and $\{V_t\}$ on $T$ and the robustness of models. This could for instance motivate the choice of the size of the SLM during training.
- Experimentation at smaller scale could be done to provide intuition on the benefits obtained when $\omega$ varies, and help illustrate how the mentioned tradeoff can be found.
- More generally, although the authors define several values that should be affected by KD in the theoretical part (e.g. $\xi$), they do not explore the effect of KD on these metrics in the experimental part.

Additionally, I noted several unnatural design choices in the experimental part, which I happen to be more knowledgeable about:
- In Figure 2, which summarizes (and introduces) the benefits of the methods, the authors use top-1 accuracy as the main LM performance metric, while it is more common to show cross-entropy or perplexity in this case. They also report log-perplexity in Table 1, instead of perplexity, which seems like a more natural choice.
- Regarding the model architecture, the chosen vocabulary size seems particularly large (256000) compared to standards in the literature for (a priori) monolingual English datasets (e.g. GPTNeo: 50k; Llama2: 32k; Llama3, which is trained on multilingual data: 150k). The choice of the tokenizer in this case could be important, as it can deeply affect the nature of the modeled distribution $\mathcal{D}$, and thus of the tradeoff offered by $\omega$.
- Regarding the optimization, my opinion is that Adam (or AdamW) would have been a more neutral choice than Adafactor, although it has been used for several models as well. The authors also do not discuss the impact of learning rate when using the KD loss, which is understandable given the cost of running such a hyperparameter search at scale.

In my opinion, another weakness of the paper lies in its lack of transparency and precision on the computational overhead of the proposed method. The authors briefly discuss this overhead, arguing that "As a rule of thumb, a forward pass constitutes 1/4th cost of a training step". This claim lacks support in my opinion, as it depends on implementation and hardware among other things, and as previous works also mention 1/3rd as an estimate (cf. https://arxiv.org/pdf/2203.15556). Moreover, their study does not take the cost of pretraining the SLM into account for the overall training time, or the memory cost of having two models in memory simultaneously. Finally, given the 12% overhead mentioned by the authors and the plot from Figure 2, it is not clear that their approach performs better for pure in-domain language modeling performance when using a FLOPS-related metric on the x-axis. These points are not prohibitive for the method, but should be discussed in the paper.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: None
soundness: 4
strengths: 1. A significant contribution of this paper is its theoretical understanding of knowledge distillation. The bias-variance trade-off in generalization bounds is particularly interesting, which effectively elucidates why even small teacher models can be helpful. However, since I am not an expert in this field, I would need to refer to the opinions of other reviewers on this part.
2. The proposed method shows good experimental results, which outperforms standard training while saving 28% training time.
summary: This paper explores the use of small language model (SLM) to provide useful supervision in the early stage of large language model (LLM) training. Specifically, the training of LLM is divided into two stages. In the first stage, it employs KD from an SLM to provide soft labels as additional training supervision. In the second stage, it resorts to the standard MLE training. Though it looks counterintuitive of employing a smaller model to teach a larger one, this paper provides theoretical justifications for such a methodology. Empirical results on language modeling also validate the utility of this approach, which saves 28% training time compared to the standard training.
weaknesses: 1. The small teacher model employed in the experiments is not small enough (1.5B teacher v.s. 2.8B student ). The authors assume that the small language model is readily available, but it is not always the case. For example, before the release of Llama-405B, the largest Llama model only has 70B parameters. It is unknown whether such a much smaller teacher model can help the training of student model. 
2. This paper lacks discussion and comparison with recent progress of efficient LLM training [1,2,3].

[1] Yao et al., Masked Structural Growth for 2x Faster Language Model Pre-training, https://arxiv.org/pdf/2305.02869  
[2] Li et al., FLM-101B: An Open LLM and How to Train It with $100K Budget, https://arxiv.org/pdf/2309.03852  
[3] Shao et al., Patch-Level Training for Large Language Models, https://arxiv.org/pdf/2407.12665

Meta-review: This paper introduces a novel approach to improve large language model (LLM) pre-training efficiency and quality by leveraging a small language model (SLM) for soft label supervision and selective data sampling. The method reduces training time and enhances model quality, supported by a theoretical framework that explains how adaptive utilization of SLM guidance balances bias and variance, with empirical validation on a 2.8B parameter LLM using the Pile dataset.

The paper was reviewed by three reviewers. All of them agree this paper has issues and may not be ready for publication at this stage. I would encourage the authors to please go through the reviews, and address them in the next iteration of this paper.

Decision: Reject