Progressive distillation induces an implicit curriculum

[Continued]

Non-terminal prediction task (Measure 3, line 424):

To further evaluate the learned representations, we measured the performance of a linear classifier when trained on top of the output embeddings on the non-terminal prediction task. The classifier trained on features derived from progressively distilled models performed competitively or exceeded the performance of classifiers trained on features from one-shot distilled models. These results confirm that progressive distillation does not compromise the long-term generalization of the learned features.

Average Non-terminal prediction accuracy:

Q2: What guidelines can help select effective teacher checkpoints for tasks like visual classification, image generation, and reinforcement learning?

Response: The heuristics and guidelines for progressive distillation can vary significantly depending on the model's loss behavior during training. Our work specifically focuses on tasks where the loss exhibits dormant phases followed by fast phase transitions. We demonstrate that leveraging teacher checkpoints from fast phase transitions can substantially accelerate student training by enabling the model to bypass the slow, dormant phase and providing guidance on the intermediate features that are most beneficial for further learning.

Extending our analysis to other diverse tasks will require building toy models that can simulate the input label relationship in these tasks. Additionally, our implicit curriculum perspective provides one possible framework through which we can measure the utility of intermediate teacher checkpoints. Expanding our framework to other perspectives and diverse domains through which larger models can facilitate more efficient and effective learning for smaller models stays an important research direction.

Q3: Does joint training of small student models with large teachers provide a similar implicit curriculum?

Response: Joint training can be considered equivalent to progressive distillation if the student model continually distills from the teacher checkpoints at every step, rather than at predefined intervals. Under this interpretation, our theory could be extended to such a setting, as the implicit curriculum would still be governed by the transitions in the teacher’s training phases.

However, for the sake of simplicity in experimentation and to facilitate ablation studies, we focus on a setup where teacher was first trained to completion, and then the student model is trained using the logits from the teacher’s checkpoints sampled at regular intervals. This approach allows us to isolate and analyze the role of each teacher checkpoint in progressive distillation in a controlled manner.

Q4: Scalability to Larger Models and Diverse Architectures

Response: Our work isn’t proposing progressive distillation as a strategy to train smaller models, instead we want to understand how a larger model can help optimization in a smaller model. Given our interpretability analysis on intermediate checkpoints, scaling to larger architectures wasn’t feasible within the scope of our study.

The algorithm of progressive distillation has been widely used in practice for efficient and improved knowledge distillation in different domains. Here are some references (many of which have been cited in our paper):

• Vision models
  • Harutyunyan et al. 2023. Supervision Complexity and its Role in Knowledge Distillation.
  • Mirzadeh et al. 2019. Improved Knowledge Distillation via Teacher Assistant.
  • Cho & Hariharan 2019. On the efficacy of knowledge distillation.
  • Beyer et al.’2021. Knowledge distillation: A good teacher is patient and consistent
  • Cao et al’23. Learning Lightweight Object Detectors via Multi-Teacher Progressive Distillation (Object detection and segmentation)
• Language models
  • Gemini Google Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.
    • Gemini Flash has been reported to be progressive (online) distilled from Gemini Pro.
  • Rezagholizadeh et al. Pro-KD: Progressive distillation by following the footsteps of the teacher

We thank the reviewer for their positive comments and suggestions about our paper. Please find our responses to your questions below.

Q1: Could progressive distillation lead to degenerate features that hinder long-term generalization?

Response: We show, by two new sets of experiments, that features learned by progressive distillation stay useful in the long run.

Setting 1: Fine-tuning experiments on BERT (trained on Wikipedia+Books, Figure 6)

On our trained BERT models on Wikipedia and Books dataset (figure 6), we observe that progressive distilled models have $0.9$ percent better performance on downstream tasks after fine-tuning. This shows that the features learned by progressive distilled models prove to be more helpful for downstream performance generalization.

Setting 2: Longer training on GPT-2 (trained on CFGs, Appendix E)

In our GPT-2 experiments, both the student and the teacher had converged by 8000 training steps on the PCFGs that we used. We want to clarify that 8000 training steps use roughly 4M training samples. In all of our experiments (Figures 29-31), progressive distillation provides acceleration in training when compared at 2000 steps, which translates to a million training samples.

On the request of the reviewer, we extended all our GPT-2 experiments to 16,000 steps. We do not observe any degradation in the performance of the progressively distilled models. Top-1 prediction accuracy:

We thank the reviewer for their positive feedback, constructive comments, and suggestions for future exploration! Please find our responses below.

Q1: Alternative interpretations, e.g. the Lottery Ticket Hypothesis and winning subnetworks.

Thank you for the very interesting and detailed comments! We would like to clarify on the agreements and differences between the comments and claims in the paper.

• Curriculum vs guided search: First, we’d like to clarify that “implicit curriculum” and “guided search” are not meant as alternative interpretations. In fact, “guided search” was the mental picture we had in mind throughout the project. What the paper shows is that implicit curriculum is a mechanism through which the guidance is provided.
  • We also want to clarify a potential confusion based on the comment: “Said differently, the discovery of progressively complex features could be the result/consequence of a guided search to efficient parameter configurations.” — The discovery of progressively complex features is in the teacher, whereas the one being guided is the student.
• Connection to LTH or network pruning: Both LTH and network pruning are based on the hypothesis that there exists a sparse subnetwork that can compute a similar function to the full network, which is about the learner network. In contrast, the implicit curriculum in our paper is about the target function. The same target function can be learned using different learner networks. For instance, the sparse parity can be learned using MLP or Transformer, both of which provides a low-degree curriculum towards sparse parity (Figure 2, Figure 12).
  • “The student benefits from getting direct signal about which features are important (winning tickets)” — To clarify, the teacher only provides supervision through the logits. For sparse parity, the logits are 2-dimensional, while the student has an input dimension of 100 and a hidden width of 100 or 1000. Hence it is impossible for the teacher to directly provide supervision on which subnetwork is the winning one, but the teacher’s logits can encode information towards an alternative function, which is what we call an implicit curriculum (e.g. a lower-degree polynomial instead of the target degree-k monomial).
• Task complexity: The reviewer commented that “calling this an ‘implicit curriculum' might be misleading because the task complexity is constant.“ We’d like to clarify that while the complexity of the target function is constant, the function induced by intermediate teacher checkpoints often does not have a constant complexity [Nakkiran et al. 2019]. As the reviewer correctly pointed out, “what is changing is the (teacher) model's internal learning progression”, and this learning progression is the source of the implicit curriculum.

References: Nakkiran et al. 2019. SGD on Neural Networks Learns Functions of Increasing Complexity

Q2: Effect of temperature

Thank you for the comment! Please note that Figure 14 in Appendix C.2 provides additional results on different choices of $\tau$, which demonstrate that setting the temperature as $\tau = 10^{-4}$ or $\tau = 10^{-20}$ indeed does not affect the conclusion.

Thank you very much also for the references on progressive stacking; this is an interesting connection and we will add a discussion in the camera ready version.

Please let us know if you have more suggestions or concerns, and we will be happy to discuss more. Thank you!

We thank the reviewer for their positive feedback and constructive comments! Please find our responses below.

Q1: Effect of model size; relative size of the teacher vs student

This is a great question. For the current paper, we’ve chosen the sizes to ensure that there is a significant performance gap (in terms of either the optimization speed or the final performance) for teachers and students trained directly on the data; the exact sizes are chosen somewhat arbitrarily as long as a gap is observed.

Below we summarize the empirical comparisons on sizes.

• For sparse parity learned with MLP, when training on the label only, there is an inverse relationship between the model width and the optimization speed, as shown in Edelman et al. 2023 and replicated in our experiments (Figure 7 in Appendix C.2). In contrast, when learning from a teacher, students of different sizes (i.e. width 100 or width 1000) are able to follow the teacher and hence replicate the teacher’s accuracy curves (Figure 2, Figure 8). We did not try further small width as we think a 500-time size difference (i.e. a width-100 student and a width-50000 teacher) is sufficiently big.
  • Additionally, Figure 15 shows similar findings for a hierarchical generalization of sparse parity.
• For sparse parity learned with Transformers, we find that increasing the number of heads or the MLP width both lead to faster training (Figure 10). The precise relation between the model size and the learning speed is an interesting direction for future exploration: the answer would be more nuanced than that of MLP, as Transformers are not permutation invariant over the positions.
• Our PCFG experiments use 4-layer Transformers with 32-head teachers, and 8-head or 16-head students. We observe similar trends for both students (Figure 5, Figure 24, Figure 25, Figure 27).
• Our natural language experiments use 12-layer Transformers with 12 heads for the teacher and 4 heads for the students. We did not experiment with other sizes due to limited compute.

Note that in all experiments, the teacher and the student share the same number of layers and differ only in some notion of “width” (the MLP width, or the number of heads). This is because we want to focus on optimization difficulty which is typically related to width; please refer to “Benefit of width in optimization” in Appendix A.1 for more details.

Understanding depth is an interesting future direction. For instance, there have been results on the depth’s effect on representation capacities [Levine et al. 2020, Merrill & Sabharwal 2022, Sanford et al. 2024] and inductive biases [Tay et al. 2021, Petty et al. 2023, Ye et al. 2024]. Moreover, we note that prior work on vision (see references for the next question) typically vary the depth across the teacher and the student.

References:

• Levine et al. 2020. The Depth-to-Width Interplay in Self-Attention.
• Merrill & Sabharwal 2022. The Parallelism Tradeoff: Limitations of Log-Precision Transformers.
• Petty et al. 2023. The Impact of Depth on Compositional Generalization in Transformer Language Models.
• Sanford et al. 2024. One-layer transformers fail to solve the induction heads task.
• Tay et al. 2021, Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers.
• Ye et al. 2024. Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process.

Q2: Can the results be extended to the vision domain?

Previous works have demonstrated that progressive distillation outperforms traditional knowledge distillation in vision tasks (cited in lines 93-99), which we highlight here:

• Harutyunyan et al. 2023. Supervision Complexity and its Role in Knowledge Distillation.
  • This is most related to our setup, where multiple teacher checkpoints are used to supervise the student.
• Mirzadeh et al. 2019. Improved Knowledge Distillation via Teacher Assistant.
• Cho & Hariharan 2019. On the efficacy of knowledge distillation.

Extending our analysis to vision domain is indeed an interesting question. In all our experiments, we have a theoretical model to simulate input label distribution (for our analysis in NLP domain, we built our theory of PCFGs). However, we do not currently have a well-suited theoretical model to map input-label functions for vision tasks, which is why we have not conducted experiments in this domain.

Thank you also for catching the typo for us. We hope our responses have addressed your concerns, and we are happy to discuss further if you have more suggestions or if there's anything unclear. Thank you!

Q4: ..more intuitive explanation for the statements in the proof of Theorem 3.2

Response: Thank you for pointing this out.

Lines 291–292: This step asserts that the intermediate teacher’s output after first stage of training can be expressed as a polynomial consisting of degree-1 monomials involving the support variables and higher odd-degree polynomials involving the input variables. The coefficients corresponding to degree-1 monomials involving the non-support variables, as well as all even-degree polynomials, are close to zero. Additionally, the higher odd-degree polynomials do not contribute to the gradients for the student model in the subsequent steps.

Lines 293–294: Due to the structured output of the intermediate teacher checkpoint, the gradients for the student’s first-layer weights exhibit higher magnitudes for the support coordinates.

Lines 295–296: As a result of the larger gradients directed toward the support coordinates, these coordinates experience faster growth in their corresponding weight parameters. To ensure that the noise in the gradients, arising from random sampling, does not outweigh the difference in gradient magnitudes between the support and non-support coordinates, a minimum number of samples is required.

Q5: For experiments in Figure 2, how do the authors explain the results that the best correlation (Middle) with in-support variable does not lead to best accuracy (Left)?

Response: It is important to note that the true label is determined by the product of the support variables. Consequently, the degree-1 monomials involving the in-support variables are the primary contributors to the error in the teacher’s performance. Therefore, a key takeaway from our work is that the noise present in the intermediate teacher’s predictions, when it has not yet achieved 100% accuracy, is not merely random when averaged across all samples. Instead, this structured noise can be effectively leveraged by progressive distillation to accelerate the training of a smaller student model.

Q6: L406, what's the formula for total variation?

Response: Total variation distance between two probability vectors p and q is given by $\sum_i |p_i- q_i|.$

We thank the reviewer for their positive comments and valuable suggestions about our paper. Please find our responses to your questions below.

Q1: Instead of masked token prediction, can the authors run experiments in challenging NLP tasks such as QA, summarization and long-form generation?

We address the question from two aspects, NLP tasks after fine-tuning, and reasoning as a type of long-form generation.

Fine-tuning on NLP tasks: We take the checkpoints trained with different methods on Wikipedia and Books dataset for 120k steps (figure 6) and fine-tune them on different GLUE tasks. We report the top-1 accuracy after fine-tuning. For each task, we append a prompt [Gao et al. 2021] to the end of each example (e.g. for sentiment task, we append ‘The sentiment is <mask>’)) and train a classification head on the <mask> token. We observe that progressively distilled models perform better than other models across all tasks.

Reasoning via long-form generation

Here, we provide preliminary experiments on training different GPT models from scratch on TinyGSM [Liu et al. 2023]. We measure the accuracy, where “correctness” requires the model to generate a python program that is both executable and returns the correct numeric solution for a grade-school math problem.

• The trained models are evaluated by pass rate, defined as the number of problems for which at least one of the python programs returned by the model is correct. In particular, we take pass@1 to be the accuracy for a single greedy generation (i.e. temperature 0) per problem, and pass@10 takes the best solution over 10 candidates generated at temperature 0.8.

Our results show that progressive distillation has better pass rates than cross entropy and one-shot distillation. This suggests that progressive distillation can improve performance of models on more difficult reasoning tasks. However, we would like to note that these experiments deserve a thorough study on its own and are not covered by our current theory. We keep a more extensive exploration to future work.

Here, the teacher was a 4 layer 768 width model, that achieved $38\%$ pass at 1 accuracy.

References:

1. Liu et al. 2023. TinyGSM: achieving >80% on GSM8k with small language models.
2. Gao et al 2021. Making Pre-trained Language Models Better Few-shot Learners.

We thank the reviewer for their helpful comments about the quality and clarity of our paper. Please find our responses below.

Q1: Impact of temperature

Thank you for the comment! The temperature is indeed an important hyperparameter for knowledge distillation. We consider the default of $\tau=1$ and the extreme of $\tau \approx 0$. Our one-shot distillation results are reported by taking the better of the two, while the progressive distillation results are reported on $\tau \approx 0$. In other words, we made sure to provide more advantage to one-shot distillation when comparing one-shot and progressive distillation. For progressive distillation alone, Figure 14 in Appendix C.2 provides additional results on different choices of $\tau$.

As mentioned in the paper, we consider $\tau \approx 0$ to remove potential regularization effects induced by soft logits, which are believed to help improve generalization by prior work [Yuan et al. 2020, Menon et al. 2021]. To the best of our knowledge, precise understanding on temperature impact’s on optimization, which is orthogonal to the main claim of our current work and is an interesting future direction.

Q2: Implicit curriculum across setups

We answer the question from two aspects:

• Across model layers: We’d appreciate it if the reviewer could clarify what this means. Here are some explanations based on our interpretations:

  • If the question is about how the implicit curriculum arises in intermediate layers of the teacher, this question could be interpreted as regarding interpretability, which we provide preliminary analyses in Figure 11 and 12 in Appendix C.2 on Transformers.
  • If the question is about how the implicit curriculum affects intermediate layers of the student, please note that 1) the teacher and the student have different architectures, and 2) the teacher only provides supervision through logits (i.e. final layer outputs), so the intermediate layers are not guaranteed to learn similar representations even if the teacher and the student had the same architecture.

• Across datasets and training objectives: Our results consider both discriminative and generative objectives, on both synthetic (i.e. sparse parity and PCFG) and real-world (i.e. Wikipedia + Books) datasets. We identify the low-degree curriculum for sparse parity, and the n-gram curriculum for both PCFG and natural languages.

Q3: Exploring more tasks and architectures

We would appreciate it if the reviewer could clarify the specific tasks and architectures of interest. Please also refer to our response to Q2 of reviewer sKbn.

Q4: Interaction of optimization algorithms and batch size

We use SGD for MLP training and Adam for Transformer training, following common recommendations in practice. Regarding batch sizes, our theory currently requires a sufficiently large batch size, but our experiments do not require so. Empirically, we do not find qualitative difference with different batch sizes. Therefore, we use batch size = 1 (i.e. online SGD) for MLPs, and use a larger batch size for Transformers to speed up training; in particular, we use batch size = 32 for Transformer on parity and batch size = 512 for Transformer on PCFG and natural languages.

We hope our response has helped clarify the questions and concerns. Please let us know if there are further comments, we would love to learn from your insights. Thank you!