Progressive distillation induces an implicit curriculum
Authors: Abhishek Panigrahi
Bingbin Liu
Sadhika Malladi
Andrej Risteski
Surbhi Goel
Venue: ICLR 2025 Oral (2025)
Keywords: knowledge distillation
feature learning
curriculum
sparse parity
PCFG
optimization
MLP
Transformer
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Progressive distillation accelerates the student model's training by providing an implicit curriculum through intermediate teacher checkpoints.

Abstract: Knowledge distillation leverages a teacher model to improve the training of a student model. A persistent challenge is that a better teacher does not always yield a better student, to which a common mitigation is to use additional supervision from several “intermediate” teachers. One empirically validated variant of this principle is progressive distillation, where the student learns from successive intermediate checkpoints of the teacher. Using sparse parity as a sandbox, we identify an implicit curriculum as one mechanism through which progressive distillation accelerates the student’s learning. This curriculum is available only through the intermediate checkpoints but not the final converged one, and imparts both empirical acceleration and a provable sample complexity benefit to the student. We then extend our investigation to Transformers trained on probabilistic context-free grammars (PCFGs) and real-world pre-training datasets (Wikipedia and Books). Through probing the teacher model, we identify an analogous implicit curriculum where the model progressively learns features that capture longer context. Our theoretical and empirical findings on sparse parity, complemented by empirical observations on more complex tasks, highlight the benefit of progressive distillation via implicit curriculum across setups.

Published as a conference paper at ICLR 2025
PROGRESSIVE DISTILLATION INDUCES AN IMPLICIT CURRICULUM
Abhishek Panigrahi* Bingbin Liu
Sadhika Malladit Andrej Risteski Surbhi Goel 8 # University of Pennsylvania
i Princeton University
Carnegie Mellon University {ap34, smalladi (ccs princeton.......m edu, (bingbini, aristeskk0cs emu edu, surbhigecis upenn. edu
ABSTRACT
Knowledge distillation leverages a teacher model to improve the training of a student model. A persistent challenge is that a better teacher does not always yield a better student, to which a common mitigation is to use additional supervision from several "intermediate" teachers. One empirically validated variant of this principle is progressive distillation, where the student learns from successive intermediate checkpoints of the teacher. Using sparse parity as sandbox, we identify an implicit curriculum as one mechanism through which progressive distillation accelerates the student's learning. This curriculum is available only through the intermediate checkpoints but not the final converged one, and imparts both empirical acceleration and a provable sample complexity benefit to the student. We then extend our investigation to Transformers trained on probabilistic context- free grammars (PCFGs) and real- real-world pre-training datasets (Wikipedia and Books). Through probing the teacher model, we identify analogous implicit curriculum where the model progressively learns features that capture longer context. Our theoretical and empirical findings on sparse parity, complemented by empirical observations on more complex tasks, highlight the benefit of progressive distillation via implicit curriculum across setups. Code is available here.
INTRODUCTION
As the cost of training state-of-the-ar models grows rapidly (Hoffmann et al., 2022), there is increased interest in using knowledge distillation (Hinton et al., 2015) to leverage existing capable models to train new models more efficiently and effectively. Knowledge distillation is effective technique to train smaller vision (Jia et al., 2021; Touvron et al., 2021; Yu et al., 2022; Lin et al., 2023) and language models (Sanh al., 2019; Gunasekar et al., 2023; Touvron al., 2023; Reid al., 2024) that permit faster inference with comparable performance. However, one curiously persistent phenomenon that better teacher does not always yield a stronger student. Prior works (Mirzadeh al., 2019; Jin et al., 2019; Jafari et al., 2021; Harutyunyan et al., 2022; Anil et al., 2018) hypothesized that this due to a capability gap between the teacher and the student. As such, they proposed progressive distillation, where the student is incrementally supervised by increasingly capable teachers. This technique has yielded strong empirical performance. One recent example is the training of Gemini- Flash from Gemini-1.5 Pro (Reid et al., 2024; Team et al., 2024): Gemini-1 1.5 Flash achieves 95% of Gemini- 1.5 Pro's performance on average and outperforms Gemini- Pro on 41 out 50 text-based long-context benchmarks, while being substantially smaller. However, little is understood about progressive distillation in terms of the optimization or generalization benefits, compared to directly learning from the data or the final teacher checkpoint (i.e., one-shot distillation).
Most prior work hypothesizes that progressive distillation enables better generalization (Mirzadeh al., 2019; Jafari et al., 2021; Harutyunyan et al., 2022). In contrast, we identify novel mechanism by which progressive distillation helps a student by accelerating its optimization (Figure 1). We define optimization acceleration as achieving improved performance with fewer training steps samples. In this paper, we use fresh training samples in each training step; hence we use training steps and samples interchangeably to measure the optimization speed.
"Equal Contribution

---

Published as a conference paper at ICLR 2025
MLP on sparse parity
Transformer on sparse parity
BERT PCFG
Top-1 Accuracy
1.0
1.0
1.0
0.8
0.8
0.9
one-shot
0.6
0.6
0.8
progressive Tear
Veav
5M 10M 15M
20M
2M
4M
6M
8M
5000 Training Steps
10000 15000
Training Steps
Training Steps
Figure 1: Progressive distillation accelerates training. Left: MLP on (100, 6)-sree parity (Definition 3. with width-50k teachers and width- students. Progressive distillation checkpoints are at 100k-step intervals, and one-shot checkpoint uses the final (20M-step) checkpoint. Middle: Transformer on (100, 6)-sparse parity, with 32-head teachers and 4-head students. Progressive distillation checkpoints are at 10k-step intervals, and the one-shot checkpoint is at 250k steps. Right: Transformers on PCFG (Section 4), with 32-head teachers and 8-head students using BERT-style masked prediction. Progressive distillation uses 8 intermediate checkpoints.
We study two tasks where learning the right features is believed to be important and show that the intermediate checkpoints provide signal towards these features. The first is learning sparse parity (Definition 3.1), which is a commonly studied setting understand the feature learning dynamics of neural networks. The second is learning probabilistic context-free grammars (PCFGs), which we use as a sandbox for capturing certain aspects of language modeling. Theory and extensive experiments in these settings support the following claims.
1. Progressive distillation accelerates student learning. Our experiments in multiple settings demonstrate that progressive distillation accelerates training compared to standard one-shot distillation and learning from the data directly (Figure 1). More specifically, for sparse parity, progressive distillation can train smaller MLP (or Transformer) at the same speed as larger MLP (or Transformer). For PCFGs, progressive distillation improves the accuracy of a smaller BERT model (Devlin et al., 2018) at masked prediction. Finally, we verify our findings on more realistic setups of training BERT on Wikipedia and Books dataset.
An implicit curriculum drives faster learning. We demonstrate theoretically and empirically that acceleration comes from implicit curriculum of easy-to casy-to-learn subtasks provided by intermediate teacher checkpoints, which is not available from the final teacher checkpoint. For sparse parity, the easy-to-learn subtasks provide supervision for the coordinates which constitute the support of the sparse parity (Section 3). As a consequence, we show progressive distillation provably improves the sample complexity for sparse parity over one-shot distillation learning directly from data (Theorem 3.2). For PCFGs, the implicit curriculum is defined in terms of learning features that increasingly capture larger n-gram contexts. Our results also provide guidance on how to select the intermediate teachers used during progressive distillation.
Related works'. One persistent surprise in knowledge distillation is that stronger teachers do not al- ways lead to stronger students. Prior works have speculated that overly large "teacher-studas gap" is the cause, and accordingly proposed to bridge this gap by introducing supervision of intermediate difficulty (Mirzadeh al., 2019; Cho & Hariharan, 2019; Harutyunyan al., 2022; Jafari al., 2021). Mirzadeh et al. (2019) used multi-step distillation involving models of intermediate sizes, and Shi et al. (2021) proposed to directly inject teacher supervision into the student's trajectory using an approxima- tion of mirror descent. Most related to our work. Harutyunyan et al. (2022) analyzed distillation for extremely wide networks and found it helpful to learn from the intermediate checkpoints of the teacher, a strategy also adopted by Jin et al. (2019). They speculated that this is because neural networks learn progressively complex functions during training (Kalimeris et al., 2019). contrast to their focus on the generalization ability of the student, we study the optimization dynamics of distillation.
is worth noting that there is also a rich body of work on understanding standard (one-shot) distillation, mostly regarding regularization effects. In particular, Menon et al. (2021) shows that learning from the teacher Icads to a tighter generalization bound when the teacher is closer to the Bayes distribution over the class labels. However, such Bayes perspective cannot explain the training
We defer detailed discussion of related work to Appendix A.
2

---

Published as a conference paper at ICLR 2025
acceleration in the feature learning tasks considered in this work, whose the Bayes distributions delta masses and hence are the same as the one-hot labels themselves. Our results fill this gap by providing an orthogonal view of implicit curriculum.
The benefit of curriculum on parity has also been explored in Abbe et al. (2024a:b), where the curriculum gradually increases the input Hamming weight. The difference though is that such curricula are defined by explicitly altering the distribution over the inputs, whereas our curriculum shows up implicitly in the teacher supervision. Moreover, our implicit curriculum emphasizes that properly chosen intermediate checkpoint, while having worse accuracy than the final checkpoint, can lead better-perfomie performing student. This can be seen plausible mechanism for weak-to weak-to-strony to-sonn generalization (Burns et al., 2023).
Outline. Section 2 describes the distillation strategies. Section 3 introduces the implicit curriculum with case study on sparse parity, presenting both empirical evidence and a provable benefit in sample complexity. Section 4 continues the empirical investigations on PCFG, and extends the observations to BERT's training on Wikipedia and Books dataset. Finally, Section 5 discusses open directions.
2 PRELIMINARIES
We now outline the distillation strategies considered in this paper and their empirical instantiation. For ease of exposition, we discuss one-dimensional label classification tasks here and generalize to requence-to-sequence functions in Section 4. Denote the teacher and student models operating on input domain X' as /T= X' - RCC and [s - TRC respectively. The outputs of a model are logits that are transformed into probability distribution over C' classes using softmax function with temperature т, denoted as p(z; T) := softmax f (x)/T). We will use PT. Ps to denote the probability distributions of the teacher and the student, and will omit the subscript denote generic model. When = 1, we omit 7 from the notation for brevity. Following Zheng & Yang (2024), we set r = for the student and vary the temperature of the teacher.
We compare two loss functions: f, where the student fs learns only from ground-trun truth labels, and (DI. where the student fs is supervised only with the logits of some teacher fr.
E(z, yi fs) = (L(e))lss())), (1: fs.
(L(PT(x;; (p)ss)
where e, is a one-hot vector whose yd entry is 1. We consider two strategies for choosing the teacher. The first is one-shot distillation, where the student learns from a fixed JT throughout the training, and the teacher is chosen as the final converged checkpoint. The second is progressive distillation, where the student learns from multiple intermediate checkpoints of the teacher's training run: Definition 2.1 ((CT, D) progresssiv distillation). Given a set of teacher checkpoints CT = {ffr, and set of training durations D. the student is trained with the logits of teacher checkpoint fr, for training length D, with E [|C+|] (1,..... "II.I.,
To simplify the presentation, the main paper tests a specific type of progressive distillation schemes, where C7 contains N equally-spaced checkpoints and the student is trained each one for T steps:
Definition 2.2 ((N, proogresssv distillation). CT contains N - 1 equally-spaced intermediate teacher checkpoints and the final teacher checkpoint. The student trained with each checkpoint for T training steps. After NT steps, the student is trained with the final teacher checkpoint.
To study the effect of each teacher checkpoint, we will also consider an extreme version of progressive distillation with N 2, where the student uses one intermediate teacher checkpoint.
Choice of temperature. We set r = 10- for sparse parity and PCFG experiments (Section 4) where the vocabulary size is smaller than 5, and (Section 4.2) whose vocabulary size is 30k. Using such small temperature makes the teacher's s outputs close to one-hot labels. This removes potential regularization effects due to the softness of the
10 for natural language experiments
We note that prior papers generally use a combination these objectives, hut we use supervision from one source in order to isolate its effects on distillation.
FFiuree 14 provides comparison temperature choices for sparse parity learning.
3

---

Published as a conference paper at ICLR 2025
labels (Yuan et al., 2020) which would otherwise be confounding factor. Moreover, the supervision with nearly one-hot labels is more representative of the setting where the student learns directly from the teacher's generations instead of the logits. This method, often described as generating synthetic data in the language modeling setting, has generally yielded small yet highly performant students (Gunasekar al., 2023; Liu al., 2024). For one-shot distillation, we report the best- rest-performing temperature among T = 1, 10-4 in the main paper and defer other results to Appendix D.8.
3 THE IMPLICIT CURRICULUM: A CASE STUDY WITH SPARSE PARITY
To elucidate the mechanism by which distillation accelerates training, we first focus on the well- studied task of learning sparse parity." Sparse parity is commonly used sandbox for understanding neural network optimization in the presence of feature learning (Barak et al., 2022: Bhattamishra et al., 2022; Morwani et al., 2023; Edelman et al., 2023; Abbe et al., 2024b).
Definition 3.1 ((d, k--spe parity task). Let S C [d] denote a fixed set of coordinates, with |S| k and k < d. Then, the sparse parity task is defined for any input x E {=1] whose label is computed as y = 1 if ILes > 0 and 2 otherwise.
We train the teacher and student models using 2-label classification, where fr and fs return logits in 2R2 The teacher and the student have the same number of layers but different sizes. We vary the model width for MLP, and vary the number of attention heads for Transformer, with a fixed per-head dimension. These choices not only affect the parameter counts, but also govern the learning speed'.
Why can larger models learn faster? A natural way to learn sparse parity with gradient descent involves first identifying the support S and subsequently computing the product of variables in the support (i.e.. Ihes x;). Empirically, the two stages of learning manifest as a long plateau period in the model's 's accuracy, followed by sharp phase transition (Figure 1, left and middle). The search for the support is what makes learning problem difficult, as it depends on the input dimension d rather than the support size (Abbe et al., 2023; Barak al., 2022). The benefit of increasing the width or the number of heads comes from providing more "parallel search queries." For MLP, prior work has shown that increasing the width accelerates training (Edelman et al., 2023), which we also observe in Figure 7 (left) in appendix. For Transformers though, we find that increasing the number of attention heads the most effective for improving the convergence speed, as opposed increasing the per-head dimension or the MLP width. A detailed comparison is provided in Appendix C.2 (Figure 10). Given this finding, we will vary the number of attention heads between the teacher and the student, while keeping the per-head dimension fixed. The number of heads hence directly controls the parameter count. This choice also aligns with the practice in open-sourced models such as the Llama series (Touvron et al., 2023).
In the following, we first empirically verify that carefully chosen intermediate teacher checkpoints constitute implicit curriculum for the student to learn from. Then, we show that this curriculum provably improves the speed of learning in the student by improving its training sample efficiency.
3.1 ACCELERATING LEARNING WITH THE IMPLICIT DEGREE CURRICULUM
The difficulty of the search problem suggests that we can accelerate student learning by providing direct supervision for what the support is (Abbe et al., 2023). We show that supplying the intermediate signal from bigger teacher model accelerates the search process for the smaller model, as described by the following set of results." 6
(R1) Intermediate teacher checkpoints constitute an implicit degree curriculum. We provide empirical evidence that the supervision from intermediate teacher checkpoints serves as an implicit curriculum supplying strong signals for certain 1 monomials, which require fewer samples to learn. In Figure 2, we report the correlation between degree-1 monomials and the prediction of the teacher logits at various checkpoints. The correlation for each monomial zj.j E [d] is computed as [Eze,, (IPT (x)]1 2,) at each checkpoint fr. Here [pT(x)) (x)]1 refers to the first output dimension of
"We also experiment with hierarchical generalization of sparse parity, which is deferred to Appendix C.3. 'In terms of the number of samples or the numher of training steps, which coincide in our experiments as we use freshly sampled batches.
'We will mark our results with (Ri) throughout the paper for easy reference.
4

---

Published as a conference paper at ICLR 2025
Teacher's accuracy behavior
Absolute Correlation
Corr. to degree-1 monomials
(2, proggesseee
1.0
0.10
1.0
Accuracy
Accuracy
0.8
1-6
0.8
candidate
0.05
candidate
Rest
candidate
0.6
0.6
0.00
0.95
1.10 Training Steps 1e7
1.00
1.05
0.95
1.00
1.05
1.10
Training Steps 1e7
Training Steps
1e6
Figure 2: Implicit curriculum for (100, 6)-spa parity. We compare 3 candidate intermediate checkpoints, labeled as (. 2 @ corresponding to 9.7M, 10.2M, and 10.8M steps, or the beginning. middle, and end of the teacher's phase transition. Left: Teacher's accuracy throughout training. Middle: During the phase transition, FT is much more strongly correlated with in-support variables Hg ,g in this case) than with off-support variables. Right: Only candidate ② (i.e., during phase transition) enables (2, - distillation to reach 100% accuracy. We use width-50k teachers and width- 100 students; Figure 8 shows similar results for width- 1000 students.
(Z
fT, which corresponds to ply = 1) p x; > 0) = - --y(y= = 2) (recall Definition 3.1). We take the absolute value as we are only concerned with the magnitude of the correlation. Importantly, these strong correlations emerge when the teacher learns the sparse parity task (i.e., during the phase transition) but diminish with continued training.
Note that the monomials need not be strictly degree-1. While our theory (Section 3.2) will only focus on degree-1 monomials for the sake of mathematical analysis, low-degree polynomials can still provide acceleration, which we also observe in practice (see Figure 9 in the Appendix for such an example). This transient low-degree supervision, available only through intermediate teacher checkpoints, may explain the superior performance of progressive distillation over one-shot distillation (Figure 1). We will confirm the provable sample complexity benefit of this implicit low degree curriculum in Section 3.2. The importance of the implicit curriculum is further strengthened by the superior performance of (2, -.oo))rrrreesvvve distillation:
(R2) Progressive distillation with a single intermediate checkpoint can outperform one-shot distillation. We consider the extreme version of progressive distillation where only single intermediate checkpoint is used (in addition to the final checkpoint). Figure 2 shows the result for (2, progrressee distillation. We consider candidates for the intermediate teacher checkpoint, occurring respectively at the beginning. middle or the end of the teacher's phase transition. Our result demonstrates that the checkpoint selection is crucial, where only the checkpoint during the phase transition is useful in accelerating training. This provides further evidence that the implicit degree curriculum the key to faster training via progressive distillation.
More complex tasks may require more intermediate checkpoints, which we discuss in more depth in Appendix C.3. Nevertheless, we find that progressive distillation can be run efficiently and effectively across tasks, and a small number of intermediate teacher checkpoints often suffice to accelerate training provided that the checkpoints are properly selected.
3.2 THE LOW-DEGREE CURRICULUM REDUCES SAMPLE COMPLEXITY
We now formalize the benefits of progressive distillation for (d, k)-sparse parity in terms of sample complexity. For the sake of mathematical analysis, we take the student fs and the teacher fr models to be layeer MLPs with ReLU activations and scalar outputs. Further, the labels y are given as +1, where (or -1) corresponds to the class dimension 1 (or 2) in Definition 1. Following previous works (Barak et al., 2022; Abbe al., 2023; Edelman al., 2023), we analyze simplified two-stage training procedure and train the model using the hinge loss: La yi fs. fr) max 1 fs (x)y) + (1 a)) max(0, fx(x)
Let's first recall the hardness of learning sparse parity. 8 For simplicity, we consider the case of MLPs of width Ō(2*) trained using online SGD. When learning from data alone, statistical query
We show similar results for width-1000 students (Figure 8) and transformers (Figure 13).
AA more detailed discussion is provided in Appendix A.1.
5

---

Published as a conference paper at ICLR 2025
(SQ, Kearns (1998)) lower bound shows that learning the support for (d, - parity requires O(d samples (Abbe et al., 2023; Edelman et al., 2023). We will show that although this lower bound also applies to one-shot distillation from a strong teacher, it can be cireumvented when learning from the implicit low-degree curriculum identified in the previous section.
Specifically, we compare the sample complexity of one-shot distillation and (2, prrogresssve distillation (Section 2). Both strategies use a well-trained final checkpoint with an error of O(c) error for an arbitrarily small E O. Progressive distillation additionally uses the teacher's intermediate checkpoint after its first phase of training, where we can provably show its predictions to have correlations at least (2(1/)) to the monomials F,, Vi E S. That is, progressive distillation first learns from the intermediate checkpoint and then switches to the final checkpoint, whereas one-shot distillation learns directly from the final checkpoint.
(R3) Progressive distillation reduces sample complexity. We formally demonstrate the sample complexity benefit of progressive distillation.
Theorem 3.2 (Informal version of Theorem B.1). Consider learning (d, (d,k))sparse parity with student model of size m = @(2²), where hides polylog factors in d, k. Suppose the teacher has loss (e) for some small > 0. Then, the total sample complexity needed for the student to reach c-loss using progressive distillation with 2 checkpoints is @(2=d& distillation requires at least Q(dd=11,e (,08)),,,)))) samples.
c 22 + +.²). However, one-shot
Proof sketch. We track the training behavior of the teacher model during its two-phase training. We show that at the end of the first phase, the teacher's predictions will have (2(/kk) correlations to degree- 1 monomials 2;, Vi E S. In contrast, the correlations are smaller for degree-1 monomials z;, Vi & S. Hence, the teacher's predictions can be written as Eies C + Eies Cg7i plus additional higher degree odd polynomials which can be controlled, with |c|| 2 Q(1/k) for e S, and |c = o(1/kd) if i & S. When training on the predictions from this intermediate teacher checkpoints, the correlation gap between in- and off-support degree-1 monomials will be reflected in the gradients of the student's weights. Namely, there is (2(1kk) gap between the support and non-support coordinates in the weight gradients. This gap allows the coordinates E S in the student's weights to grow quickly with only Olk lo)(m)) samples.
On the other hand, for a teacher that has loss Oc)). a similar argument can show that the separation gap between the correlations of the teacher's 's predictions to degree-1 monomials on support and outside support can be at most (c). So, harnessing this gap will require a sample size of at least ((c=2) by concentration inequalities. Learning directly from the labels will require 1(d samples from the SQ lower bound as discussed above. This gives the sample complexity differences between one-shot and progressive distillation. The full proof is provided in Appendix B.
Remark. One gap between our theory and experiments is that our analysis applies to large-batch SGD with small gradient noise, whereas the experiments use online SGD with batch size 1. Bridging this gap, such by adapting the analyses in Abbe et al. (2023) on Gaussian data, is an interesting future direction.
4 IMPLICIT CURRICULUM WITH PCFGS AND NATURAL LANGUAGE
In this section, we empirically show that an implicit curriculum emerges generally, both when learning on probabilistic context-free grammars (PCFGs) and when performing natural language modeling tasks on the Wikipedia and Books datasets. We focus on BERT models (Devlin et al., 2018)', and discuss experiments GPT-2 (Radford et al., 2019) in Appendix E.
The masked prediction task. Our experiments will be based on BERT models trained to perform masked prediction, which requires filling in masked-out tokens in input sequence and excels at feature learning in natural languages (Hewitt & Manning, 2019; Tenney et al., 2019; Li et al., 2022).
Definition 4.1 (Masked prediction task with mask rate p). Let v denote the vocabulary that contains special token [mask], and let h denote an arbitrary sequence length. Given a sequence r",
See Appendix D.5.1 for primer on BERT.

---

Published as a conference paper at ICLR 2025
Figure 3: An example of PCFG tree T(x) that gen- erates ="The cat ran away". "The cat" is an ex- ample of level-2 span, and "cat" is as boundary to- ken for the spans of both the level-1 non-terminal Noun and the level-2 non-terminal Noun Phrase.
Start (S)
Level
Noun Phrase
Verb Phrasc
Level
Determinant
Noun
Verb
Adverb
Level
The
away
sample set of masked positions M E [h] following P(i -M Vi [h]. Create masked input X,M from x by replacing tokens at positions in M with [mask], a random token from X, or kept unchanged with probabilities 80%, 10%, 10% respectively. Then, the masked prediction objective is the cross-entropy of the model's predictions at positions E M on input X.M-
modeling, we need to generalize the definition of the teacher fT and student fs from Section 2 accordingly, denoted as fr e + RaxC and (t)
Since we are performing pruuence-to-sqqeeee
denote the teacher's output fs v² + 1x We will use PT (x: T) softmax([f((x)))//))
distribution on the ith position; similarly for p As before, we omit when = We use the following loss functions for the masked prediction task (Definition 4.1):
e(x: fs) = EM IM
Σ
KL(e,,,lpg (x.M)).
(3)
TEM
(x fT) EM [M
Σ
KL(P)) (x.MiT) (x,M),
(4)
REM
where ey is one-hot vector whose You entry is 1.
We train BERT models with f, fDL and report the average top-1 accuracy on the masked tokens. As discussed in Section 3.1, the teacher and student have the same depth (4 layers) but differ in the number attention heads, with 32 heads for the teacher and 8 heads for the student. Each attention head has dimension 8, so the teacher has width 256 and the student has width 64. All hyperparameter details are Appendix D.6.
4.1 7-GRAM CURRICULUM PCFGS
We first consider probabilistic context free grammars (PCFGs), which commonly used emulate the structure of natural language and thus provide mechanistic insights into language models (Zhao al., 2023; Allen-Zhu & Li, 2023a). A PCFG generates sentences following tree structure; Figure shows an example for the sentence "The cal ran away." More precisely, a PCFG G = (N, R. P, v) is defined by of non-terminals N. rules R
the non-terminals, probability distribution P over R, and a vocabulary (terminals) A sentence X is associated with generation tree T(x), whose intermediate nodes are non-terminals in N,, leaf nodes are terminals in v, and edges are defined by rules sampled from R according to P. A formal definition of PCFG is provided in Appendix D.I. Our choices of PCFGs are taken from Allen-Zhi have the same distance to the root. Experiments in the main paper are based on the PCFG cfg generated by depth-7 trees, and results on other PCFGs are deferred to Appendix D.4.
Li (2023a), where all leaves in the same tree
4.1.1 PROGRESS MEASURES OF IMPLICIT CURRICULUM
Unlike experiments parity, what constitutes feature is less straightforward for PCFG. We will use three progress measures to quantify the implicit curriculum for masked language modeling on PCFGs, based on r-gram statistics and non-terminal prediction.
Measures that use n-gram statistics will measure the dependence of the model's predictions on tokens in the neighboring contexts, defined as follows:
Definition 4.2 (n-gram neighboring context). For a h-length sentence æ E a and for E [h], define the n-gram neighhoring context around the 5th token the set of tokens positions
7

---

Published as a conference paper at ICLR 2025
Teacher's loss behavior
Median (over (x. i))
Mrobust for n-grams
Final Top-1 Accuracy
(2. T-progressive
Validation loss
1.0
①
10-
0.90
0.5
10-3
C
0.85
one-shop
G1
0.0
0.6
4
8
0.4
Training Steps (x10²)
Training Steps (x10²)
Teacher Checkpoint(x109)
Figure 4: BERT on the PCFG cfg3b. Left: A 32-head teacher's loss exhibits three distinct phases: L an initial phase with little change, ② a middle phase with rapid drop, and 3 a final plateauing phase until the end of training. The triangles mark the selected checkpoints for progressive distillation, with the first teacher checkpoint (denoted by C1) located at the middle of phase ②. Middle: Mrnhutt across training, which peaks at C₁. The model gets more robust to shorter n-gram perturbation as training progresses. The median is taken over the input sequences. Right: 8-head student's final accuracy with (2, pprogressive distillation after 4000 total training steps. The z-axis marks the choice of the first teacher checkpoint. is grid- grid-searched over {500, 1000, 2000). The best performance is obtained by choosing Cj. Although results in the plots are for single training run of the teacher, similar behaviors occur robustly across random seeds.
within (n - 1)/2 distance from i, denote as n--raa(()) = (i max(i 1)/21,0 SjS min(i + - 1)/
In the example of Figure 3, for the word "cat", its 3-gram neighboring context consists of words "The" and "ran", and its 5- 5-gram neighboring context additionally includes the word "away." The choice of n-grams is inspired by results Zhao al. (2023), which show that BERT model can solve masked prediction by implementing dynamic programming algorithm that builds hierarchically on increasingly larger n-gram neighboring context spans (Definition 4.2). model that primarily uses short n-gram neighboring context will be largely affected if the tokens within the context are perturbed during evaluation. This motivates us to consider two n-gram based measures.
Measure 1: Robustness to removing 7-gram context. Our first progress measure of feature learning checks how the model's prediction changes when the 7-gram context present or absent. For each masked position i, measure the total variation (TV) distance between the probability distributions when masking out only the current token, and when masking out all the tokens in n-gram(() (i). i.c. the neighhoring n-gram context centered é. Recall that X.M denotes masked version of x with masked set M (Definition 4. I), and that () denotes model's output probability distribution at the les position. Then, our first measure is defined as
Mrobus (f, x, i, n) = TV (x4g)) (x), (nn-gra(()))) ))-
(5)
We report median of Meobuu (f, x, i, n) over randomly sampled x and 10, A larger Mosus (f,, x, i, indicates the model heavily depends on neighboring n-gram context tokens for the masked prediction.
Measure 2: Closeness between full and n-gram predictions. Our second progress measure examines the change in predictions when the model is given the full sequence versus only local n-gram window:
Melose x, i, TV ( (x))))))))))))))),
(6)
where Xx)gnam(X) (+) denotes the n-gram context centered position i, minus the position i itself. We report median of Metose x, i, nt) over randomly sampled x and i. A large Mekose (f,x, i, indicates that the model utilizes contexts outside a n-gram window in its predictions.
Measure 3: Non-terminal prediction. Finally, also measure how well the model outputs encode
the features of the underlying PCFG by checking the accuracy at predicting non-terminals (Allen-Zht & Li, 2023b). The predictions are given by linear classifier on top of the output embeddings.
Definition 4.3 (PCFG non-terminal prediction task). Define the span of a non-terminal n as the set of terminals within the subtree rooted n. denoted by span(n). The (right) boundary of span(n) refers
Our observations stay the same for other percentiles.
8

---

Published as a conference paper at ICLR 2025
to the rightmost position within span(n). We say a non-terminal is of level if it is at distance i from the root. Then, level-i non-terminal prediction task aims to predict n at boundary of span(n(*).
As an example, in Figure 3, the level-2 non-terminal prediction task aims to predict the non-terminals Noun Phrase and Verb Phrase at words "cat" and "away" respectively. More details are provided in Appendix D.3.
4.1.2 EMPIRICAL VERIFICATION OF THE n-GRAM CURRICULUM
Similar to Section 3.1, we will start with examining the training dynamics of the teacher model. We observe a phase transition period akin to that of sparse parity, during which we identify an inflection point concerning Mobus and Melose This inflection point proves to be a crucial intermediate checkpoint. We then demonstrate that progressive distillation improves feature learning in the student model, substantiated by the three measures defined in Section 4.1.1.
For training dynamics, we observe 3 distinct phases of training in the teacher's loss (Figure 4 left): 1) an initial phase where the loss doesn't change much for the first 5% of training; 2) a rapid loss drop phase in the next RJ 20% of training: and 3) final phase of slow loss drop till end of training. In particular, the rapid loss drop phase reminiscent of the phase transition in sparse parity (Section 3). Moreover, we identify an inflection point (marked by C1) during the second phase: before the inflection point, the robust loss Moobus increases (Figure 4 middle), and the loss Melose stays high (Figure 22 left); after the inflection point, both Mrobus and Melose start drop rapidly, suggesting that the model learns to utilize longer contexts as opposed to short neighboring n-grams.
(R4) The inflection point is best for (2, prrogressivve distillation. We study the importance of each teacher checkpoint by comparing the performance of (2, (2.T") p-rogresssiv distillation, where the student learns from a single intermediate checkpoint in addition to the final checkpoint. The value of T is grid-searched (more details in Appendix D.6). For the choice of the intermediate checkpoint, Figure 4 shows that the best intermediate checkpoint is the one at the inflection point (at 1000 training steps), which we denote as Cj. Note that the inflection point, the teacher has the highest reliance on shorter n- n-grams (c.g. for = 3), which are analogous to the low-degree monomials in Section 3 and serves as intermediate tasks that are likely easier learn. Hence, C1 being the optimal checkpoint choice further strengthens our hypothesis that an implicit curriculum is the key to the acceleration enabled by progressive distillation.
Following (R4), we will choose the checkpoints for progressive distillation training steps that are multiples of that of C₁, i.e. at steps {i 10 As shown in Figure (right), progressive distillation helps the student learn faster than both one-shot distillation and cross entropy training. Furthermore, progressive distillation leads to improved feature learning.
(R5) Progressive distillation improves feature learning on PCFG. Progressive distillation improves over one-shot or no distillation over all 3 measures mentioned in Section 4.1.1. 1. As shown in Figure 5, progressive distillation makes the student better utilize long contexts rather than local n-gram windows, evidenced by lower Mrobust and Melose The student can also better predict the non-terminals, suggesting better structural learning of the underlying PCFG.
Median (overtx, m))
Mzose for n-grams
M rotust for n-grams
Non-terminal prediction
100
10-1
10
Accuracy
10-3
10 3
80
one shot
CE
progressive
3
5
1
n-gram context
n-gram context
Non-terminal level
Figure 5: Comparisons on a 8-attention head BERT model. (Left) Melose for different n n-grams. Progressive distillation has lower Melsse with longer context. (Middle) M for different n-grams. Progressive distillation has lower Mromust for all n-gram contexts. (Right) Probe performance predict the non-terminals (NTs) (Definition 4.3). Progressive distilled student performs better when probed for higher level non- ann-terminals in its contextual embeddings.

---

Published as a conference paper at ICLR 2025
4.2 BEYOND SYNTHETIC SETUPS: IMPLICIT CURRICULUM IN NATURAL LANGUAGES
We conduct experiments on BERT training (Devlin et al., 2018) on Wikipedia and Books (details in Appendix F). The teacher and student both have 12 layers, with 12 and 4 attention heads per-layer respectively. Each attention head is of dimension 64, corresponding to a width-768 teacher and a width-256 student. Similar PCFG, the teacher's loss exhibits distinct phases (Figure 6 left), with an inflection point marking the change in Mosust (Figure 6 middle). The inflection point can hence provide an implicit currieulum towards easier-to-learn local n-grams. Finally, progressive distillation helps the student achieve better accuracy at masked language prediction (Figure 6 right). In Table 3 (appendix), we further show that progressive distilled models achieve an average 0.9 performance improvement on downstream tasks after fine-tuning,
Connections to related works. Our results align with those of Chen et al. (2023), who observed a phase transition in loss when training BERT on real-world language data corresponding to the model learning syntax rules of language. Comparable findings were also reported in a concurrent work on matrix completion (Gopalani et al., 2024). For auto-regressive models, prior work has discussed the emergence of n-gram induction heads which indicate phases in which the model learns to perform in-context learning (Akyürek al., 2024; Quirke et al., 2023; Olsson al., 2022). We observe similar behavior on PCFGs and Wikipedia datasets, quantify the phase change using n-gram context dependence, and further leverage these transitions to accelerate training of smaller student model.
BERT's loss behavior
Median (over (x. .)
Mrobuct for n-grams
Student BERT
Validation loss
100
Top-1 Accuracy
0.5
10
CE progressive
10-2
0.4
104
10"
16 32 64
24000
60000
120000
Training Steps (log scale)
Training Steps (x10²) (log-scale)
Training Steps
Figure 6: BERT on Wikipedia and Books. Left to right: (a) Similar to our experiments on PCFG (Figure 4), we observe three distinct phases in the loss behavior of 12-head teacher. The rapid loss drop phase signifies a transition phase for the model. The triangles mark the selected checkpoints for progressive distillation, with the first teacher checkpoint roughly picked in the middle of the second phase (C1). (b) We observe Mrobust peaks C1. and the model gets more robust to shorter n-gram context masking, as training progresses. (c) A 4-head student achieves better top-1 accuracy on masked prediction objective with progressive distillation.
5 DISCUSSIONS
We have shown that progressive distillation can improve the student's 's feature learning via an implicit curriculum provided by the intermediate checkpoints. We discuss limitations and potential future directions below, and provide preliminary results for some of them in the appendix (see Appendix A).
Impact of temperature. The teacher temperature is important hyperparameter in knowledge distillation, where varying T can sometimes lead to a greater performance gain than changing the distillation method (Touvron al., 2021; Harutyunyan et al., 2022). Our results are consistent with these prior findings. However, our experiments use limited temperature choices, i.c. the default (r 1.0) and low temperature (7 10 or 10 10-20). A more precise understanding of temperature, especially its impact on optimization, is an interesting direction for future work.
Distillation via generations. Another related distillation setting is training smaller (language) models using the generations of larger models, which has been shown to greatly improve various abilities (Liu et al., 2024; Yue et al., 2023; Yu et al., 2023; Luo et al., 2023; Chaudhary, 2023; Taori et al., 2023; Zheng et al., 2023). There are two differences between our experiments and these generation- approaches. First, the supervision in our experiments are distributions (over classes or the vocabulary). while generations are samples from distributions. Our experiments with low or zero temperature provide positive evidence towards bridging this gap, but the precise effect remains to be explored. More importantly, given input, there is unique supervision in our settings, whereas there could be multiple generations given by multiple steps of unrolling of the teacher. Extending our framework to these generative setting will be an important direction for future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
AP and SM acknowledge funding from NSF, ONR, Simons Foundation, and DARPA. AR and BL are supported in part by NSF awards IIS-221 1907, CCF-2238523, IIS-2403275, an Amazon Research Award, a Google Research Scholar Award, and an OpenAl Superalignment Fast Grant. BL also acknowledges the support of the Tang Family Presidential Fellowship. SG was supported in part by Microsoft Accelerating Foundation Models Grant and an OpenAl Superalignment Fast Grant.
REFERENCES
E. Abbe, Enric Boix-Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: eeap complexity and saddle-to-saddle dynamics. Annual Conference Computational Learning Theory, 2023. doi: 10.483500arXiv.2222..000
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. arXiv preprint arXiv: 2202.08658, 2022.
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum. Journal of Machine Learning Research, 25(331):1-58, 2024a.
Emmanuel Abbe, Elisabetta Cornacchia, and Aryo Lotfi. Provable advantage of curriculum learning on parity targets with mixed inputs. Advances in Neural Information Processing Systems, 36, 2024b.
Ekin Akyürek, Bailin Wang. Yoon Kim, and Jacob Andreas. In-context language learning: Architec- tures and algorithms. arXiv preprint arXiv:2401.12973, 2024.
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self- distillation in deep learning. The Eleventh International Conference on Learning Representations, 2023a. URL https: openreview net/ Eorum? (eprummii--ine www.:::dd--uuu00000e FXGA.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXi::2305. /3673, 2023b.
Rohan Amil, Gabriel Pereyra, Alexandre Passos, Róbert Ormándi, George E. Dahl, and Geoffrey Hinton. Large scale distributed neural network training through online distillation. International Conference on Learning Representations, 2018.
Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. In NeuriPS, 2022. URL http:// /tapt::://ppprs..... papprspppeses// nips cccpppperfilleeeeeceeeeeeeeeeeee http:/ppprrrflleeppprprrrr////// http//ppppprr22//2 //ash// 884ba 653 9211 707 637276277777 408 7bde
html.
S. Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. Annual Meeting of the Association for Computational Linguistics, 2022. doi: 2.8550//rXiv.2211..223
Blake Bordelon, Alexander B. Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024a. https: httttt::///pprreiieet wwwpnrvvinneeeeeeeeeeeeeee net forum?id- Rc.
Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. arXiv preprint arXiv: 2405.15712 2024b.
Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschenbrenner, Yining Chen, Adrien Ecoffet, Manas Joglekar, Jan Leike, Ilya Sutskever, and Jeff Wu. Weak-to- strong generalization: Eliciting strong capabilities with weak supervision. arXiv preprint arXiv: 23 2023.
Sahil Chaudhary. Code alpaca: An instruction -oollwiion-folooiingg llama model for code generation. GitHub repository, 2023.
11

---

Published as a conference paper at ICLR 2025
Angelica Chen, Ravid Schwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in mlms. arXiv preprint 1xXiv:2309.07311. 2023.
Lénaïc Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal (eds.), Proceedings Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp.
PMLR, 09-12 Jul 2020. URL https::///// http:s::/ppcceeeeeeeeee mlr. press/v125/
chi z html.
Lénaïc Chizat and Francis R. Bach. On the global convergence of gradient descent for vee-paaraettriized models using optimal transport. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Ad- vances in Neural Information Processing Systems 31: Annual Conference Neural Informa- fion Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 3040-3050, 2018. URL https::/proceedinggg weweiiinnnnneeuissss.eemm nwwrips cc/paper/2018 //aas// alafc58c6ca$5 oow.eecaa55000077299111101616111 raw.:tdd77--b-baaattttaaaaa
Abstract .tmll.
Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. arXiv preprint arXiv: 1910.01348, 2019.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason Lec. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2024a.
Alex Damian, Loucas Pillaud- -iviie,, Jason D. Lee, and Joan Bruna. computational-statiiiinn gaps in gaussian single-index models. Annual Conference Computational Learning Theory. 2024b. doi: 10. 1144550/arXi.22403..00
Tri Dao. Govinda M Kamath, Vasilis Syrgkanis, and Lester Mackey. Knowledge distillation as semiparametric inference, 2021.
Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint (xXiv:881000880) 2018.
Yihe Dong, Jean- Jean-Bapiin Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: pure attention loses rank doubly exponentially with depth. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 2793-2803. PMLR, 2021. URL http::ttp//. / proceedicceeeeeeeeeeeeeeppe www. www.pprssseee..ee....ee / vv19 / dong2 html.
Simon Du and Wei Hu. Width provably matters in optimization for deep linear neural networks. In International Conference Machine Learning, pp. 1655-1664, PMLR, 2019.
Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang, Pareto frontiers in neural feature learning: Data, compute, width, and luck. arXiv preprint arXiv:2309.0330 2023,
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks high dimension. arXiv preprint arXiv: 1904. 12191, 2019.
Margalit Glasgow. SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: case study in the XOR problem. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: http//pppeevviiieeee net forum? d=HHoJ]1 6.
Pulkit Gopalani, Ekdeep Singh Lubana, and Wei Hu. How do transformers fill in the blanks? case study on matrix completion. In ICML 2024 Workshop on Mechanistic Interpretability, 2024. URL
https: /tttps://ppenrewe
net www.nnee//oornne forum wwt.mmmddidnnnnnnneemmmmm
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giomo, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, Adil Salim, Shital Shah. Harkirat Singh Behl, Xin Wang, Sébastien Bubeck, Ronen Eldan, Adam Tauman Kalai, Yin Tat Lee, and Yuanzhi Li. Textbooks are all you need. arXiv preprint arXiv: 2306. 1644, 2023.
12

---

Published as a conference paper at ICLR 2025
Hrayr Harutyunyan, Ankit Singh Rawat, Aditya Krishna Menon, Seungyeon Kim, and Sanjiv Kumar. Supervision complexity and its role in knowledge distillation. In The Eleventh International Conference Learning Representations, 2022.
John Hewitt and Christopher D. Manning. A structural probe for finding syntax in word repre- sentations. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Techmologies, Volume (Long and Shorl Papers), pp. 4129-4138, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-114 URL https: /aaaacaaanthooooo wwg./y..orggi.. Ntgg/N19-- 1419.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv: 1503.02531, 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. arXiv preprint arXiv: 2203.15556, 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. arXiv preprint arXiv: 10540, 2020.
Daniel Hsu. Dimension lower bounds for linear approaches to function approximation. 2021. URL https:/ www. . cs col u a -du/-d pttp://pparrs///a/ dttp:///hhhapppppr///ddiiaaaeaae mens on- cenn--rrummnte pdf.
A. Jafari, Mehdi Rezagholizadeh, Pranav Sharma, and A. Ghodsi. Annealing knowledge distillation. Conference of the European Chapter of the Association for Computational Linguistics, 2021. doi: 10.18833/v1/2021. mai.m.....clmmmi.112222.
Ding Jia, Kai Han, Yunhe Wang, Yehui Tang, Jianyuan Guo, Chao Zhang, and D. Tao. Learning efficient vision transformers via tine-graine manifold distillation. Neural Information Processing Systems, 2021.
Xiao Jin, Baoyun Peng, Yichao Wu, Yu Liu, Jiaheng Liu, Ding Liang, Junjie Yan, and Xiaolin Hu. Knowledge distillation via route constrained optimization. IEEE International Conference on Computer Vision, 2019. doi: 10. 10991CC.222119.000444
Dan Jurafsky. Speech language processing. Pearson Education India, 2000.
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin L. Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. SGD neural networks learns functions of increasing complexity. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 3491-3501 2019. URL htwps: http::///cccccceeeeeeeeoo neurips cc/ /appp/r///// 2019 9019/8 htt./hhh/aatt 4 2019/hasshh
neeriiirsinnnnnnnnnnsnee......ss 87 7cc.eecc776.aaa55ec5552222 wbs.ettttteeaaattttttaa......... htmi.
Michael Kearns. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM),
45(6))983-1006, 1998.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 141.::1412.6980, 2014.
Huayang Li, Yixuan Su, Deng Cai, Yan Wang, and Lemao Liu. A survey on retrieval-augmented text generation. arXiv preprint arXiv: 2202.0/110, 2022.
Han Lin, Guangxing Han, Jiawei Ma, Shiyuan Huang, Xudong Lin, and Shih-Fu Chang. Supervised
masked knowledge distillation for few-shot transformers. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 19649-19659, June 2023.
Bingbin Liu, J. Ash, Surbhi Goel, A. Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. International Conference On Learning Representations, 2022. doi: 10.48550/arXiv. 2210. 10749.
13

---

Published as a conference paper at ICLR 2025
Ruibo Liu, Jerry Wei, Fangyu Liu, Chenglei Si, Yanzhe Zhang, Jinmeng Rao, Steven Zheng, Daiyi Peng. Diyi Yang, Denny Zhou, et al. Best practices and lessons learned on synthetic data for language models. arXiv preprint arXtv:2404.07503, 2024.
lyya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 000.v33608.03883, 2016.
Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qingwei Lin, Shifeng Chen, and Dongmei Zhang. Wizardmath: Empowering mathematical reasoning for large language models via reinforced evol-instruct. arXiv preprint arXiv:2308.00 09583, 2023.
Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, Seungyeon Kim, and Sanjiv Kumar. A statistical perspective on distillation. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 7632-7642. PMIR, 18-24 Jul 2021. URL https: proceed i ngs. mlr w....pres///ee / http///1199mmmmm/1/11111 ahtml..
William Merrill, Ashish Sabharwal, and Noah Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-856, 2022.
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and H. Ghasemzadch. Improved knowledge distillation via teacher assistant. AAAI Conference on Artificial Intelligence, 2019. doi: 10.1609/AAAL V34104.5963
Hossein Mobahi, Mehrdad Farajtabar, and Peter Bartlett. Self-distillation amplifies regularization in hilbert space. Advances in Neural Information Processing Systems, 33:3351-3361, 2020.
Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu, Rosie Zhao, and Sham Kakade. Feature emergence via margin maximization: case studies in algebraic tasks. arXiv preprint arXiv: 23/1.07568, 2023.
Vaishnavh Nagarajan, Aditya Krishna Menon, Srinadh Bhojanapalli, Hossein Mobahi, and Sanjiv Kumar. On student-teacher deviations in distillation: does it pay to disobey?, 2024.
Eshaan Nichani, Yu Bai, and Jason D Lee. Identifying good directions to escape the ntk regime and efficiently learn low-degree plus sparse polynomials. Advances in Neural Information Processing Systems, 04444588448111 2022.
Ryan O' D'Donnell. Analysis of boolean functions. Cambridge University Press, 2014.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint ar arXiv:2209. 1895, 2022.
Lucia Quirke, Lovis Heindrich, Wes Gurnee, and Neel Nanda. Training dynamics of contextual n-grams in language models. arXiv preprint 1exiv:2311.00866 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint 27X1v:2403.05530 2024.
Yi Ren, Shangmin Guo, and Danica Sutherland. Better supervisory signals by observing learn- ing paths. In International Conference on Learning Representations, 2022. URL https: //ooprrevvies n forum? f
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 19: 19.00/ 2019.
14

---

Published as a conference paper at ICLR 2025
Wenxian Shi, Yuxuan Song. Hao Zhou, Bohan Li, and Lei Li. Follow your path: progressive method for knowledge distillation, 2021.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory, 65(2):742-769, 2018.
Daniel Soudry and Elad Hoffer. Exponentially vanishing sub-optimal local minima in multilayer neural networks. arXiv preprint arXiv: 1702.05777, 2017.
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Alpaca: A strong, replicable instruction-follomim model. Stanford Center for Research on Foundation Models. https://crfm.. stanford. 2ttt//022//33/1//appppa huml, 36)77, 2023.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rívière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery. Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Clément Crepy, Daniel Cer, Daphne Ippolito. David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Christeee Christian Muraru, Grigory Rozhdestvenskiy, Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross Mcllroy. Ruibo Liu, Ryan Mullins. Samuel L. Smith, Sebastian Borgeaud, Sertan Girgin. Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojeiech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong. Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy. Gemma: Open models based gemini research and technology. arXiv preprint arXiv: 2403.08295, 2024.
Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher Choquette-
lan Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nip pipeline. ACL., 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers and distillation through attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 18 24-24JB Jul 2021. URL https: httthttts///prreeeeee /ttp///eeeddiieeeeee/ee www.ceedeingg...... mlr .wr.esssss....................ee http///tttttttvv...11. htmll.l...mmmmmm.mmmmmmmmmm..... html.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro. Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. / 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv: 2203003666, 2022.
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok. Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint ar0iv:2309. / /2284, 2023.
15

---

Published as a conference paper at ICLR 2025
Shixing Yu, Tianlong Chen, Jiayi Shen, Huan Yuan, Jianchao Tan, Sen Yang, Ji Liu, and Zhangyang Wang. Unified visual transformer compression. International Conference on Learning Representa- tions, 2022. doi: 14485500/rXX1.22233000222.
Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng. Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 3903-3911, 2020.
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint MrXiv:2309.05653 2023.
Haoyu Zhao, A. Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? Conference Empirical Methods in Natural Language Processing, 2023. doi: 10:485550/rXxi.2233.8811
Kaixiang Zheng and En-Hui Yang. Knowledge distillation based on transformed teacher matching. arXiv preprint arviv: 2402.1 2024.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin. Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang. Joseph E. Gonzalez, and Ion Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
Helong Zhou, Liangchen Song, Jiajie Chen, Ye Zhou, Guoli Wang, Junsong Yuan, and Qian Zhang. Rethinking soft labels for knowledge distillation: A bias-variance tradeoff perspective. In Interna- tional Conference on Learning Representations, 2020.
16

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
A Overview of the appendix
17
A.1 Additional related works
18
B Proofs of results in Section 3.2
19
B.1 Analysis for the teacher Analysis the student
20
27
C Results on sparse parity and its generalization
28
C.I Additional results sparse parity with MLP
28
C.2 Learning with Transformers: parallel search with attention heads
28
C.3 A hierarchical generalization of sparse parity
34
D Extensive study on PCFGs
37
D.1 A formal description of PCFGs
37
Variants progressive distillation
37
D.3 Details Non-terminal prediction with Multi-heas linear probing
38
D.4 Details the synthetic PCFGs
38
D.5 Extensive experiments BERT
39
D.6 Hyperparameter details
40
D.7 Additional Curriculum probing on the teacher's checkpoints
41
D.8 Ablations with hyperparameters
41
E Autorcgressive training with GPT2
44
E.1 Observations
45
F Details Wikipedia + Books experiments
49
A OVERVIEW OF THE APPENDIX
The appendix provides omitted proofs and additional empirical explorations, which we outline below.
Omitted proofs We will start with the proof of Theorem 3.2 in Appendix B. The main idea is to show that the teacher can develop stronger correlation to in-support variables than to off support variables, which can then be utilized by the students reduce sample complexity.
experiments with MLP (Ap- pendix C. and Transformers (Appendix C.2), as well as results on learning a hierarchical extension of sparse parity (Appendix C.3). For Transformer experiments, dimensions of the architecture, such as MLP width and number of attention heads, affects the search of support for sparse parity. We discuss the effect of temperature in Figure 14. For the hierarchical extension of sparse parity, we show that the implicit curriculum occurs in different phases, which suggests a natural choice for number of intermediate checkpoints used in progressive distillation.
Additional empirical results sparse parity We present
study how scaling along different
Masked prediction on PCFGs In Appendix D.5, we provide a formal definition of probablistic context-free grammar (PCFG) and introduce the PCFGs that we use from Allen-Zhu & (2023b). We then provide details of our experimental setup and conduct extensive ablation studies on training a BERT model using the masked prediction task with PCFG data. We experiment with variants of progressive distillation and confirm that they lead to improved performance PCFGs, measured
17

---

Published as a conference paper at ICLR 2025
by accuracy and the three progress measures introduced in Section 4.1.1. Furthermore, investigate the effect of temperature, masking rate, and PCFG variation in Appendix D.8.
Next-token prediction on PCFGs In Appendix E, we conduct next-token prediction experiments using GPT-2 models on PCFG "cfg3f", i.e. the most complex PCFG in Allen-Zhu & Li (2023a). We characterize conditions under which progressive distillation provides significant gains.
A.1 ADDITIONAL RELATED WORKS
Understanding knowledge distillation There have been many works dedicated to understanding the effectiveness of knowledge distillation (Hinton et al., 2015; Mobahi et al., 2020; Menon et al., 2021; Dao et al., 2021; Nagarajan et al., 2024). For classification tasks, which are the focus of most knowledge distillation works, one intuitive explanation is that the teacher output provides a distribution over the class labels, which is more informative than the one-hot data labels. Menon et al. (2021) formalizes this intuition and shows that a teacher that provides the Bayes class probabilities leads to tighter generalization gap. Motivated by their result and the observation that high-accuracy teacher can be poorly calibrated, Ren et al. (2022) proposes to supervise the student using a moving average of the teacher across the training trajectory. While Ren et al. (2022) uses information of trajectory, their student learns from a fixed target throughout training, which is a major difference from progressive distillation. The teacher supervision also provides regularization benefits, such as controlling the biriinari tradeoff (Zhou et al., 2020), encouraging sparsity (Mobahi et al., 2020), or as a form of label smoothing (Yuan et al., 2020).
Learning sparse parity There are well established hardness results for learning sparse parity. When given access to labels only, learning (d, ks-sparse parity with gradients from finite samples is example of learning with statistical queries (SQ) (Kearns, 1998), for which a Rid') SQ computational lower bound applies (Edelman et al., 2023). When learning with a fully-comnected network (MLP), these parallel queries correspond to a combination of model width (i.e. neurons) and training steps, 11 and hence the SQ lower bound implies a fundamental trade-off between the width, the number of training steps, and the number of samples (Edelman et al., 2023). In particular, given the same number of training steps, narrower models require more samples to learn parity.
Feature learning In this work, we use feature learning to refer a learning process that recovers a fow-dimensional "feature" which helps reduce sample complexity. Sparse parity is task that can benefit from feature learning, where the feature is the support. For the special case of A: = 2, Glasgow (2024) shows that feature learning using jointly-optimized 2-layer neural network can reduce the sample complexity from (dd2) (corresponding to learning with NTK (Wei et al., 2019; Ghorbani al., 2019)) to (dpoly log d). Sparse parity is an example of a single-/multi-indes function, where the label is determined by a have also been studied Gaussian inputs (Nichani et al., 2022: Abbe al., 2022; 2023; Damian et al., 2024a;b) and have known separation between neural networks (Abbe et al., 2022; 2023) and non-feature-eemming kernel methods (Hsu, 2021).
projection of the data. These functions
Benefit of width in optimization Prior work has shown that width plays an important role in the optimization difficulty, where wider networks are more optimized easily. Du & Hu (2019) shows that sufficient width is necessary for the optimization on deep linear networks. Multiple works show that overparameterization leads to favorable optimization landscape, such as fewer sub-optimal local minima (Soudry & Hoffer, 2017; Soltanolkotabi et al., 2018) or guaranteed convergence at the limit (Chizat& & Bach, 2018: 2020). Wider models also exhibit faster decaying loss empirically (Yang et al., 2022; Bordelon et al., 2024a). Most related to our focus learning sparse parity, Edelman et al. (2023) relates the width to the number of parallel statistical queries (SQs). Combined with sparse parity's 's SQ lower bound, their result implies trade-off where larger width requires fewer optimization steps. Our work also acknowledges the benefit of width in optimization, but takes a different perspective by demonstrating that smaller student can inherit the optimization benefit when learning from a higher-width teacher. Moreover, we consider the number of attention heads as another
Mooe precisely, it is a combination of width and steps, as well as the batch size which affects the precision of the stochastic gradient. We omit the impact of hatch size here since we keep the batch size unchanged in the experiments.
18

---

Published as a conference paper at ICLR 2025
Algorithm 2- 2-stage training
Require: Stage lengths: Ti, T2, learning rates 7)1, 72, batch size B1, B2, weight decay A2. for [0, Ti] and E [m do
Sample B -smmpee {(x(1, y()))
Update the weights as w end for for t [0, and all E [m] do Sample Ba-samples {(x(1), ((x())y(11)) R2 Update the outer layer weights as
w (t-1)
Vws (Eecos (x, + A ||20;11)
(x))
y
(1+TT) t ai (1+Ti-1) 7e
Va, (L + )2a
x... (19))
end for
scaling dimension for Transformers, where the intuition is similar to having more "paths" (Dong et al., 2021). There have been results on studying the limiting output distribution as the number of attention heads goes infinity (Hron et al., 2020; Bordelon al., 2024b), though to our knowledge, there are no quantitative descriptions for finite number of heads.
B PROOFS OF RESULTS SECTION 3.2
We provide the formal version of Theorem 3.2 in this section. Recall that the teacher model is defined as
[T(x) Σ aso ((wss x) + b.).
The student model is similarly defined as
m
fs(x)= = _ à ((r0, x) + +b). 1=1
Setup We assume the data points sampled random from ({+1))d) Without loss of generality, let the target k-sparse parity function be = x 22 ... It Symmetric initialization: Following (Barak al., 2022), use the following symmetric initialization: for each < m/2,
w; ),{(((1))"), b; ~ - 1 I),, ~ u({+l/m)) Wi+m/2 = Wis bitm/2 = by, aitm/2 /2 = -
Following prior work (Barak al., 2022: Abbe et al., 2023; 2024b), we adopt two-stage batch gradient descent training. where we first train the first-layer weights [wi.... Wm. J. keeping the output weights {a:) fixed. In the second stage of training, we the output weights {a: (m while keeping others fixed. We keep the biases {b: mm fixed throughout training. Similar strategy for training the student model as well. The teacher trained with hinge loss, given by f(x, y) max(0, fx))) The student is trained with &DL yi fs. fT) max(0, (x))).
Two-stage training:
The training process is summarized in Algorithm 1.
Sample complexity benefits with progressive distillation for the student Our result is that progressive distillation provably reduces the sample complexity compared to (one-shot) distillation distillation. The key is establish separation between the correlations with in-support and off-support variables, which happens with high probability as formalized in Corollary B.6. Under such event, show:
19

---

Published as a conference paper at ICLR 2025
Corollary B.6: conditions satisfied by the teacher after first phase
Wh..p. the output the teacher after the first phase satisfies the following condition for all
Ex.y ((1) (x) (
2 1)(
if i E [+]]
((1) Ex.y JT (x) Ma
S o( if $ [R]].
Theorem B.1 (Sample complexity benefits with progressive distillation). Suppose the teacher model has been trained with 2-stage training in Algorithm 1, which satisfies the conditions in Corollary B.6 at the end of first stage and achieves loss Odd for some constant e 2 1 at the end of the second stage. Suppose we train a student model Is of size in = =0(2k) k) using the following two strategies:
1. Progressive distillation: Train for the first Ti steps w.rd the teacher's 's logits at Ti checkpoint. Then, train with the final teacher checkpoint in the second stage.
2. Distillation: Train with the final teacher checkpoint throughout training.
Then,
1. Under progressive distillation, the total sample complexity to reach a loss of c with probabil-
ity 1 8 is
@(k² log(dm// + 2f log(k/8)).
2. The necessary sample complexity under distillation is at least ((dmin(2ck-11)
The proof consists of two parts: 1) showing that the teacher develops strong correlation with the in-support variables after the first stage training (Lemma B.2, Corollary B.6), and 2) showing that given the support, the second phase of training converges quickly (Corollary B.8). These two helper lemmas are proven in Appendix 1.1.1 (first stage) and Appendix B.1.2 (second stage). The proof of Theorem B.1 is given in Appendix B.2.
Notations Before stating the proofs, we provide a list of necessary notations.
At any training step t, f/ will refer the teacher's output that step. Its parameters are referred
to as ((0) = (a{") w. (t) (,00),m the student fs are defined similarly.
Ji=1- The loss for / is denoted by L(S()) Lgto Notations for
Given a set X denotes the Fourier function on 3, where X (x) Iiss Ty We are particularly interested in 5 = S, i.e. the support of the sparse parity.
Maj: + 1 1 represents the majority function. On any x, Maj returns the sign of Σd Di=1 X,-
G for i > represents its ith fourier coefficient, i.e. Exz,y Maj(x)xs(( for any S E 10, 1)dd with = = C when even, and Ç, = ((i-1/3/(5) when is odd (O' Donnnlll, 2014).
To denotes the error tolerance in the gradient estimate due to mini-batch gradient estimation: let the population gradient and ğ be the estimated gradient with a few examples, Ty defined such
that [ - g0 S Tg. A -rror gradient estimate can be obtained using a batch size of ((1/>²))
B. ANALYSIS FOR THE TEACHER
B.1.1 FIRST STAGE ANALYSIS FOR THE TEACHER
First, we show that with an appropriate learning rate, the magnitude the weights wi, on coordinates
i increases to 26 while the coordinates & stay O (11) small.
Lemma B.2 (Single step gradient descent, adapted from Claims
Barak al. (2022)). Fix
Tg. > 0. Ser Ti as 1. Suppose the batch size B 2 2(r log(md For learning rate MiCz-1 and \ı 1, the following conditions hold true for all neurons E [m] at the end of first stage of training w.p. least -8
20

---

Published as a conference paper at ICLR 2025
1.
(1)
sigu(a, Ca- 1 1o (w.
S
1 all j E [k].
Ksx To
24
2.
Wij
(1)
1 150-1
sign(a,
(i (w)) (1)
for all >
2h
Proof. The proof follows that of (Barak et al., 2022), which we outline here for completeness. The proof has two major components: First, the magnitude of the population gradient at initialization reveals the support of the sparse parity. Second, the batch gradient and the population gradient can be made sufficiently close given a sufficiently large batch size. We will explain each step below.
Claim B.3. At initialization, the population gradient of the weight vector in neuron is given by Ex.y V Wwy ((x, yi f -Ex,y Vw., S7 (x)y, which can be split across the coordinates as
Ex, way J (000 (x)y
(0) (4-IXIINU (w) for all j E
1 (0) Ex,y V my 1 (x)y - -- (1x1 for Jorall)) S
Thus, the gradient of the weight coordinates Wiy for any neuron i and E S has magnitude Ks-all while the gradients of the weight coordinates th for any neuron i and & S has magnitude [[++1] The gap between the gradient in support and out of support is given by [Ca- K4+1 (Lemma 2 in Barak et al. (2022)).
()-
The second component involves applying hoeffding's inequality show the gap between sample and population gradient.
Claim B.4. Fix 6, Tg > O. For all i, for a randomly sampled batch of size Bi, {(xx ) with probability at least - 8.
Ex Exxr)(((1))) V way 1 (x) (( Veeg 1 (x)) S Tgy Tg.
provided B₁ 2 A(TTT² log(md/0)).
Because want the noise Tg be smaller than the magnitude of the true gradients for the coordinates in the support S, we want Ty to be smaller than ((x- We set this to get favorable condition for second phase of training (see Lemma B.7).
On the other hand, we show that after the first phase, the output of the network has positive correlations the individual variables in the support of the label function, and thus the checkpoint after the first phase can be used to speed up training of future models.
Lemma B.5 (Correlation with in-support variables). Under the event that the conditions in Lemma B.2 are satisfied by each neuron, which occurs with probubility at least 1 8 w.r.t. the randomness of initialization and sampling, the output of the model after the first phase satisfies the following conditions:
ak 1. Ex.g (1)) (x)x; 2 1 + (GT - (m- for all
2. Ex,y ((1) (x)x: < ((kk))-1 for all i & S.
3. 1 (1) (x)xs(x) s (4 for all S with even (S].
4. ( = Bx (x)]² < Od/k))
Proof. Consider a neuron i E [m/2] and its symmetric counterpart i m/2. W.L.O.G., we assume sign(w(()) (agn(a;") (")(((()))( (G-1) for all j [k)]. and sign(a;") 1. Recall that k is assumed be even,
21

---

Published as a conference paper at ICLR 2025
hence sign( (X[x) (m) 1. Then, the condition in Lemma B.2 can be simplified
Wij
(1)
2k
Vijs
W (1)
4tm/2.j
2k: Vays for all j [k],
(1) 1 Sk+1 sign( wij
w (1)
CR+1
sign(w?y (0)) + Vij. for all j 2 k
2h K5--1
Vijs
i+m/23
2h
where Viy satisfies the following conditions.
Tg
loyy
for all j [k]]
Kh-11
Tg 1kCk-1 for all k:
Then, the sum of the output of the neurons i and + m/2 on an input (ignoring the magnitude of a,) is given by
k
d
(x)) =
Ck+1
sign(ur (0) )o, + (w,, x) + bi
Jyj
+
2k
2k
j=1
j=++1
k:
d SK+1
2
Σ
+ (0,, x) + b, sign( w 10) 1 Jzz
zzj
2k
2k S4-11
1=1
j=k+1
and
m/2
m/2 1
1 (x) Σ Mx
Σ m
1. In-support correlations: We are interested in the correlation of this function a variable x.
for u E S. We argue for u = 1, as the similar argument applies for others. Thus, we are interested in
k
Ex,y (F, = Ex
1
1
Sk+1
(0) )ay sign(w Jas + (0,, x) + bi
+
x
2k
2k
Ck-1
J=1
j=k+1
d Σ
Ck+1
(0) sign w )oi + (0i, x) + bi In
2k
2k
|5--1
1=1
j=b+1
(7)
We focus on the first term; argument for the second term is similar. First of all, we can ignore (v;, incurring an error of O(Tyd ((x-1l ).
1
1
Sk+11
Ex
Iz
+
sign (m) Ja,, + by
21
2k
2k
(s-1l
j=1
j j=
&
E Ex.ga +10
1
+
1
Σ
T.
+
1 Cx+1
Σ
sign (u:; )or, + bi
2k
2h
21
j=2
j=
b
1
1
1 GE+1
sign (w: (0) Jz + by
-Exys1==10
+
Izj +
2k
2k 1=2
2k (E-11
j=k+1
1 F Ex,y I
C1+1
Σ
sign(w Jary be 20)
J.j
2h
2h
2 [K-1
j=k+1
22

---

Published as a conference paper at ICLR 2025
The final step follows from the observation that the argument of a in the first term is È higher than the argument of a in the second term. This implies that when the first term is non-zero, it's at least 2E higher than the second term. Hence, we lower bound by considering term is non-zero.
scenario where the first
Continuing, can further split the indicator function into cases when each term in the argument of the indicator function is positive.
d
Ex
Ex.g
I
Ck+1
Σ
sign uu )r + bi x1
2k
2k [A--1
1=1
j=
d
Ex
Ex
Σ
Ti +
CK+1
(0) sign(w Jrú 20)
2k
2h j=2
2h
15-1
j=k+1
1 E
Ex, Il
Tyy 2000
Ij 2 26)
2k
j=2
j=k+1
I (b, 20.
8k
From Equation (7), we then have
1 I(; 2 0) + O (4- 4k:
2
As by has been kept at random initialization and thus is random variable selected from the set {-1+ E
}. with probability 3. (b; 2 0). This implies, w.p. atleast 1/2 w.r.t. a neuron's bias initialization, (xx (xx))) 2 Jh OT, I The final bound comes from the fact that Ex (x)x=111111)))))) Ex,y m E (F()
1/22), O(m-1/2), where the ()))) (x)=1 2 ak + O(Tyd (ondkka-l--- + O(m
error term is bounded using Hoeffding inequality.
Out-of-support correlations:
Similar to the Equation (7), we have for u & S,
(FF))))))((x) Ex,yo
1
+
1
Sk+1
Σ
sign ( 00)
Jæy + ((,,,)) by Ir
2k
2k Ck- 1
1++1
d Σ
Ck+1
(0) sign(w: )az + (vi,)) + b; au.
xj
2k
2k Ck-
3=1
j=k+1
(8)
However, we observe that the influence of x, in each of the terms is bounded by 1 Ct.11
& [a-ill Consider
the first term; the argument for the second term is similar. We can again ignore (v., x) incurring an error of OTT (G-yd(Ca111 (1-11-
23

---

Published as a conference paper at ICLR 2025
k
Ex
2ij +
Ck+1
Σ
sign wij (00) )ay + bs
Iu
2k
2k
Ck-1 g=
j=1
E Ex,,, =110
1 CK+1
sign(w (0))
+ 2h: j=1
Ii +
Ck+1
sign(ur (0) )oy + bb
2k Sk-1
2h 15k-11
j=k+ + 1
-xxyxx==-00
1 Ck+1
sign(w (w
(0)
Σ
S1++1
E
sign (uoj (0) )ry + bi
2k |Sk-1l
2k j=1
2k (s--1l
jmmt Iddijj
((x) Ck+1 Ex
sign(w/u) (0) I
Ck+1
sign(w )ay b; 2
S5-11
2k
2k
j=b+1-ddj
where C(x) E {1,2) denotes a function that depends on x. The final step follows from a first order taylor expansion of a. The magnitude can hence be bounded by 1 |Ce+1] k I[e--l" This can be bounded
by kd (section 5.3, O' O'Donell (2014)). The final bound comes from the fact that Ex (x)x... = Ex. m (T S ((hd))-1).
3. Correlations to support of an even size:
The function (f(1)); is given by
(xx) =0
1
1
S4+1
(0) Jay (0,, x) bi sign(w!y
3j
2k
2k
Sk-1
j=1
j=
1
1
Gh+11
sign(w!) (0) Joz + (0,, x)) + be
2k
2k 1Sk-1l
j=1
j=k+1
1
1 Sh+1
(on 00) sign( wiy Jay + bb
33 +
2h
2k
Sh-11
1=1
j=k+1
k
1
Sh-m1
Σ
sign (u (0) ja, + b. + O(TTd (5-1a1))
zj
2k
2k
Sh-
j=1
j=b+1
[ OTy (
One can observe that (x) is symmetric function and so an odd function. Thus, (x)
(exercise 1.8, O' O'Donelll (2014)) and so, Ex,y (T+ ()=))))))(())))))))))))) (x)x)xs(x) = OTyed (
4. Output norm: Focusing on function (f
((())) 1/ Ex.u(ff)) (xx))
Ex.y (o((w!)) (1) ; x) + +b)) - (( 2 x) + bi)
S Ex.y min (100 || + 12 w 2 + b2 ||x|| = o (²22)
d
1
The intermediate step uses Cauchy-Schward inequality, and the final step uses the values of
wiy w i+mm/255" As fx)
(1) (1)
1 ( (xx), have || < 2 Dm/² ((
(1)
24

---

Published as a conference paper at ICLR 2025
Corollary B.6. Under the event that the conditions in Lemma B.2 are satisfied by each neuron, which occurs with probability at least 1 8 w.r.t. the randonness of initialization and sampling. the output of the model after the first phase can be given as:
f (x)
Σ j=1
E
cxy+
Σ
csXs(x)
Σ
csxs(x),
j=k+1
$C4:51
SCI4):1|192=0
where
lc 2 (((kkk)
for all 1 S S k, for all j >
1c & (((kd)
lesl S O (4- 11 1 |cs| - (ddk))
for all C [d] with SS|% = 0,
for all S C [d] with S1|% 1.
As such, the following correlations hold true for all i.
Ex (xx) (x). Maj(x)s; = 2 + O(Tgd//3 G
batch size Bı is set 2 21132 such that Ty S 1(k-1dd-- (Sk-1)), then the following holds for all i.
Ex,yJT ((2) (x) Maj(x))a 2 2(
ifi E [k]
Ex.g (2))1(())))) (x) Maj(x))x
s o(k
if $ [K],
Proof. The form of p follows from the fourier coefficient analysis in Lemma B.5. Now, we can use the formulation to derive
E
(x)
Maj(x)a:
d
Ex.y
Cjogy Maj(x) ds + Ex.y
Σ
MsMax())xs(() - as
j=l
1
Ex,y
csxs(x) Maj(x)æ;
-Ex,y
cj Maj(x) 24 + Ex,y
Σ
Msjax())xs(x) 2z
/=1
=c Ex Maj(x) + Ex,y
Σ
Maj))()) + Ex,y
cxMax())xs(x) SC14
Σ
2 G. + Ex,y
(xMaj(x)xs(x) - x4. SC/2)
Σ
The second step removes Ex,y Escidnssiz2 (sxs(x) Maj(x)a;; because Maj(x) odd function, and so Ex,, Maj())))) will be 0 for odd sized S. Similar argument holds for removing Ex, Σ Engze Maj in the final step. We finish the proof by bounding Fz.
Ex,y es
25

---

Published as a conference paper at ICLR 2025
As |cs| 1 (k- for all S with |S|%2 0, can bound it as
Ex,y
Σ
(sMax(x)xs(X) 3i
SC/dH|/1/220
<O(Tgd1SA-117
Ex Majx)
SCId)-181%200
5O(Tgd (SK-
1|-
(ajx))s(
SC
50(Tgd [Sk-1 1)
Σ
1x Majx))((
5Cld
SS 1/3
-O(Ty |sc-11 1
+
(15)
SC14]
- 1
Here the pre-final step follows from the bounds on the Fourier coefficients of Maj outlined in Appendix B. Finally, we set B1 2 2((T-)) is set such that Tg S (0-1 This makes OlTgd5/3 (S-il=") o(1/k). Hence, with appropriate batch size B1.
1 Ex.yy (x) Maj(x)a ~ o1//k)) The proof follows from the magnitude of c, derived above.
B.1.2 SECOND STAGE ANALYSIS FOR THE TEACHER
Lemma B.7 (Second stage Training, cf. Theorem 4 in (Barak et al., 2022)). Fix e,8 Suppose m 2 Q(2*k k kog(k/))), 2 Ω (Is+ log(kd/e)). Furthermore, suppose B1 2 20( 2 loo (k)
the weights satisfy the conditions in Lemma B.2 with Ty 0((c-11--1) after the first phase. Then after T2 13mmd213// steps of training with batch size B2 1, learning rate 112 = -k. //d.5//d//mmm m 1)) and decay x2 0, we have with expectation over the randomness of the initialization and the sampling of the batches:
min E [g (x, v)] S E. (2
Thus, the minimal sample complexity to reach a loss of is given by
T Bi + x Bg = 1 2² log(kd/e)) +m
=6 log(dk/e /E) + + d log(k/)))
Corollary B.8. Under the conditions outlined in Lemma B.7, after Tr steps of training in the second phase, if denote the time step which the model achieves the minimum loss, i.e. arg min I. [T.ges / (x, U}}, then
E 1 (x)
S s, for all E [d].
The proof follows from the fact that the correlation along = Ilis 2; is large (2 1 - e as hinge loss is below e), the correlations along other Fourier basis functions will be small. Hence, depending on how saturated the model is, the signal along the support elements are small.
We will slightly modified version of Lemma B.7 with higher sample complexity in the first phase, to ensure the stronger conditions of Corollary B.6 hold true as well. This will be necessary to get improved signal to teach smaller student.12 12
le haven't 't optimized the error bounds in Corollary B.6. Our sample complexity bounds are likely loose in Corollary B.9
26

---

Published as a conference paper at ICLR 2025
d 2 Q (K+ log(kd/e)). Furthermore, suppose B1 2 Q(Kk- /2 log(kd/e)) s.t. the weights satisfy the conditions in Corollary B.6 with Ty - - d 5/3) after the first phase. Then after T2 = Mmd /e²) steps of training with batch size B2 = L, learning rate T22 4k2 /(dy m 1)) and decay X2 0, we have with expectation over the randomness of the initialization and the sampling of the batches:
Corollary B.9 (Modified Version of Lemma B.7). Fix e, > O. Suppose m 2
min E [Lgto (x. v)] S. TE(Ta]
Thus, the minimal sample complexity to reach a loss of is given by
T1 x B1 T2
B2 1(((--11 /10/3/² h22 log(kd/c)) @mmd2 (2)
0 log(dk/e) + dt log(k/00)
B.2 ANALYSIS FOR THE STUDENT
Proof of Theorem B. I. We will first prove the sample complexity upper bound for progressive distil- lation, followed by a sample complexity lower bound for distillation.
Sample complexity for Progressive distillation: Under progressive distillation, the label is given by ((T)) for the first Ti steps. We will follow similar steps as Lemma B.2, where the label replaced by (2)) Claim B.3 changes, while Claim B.4 stays the same. We will showcase the change in Claim B.3 here.
At initialization, the population gradient of the weight vector in neuron i at coordinate j is given by
Ex.0
ton LDL (x yi
(0) fr)
=-Ex,p (0)
(x) (1)) (x)
= az Bx,y
(un
(0) x) + i 2 0] ( (x)((())
1
=aa,Bay Bx.
+
Maj(
xx) 1 (x)(s
2
2
-a Ex.v
1
(Ti)
(x)xr- a. 2 Ex,g Maj( ( x) (x)x;;
where the relation between (20)
(0)
x) + 2 of and Ma x) follows because of b. <
at
initialization. From Corollary B.6.
Ex
x
DL (x, y: 16) fr) 2 1)k-)),
if j E [R]],
Ex.y (x
(x, y: 13 (0) .T) s <o(k)))
if j & [A]]
Thus, a fourier gap exists between the population gradients in-support and out-of-support coordinates in the gradients. We can then apply Claim B.4 to show that a finite batch size of B1 s2kk2 log log(dm//)) is sufficient to maintain this gap between the coordinates in support and out of support. Thus, the change in the necessary sample complexity comes from the reduced sample complexity in the first phase. The proof for the second phase training is exactly equal the proof for the teacher in Theorem B.I.
Sample complexity for Distillation: On the other hand, for the teacher checkpoint with loss O(d-s) the correlation to the monomial terms in the support is bounded by (dd(-)) (by Corollary B.8). If we want to learn from the correlations to the support, we need the number of samples to be at least ((d2)) as the gradient noise needs to be lower than O(d (by Claim B.4). To learn the support from the true label, we need the number of samples to be at least ((d ), by the following result:
Lemma B.10 (Width-optimization trade-off, cf. Proposition in (Edelman et al., 2023)). For O. 8r2 gradient noise Tg > O, and model width m > 0, ifT < 1 (d) m. then there exists a (d, k)-parse parity such that w.p. at least 1 - 5 over the randomness of initialization and samples, the loss is lower bounded as ( 2 1 Tg for all 1. (1...T).
27

---

Published as a conference paper at ICLR 2025
MLP on sparse parity
Transformer on sparse parity
1.0
1.0
Top-1 Accuracy
Width
Heads
100
0.8
1000
10000
0.6
16
0.0
0.4
32
SM
10M
15M
20M
0.2
2M
4M
6M
BM
Training Steps
Training Steps
Figure 7: Larger models learn sparse parity faster. A larger model has more width (MLP, left) more attention heads (Transformers, right). The results are for (100, 6)-partty, aggregated over runs for each setup.
This result implies that for fixed batch size (and hence a fixed Tg); we either require bigger width, or more number of gradient steps (which translates to sample complexity since we are using fresh samples each batch). Hence, for the model to learn the support from a combination of the two components, needs a sample complexity at least ((dminn2e.k)))) /m).
C RESULTS ON SPARSE PARITY AND ITS GENERALIZATION
C.1 ADDITIONAL RESULTS ON SPARSE PARITY WITH MLP
We take both the teacher and student models to be 1-hidden-layer MLPs with ReLU activations. The teacher has a hidden width of 5 x 10a, and the students are of widths 102 or 10 All models are trained using SGD with batch size 1 for 20M steps on sparse parity data with n = 100 and A 6 (Definition 3.1). The support set to be the first 6 coordinates of the input vector without loss of generality. The learning rate is searched over {10 2,5 x 10 10 $). Evaluation is based held-out set consisting of 4096 examples, and we report the average across 3 different training seeds. For one-shot distillation, which point the teacher has fully saturated. For progressive distillation, we use N = 200 equally spaced teacher checkpoints that are 0. steps apart.
use the teacher checkpoint the end of training (20M checkpoint),
Teacher's accuracy behavior
Absolute Correlation
Corr. degree-1 monomials
(2, 1M) progressive
0.10
1.0
Accuracy
Accuracy
0.8
1-6
0.8
candidate
0.05
candidate
Rest
candidate
0.6
0.6
0.00
0.95
1.00
1.05
1.10
0.95
1.00
1.05
1.10
0
2
Training Steps
1e7
Training Steps
1e7
Training Steps le6
Figure Repeated experiments from Figure 2 for student of width 1000.
C.2 LEARNING WITH TRANSFORMERS: PARALLEL SEARCH WITH ATTENTION HEADD
The benefit of progressive distillation and the implicit curriculum is not specific to MLP. This section presents similar results with Transformers (Vaswani et al., 2017). The d-dimensional input vector is now treated as a length-d sequence, and the label is predicted using the last token's 's output. We fix the support S to be the first 6 coordinates of the sequence. Note that unlike MIP,, Transformer's learning not
the location of S due to the causal mask. Nevertheless, given the same S, the comparison on learning speed is still meaningful.
For Transformers. the parallel queries come from both the MLP width and also the number of attention heads. To illustrate this, consider the following two solutions (which we formalize in Appendix C.2.1) to sparse parity: The first solution uses attention locate the support and then uses MLP to compute the product of the suppor variables. The second solution copies over all variables to the final position, whose MLP is then responsible for both identifying the support and computing the product.
28

---

Published as a conference paper at ICLR 2025
Teacher's accuracy behavior
(2.0.MM)-rrorressve
1.00
Accuracy
Accuracy
candidate
0.75
0.8
candidate 2
candidate
0.6
0.50
7
8
2
Training Steps
le6
Steps
1e6
Absolute Correlation
Corr. degree-1 monomials
Absolute Correl lation
Corr. to degree-2 monomials
Absolute Correlation
Corr. to degree-3 monomials
3
0.04
0.04
0.04
16
16
1-6
0.02
0.02
0.02
0.00
0.00
0.00
Training Steps 1e6
Training Steps
1e6
Training Steps
1e6
Figure 9: Repeated experiments from Figure 2 but for different teacher. For (2, 0.1M) progressive distillation, the checkpoint that lies in the middle of the second phase accelerates training the most. In Figure 2, we used teacher that was trained with learning rate 5 x 10-3, The correlation plot for the teacher to degree-1 monomials had a clear gap for degree-1 monomials in-support and out-of-support at the middle of the second phase (indicated by candidate 2). However, for teacher that is trained with a higher learning rate 10-2, we didn't 't find such a clean gap in correlations for degree-I I monomials. On the other hand, correlations degree-2 and degree-3 monomials showed a clean gap between in-support and off-support variables at the middle of the phase transition. Hence, the student needn't 't learn only from degree-1 monomials get training acceleration, any low degree monomials suffice to teach the student about the support. Rest for degree-2 monomials refers to all monomials the form T,Iy where atleast one of i,j & S. Similar definition for degree-3 monomials.
29

---

Published as a conference paper at ICLR 2025
The second solution is less interesting it reduces to MLP, focus the first solution in the following, which utilizes the attention mechanism unique to Transformers.
(R6) More attention heads helps with the search for support Our experiments are based on 2-layer Transformers with 8 dimensions per attention head. As shown Figure (right), increasing the number of heads makes learning faster. There are clear phase transitions similar to the MLP case.
Ablation with other ways to vary the model size Most Transformer experiments in this work keep the per-head dimension to be fixed and vary the number of attention heads between the teacher and the student. The MLP input dimension is the sum of the attention head dimensions, so a student with fewer heads will have smaller MLP than the teacher, which is preferable in terms of efficiency. Fixing the per-head dimension is a widely adopted setup in practice, such as the Llama series (Touvron et al., 2023). We additionally consider particular, we vary the number of heads, while 1) fixing the hidden dimension (i.e. the total dimension of all heads concatenated) be 256, or 2) fixing the dimension of each head be 256 and averaging the output from each head, in which case the hidden dimension is also 256. These two setups are less common practice but nevertheless serves as complementary evidence: the performance difference comes solely from the number of attention heads, the MLP dimension is kept the same. As shown Figure 10 (b,c), increasing the number of attention also increases the training speed in these two setups.
other ways vary the model size. In
1.0
1.0
head
Acuraayy
Accuracy
0.8
beads
hends
0.6
hoads
headh
0.6
heads
Broodt
0k
50k
100k
150k Steps
200k
0.4 0k
50k
100k
150k Steps
200k
5.c
3306
150c
200L
(c) Dimension 256, averaging heads Figure 10: Increasing the number of attention heads speeds up training. Each plot compares the accuracy throughout training for 2-layer models with various heads, while fixing: (a) the per-head dimension to 8; (b) the MLP hidden dimension to 256; (c) both the per-head and MLP hidden dimension 256, by averaging (rather than concatenating) the heads. We report runs with the learning rate that has the highest mean accuracy and break tie with training speeds. The shadows show the variances of the runs.
(a) Per-head dimension =
(b) Hidden dimension 256
0.5
zttn
zltn ver
6tth
ottn ver
41th ver
4tth ver
1tth
0.3
6tth ver
1tth var
11th var
1thh ver
Absention
1tth ver
149 wihh off pportttttt ables
sccuracy (scaled
0.2
accuraiy (500eet
anhh off support vor ables
0.2
0.1
0.0
0.0
5000C
100000
150000
200000
250000
50000
100000
150000
200000
250000
Stop
Stop
Figure 11: In-support attention growth co-occurs with accuracy increase Attention on individual coordinates on or off the support of the sparse parity, taking the median of 1024 random binary input sequences. The shade highlights the teacher's phase transition period. The model accuracy is marked by the gray dashed line, with scale adjusted for better display. The two subfigures show the same type of results but with different randomness seeds.
Ablation with 2-shot distillation We repeat the 2-shot distillation ablation for MLP. We first confirm that the low-degree curriculum described in Section 3.1 is also observed in Transformers. As
We use layers since I-layer Transformers are hard train empirically, despite being representationally sufficient solve sparse parity.
30

---

Published as a conference paper at ICLR 2025
Acuuacy (sckee)
0.000
hosuppe Enhiss
Ofsansent (Khees
0.255
50.000
0.006
0.10
0.024
0.03
0.002
0.00
0.000
25320
50036
15005
100000
125000
195000
179000
53000
100030
150000
200000
Staso
0.210
0.012
tolics
0.000
0.020
ountort Solves
0.000
0.004
0.004
0.02
0.022
6.000
0.000
20003
+2000 50000 100000
220000
L400C0
20:00
73200 303200 123000
130000
175030
Figure 12: Low-degree curriculum in Transformers on (100, 6)-saasee parity. The r-axis shows the training steps, and y-axis shows the 2-layer 32-head teacher's correlation with in-support (orange lines) vs off-support (blue lines, aggregated into mean and standard deviation) degree-1 monomials. The black dotted lines mark the accuracy, scaled for better display. The correlation values calculated using 100k randomly drawn sequences. The 4 subplots correspond to models trained using 4 random seeds.
1.0
0.9
progressive
Accuracy
progressive (2-shot) one-shot
teacher
0.5
0.4
0k
50k
100k
150k
Steps
Figure 13: 2-shot progressive distillation with Transformers: Compared to cross-entropy training or one-shot distillation, transformers learn faster with progressive distillation, where the intermediate checkpoints are taken either at regular 10k intervals ("progressive"), during the phase transition ("progressive (2-shot)"). The two vertical lines show the teacher training steps at which the two checkpoints for 2-shot distillation are chosen. We set the teacher temperature be 10-4 for progressive distillation, and r for one-shot distillation.
shown in Figure 12, the 2-layer 32-head teacher model exhibits significantly higher correlation with the in-support monomials (i.e. {aihes his than with off-support monomials during the phase transition. Then, we show in Figure 13 that using as few as intermediate checkpoint suffices to significantly speeds up the training of the student.
Ablation with various temperatures As mentioned in Section 2, our progressive distillation results use low temperature in order to remove potential favorable regularization effects from soft labels. We chose a temperature of 10 for sparse parity, where the output dimension is 2. In Figure 14, we empirically confirm that setting the temperature to be below 0.01 is sufficient to get results that are qualitatively similar to using т = 0 (i.e. taking the argmax). Note that using a higher temperature such = can make learning slower despite potentially having more regularization effects from softer labels. We leave understanding the exact effect of temperature to future work.
Note that the upper right subplot in Figure 12 has second correlation spike with the in-support variables. However, supervising with this second checkpoint does not provide acceleration. This suggests that there might be mechanisms other than the low-degree curriculum play.
31

---

Published as a conference paper at ICLR 2025
1.0
0.9
0.8
0.7
temp
temp 0.01
0.6
temp le-4
temp
teacher
100k
150k
200k
Stops
Figure 14: The benefit of progressive distillation holds with hard labels, as shown by comparing 2-shot progressive distillation with different temperatures. The two gray vertical lines mark the training steps at which the teacher checkpoints are taken.
C.2.1 TWO TRANSFORMER SOLUTIONS FOR SPARSE PARITY (PROPOSITION C.I AND PROPOSITION C.3)
We consider simplified version of Transformer block, without the residual connection or the layernorm:
fulock Sulpi factas
where
Sastn (X; Wo, WK, Wv) = CausalAttn(XW XxW,,
with Wo, WK, Wv being the query, key, value matrices, and fulp (z; (Wi, by bree)()) is L-layer MLP that recursively apply ( (x) o o f (ax) b(+1) position-wise. or is the relu function for I E [T. 1), and is the identity function for 1 L.
Proposition C.1 (Attention support selection). (d, k)--pprrss parity can be solved by a I-layer Transformer with a 2-layer MLP, whose attention weights satisfy ai, x exp(cili E S]) for some large constant
The MLP has hidden dimension 4(k + 1), Loe norm bounded by 4k(k + 1).
Proof. The idea is that the attention selects the in-support variables, and the MLP computes the product of these variables.
To select the in-support variables, we want the attention weight az xx exp(c1 [i E S)). for some large constant c. This can be achieved by having the projection matrices WQ. WK focus only on the position and ignore the tokens. In particular, let z denote the input sequence, and let the embedding of token be x; vs, + pi. where (vo, di embeddings for the binary token 0 1, and p; is the position encoding for position i. Take (vo }. {pth: such that Us, Lp:. Let c 0 be large enough constant. Choose Wo. WK such that for any E [d]. x W W WK = P. W WKPs
1 S].
This ensures that a4, x exp(c1 [i E S)).
Then, the role of attention is to average over the in-support tokens. For simplicity of exposition, let's take for now (i.e. using saturated attention (Merrill al., 2022)), that ag,, 1[tEs], that is, the attention weights at the last position average over the in-support variables. Take Wy to vector, such that Wy ignores the positional information and the input token 0, and only preserves the input token 1, i.c. Wyp; = O, Vi E [d], Wy = O, and Wy =
Next, the MLP needs compute the parity function over the in-support variables. The input to the MLP is hence proportional to (Eies 1[2; = 1]) and the size of the set of inputs is l + 1. To determine the size of the MIP, we use the following lemma:
Lemma C.2 (ID discrete function interpolation with an MLP (Lemma in Liu al. (2022))), Let X be a finite subset of R. such that [x] < B₁ for all x E X, and jx x'1 2 for all x & x' E X. Let + d be such that [If())) < for all x E X. Then, there 2-layer ReLU network for which
Smip( u (:0mp) f(x)
Ve X,
14 A/4.
32

---

Published as a conference paper at ICLR 2025
The inner dimension is d' 1 and the weights satisfy
4
4Bx +2,
|| 00 S
|161l120 S
11 S By: b2 = O.
A
Setting B 1, By 1, and A [ST the parity function these k + input values can be approximated by a 2-layer MLP with inner dimension 4(k: + 1), with norm bounded by 4k(k + 1).
As concrete example,
way satisfy the requirements above is set the attention weights to 0 = Wv e == [1,0, 0, 0,
0). on = -2 == [0, 0j. Set ps Pa eg for i E S, and p; e4 for
0
0
i g S. Set We Wx
0 for some sufficiently large e > O.
0
Proposition C.3 (No attention selection). There exists a I-layer Transformer with 3-layer MLP that computes k-sparse parity, whose attention weights satisfy ai,d Consequently, the MLP computes the sparse parity function given the full set of variables.
Proof. The idea is for the uniform attention to copy all tokens to the last position. However, unlike in Proposition C.I, the attention needs to copy the tokens into length-d embedding vector, need to preserve the position information in this embedding vector. We need to generalize Lemma C.2 accordingly to handle multi-dimension
inputs:
Lemma C.4 (General discrete function interpolation with an MLP; (Lemma 2 in Liu et al. (2022))). Let X be a finite subset of TRd, such that || ||= B2 for all X, and || x'1l A for all + x' E X. Let : X - Rdous be such that (x)]l S B, ReLU network for which
B, for all X. Then, there is 3-layer
fmip(æ + E; mmlp) = f(x)
E X, 1C| S A/4.
Letting X, denote the set of unique values in coordinate the inner MLP dimensions are as follows:
di 4
Σ 1&
d2 = 12
The weights satisfy
4 A
4B, 1bb1 <
W lax S
+2,
||2 S 1 |lb2llo0 dnn,
| S By: by 0.
Then, the MLP at the last position computes the sparse parity over the k coordinates while ignoring the others. Hence the effective input set is | = 2* Setting BN = 1, B, = 1, and = there exists a 3-layer MLP with width 2 and norm bound 2 by Lemma C.4.
Preliminary interpretability analysis: Transformer does utilize attention in practice We observe that the model focuses attention on relevant tokens and that the amount of attention weights put on the support is tightly correlated with the accuracy, which suggests that the model indeed utilizes the attention mechanism in learning sparse parity.
Specifically, Figure 11 shows the results on 2-layer 16-head GPT-2 models. The attention weights for the final position, whose logits used for computing the binary parity label for the entire sequence. We track the attention weights along length-2 paths from the first and the second layer. For example, for a single-head model, let a
(1) E d denote the lay- attention vector at the
ieh position; then, the on-support attention for given sample is computed as (a (2)) v("),, where [v EjET [[] is the total amount of first-layer attention weights that the 4eh position puts on the support T. For multi-head models, a! A d-1 is defined the sum of attention vectors from all heads, and the rest is computed similarly.
33

---

Published as a conference paper at ICLR 2025
1.0
1.0
progressive ane-shot teacher
Accuracy
Accuracy
0.5
0.5
progressive
one-shot
1 Training Steps le7
2
Training Steps 1e7
Width-100 student.
(b) Width- 1r0th-10000.om student.
Figure 15: 8-way classification using hierarchical decision tree of depth 3, with each node repre- sented by 5-sparse parity. Progressive distillation helps student learn faster from width-50k teacher, compared to one-shot distillation from the final checkpoint.
Level
sigu (
sign(etasttotta))
Level
An illustration of hierarchical data generation, for 3-level tree with 3 variables per feature. A feature corresponds to a tree node, each marked by a rectangle. The product of the binary variables in feature determines which child take: the left child is chosen if the product evaluates to 1, and the right child is chosen if the product is +1. The final label for an example is decided based on the tree leaf reached.
Figure 16:
C.3 A HIERARCHICAL GENERALIZATION OF SPARSE PARITY
This section considers an extension of sparse parity, where the labels are given by a decision tree. Sparse parity be considered special case with tree depth 1.
The input x is boolean vector picked uniformly at random from the d-dimensional hypercube (+1)d, and the label y E [K] where K := 2D for some fixed D N. The underlying labeling function for y follows binary decision tree of depth D, whose leaves correspond to class labels. The branching at a node depends on sparse parity problem. An example visualization is provided in Figure 16.
Definition:
More formally, the nodes the decision tree are represented by a set of sparse parity problems S {T1, Tz, ... Tx where T, is determined by product of subset of size k variables selected from the dimensions of the input x (e.g. 21 ... Ty for = 5). An input x belongs to the class i [K] iff
D II
I (T (o (x) 0) > 0, where
11, { -1,
ifi
cli,))
otherwise
Here, u D denote the features that on the path joining the root of the decision tree to the leaf representing the label i. An example is given in Figure 16.
34

---

Published as a conference paper at ICLR 2025
Absolute Correlation
Corr. to degree-1 monomials
Corr. to degree-1 monomials
Corr. to degree-1 monomials
0.03
0.03
0.02
18-20
Rest
Rest
Rest
0.02
0.02
0.01
0.01
0.01
0.00
0.00
0.00
1e6
Training Steps
1e6
1e6
Absolute Correlation
Corr. to degree-2 2 monomials
Corr. to 2 monomials
Corr. to degree 2 monomials
15
0.02
5030
0.015
16:20 20
0.02
Rest
Rast
Resi
0.010
0.01
0.01
0.005
0.00
0.00
0.000
4
6
2
1e6
Training Steps
1e6
1e6
Figure 17: Setting: 8-way classification using hierarchical decision tree of depth 3, with each node represented by 5-sparse parity. The relevant features for class .1...15, my -.., -1.0---0000 at tree levels 3, 2, and 1 respectively (Figure 16). The irrele- vant features are I36, .. , F1000 Here we plot the magnitude of correlation to degree-1 monomials Ex.(Pr(x))]= for each in the relevant feature groups for class Because the degree-1 monomials show noisy correlations, we also report the magnitude of correlation to degree-2 monomials Ex,y(PT (x)]i Jyy for each j the relevant feature groups for class For degree-2 monomials, rest refers to correlation to monomials of the form 2;2r where atleast one variable is outside support variables (z36,-. ... 2 00). The correlations to degree-1 (or 2) monomials on the relevant features spike at different training steps.
are
Teacher behavior
po--)progressieetree 1)
(3. 2M--progressive(r=100
1.0
1.0
1.0
Accuracy
Accuracy
Accuracy
1W0 31 SMIS
0.5
Candidates for (3, TI- progressive
0.5
0.5
4.29 41-44--88 43NB4)-5
3.5 5.3
10 Training Steps(x109)
3
5
Steps 1e6
Steps
1e6
Figure 18: Setting:
8-way classification using a hierarchical decision tree of depth 3, with each node represented by 5-sparse parity. (3, 2.) pprogressive distillation from 3 checkpoints
1000 width student; 2 intermediate teacher checkpoints are used each for 2M steps, and then the final checkpoint is used till end of training. Observations: (a) Teacher shows phase transition in accuracy during training. 6 candidate checkpoints for (3, -progressee distillation have been marked, out of which 2 are selected in each setting. The checkpoint at 6MM lies outside the phase transition of the teacher. (b): We show the behavior of a few representative settings. Two main observations: (1) Selecting only a single checkpoint during the phase transition of the teacher is sub-optimal, as shown by plots that contain 6.M checkpoint as an intermediate checkpoint, (2) 2 checkpoints during the stage transition suffice train the student 100% accuracy, however the performance can heavily depend on their selection. Figure 17 shows that the teacher learns the low-level features at 4.5M checkpoint, making it crucial for distillation. (c): Even with extremely low temperature, the benefit of the phase transition checkpoint persists, suggesting that the monomial curriculum, not regularization, is the key to the success progressive distillation.
Experiment Setup: In this section, focus 8- way classification, where the data is generated by a tree of depth 3. Each feature in is given by a product of 5 variables. We keep the variables distinct in each feature, i.e., T 21 222 Us. T2 2637 Tin and so on.
35

---

Published as a conference paper at ICLR 2025
Teacher behavior
2M) po-grrrgrrssiee =1
(3. 10
1.0
1.0
1.0
Accuracy
Accuracy
Accuracy
31931-420 33993145M13
0.5
Candidates for [3, TJ-
0.5
3,5N
0.5
presfet
4.38-41-41 SNI
progressive
2,5HEL1 5:00000
3.54:21-4550
3
5
0
1.N421-4.0000
Training Steps(x10s)
Steps 1e6
Steps
1c6
Figure 19: Same experiments as Figure 18 for a width- 100 student.
Experiments and Observations: We conduct similar experiments as our sparse parity experiments. In Figure 15, we show that progressive distillation helps train smaller student as fast as the teacher, and even reach 100% accuracy.
Low-degree curriculum: We show the correlations of the teacher's logits for particular label and its relevant features in Figure 17. We observe similar spikes in the degree-1 monomials involving the support of the features. However, because there are multiple features defining a label class, with features at level being shared among multiple labels, the spikes appear in the degree- monomials of the features. As such, single teacher checkpoint won't t give information of entire support to student to learn from.
a difference in the time-frames which
Effectiveness of (3, prrogressivv distillation: We consider progressive distillation with 3 check- points, where the student only uses intermediate teacher checkpoint addition to the final one. We show in Figure 18 that there exists (3, prrogresse distillation that can help train student successfully. Furthermore, we demonstrate that these two intermediate checkpoints must be posi- tioned within the phase transition achieve 100% accuracy in training the student. This supports the hypothesis that low-degree curriculum crucial for progressive distillation since the correlations with degree-1 monomials are high only during the phase transition period. Additionally, distillation strategy with only a single intermediate checkpoint and the final checkpoint is insufficient for the student to achieve 100% accuracy, which aligns with our observation that degree-1 monomials for different features emerge at different steps. However, we also note that even within the phase transition, the optimal selection of the two checkpoints can significantly impact the student model's performance.
find that
36

---

Published as a conference paper at ICLR 2025
D EXTENSIVE STUDY ON PCFGS
D.1 A FORMAL DESCRIPTION OF PCFGS
We study progressive distillation using probabilistic context free grammar (PCFG). Compared to sparse parity and hierarchical data, PCFG is realistic proxy for natural languages and has been commonly used as a sandbox for mechanistically understanding the training of language models (Zhao et al., 2023; Allen-Zhu & Li, 2023b). A PCFG consists of set of non-terminals (NTs) and grammar rules involving the non con-terminals that specify the generation process of a sentence. For example, for the sentence The cat ran away, the grammatical structure dictates words the, cat, ran, away as dete rmi nant, noun, verb, and adverb. ran and away together represent a verb phrase, and the, cat together represent a noun phrase (see Figure 3). For a language model to generate grammatically correct sentences, it needs to learn the underlying grammatical rules.
probabilistic context-free grammar (PCFG) is defined as 4-tuple G (W, c, R, P), where
N is the of non-terminals, which can be considered as internal nodes of parse tree. There is special non-terminal S, known as the start symbol.
[u] the set of all possible words, corresponding to parse tree leaves.
R denotes a set of rules. For all A, B, C e N,, there is rule A BC' in R. Furthermore, there [u].
are rules A + for all A E
P specifies the probability of each rule to be used in the generation process. For rule r E R. if P(r) = then the rule is an invalid nule under the generation process. Furthermore, for each non-terminal A E N, on all rules r E R of the form A - L P(r) = 1. We denote R(A) as the set of all non-zero rules from A.
A concrete example of PCFGs is to model grammars of natural languages (Jurafsky, 2000). In this case, language tokens form the vocabulary of PCFG, while parts of speech such nouns, verbs noun phrases, verb phrases form the non-terminals. Rules like noun phrases being composed of determinant and noun form the core of such PCFG, while the probability of each rule is determined by their occurrences across sentences in the language.
Data generation from PCFG Given a PCFG g = (W, u, fashion as follows: we start with $1 ROOT at step and maintain string s, ([o] U WN) at step f. At step f, if all characters in belong to [u]. the generation process terminates, and is the resulting string. Otherwise, for each character
P). U.N,RRP))
string is generated in recursive
st, E N, we sample a rule r E R of the form A with probability P(r) and replace A by characters given by r(A).
Tracking n-grams As outlined in Section 4, we track the behavior of trained models by measuring the behavior of their output on the neighboring n-gram context. the context PCFGs and masked language modeling for BERT, Zhao et al. (2023) theoretically demonstrate that one of the optimal algorithms for predicting masked tokens is dynamic programming algorithm based on the inside-outside algorithm (textbook reference: Jurafsky (2000)). This algorithm computes "inside probabilities" for spans of tokens of various lengths, representing pairwise token dependencies within those spans. For example, in the setting of Figure the inside probability for the span "The cat" indicates the likelihood that these two tokens co-occur. The dynamic programming approach calculates these inside probabilities hierarchically, with smaller spans forming the basis for larger spans. The model's s performance ultimately depends on how accurately it represents span probabilities across different lengths. For instance, the token "cat" is masked in the sentence "The cat ran away". the success of the model depends on the representation of the likelihood of the spans "The cat", "cat ran", "The cat ran", and "The cat ran away". We denote the neighboring tokens in the n-gram window span of token as its 7--gram context.
D.2 VARIANTS OF PROGRESSIVE DISTILLATION
Comparisons at different lengths We follow common practices for training self-attention models for both one-shot distillation and progressive distillation. We use Adam optimizer (Kingma & Ba, 2014), 512 batch size training (to imitate large batch training), and cosine learning rate schedule
37

---

Published as a conference paper at ICLR 2025
(Loshchilov & Hutter, 2016) which is generally used train large language models. As cosine learning rate depends on the total training horizon, in order to show that progressive distillation converges faster than one-shot distillation, we compare the two algorithms by varying the number of training samples for the student. That is, we train the teacher model with 4 x 10 training samples (equal to 8000 steps), and compare the two algorithms for student model {1,2,4, 4, 10 training samples (equal to {2000, 4000, 8000, 16000} steps).
Progressive Distillation choices Because are considering comparisons different training lengths for the student, we have to consider a more general version of progressive distillation introduced in Definition 2.2. In Definition 2.2, progressive distillation defined by two parameters, (a) number of teacher checkpoints (N) for supervision, and (b) training steps per checkpoint. We define our selection criteria for the N checkpoints later. However, after selecting the N checkpoints, we have the following two variants of progressive distillation.
1. N-shot Equal-split distillation: Here, we simply split the entire student's training length into equal intervals, where the student is supervised by the ith teacher checkpoint interval E [N].
2. --shot Kil--Equa--plit distillation: Here K E (0, 1], and To refers to the total training length of the teacher. The idea is to decide the allocation on the basis of the training length of the teacher, instead of the training length for the student. We train the student under the supervision of each checkpoint for <Ta/N training steps. Teacher checkpoints that fail to fit into the student's s supervision schedule are ignored (corresponding to large k), and the final checkpoint is kept till the end of training if the student is trained for longer than kTop We can view as the amount of "speed up"; ; for instance, with i - 0. Our experiments (Appendix D.2) suggest that x = 1/2 is a reasonable rule of thumb that can help the student learn faster than the teacher at any given training length.
recover one-shot distillation
In the main paper, in Figures 5 and have reported performance PCFG and Wikipedia for N-shot To- To-Equal-split distillation as progressive distillation. We conduct more ablation studies on k in Appendix D.2. We keep the exploration of optimal strategies of progressive distillation future work.
Selection criteria for N teacher checkpoints: While there are multiple ways in which one can pick the reference checkpoints to train the student model, we use simple strategy which is sufficient to demonstrate the benefit of progressive distillation. Similar to our observation of transition phase for parity in Section 3, we search for transition phases in the loss behavior of the teacher and select the first teacher checkpoint roughly the middle of the transition phase. The rest are picked at multiples of this initial checkpoint.
D.3 DETAILS ON NON-TERMINAL PREDICTION WITH MULTI-HEAD LINEAR PROBING
Following Allen-Zhu Li (2023b), we train a position-based linear attention on the model's 's em- beddings to predict the non-terminals each level of underlying PCFG. We consider set of linear functions f. Rd +> RW,, where [H] and H is the number of "heads" in the linear attention model. If e1. ... e denote the model's output embeddings for a sequence Φ1, L. then the prediction of the model at each index i E [z]] is given by
G,(x)= = Σ
Wri-sff(ee).
ME(L)
expl (P Pk.)))
Wrii = EveL expp (Prrr, Pw.r))
for trainable parameters Psr E Rdd We train the parameters with logistic regression on 51200 examples and test on a validation set of 1024 examples.
D.4 DETAILS ON THE SYNTHETIC PCFGS
We use 5 synthetic PCFGs considered by Allen-Zi Shen-Zhuu & Li (2023a) (please see Figure 20 for the rules involved in the PCFGs). These PCFGs differ in difficulty, based the number of rules
38

---

Published as a conference paper at ICLR 2025
per non-terminal and the ambiguities in the rules per non-terminal. Under PCFG, each string is generated by generation trees of depth 7. We give differences in the PCFGs, as outlined by Allen-Zhi & Li (2023a) below.
In cfg3b, the PCFG is constructed such that the degree [R(A)] = 2 for every non-terminal A. In any generation rule, consecutive pairs of symbols on the generated symbols are distinct. The 25%, 50%, 75%, and 95% percentile string lengths generated by the PCFG are 251, 278, 308, 342 respectively.
In cfg3i, [R(A)] = 2 for every non-terminal A. However, the consecutive pairs of symbols needn't 't be distinct in generation rules. he 25%, 50%, 75%, and 95% percentile string lengths generated by the PCFG are 276, 307, 340, 386 respectively.
efg3h, [R(A)] E {2 for every non-terminal A. he 25%, 50%, 75%, and 95% percentile string lengths generated by the PCFG are 202, 238, 270, 300 respectively.
In cfg3g, [R(A)] 3 for every non-terminal A. he 25%, 50%, 75%, and 95% percentile string lengths generated by the PCFG are 212, 258, 294, 341 respectively.
In efg3f, (R()) E f3, 4] for every non-terminal A. he 25%, 50%, 75%, and 95% percentile string lengths generated by the PCFG are 191, 247, 302, 364 respectively.
220 270
221 (1000022011
22] (1823
273 21
sample from cfg3b:
131 20
121
221-3020
-1-33
221-13312
3133129000 121 3313 220022000111 HETELLU
221-118
221-1323
121-11111
16l-415:15
140-11 -14 15
721 .18
SPEP 1191 10 101 1o 10111111111111
2PP11171111110 31377 798991991
173-444 23:05
271-05-35
111--42-2 54
173-415-333
sample from cfg3i:
180-14.222
121-433 15-33
31 11 2012 22102122300 190001 1111212
20015
171
-1213221 10
20-22
14- 221-22224
221 2211
113-879
50100 71
13
L3:33
sample from cfg3h:
1710897
191
221-870
71n212
241--2221 13
121-013 21
STEE 1w1.031113331110
THE 2121111111 www.
11-135
51
141-1012
cfg3b
41-053
cfg3i
01-23
221087
341-113 21
31313000 11333111 3101
11-23
91-02
31-232
15--1031133
sample from cfg3g:
R1-12
131-07
121-1333311 13
cfg3h
2011 2211 1122231 223312
11-123
12--03
1122
11:e85
31112201 15220223131332111 www. aw aes.aaaaaaaaa......... 1a
7ix23
1200337
210232511112
51.75
a sample from cfg3f:
cfg3g
0i-331
31-223
51-25
11:537
3271991992
191911 1.2 201000001110000
20111129 2022
10
531 3133120
2011 211111111111111111111111
cfg3f
Bi-233
16-125
2.11382617228 12232222222222 38221 12121 2012212212222181222222 121132222222
31-13
Figure 20: The synthetic PCFGs considered from Allen-Zhu & Li (2023b). Vocabulary is 21, in each setting. More details on the differences between the PCFGs are in Appendix D.4.
D.5 EXTENSIVE EXPERIMENTS ON BERT
We first give some details on the architecture of BERT and its pre-training loss function.
D.5.1 A PRIMER ON BERT
BERT (Devlin et al., 2018) is an encoder-only transformer that is trained with masked language modeling (MLM) (Figure 21). In encoder-only architecture, the contextual information are shared across the tokens using bidirectional self-attention layers. During pre-training the model is trained with ML loss, that perturbs certain fraction of the tokens in the input random and the model is
39

---

Published as a conference paper at ICLR 2025
trained to predict the original tokens positions of the perturbed tokens. The pre-training recipe follows a 80- 10-10 principle, where tokens at 80% of the perturbed positions are replaced by a special replaced by random tokens from
(mask) token, while tokens at 10% of the perturbed positions the vocabulary, while remaining positions are filled with the original tokens themselves. We stick to this principle, while creating data for training from different PCFGs.
Model architecture considered: We train depth-4 BERT models with {8, 16, 32} attention heads. each of which operates on 8 dimensions, using a 30% masking rate. The head dimension is fixed to 8, with the corresponding width of the 4 models being (64, 128, 256) respectively.
ran 0.4 danced 0.1
cried 0.01
Bi-directional attention
BERT
The
<mask>
away
Figure 21: An informal representation of BERT (Devlin et al., 2018). The model uses bidirectional attention layers share contextual information across the tokens. During pre-training, few of input to- kens are replaced by special < mask > tokens, and the model is trained to predict the masked tokens.
D.5.2 DATA GENERATIONS
Data for masked language modeling: We generate 8 x 106 random sequences for each PCFG. We follow Devlin et al. (2018) to create masked input sequences and output labels, i.e. for each sampled sequence we mask 5% of tokens for input and the labels are given by the tokens in the masked positions of the original sequence. We also follow the 80-10-10 10 principle, where for input, the tokens in 80% of the masked positions are represented by special mask token [mask], while 10% of the masked positions represented by randomly sampled token from the vocabulary and the remaining 10% are represented by tokens from the original sequence.
D.6 HYPERPARAMETER DETAILS
We use batch size of 512 in each setting. We Adam (Kingma & Ba, 2014) optimizer with 0 weight decay, 91, 32 = (0.9, 0.95). We use cosine decay learning rate. We extensively tune the learning rate in the grid (10-2 7.5 10-3, 3,3,5 x 10-3 2.5 10-3 10-3) 3) in each setting. We train the teacher on 4 x 10 training samples (equal to x 10 steps).
Distillation experiments at different training horizons: To thoroughly compare the sample evaluate both algorithms using
complexity requirements one-shot and progressive distillation, smaller student model across various training sample sizes. The smaller student is trained with 21,2,4 4, 8} x 10 training samples (equal (2 103, 4 10 8 x 10 16 103} training steps) and the performance is compared in each horizon. For example, Figure (right) plot contains 4 distinct points for each method which represents the performance of smaller model under the different training steps (sample sizes).
Training split for (2, -roggeeross distillation for PCFGs: We report the performance in Figure for 4000 training steps. We find the best training time split T between the intermediate checkpoint and the final checkpoint in the grid (500, 1000, 15000, 2000), i.e. the student trained
40

---

Published as a conference paper at ICLR 2025
Meuse for n-grams
prediction different levels
100
Median (over
-0 -1
Accuracy
4
2
3
10-3
60
0.6
2
4 8
0.6 1
2
4
Steps of Training(x10") (log-scale)
Steps of Training(x10") (log-scale)
Figure 22: We conduct additional probing experiments the teacher's (4 layer, 32 attention head BERT) logits during training to indicate curriculum learning. (left) TV distance between model's predictions with full context and context with only 7--rrmm tokens (Melsse). We observe that the teacher's s logits get closer to higher n-gram context predictions, and the inflection appears at the middle of the second phase (our first selected checkpoint for progressive distillation) (right) Performance of linear classifier probe on teacher's 's intermediate checkpoints to predict the non- tertenemills at different levels of the PCFG generation tree. We observe that the probe's 's performance is 95% of the final probe performance by the middle of the second phase, indicating the model has almost learned the underlying PCFG features by this time.
with the logits of the first intermediate teacher checkpoint till step T and then the teacher is switched to the final teacher checkpoint.
Low-temperature distillation: We focus distillation with small temperature of 10-4 4 (in Equation (1)), for the following reasons. First, as discussed in Section 3, it removes any potential regularization effects induced by soft labels. Moreover, using such small temperature corresponds to training with the top-1 predictions of the teacher model, which is more memory-efficient compared
training with the full teacher logits, especially when the vocabulary size is large.
D.7 ADDITIONAL CURRICULUM PROBING ON THE TEACHER'S S CHECKPOINTS
In this section, we study the performance of different progressive distillation variants and compare them to one-shot distillation. As per our experiments in Figure 4, we use the 8 teacher checkpoints selected for supervision. In Figure 24, we compare one-shot distillation the two variants of progressive distillation, i.e. 8-shot Equal-split and 8-shot To -Equal-split distillation. We observe that both variants of progressive distillation help the student learn faster than one-shot distillation, and the gap diminishes the students are trained for longer. The optimal strategy for progressive distillation depends on the training budget for the student. For training steps lower than the teacher's Tu budget, To Eqqul--plit distillation slightly performs better than Tn-Equal-split distillation, which changes as 2
2
To. we train longer. To keep things simple, we focus on 2 -E-qual-split progressive distillation in all of our subsequent experiments.
D.8 ABLATIONS WITH HYPERPARAMETERS
Ablation with temperature Here, we compare progressive distillation and one-shot distillation at temperature and temperature 10-4 (representing hard label supervision) (Figure 25). We observe that progressive distillation at temperature 10 performs better than one-shot distillation both temperatures. However, progressive distillation at temperature can perform worse than one-shot distillation for stronger student. We keep explorations the effect of temperature on the algorithms as future work.
Ablation with mask rate In Figure 26, compare progressive distillation with one-shot distilla- tion at different masking rates. We observe that at all masking rates, progressive distillation performs better than one-shot distillation.
41

---

Published as a conference paper at ICLR 2025
1.0
Validation Loss
1.0
Top-1 Accuracy
4
0.8
8
16
0.5
16
0.6
32
32
0.0
0.4
100
1000
8000
100
1000
8000
Training Steps
Training Steps
Figure 23: Comparison of BERT's training behavior on cfg 3b with varying numbers of attention heads (where the embedding dimension scales linearly with the number of attention heads) over 8 x 10 training steps. The x-axis represents the number of training steps and is in log scale. Larger BERT models show an earlier and more pronounced drop in loss/increase in accuracy compared to smaller models. For reference, each training curve is annotated at the point where the model reaches 80% its performance at the final step.
0.5
95
To Equal-split progressive
0.4
Equal-split progressive
90
0.3
Equal-spiit progressive
85
To Equal-spiit progressive
Top-1
Is Equal-spiit progressive
0.2
80
Equal-split progressive
5000
Number Training 15000
10000
5000
Number Training 15000
10000
Accuracy
95
0.4
To Equal-split progressive
Equal-split progressive
90
Logg
0.3
Equal-spiit progressive
To Equal spiit progressive
85
1 Equal-split progressive Equal-split progressive
80
0.1
5000
10000 15000
5000
10000
15000
Number of Training Steps
Number of Training Steps
Figure 24: Experiments on BERT (Left to right/top to bottom): (a), (b) show the comparisons for an 8-attention head student, (c), (d) show the comparisons for 16-attention head student. We observe differences between the different variants of progressive distillation at different training steps. For training steps lower than the teacher's (marked by To), To/ /2-Equal-split progressive distillation is better, implying that for shorter training, shouldn't i try to fit all the teacher's checkpoints. The trend reverses as training sample budget approaches To and beyond.
42

---

Published as a conference paper at ICLR 2025
Accuracy
90
Top-1 1p-1AA Accuracy
one [t= 1.0)
one-shot (t= 1.0)
one-shot (r 10
one-shot = 10-
80
Progressive (T = 1.0)
Progressive (T 1.0)
Progressive (T 10-10-1
Progressive (r (0=10-1 10
70 2000
8000
70
2000
Number Training Steps
Number of Training Steps
(a) Number of heads=8
(b) Number of heads=16
Figure 25: The experiments above compare progressive distillation and one-shot distillation at temperature 1 and 10-4 (representing hard label supervision) for PCFGs cfg3b masking rate 30% using a BERT model with 8 attention heads (left)/ 16 attention heads (right), per head dimension and 4 layers. We observe that progressive distillation with hard labels performs better than one-shot distillation at temperatures and 10-4 However, progressive distillation at temperature 1 can perform worse than one-shof distillation for stronger student. We keep explorations on the effect of temperature the algorithms as future work. Here, as progressive distillation, where To = 8000 is the total number steps used for teacher training.
2 Te-Equal-split progressive distillation
Steps = 2000
Steps = 4000
Steps = 8000
Top-1 Accuracy
100
4.8%
2.0%
80
7.3%
11.0%
4.5%
1.2%
one-shot Progressive
60
0.15
0.3
0.5
0.15
0.3
0.5
0.15
0.3
0.5
Mask rate
Figure 26: The experiments above compare progressive distillation and one-shot distillation for PCFGs cfg3b at different masking rates using BERT model with 8 attention heads, per head dimension 8, and 4 layers. The relative gap between the performance progressive distillation and one-shot distillation have been reported on the bar plots. We observe that progressive distillation performs better than one-shot distillation at all masking rates, with the gap diminishing with the number of training steps. Here, we use - -e-Equal-split progressive distillation as progressive distillation, where To 8000 is the total number steps used for teacher training.
Ta
43

---

brief introduction into GPT models: GPT models are trained with the auto-regressive loss. The teacher and student models operate on sequences of input domain fr x g and fs 2 - C where the input sequence length h can be arbitrary. Denote the length-h input Published as a conference paper at ICLR 2025
Ablation with difficulty of PCFG In Figure 27, compare progressive distillation with one-shot distillation with increasing difficulty of the underlying PCFG. The benefit of progressive distillation over one-shot distillation is influenced by the model's capacity and the specific PCFG being trained.
Steps
2000
Steps 4000
Steps 2,0%
8000
Top-1 Accuracy
7.8%
5.2%
0.9%
80
1.7%
one-shot
Progressive
60
cfg3b cfg3i cfg3h
cfg3b cfg3i cfg3h
cfg3b cfg3i cfg3h
PCFG
(a) Number of attention heads=8
Steps
2000
4000
Steps
8000
Top-1 Accuracy
100
4.0%
0.4%
1.0%
80
2.7%
1.0%
0.1%
one-shot Progressive
60
cfg3b cfg3i cfg3h
cfg3b
cfg3i cfg3h PCFG
cfg3b cfg3i cfg3h
(b) Number of attention heads=16
The experiments above compare progressive distillation and one-shot distillation for PCFGs cfg3b, cfg3h, and o ofg3i at masking rate 30% using BERT models with 8/16 attention heads, per head dimension 8, and layers. The relative gap between the performance of progressive distillation and one-shot distillation have been reported on the bar plots. The benefit of progressive distillation over one-shot distillation is influenced by the model's capacity and the specific PCFG being trained. For instance, on cfg3i, the student model can only achieve top-1 accuracy of 75%. Progressive distillation reaches this within 2000 steps but fails improve further, resulting in minimal gains over one-shot distillation when compared with cfg3b. The comparisons are To E-Equll-Sppit Progressive Distillation as Progressive at temperature = 10 Here, we use 2 Distillation, where To 8000 is the total number steps used for teacher training. Our teacher is BERT model with 32 attention heads, per head dimension 8, and 4 layers, which doesn't train on cfg3g and cfg3f. f, henee we don't report the performance of the student on these PCFGs.
Figure 27:
E AUTOREGRESSIVE TRAINING WITH GPT2
Setting: Similar to experiments on BERT, we train GPT2 models of depth 4 with {8, 16, 32} attention heads, while keeping the dimension per attention head fixed
sequence as : [..,...,,. ,.,). and denote X,y as the subsequence [...... .ry] (i.e. the indexing inclusive on both ends). The cross entropy loss for next-token prediction training on x is given by
1 xl(ex.)ls(xx==))).
h
where Ex; denotes a one-hot vector with in æ,th coordinate. We take a different approach, where we compare the algorithms at different difficult levels, by training subset of tokens in each sequence. The subsets that we consider are the boundary tokens at different levels of PCFG generation (recall Figure 3).
44

---

Published as a conference paper at ICLR 2025
Formally, if (z) represents the set of level-f boundary tokens, then we define the cross entropy loss and the distillation loss corresponding to boundary tokens at any level the PCFG as
(x:
Σ
KL(e,, (xxxx-1)s((
(9)
(()
(
10 xx Js.
(10)
(c()(z)) (c))
There are few remarks that need to be made about the above loss function. First, note that the subsets satisfy the condition (ar) S (((2) for all & 2 62- Hence, the loss [,(2) includes loss ((4) for all f 2 le and losses E (e, fDL). Second, [ will average the losses at all tokens, which is the standard auto-regressive loss used in practice train large language models.
We focus on 3 f that has 6 levels in the generation process, and we report the behavior of the models when trained with losses (2), ((3), (1))),1 (4),, with L E (DL). We focus on Ta --Euul--splt progressive distillation.
2
Definitions for Mobust and Melse- Similar experiments BERT, track the change in the model's predictions with and without the n-gram context tokens. However, as the model is trained autoregressively, need to change definitions of Myobust and Melose from Equations (5) and (6), as well as the definition of n-grams.
For h length sentence x and for i E [b], we define the n-gram neighboring context around the den token as the set of tokens at positions within n - distance to the left from i, i.e. the set { for i
For Meosse on teacher fr and ngram length n, we measure the TV distance between the model's probability distributions of the model any position when all the tokens at positions 1, 2....... are available, and when only the tokens in the neighboring n-gram context window are available (i.e. 1)
positions n n 2, Melos (fT, i, P
(11)
For Mrouust on teacher fr and an n-gram length n, we measure the total variation (TV) distance between the model's probability distributions any position i, considering two scenarios: where all tokens at positions 1, ...... 1 are available, and another where the tokens within the n-gram masked. However, since the attention mechanism in GPT requires token at
context window position i - 1 before it can predict 2; and we don't have a special token to replace the masked tokens, we cannot remove that specific token from the context. Therefore, keep the token at position 1 intact while masking the other tokens within the n-gram context window. We refer to this modified approach as "skip n-gram."
Mrodust ( (ST, i
TV (x(a)) (-+1a))))
(12)
E. OBSERVATIONS
Teacher's behavior during training Figure 29 shows the loss behavior of a teacher run. We observe 2 distinct phases of training: rapid loss drop phase in the first 10% of training, and final phase of slow loss drop till end of training. In Figure 28, we compare the training accuracy behavior across models of different sizes. At log scale, we observe very small dormant phase in the training behavior at the start of training. Larger models transition to the rapid loss drop phase faster than smaller models and also show more prominent change in this phase.
Teacher's checkpoint selection for progressive distillation As outlined in the previous section, we select the first supervision checkpoint at roughly the middle of the first phase (1/20,, fraction of training), and the other checkpoints are selected at (i/20) 2h fractions of training.
Similar inflection points loss as BERT and an implicit curriculum:
others simply masked out during attention score computation avoid shifts in position embeddings.
45

---

Published as a conference paper at ICLR 2025
Top-1 Accuracy
0.8
0.8
0.8
17
87h
0.7
8
16
0.6
0.6
0.6
100
1000
8000
100
1000
8000
100
1000
8000
Steps Training
Figure 28: (left to right) Models are trained with the cross entropy loss e(4), g(2) respectively. Here, we compare GPT's training behavior with cross entropy loss cfg3f with varying numbers of attention heads (where the embedding dimension scales linearly with the number of attention heads) over x 10x training steps. Larger models show an earlier and more pronounced increase in performance compared to smaller models. For reference, each training curve is annotated at the point where the model reaches 80% of its performance at the final step.
Median over ((x. i))
Melaee for n-grams
Median (over (x, i))
Mrobust for skip n-grams
Eval Loss
0.6
10-
0.4
10-2
10-2
0.2
10-3
0.010=
10a
10a Steps of Training
10 Fraction of Training
Fraction of Training
(b) Median of Maww (JT, x, (c) Median of Motau(JT x, i,
over x, for different n cfg3f with cross entropy loss: g(3), We observe two distinct phases; (2) rapid drop in loss phase, and (3) slow drop loss till end of training. The rapid loss drop phase signifies a transition phase for the model, similar to one we observed for hierarchical boolean data (Section 3). All selected checkpoints for progressive distillation are marked by triangles. The first teacher checkpoint is roughly picked at the center the second phase. The rest of the checkpoints are picked at training steps that are multiples of the first one. (b) and (c) show inflection points in the teacher's predictions with full context and with/without n-gram contexts the selected checkpoint.
(a) Loss behavior of teacher model over x, for different n
Figure 29: Experiments on GPT: Behavior of teacher model when trained
We observe inflection points in the model's behaviors at the first selected checkpoint. Similar to our observations on BERT, we observe a curriculum on the reliance of the model's predictions on 3-gram predictions (Figure 29). Hence, we check whether progressive distillation can help train smaller model faser.
(R7) Progressive distillation helps train smaller model faster In Figures 30 and 31, compare one-shot distillation To -Equal--spit progressive distillation. We observe that progressive distillation help the student learn faster than one-shot distillation, and the gap diminishes as the students are trained for longer. However, the gap between progressive distillation and distillation decreases as more tokens are involved in the loss function i.c. the gap is smaller for loss L compared to loss L (4) We conjecture that auto-regressive training with all tokens involved provides strong curriculum for the model to learn the structure of the language. We keep a thorough study of this analysis to future work.
Does longer training lead to degradation in performance for progressive distillation?
One primary concern is whether progressive distillation can lead to degradation in the student model's performance in longer steps. We do not observe any degradation in the performance of the progressively distilled models in Table which shows that progressive distillation stays useful in the long run.
We extended all our GPT-2 experiments in Figures 30 and 31 16,000
To further evaluate the learned representations, we measured the performance of a linear classifier when trained top of the output embeddings the non-terminal prediction task (Definition 4.3).
46

---

Published as a conference paper at ICLR 2025
Top-2 1. Accuracy
Top-1 Accuracy
90
Top-1 Accuracy
90
85
one-shot
85
one-shot
one-shot
80
progressive
progressive
progressive
2000
4000 6000
8000
2000 4000 6000
8000
2000
4000
6000
8000
Number of Training steps
Number training steps
Number of Training steps
(a) Loss: g(
(b) Loss: e(3)
(c) Loss: Eo
Figure 30: Experiments GPT (Left to right) for an 8- attention head model at different losses. Here, progressive distillation refers to To -Equal-split progressive distillation. We observe that progressive distillation outperforms one-shot distillation all training sample budgets, with the gap diminishing with increasing training sample budget. The gap between progressive distillation and distillation decreases as the number of tokens involved in loss function increases i.e. the gap
is smaller for loss 7(22 compared to loss (4)
92
Top Accuracy
Top Accuracy
Top-1 Accuracy
85
Name one-shot progressive
one-shot progressive
83
one-shot
progressive
2000
4000
6000
8000
2000
4000
6000
8000
2000
4000
6000
8000
Number Training steps
Number of Training steps
Number Training steps
(a) Loss: (
(b) Loss: el
(c) Loss: E
Figure 31: Experiments on GPT (Left to right) for a 16-attention head model at different losses. Here, progressive distillation refers to 2. -qquall-spl progressive distillation. We observe that progressive distillation outperforms one-shot distillation at all training sample budgets, with the gap diminishing with increasing training sample budget. The gap between progressive distillation and distillation decreases as the number of tokens involved in the loss function increases i.e. the gap is smaller for loss _ compared loss 10a
Steps = 500
Steps = 1000
Steps 2000
Top-1 Accuracy
3.4%
2.4%
0.4%
1.2%
0%
4.4%
-0 0%
-0,3%
1.3%
CE
one-shot
progressive
cfg3h cfg3g cfg3f PCFG
cfg3h cfg3g cfg3f PCFG
cfg3h cfg3g cfg PCFG
Figure 32: The experiments above compare progressive distillation and one-shot distillation for PCFGs cfg3h, efg3g, and ff on GPT models with 8 attention heads (each head having dimension of 8) and 4 layers. The models performance gap between progressive and one-shot distillation is presented in the bar plots. Notably, the advantage of progressive distillation over one-shot distillation depends on the specific PCFG being trained. For example, with cfg3f, the student model can achieve beyond 90% top-1 accuracy, and progressive distillation allows it to reach this more quickly. contrast, for cfg3g, the student model's 's top-1 1 accuracy plateaus at 84%, and after 500 steps, progressive distillation shows only marginal gains over one-shot distillation. All comparisons were made at a temperature - 10-4 Here, progressive distillation refers to Tu-Equal- progressive distillation, where To S000 denotes the total number of steps for teacher training. The teacher model is GPT with 32 attention heads, each with dimension and 4 layers.
trained using the distillation loss T The relative
47

---

Published as a conference paper at ICLR 2025
Table I: Progressive distillation remains beneficial during extended training. We train GPT-2 student models with 8 and 16 attention heads for longer. We evaluate their performance using top-1 accuracy on the next token prediction task from their respective pre- pre-training objectives. Our results show that progressively distilled models do not exhibit any performance degradation compared to those trained with one-shot distillation.
Number of attention
Training loss
Steps 8000
Steps 16000
heads of student model
Progressive
One-shot
Progressive
One-shot
D1.
(3)
94.3m..
93.4...
94.40.
93.1wa
8
p(4)
90.5...)
90.7
91.0pa
90.4m.
DL.
((5)
DD
85.1
84.9(an)
84.8
(3)
DL
95.1 aam
95.0ma
95.1
94.9ma
16
((4)
91.7(0.1)
91.7
91.8,a
91.71000
DL
(5)
85.3.0.0:
85.3
85.3
85.3
DL.
The classifier trained on embeddings from progressively distilled models performed competitively exceeded the performance of classifiers trained embeddings from one-shot distilled models (Table 2). These results confirm that progressive distillation does not compromise the long-term generalization of the learned embeddings.
Table 2: Progressive distillation leads to favorable linear probe results on non-terminal predic- tion (Definition 4.3), compared to one-shot distilled models.
Number attention heads of student model
Training loss
Steps 8000
Steps 16000
Progressive
One-shot
Progressive
One-shot
p(3)
79.40.4
78.8
79.5
78.3
DL
8
((4)
DD
79.6mm
79.6Gzz
80.2aan
79.3
((5)
80.8/o.4
80.7
80.2(02)
80.
(13)
81.0ka.n
80.4ka.)
80.9(22)
80.4:0.0
DL
16
((4)
81.9t..
81.8(a2)
81.9
81.7
DL.
((5)
82. (
82.1
82.2(a.)
82.1
DL
48

---

Published as a conference paper at ICLR 2025
Table 3: Performance comparison on downstream tasks. We pick the checkpoints from Figure that have been trained with 120 10 training steps on each individual method and fine-tume them further with Adam optimization on different downstream tasks. Models trained with progressive distillation outperform others by an average of 0.9 points across all downstream tasks.
Method
Sentiment
NLI
Topic Classification
Average
SST-2
IMDB
QNLI
MNLI
SNLI
AG-NEWS
Progressive
86.9
90.7
84.4
74,8
86.0
94.1 (0.1)
86.2
One-shot
86.3 10)))
90.4 (0.1)
82.6 10.1)
72.8 (0.3)
85.7 (0))
93.9 (0.11
85.3
CE
85.5
90.6
83.1 (0.))
73.5
85.3
94.0 (0.1)
85.3
Teacher
87.9 (0.0)
91.4 (0.2)
86.2 (0.3)
78.4 (0.3)
87.9 (02)
94.4 (0.)
87.7
F DETAILS ON WIKIPEDIA + BOOKS EXPERIMENTS
We the same hyperparameters for Adam training experiments BERT and PCFG in Appendix D.6. However, we fix the peak learning rate to 10-1 (Devlin et al., 2018) each case to minimize computation costs.
Fine-tuning on downstream tasks: We further report the performance of the trained student models (in Figure pick the checkpoints that have been trained with 120 10 training steps on each individual method) on various downstream tasks after fine-tuning in Table 3. We train the models with Adam optimizer with epochs and cosine learning rate, and vary the peak learning rate in 21,2,5 10-5. Our results indicate that models trained with progressive distillation outperform others by average of 0.9 points across all downstream tasks. This suggests that progressive distillation enables the student model to acquire generalizable features from the teacher during pre- training, which can also help improved generalization on different downstream tasks after fine-tuning.
49


Reviews:
[{"review_id": "gxwbHHXl48", "reviewer": "Reviewer_oRqc", "summary": "Knowledge distillation is a widely used technique for training smaller \"student\" models by leveraging the knowledge captured by larger, pre-trained \"teacher\" models. This paper focuses on progressive distillation, where a student model is trained on a series of intermediate checkpoints from the teacher's training process, as opposed to one-shot distillation, which relies solely on the final teacher checkpoint.\n\nThe authors provide both empirical and theoretical evidence that progressive distillation accelerates student learning and leads to better generalization compared to one-shot distillation. They attribute this improvement to an \"implicit curriculum\" embedded within the intermediate teacher checkpoints. This curriculum emphasizes easier-to-learn aspects of the task, facilitating the student's learning process.\nTwo main distillation strategies are compared: one-shot distillation and progressive distillation.  One-shot distillation involves training the student with a fixed, converged teacher model.  In contrast, progressive distillation utilizes multiple checkpoints from the teacher's training trajectory.  The authors also explore a variant of progressive distillation where only a single, carefully selected intermediate checkpoint is used.    \n\nThe authors delve deeper into the mechanisms behind progressive distillation by utilizing the sparse parity task, a well-established benchmark for studying feature learning in neural networks.  They employ teacher and student models with identical architectures but varying sizes.  For MLPs, the model width is adjusted, while for Transformers, the number of attention heads is modified.  Increasing either parameter effectively increases the number of \"parallel search queries\" the model can perform during training.    \n\nA key finding is that the intermediate teacher checkpoints implicitly provide a \"degree curriculum.\"  This means that the checkpoints guide the student model to learn features in a progressive manner, starting with simpler, lower-degree features. Remarkably, progressive distillation using just a single, well-chosen intermediate checkpoint can surpass the performance of one-shot distillation.  Furthermore, the authors demonstrate that progressive distillation reduces the sample complexity of learning the sparse parity task.    \n\n\nThe authors conclude by highlighting the effectiveness of progressive distillation in improving feature learning through the implicit curriculum present in intermediate teacher checkpoints.  They discuss the importance of teacher temperature as a hyperparameter in knowledge distillation and acknowledge limitations in their exploration of this aspect.  Further investigation into the precise role of temperature, particularly its impact on optimization, is proposed as a promising direction for future research.    \n\nAnother avenue for future work is extending progressive distillation to the setting where student models are trained on synthetic data generated by teacher models.  This approach has shown significant potential in recent studies. The authors identify key differences between their current work and generation-based methods, primarily in the type and quantity of supervision provided.  Bridging this gap and developing a unified framework for progressive distillation in various training scenarios is an important challenge for the field.", "questions": "No questions beyond suggestions above:", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality - The paper presents a novel perspective on progressive distillation by identifying and formalizing the concept of an \"implicit curriculum.\" While prior work has explored progressive distillation empirically, this study delves deeper into the underlying mechanisms and provides theoretical grounding for its efficacy. The connection between intermediate teacher checkpoints and an implicit curriculum is a fresh insight that contributes to a better understanding of knowledge distillation.    \n\nQuality - The research is technically sound, employing rigorous methodology and comprehensive experiments. The authors combine theoretical analysis with empirical validation, drawing on diverse tasks like sparse parity, PCFGs, and natural language modeling. The use of multiple progress measures to quantify the implicit curriculum further strengthens the quality of their analysis. The study is well-designed, and the results are convincing.    \n\nClarity -The paper is clearly written and well-organized. The authors present their ideas in a logical progression, starting with a clear motivation and gradually building up their analysis. The concepts are well-explained, and the figures effectively illustrate key findings. The paper is accessible to readers with a background in knowledge distillation and deep learning.\n\n\nSignificance - The findings have significant implications for the field of knowledge distillation. By elucidating the role of an implicit curriculum in progressive distillation, the study provides valuable insights for designing more effective distillation strategies. The theoretical results on sample complexity offer a deeper understanding of the optimization benefits of progressive distillation. The practical implications of this work are substantial, particularly for training efficient and capable models in resource-constrained settings.", "weaknesses": "Some possible improvements I can see are the following :-\n\na) more investigation of impact of temeprature on knowledge distillation. this seems to be a bit missing in the main sections of the paper\nb) analysis of how implicit curriculum learning varies across model layers, across datasets, trainign objecives etc\n\nc) exploring more tasks and architectures\n\nd ) explore interaction of optimization algorithms, batch size etc with curriculum learning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "td7I1pRYtP", "reviewer": "Reviewer_ktaG", "summary": "This paper provides theoretical and empirical analysis on how progressive distillation can be helpful to speed up knowledge distillation. Specifically, the author studies a toy use case in sparse parity with synthetic data to show both mathematically and empirically how using intermediate teacher checkpoints can assist student models to learn faster. Then they run progressive distillation experiments in PCFGs and BERT masked token prediction using Wikipedia and Books data to further verify their findings.", "questions": "1. For experiments in Figure 2, how do you explain the results that the best correlation (Middle) with in-support variable does not lead best accuracy (Left)?\n2. L406, what's the formula for total variation?\n3. How do we pick the most appropriate teacher checkpoints in other more complicated tasks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written, and despite the complexity of the narrative, it is generally easy to follow, and I enjoyed the reading.\n2. Though only applied to a simple use case, the mathematical analysis does provide useful insight about sample efficiency of progressive distillation.\n3. The metrics selected in the analysis such as $\\mathcal{M}_{robust}$ is quite useful to understand the feature learning aspect of the method.\n4. The authors run experiments in various settings including three datasets, MLPs, BERT and GPT-2 (in the appendix) to show the gains can be generalized.", "weaknesses": "Despite the strength, I think the paper can be improved.\n1. I understand the necessity to use a toy use case (sparse parity) to show rigorous mathematical analysis, but the following experiments can be more practical in order to provide stronger empirical evidence of the effectiveness of progressive distillation.\n- Instead of masked token prediction, can run experiments in challenging NLP tasks such as QA, summarization and long-form generation.\n- Can also experiment with more recent LLMs - GPT-2 in the appendix is a good start, but it is still a very outdated model.\n2. Some writing (particularly Theorem 3.2) can be made even clearer. For example, L291-292, I didn't quite follow where higher degree odd polynomials come from. L296-297 \"This gap allows .. to grow quickly with only ... samples.\" This statement isn't clear if I don't read the entire proof in the appendix. Please consider writing it in a more intuitive way in the main text.\n3. Probably unfair as this is more of an analysis paper, but the overall contribution appears to be limited considering the scope of its experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kO7dJT4J19", "reviewer": "Reviewer_qfzz", "summary": "The paper demonstrates how progressive distillation helps in better training of the student models by inducing an implicit curriculum. The experimental results demonstrate that progressive distillation results in a higher sample efficiency in all cases as well as a higher performance in some cases, as compared to vanilla knowledge distillation (referred to as \"one-shot distillation\" in the paper). The results also show that progressive distillation induces an implicit curriculum wherein the intermediate checkpoints provide a stronger learning signal and act as \"stepping stones\" for the student models during the learning process. These results are validated by experimenting across two different model architectures and 3 different tasks. In the case of the sparse parity task, authors also provide a theoretical proof of how progressive distillation with 2 checkpoints (one intermediate and one final) leads to a better sample complexity as compared to one-shot distillation.", "questions": "* How were the sizes of student models decided? Can the authors show some results on one task (preferably the natural language modeling OR the PCFG tasks) for what happens when the sizes of student models are varied while the size of the teacher size is kept constant?\n* Can the authors run similar experiments on a simple multiclass classification task in the visual domain? Even results on something simple such as multiclass classification in CIFAR-100 with an ImageNet pretrained ResNet-18 finetuned on CIFAR-100 as the teacher, and a smaller CNN as the student would be interesting to see.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The paper is well-written and easy to read.\n* The paper includes results on tasks across different complexity levels - going from a toy setting of sparse parity to PCFGs and then to a non-synthetic task of natural language modeling.\n* Authors also run experiments across multiple model architectures, name MLPs and transformers of different sizes.\n* The induced curriculum is discussed from a human interpretability point of view (i.e. showing the correlation between degree 1 monomials and the logits of the intermediate teacher checkpoint in the sparse parity task, and drawing similar analogy in the PCFG task).\n* The paper (more specifically; the appendix) includes further extensive experimentation on the settings discussed in the main paper.", "weaknesses": "* There is a typo in Definition 4.3: I believe it should be \"boundary of span(n^{(i)})\" instead of boundary of n^{(i)}\n* Discussion about how the relative sizes of teacher and student models were decided is missing. It would be interesting to see a study of how the performance is affected w.r.t the size of the student models\n* Empirical analysis on tasks in the vision domain and with other model architectures such as CNNs and recurrent networks would strengthen the paper significantly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BuXrPiMRIV", "reviewer": "Reviewer_9DA8", "summary": "The paper investigates how progressive distillation can accelerate the training of neural networks by leveraging an implicit curriculum provided in the form of intermediate teacher checkpoints. The authors conduct experiments on tasks like sparse parity, probabilistic context-free grammars (PCFGs), and real-world datasets (\"Wikipedia and Books\") using models like MLPs, Transformers, and BERT. The main claimed findings are:\n\n1. Progressive distillation accelerates student learning compared to one-shot distillation.\n2. An implicit curriculum emerges from intermediate teacher checkpoints, guiding the student through increasingly complex subtasks.\n3. This implicit curriculum results in an observable empirical acceleration, and stems from provable benefits.", "questions": "The only perhaps unanswered question in this paper is: are simpler features learned first because they're easier, or because they're sufficient for initial loss reduction? Does the progression represent increasing task difficulty or natural optimization paths as the student discovers incrementally winning subnetworks?\n\nAs discussed above, it would be interesting to at least provide an analysis of how progressive distillation differs from or extends the Lottery Ticket Hypothesis (and the variants of progressive stacking/learning subnetworks etc). Would be (may be) interesting to discuss whether there could be experiments to show if the observed acceleration is due to narrowing the search space via subnetworks and how this impacts the interpretation of an implicit curriculum.\n\n**Future experiments to isolate effects/figure out the right interpretations**\n\nBelow are just some ideas, probably mostly for future work. \n\nIt would be interesting to evaluate (in the experiments of Figure2), the corresponding sizes of the winning subnetworks as per the lottery ticket hypothesis. \n\nOne could also analyze the sparsity patterns of successful networks at different checkpoints, compare with explicit lottery ticket pruning approaches, then try to isolate whether the benefits come from parameter space guidance or feature complexity. \n\n**Quick check of the literature**   \nThe following is (slightly) related from a quick search, mainly for your interest:\nhttps://arxiv.org/pdf/2402.05913 too recent of course. \nOlder work on progressive stacking: https://proceedings.mlr.press/v202/j-reddi23a.html  also only slightly related\nBut more generally, it might be useful to reprobe the related work, with the lens of using winning subnetworks to accelearate learning, perhaps checking how they form and how they help accelerating learning (whether there's an increase in task complexity ...)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "I liked the spirit of this paper, which I found complete, well-written, carefully designed and executed. Overall, I enjoyed the following strengths:\n\n- **Experimental completeness:** The paper is generous in providing extensive empirical evidence across synthetic tasks (sparse parity, PCFGs) and real-world datasets (Wikipedia, Books).\n\n- **Theoretical Analysis:** Offers mathematical proofs demonstrating sample complexity benefits in specific cases like sparse parity learning. \n\n- **Memorable Observations:** Beyond the idea that checkpointed learning leads to faster optimization, it identifies phase transition in the teacher's training where intermediate checkpoints can be most beneficial to the student. These correspond to acquring new skills of increasing complexity.\n\n- **Impact and Practical and actionable implications**", "weaknesses": "- **Alternative intepretations, e.g. winning subnetworks and the Lottery Ticket Hypothesis:**\n\nThe first thing I thought about, as a possible intepretation of the empirical (and theoretical) findings of the paper is the lottery ticket hypothesis (LTH), which could alternatively explain the benefits observed in progressive distillation. The hypothesis posits that within a large neural network, there exist smaller subnetworks (winning tickets) that can learn the tasks efficiently. Searching for these subnetworks takes long, but once they are identified, training only on them makes the learning faster. \n\nSo while the paper frames the intermediate checkpoints as providing an implicit curriculum of increasing task complexity, this could be reinterpreted as the teacher progressively revealing parts of the winning subnetwork, narrowing the student's search space at each step. Hence a crucial question is: could the novelty in using intermediate checkpoints be just an operationalization of the lottery ticket hypothesis in a distillation context (i.e. Guided Search) rather than a fundamentally new concept (implicit curriculum)? \nThere, the student model is being guided to explore progressively smaller regions of the solution space, as per LTH. These correspond to learning features of increasing complexity as the paper points out, but that's because that's how networks progress in their learning of the best parameters to solve a problem. By emphasizing the curriculum aspect, the paper might divert the \"reader\" from other factors contributing to accelerated learning, such as the inherent properties of the optimization landscape or the effects of network pruning (implicit or explicit).\n\nAnd both interpretations would be aligned with the empirical and theoretical findings. \n\nFirst, the fact that not all checkpoints provide good performance could be  due to the model learning more and more complex skills as it discovers the winning tickets, and not because the tasks are more and more complex (curriculum). Matter of fact, the training data is not going from simple tasks to complex ones, what goes from simple to complex is what the model has learned, not what is was trained on. Said differently, the discovery of progressively complex features could be the result/consequence of a guided search to efficient parameter configurations.  \n\nSecond, looking at the sample complexity improvement that the authors prove:\n\"the total sample complexity needed for the student to reach ϵ-loss using progressive distillation with 2 checkpoints is Õ(2^k d^2 ϵ^−2 + k^3). However, one-shot distillation requires at least Ω(d^k−1, ϵ^−2) samples\"\ncould perhaps be reinterpreted as: The student benefits from getting direct signal about which features are important (winning tickets), rather than having to discover them from scratch like in one-shot distillation.\nSame goes for the theorem about intermediate checkpoints having stronger correlations with degree-1 monomials: could be read as a curriculum of increasing task complexity, or about how the teacher naturally learns (simple correlations first) and why intermediate checkpoints help (they provide clearer signal about important features).\n\nSeen this way, calling this an \"implicit curriculum\" might be misleading because the task complexity is constant, and what is changing is the model's internal learning progression. \n\n-  **Other minors**\n\nThe specific temperature values, as well as the choice of dramatically different temperatures (10^-4 vs 10^-20) based on vocabulary, could benefit from more rigorous exploration. Ok, the authors still acknowledge this as a limitation they defer to future work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gxi4jcKQlU", "reviewer": "Reviewer_sKbn", "summary": "Traditional knowledge distillation relies on a single, powerful teacher model to train a smaller student model. However, it has been observed that stronger teachers do not always yield better-performing students. To address this issue, the authors explore progressive distillation, where the student model is incrementally guided by a sequence of increasingly capable teacher checkpoints. Through theoretical analysis and extensive experiments on tasks such as sparse parity and probabilistic context-free grammars (PCFGs), the paper demonstrates that progressive distillation accelerates the optimization process, achieving improved performance with fewer training steps compared to one-shot distillation and direct learning from data.", "questions": "* **Teacher Checkpoint Selection Across Domains**: What guidelines or heuristics can assist in selecting intermediate teacher checkpoints effectively for diverse tasks such as visual classification, image generation, or reinforcement learning?\n* **Joint Training of Student Models:** What are the effects of training small student models jointly with large teacher models, instead of using progressive distillation with intermediate checkpoints? Does this approach provide a similar implicit curriculum that benefits the student’s learning?\n\nI am willing to improve my score if the authors demonstrate that progressive distillation remains beneficial compared to one-shot distillation in the long run for large models, such as GPT-2.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Strengths\n* **Intuitive Motivation and Theoretical Foundation:** The paper is well-motivated, addressing a significant challenge in the effective distillation from large to small models. The theoretical underpinnings are well grounded, with rigorous proofs using the sparse parity example demonstrating why progressive distillation is effective.\n* **Empirical Validation Across Diverse Tasks:** The authors conduct comprehensive experiments on both synthetic tasks (sparse parity and PCFGs) and realistic settings (training BERT on Wikipedia and Books datasets). This breadth of evaluation underscores the generalizability of their findings.", "weaknesses": "**Weaknesses**\n\n- **Scalability to Larger Models and Diverse Architectures**: While the experiments include models like BERT, the paper does not thoroughly investigate how progressive distillation scales with increasingly large models (in the main paper).\n\n- **Effectiveness in GPT Models**: In Appendix E, the paper examines autoregressive training with GPT-2. Although progressive distillation accelerates learning speed, the performance plateaus around 8,000 training steps. It remains unclear why this convergence occurs and what the outcomes would be if training continued beyond 8,000 steps (e.g., up to 20,000 steps).\n\n- **Potential for Degenerate Features**: While progressive distillation leverages an implicit curriculum to identify key patterns, there is a concern that this curriculum might lead to degenerate features that could hinder long-term generalization. This issue could be further investigated by extending the training duration in GPT-2 to ensure that no negative consequences arise from prolonged training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)