MiniPLM: Knowledge Distillation for Pre-training Language Models

Dear Reviewer,

Thank you for your detailed and thoughtful comments! We are currently working on our response and noticed that the reference [1] mentioned in your comments appears to be missing. We believe this reference would be highly valuable for our response and for further improving our paper. We would greatly appreciate it if the reviewer could kindly provide the reference.

Thanks,

Authors of Submission3739

We sincerely thank the reviewer for recognizing the novelty of our approach, the clarity of our writing, the strong experimental results, and the importance of the problem we address.

How knowledge distillation is used (weakness #1)

Knowledge distillation is used for pre-training in MiniPLM because, as explained in Section 2.1, the data sampling process in MiniPLM is essentially derived from minimizing the reverse KLD between the student and teacher LMs, a widely used objective in KD for training LMs [1,2,3]. While effective, directly minimizing reverse KLD is challenging to apply during the pre-training stage due to its high computational overhead for optimization. To address this issue, we propose MiniPLM based on Difference Sampling, which is derived from the optimization of reverse KLD with approximations supported by theory (Proposition 2.1) and empirical results (Appendix D.4). MiniPLM offers a more effective, efficient, and flexible solution for KD in pre-training LMs compared to existing KD methods [4,5] which are primarily used for fine-tuning. This motivation underpins the title of our paper: “Knowledge Distillation for Pre-training Language Models.”

[1] MiniLLM: Knowledge Distillation of Large Language Models. 2024. In ICLR.

[2] On-policy distillation of language models: Learning from self-generated mistakes. 2024. In ICLR.

[3] DistiLLM: Towards Streamlined Distillation for Large Language Models. 2024. In ICML.

[4] Vicuna: An open-source chatbot impressing gpt-4 with 90% chatgpt quality. 2023.

[5] Compact Language Models via Pruning and Knowledge Distillation. 2024. In NeurIPS.

The instances with negative scores (weakness #2)

We thank the reviewer for pointing out the “overthinking” phenomenon of PLMs. During our response, we noticed that the reference [1] mentioned in your comments appears to be missing. We believe this reference would be highly valuable for our response and for further improving our paper. Without the reference, it is hard for us to determine how the instances causing the "overthinking" phenomenon look like, which is critical for us to explore how these instances affect the pre-training of LMs. We would greatly appreciate it if the reviewer could kindly provide the reference.

How $p_{\text{ref}}$ can mimic the behavior of $q_{\theta}$ (weakness #2)

We use $p_{\text{ref}}$ to approximate $q_{\theta}$ for reasons of computational efficiency. While using $q_\theta$ directly might yield better performance in MiniPLM, as suggested by Figure 8, where larger $p_{\text{ref}}$ models lead to improved results, we find that smaller $p_{\text{ref}}$ models already provide strong performance.

To further support this claim, we compute the correlations between $\log \frac{p(x)}{p_{\text{ref}}(x)}$ (using 45M, 104M, and 200M models for $p_{\text{ref}}$) and $\log \frac{p(x)}{q_{\theta}(x)}$ (using the 500M student model for $q_{\theta}$) on a holdout set of the Pile. Additionally, we calculate the sampling accuracy, which measures how many instances are correctly classified as selected or discarded when using $\log \frac{p(x)}{p_{\text{ref}}(x)}$ with a sampling threshold of 0.5. The ground truth is determined by sampling based on $\log \frac{p(x)}{q_{\theta}(x)}$. As the following table shows, using $p_{\text{ref}}$ achieves high correlations with $q_{\theta}$ and high data sampling accuracy. We have added these results to Appendix D.4 in the revised paper.

Similar strategies, where small models approximate probability differences for larger models, have been successfully employed in quite a few prior works [1,2,3,4].

[1] Selection via Proxy: Efficient Data Selection for Deep Learning. 2020. In ICLR.

[2] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pre-training. 2023. In NeurIPS.

[3] Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models. 2024. arxiv pre-print.

[4] Irreducible Curriculum for Language Model Pre-training. 2024. arxiv pre-print.

We highly appreciate the reviewer’s insightful feedback, which clearly helped us improve our paper.

In our response, we have addressed key concerns, including how knowledge distillation is applied during pre-training, how effectively $p_{\text{ref}}$ mimics the behavior of $q_{\theta}$, why MiniPLM outperforms MiniLLM, and the relatively low complexity of MiniPLM compared to other pre-training or knowledge distillation methods.

With the deadline approaching, we sincerely hope the reviewer find our response useful and update the scores if the concerns have been resolved. We also hope the reviewer could kindly provided the reference [1] in the comments for us to learn more about the “overthinking” problem of PLMs and its relationship to our methods. We remain open to further discussions and are happy to address any additional questions.

We thank the reviewer for acknowledging our well-motivated approach, the elegant design of Difference Sampling, the efficient offline computation strategy, and the comprehensive experiments conducted across model scales.

Analysis of finite-sample scenarios (weakness #1)

In Section 3.4, we explore the performance of MiniPLM in scenarios where the number of original pre-training tokens (before applying Difference Sampling) is limited to 50B. The results demonstrate that MiniPLM significantly improves data utilization under finite-sample conditions.

Larger teacher models (weakness #2 and question #3)

We run an additional experiment using the 4B Qwen-1.5 model as the teacher and the 500M model as the student. The results confirm that MiniPLM continues to perform well with a 4B teacher LM.

Following the trend in Figure 6, the 10B model size may not be optimal for training the student LMs with 200M, 500M, and 1.2B parameters, as the large model size gap can hinder effective knowledge transfer. This phenomenon is consistent with findings in prior knowledge distillation studies [1,2]. We plan to explore distilling the 10B model to larger student LMs (e.g., 4B and 7B) or incorporating the teacher assistants [1] to bridge the model size gap in future work.

[1] Improved Knowledge Distillation via Teacher Assistant. 2020. In AAAI.

[2] MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers. 2020. In ICML.

The cross-family experiments (weakness #3)

As discussed in Section 3.3, we distill the knowledge from the Qwen model to Llama and Mamba models, which covers the most recently used architectures (transformer and state-space models) for language model development. Additional details of these experiments have been included in Appendix B.

Alternative sampling strategies (weakness #4)

We compare Difference Sampling with two alternative sampling strategies: sampling based on (1) large model’s preference: $\text{topk}\left[\log p(x)|x\in D\right]$ and (2) the sample difficulty reflected by the reference model: $\text{topk}\left[-\log p_{\text{ref}}(x)|x\in D\right]$.

The results show that considering both the signals of sample difficulty and teacher’s preference, as in Difference Sampling, is essential to achieve high performance.

The selection criteria of the reference model (weakness #5)

We provide empirical results on the choice of reference model in Figure 8. Basically, the closer the reference model size is to that of the student model, the better the performance. Notably, a reference model with 1/5 of the parameters of the student model achieves strong results, outperforming Vanilla KD. Users can adjust the reference model size to balance computational cost and final performance based on their specific requirements.

Theoretical justification for using $\log \frac{p(x)}{p_{\text{ref}}(x)}$ (question #1)

Theoretically, the metric $\log \frac{p(x)}{p_\text{ref}(x)}$ approximates to minimizing the reverse KLD between the student and teacher LMs, as described in Eq. (1). This is because minimizing $E_{x\sim q_{\theta}}\log \frac{q_{\theta}(x)}{p(x)}$, can be effectively achieved using the Best-of-N algorithm, which selects examples with high $\log \frac{p(x)}{q_{\theta}(x)}$scores from instances sampled by $q_{\theta}$. The $q_{\theta}$-sampled instances can be replaced with the original pre-training corpus, which is theoretically supported by Proposition 2.1 when the dataset size is sufficiently large. To improve efficiency, $q_{\theta}$ in $\log \frac{p(x)}{q_{\theta}(x)}$ is approximated with $p_{\text{ref}}$ , leading to the final metric $\log \frac{p(x)}{p_{\text{ref}}(x)}$.

The comparison to other small language models (weakness #4)

We did not compare our model to MobileLLM because the data, model architecture, and training computation differ significantly. MobileLLM employs extensive model architecture optimizations, including searching for the optimal model width and depth as well as implementing weight-sharing strategies. Furthermore, it requires $2.1\times10^{21}$ FLOPs to train the 350M model on 10T tokens, which is 10x more computation than that we can afford.

Since the purpose of our work is to explore improved knowledge distillation techniques for pre-training language models, we keep all other training configurations, like pre-training data, model architecture, and optimization strategies at conventional settings to ensure a fair and straightforward comparison.

In the following, we show the comparison of our 500M model to the MobileLLM’s associated baselines trained on the Pile corpus with the similar model sizes and comparable training computation to ours. MiniPLM outperforms models trained with 7.0x and 2.5x more FLOPs (Bloom-560M, OPT-350M), and achieves results comparable to Pythia-410M, trained with 5.5x more FLOPs.

We also want to point out that there are quite a few existing works showing that the model architecture optimization is complementary to knowledge distillation or data refinement [1,2,3]. We leave further exploration on combining these techniques to future work.

[1] Compact Language Models via Pruning and Knowledge Distillation. 2024. In NeurIPS.

[2] Sheared llama: accelerating language model pre-training via structured pruning. 2024. In ICLR.

[3] Structured Pruning Learns Compact and Accurate Models. 2022. In ACL.

The experimental details of cross model family distillation (the question)

We have included the relevant experimental details in Appendix B of the revised paper.

We sincerely thank the reviewer for acknowledging the novelty of our Difference Sampling method, its significant storage efficiency, and the clear computational gains it provides. We also greatly appreciate the reviewer’s recognition of the solid theoretical justification and detailed analysis supporting our approach.

The effectiveness of knowledge transfer (weakness #1)

We thank the reviewer for the thoughtful feedback. Our teacher model, the official Qwen-1.8B, is based on a non-public large-scale pre-training corpus, making a direct comparison with our pre-trained models inherently unfair.

To demonstrate the effectiveness of knowledge transfer in a fair setting, we eliminate the impact caused by the differences in pre-training corpora and distill a 500M LM from a 1.2B teacher LM pre-trained on our data (Due to time constraints during the rebuttal period, we could not pre-train a 1.8B teacher LM on our data. Therefore, we used the previously pre-trained 1.2B model as the teacher LM.). In this controlled setting, the 500M student model achieves performance comparable to the teacher LM. The results, shown in the table below, demonstrate that MiniPLM effectively preserves most of the knowledge from the teacher LM.

Larger model gap for distillation (weakness #2)

Due to limit of computational resources, we are unable to complete the distillation of Llama3-8B to 1B within the rebuttal period, as this requires pre-training 1B models from scratch. As a supplement to the experimental results achievable within our computational budget, we provide results for distilling a 4B Qwen-1.5 model into a 500M model, which similarly represents an 8x model size reduction.

Controlled experiments for Difference Sampling (weakness #3)

Both SeqKD and our method use pre-training token lengths of 1024, ensuring a controlled and fair comparison. As described in lines 295–297, we use 768 tokens as the prompt and have the teacher generate the remaining 256 tokens, forming a total length of 1024. To further analyze Difference Sampling, we compare it with two alternative sampling strategies: sampling based on (1) large model’s preference: $\text{topk}\left[\log p(x)|x\in D\right]$ and (2) the sample difficulty reflected by the reference model: $\text{topk}\left[-\log p_{\text{ref}}(x)|x\in D\right]$.

The results show that combining both signals, as implemented in Difference Sampling, is essential to achieve a high performance.

We sincerely thank the reviewer for recognizing the theoretical foundations, clarity of contributions and experiments, and the improved performance and data efficiency demonstrated by our method. The reviewer’s positive feedback and lack of significant criticism are highly encouraging.

The discussion on resource consumption (weakness #1)

We have added a detailed discussion on resource consumption in Appendix C. To summarize, we train a reference model and conduct teacher and reference model inference on the pre-training corpus, which are performed only once offline to distill the teacher LM’s knowledge to multiple students. This computation can be further reduced by 95% using a proxy model for Difference Sampling as described in Appendix E. The student model's training-time computation remains identical to that of Pre-Train w/o KD.

The complexity of the training process (weakness #2)

To ensure reproducibility, we will open-source all our code, model, and data used in our work. We have uploaded our code to the supplementary material. The overall training process is straightforward:

• Model inference is performed prior to pre-training and is well-supported by existing open-source libraries [1,2].
• The data sampling process is performed on CPUs, requiring only ~60 lines of code.
• The pre-training framework for both the student and reference LM is nearly identical to that of Pre-Train w/o KD [2,3], with only minor adjustments to the data loading path.

A particularly noteworthy advantage of MiniPLM is that, once $\log \frac{p(x)}{p_{\text{ref}}(x)}$ values are computed for a given corpus, the teacher LM’s knowledge can be distilled into any future LMs without requiring access to either the teacher or reference model. This makes the complexity of applying MiniPLM nearly equivalent to that of Pre-Train w/o KD.

[1] DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale. 2022. arxiv pre-print.

[2] Transformers: State-of-the-Art Natural Language Processing. 2020. In ACL.

[3] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arxiv pre-print.

We sincerely thank the reviewer for acknowledging the strengths of our work, including the innovation of MiniPLM with its Difference Sampling method, the extensive experimental validation, and the practical benefits in computation reduction and performance scalability.

Explanation of Difference Sampling (weakness #1 and question #1)

We thank the reviewer for the constructive suggestion. In Appendix F, we provide the intuitive examples from our pre-training corpus together with a detailed discussion to further clarify the Difference Sampling mechanism, and add a cross reference to this information in lines 238-240 of the revised manuscript.

Distillation from closed-source models (weakness #2 and question #2)

We discussed this limitation in Section 5 and provided a simple solution at the expense of a large number of API calls. Furthermore, we have explicitly clarified in lines 77–79 of the revised manuscript that our work focuses on knowledge distillation in scenarios where the teacher LM's probabilities are accessible, such as in the open-source or white-box settings.

Transferability and generalizability of MiniPLM (weakness #3 and question #3)

MiniPLM is designed to impose minimal constraints on data and models, enabling its application across a wide range of language models trained with cross-entropy loss, just like other works on knowledge distillation [1,2,3]. In addition, in Table 3, we show that MiniPLM generalizes to the setting of knowledge distillation across model families. Since MiniPLM operates during pre-training, the distilled model retains compatibility with downstream tasks, which are typically handled by conventionally pre-trained models without knowledge distillation.

[1] Compact Language Models via Pruning and Knowledge Distillation. 2024. In NeurIPS.

[2] MiniLLM: Knowledge Distillation of Large Language Models. 2024. In ICLR.

[3] On-policy distillation of language models: Learning from self-generated mistakes. 2024. In ICLR.

The discussion of MiniPLM’s limitation (weakness #4 and question #2)

One possible limitation of MiniPLM is that Difference Sampling samples data points individually. This means that if the original corpus contains too many duplicates, Difference Sampling will also select duplicated samples. Therefore, MiniPLM is better suited for deduplicated data, like many recent pre-training corpora [1,2,3,4].

[1] The Pile: An 800GB Dataset of Diverse Text for Language Modeling.

[2] The RefinedWeb Dataset for Falcon LLM: Outperforming Curated Corpora with Web Data, and Web Data Only.

[3] SemDeDup: Data-efficient learning at web-scale through semantic deduplication. 2023. In ICLR Workshop.

[4] D4: Improving LLM Pretraining via Document De-Duplication and Diversification. 2023. In NeurIPS.