Elamur Research
☰
home / papers / 2024 / DistiLLM: Towards Streamlined Distillation for Large Language Models

DistiLLM: Towards Streamlined Distillation for Large Language Models

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2402.03898

Summary

DistiLLM tackles knowledge distillation (KD) for auto-regressive language models, where a large teacher LLM is compressed into a smaller student. The authors identify two weaknesses in existing approaches. First, nobody agrees on which loss to use: standard KLD tends to cause mode averaging or collapse, and the best alternative divergence seems to change from task to task. Second, methods that train on student-generated outputs (SGOs) are very slow, because they generate new SGOs at every step. DistiLLM fixes both. It introduces a skew KLD objective, in forward (SKL) and reverse (SRKL) variants, with theoretical backing for stable gradients and small approximation error. It also adds an adaptive off-policy scheme: SGOs go into a replay buffer, and a scheduler raises the probability of using them whenever validation loss gets worse. The method was tested on instruction-following and summarization with GPT-2 and LLaMA teacher–student pairs. Students reach state-of-the-art quality and train 2.5–4.3× faster than MiniLLM and GKD. This makes small, capable LLMs cheaper to build.

Key points

  • The two components are the skew KLD loss (SKL/SRKL) and an adaptive SGO scheduler with a replay buffer. The share of SGOs used rises when validation loss increases and the buffered data is reused off-policy. Code is at github.com/jongwooko/distillm (ICML 2024).
  • The paper reports that in on-policy methods, producing SGOs takes up to 80% of total training time. DistiLLM's training speedup over MiniLLM and GKD is 2.5–4.3×.
  • Teacher–student pairs include GPT-2 XL (1.5B) → GPT-2 (0.1B) and LLaMA-7B → LLaMA-3B. Evaluation covers Dolly Eval, Self-Instruct and Super-Natural.
  • When the GPT-2 student starts from a model with no prior fine-tuning, DistiLLM (SRKL) still scores about 26.1 on Dolly Eval. MiniLLM scores about 22.6 and GKD about 23.2 in the same setting, so the method works in a single distillation stage.
  • The SGO probability the adaptive scheduler arrives at is within 0.2 (SKL) or 0.1 (SRKL) of the best hand-tuned value, and in most cases it scores better than that value.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: Natural language processing · AI & machine learning

Cited by

107 papers of our library cite it (of the 140673 whose bibliographies we read).

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2402.03898 · arXiv 2402.03898

@misc{ko2024distillm,
  title = {{DistiLLM: Towards Streamlined Distillation for Large Language Models}},
  author = {Ko, Jongwoo and Kim, Sungnyun and Chen, Tianyi and Yun, Se-Young},
  year = {2024},
  doi = {10.48550/arxiv.2402.03898},
  eprint = {2402.03898},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.03898}
}

References

53 references, 14 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. //www.databricks.com/blog/2023/04/12/ dolly-first-open-commercially-viable-\\ instruction-tuned-llm. find
  2. Agarwal, R., Vieillard, N., Stanczyk, P., Ramos, S., Geist, M., and Bachem, O. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/ forum?id=3zKtaqxLhW. ↗ our card
  3. Fedus, W., Ramachandran, P., Agarwal, R., Bengio, Y., Larochelle, H., Rowland, M., and Dabney, W. Revisiting fundamentals of experience replay. In International Conference on Machine Learning, pp. 3061–3071. PMLR, 2020. find
  4. Almazrouei, E., Alobeidli, H., Alshamsi, A., Cappelli, A., Cojocaru, R., Alhammadi, M., Daniele, M., Heslow, D., Launay, J., Malartic, Q., Noune, B., Pannier, B., and Penedo, G. The falcon series of language models: Towards open frontier models. https://arxiv.org/abs/2311.16867, 2023. ↗ our card
  5. Geng, X. and Liu, H. Openllama: An open reproduction of llama, May 2023. URL https://github.com/ openlm-research/open\_llama. ↗
  6. Gliwa, B., Mochol, I., Biesek, M., and Wawer, A. Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. In Proceedings of the 2nd Workshop on New Frontiers in Summarization, pp. 70–79, 2019. find
  7. Anil, R., Dai, A. M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z., et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403, 2023. find our card
  8. Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Openwebtext corpus, 2019. find
  9. Arora, K., El Asri, L., Bahuleyan, H., and Cheung, J. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. In Muresan, S., Nakov, P., and Villavicencio, A. (eds.), Findings of the Association for Computational Linguistics: ACL 2022, pp. 700–710,… find
  10. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877–1901, 2020. find
  11. Gu, Y., Dong, L., Wei, F., and Huang, M. MiniLLM: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum? id=5h0qf7IBZZ. ↗
  12. Hinton, G. E., Vinyals, O., and Dean, J. Distilling the knowledge in a neural network. ArXiv, abs/1503.02531, 2015. URL https://api.semanticscholar. org/CorpusID:7200347. ↗ our card
  13. Honovich, O., Scialom, T., Levy, O., and Schick, T. Unnatural instructions: Tuning language models with (almost) no human labor. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 14409–14428, Toronto, Canada, July 2023. Associatio… find our card
  14. Cettolo, M., Federico, M., Bentivogli, L., Niehues, J., Stüker, S., Sudoh, K., Yoshino, K., and Federmann, C. Overview of the IWSLT 2017 evaluation campaign. In Proceedings of the 14th International Conference on Spoken Language Translation, pp. 2–14, Tokyo, Japan, December 14-15 2017. Internationa… find our card
  15. Hsieh, C.-Y., Li, C.-L., Yeh, C.-K., Nakhost, H., Fujii, Y., Ratner, A., Krishna, R., Lee, C.-Y., and Pfister, T. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. arXiv preprint arXiv:2305.02301, 2023. find our card
  16. Chiang, W.-L., Li, Z., Lin, Z., Sheng, Y., Wu, Z., Zhang, H., Zheng, L., Zhuang, S., Zhuang, Y., Gonzalez, J. E., Stoica, I., and Xing, E. P. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https://lmsys.org/blog/ 2023-03-30-vicuna/. ↗
  17. Hu, E. J., yelong shen, Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https: //openreview.net/forum?id=nZeVKeeFYf9. find
  18. Ji, H., Ke, P., Hu, Z., Zhang, R., and Huang, M. Tailoring language generation models under total variation distance. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview. net/forum?id=VELL0PlWfc. find
  19. Conover, M., Hayes, M., Mathur, A., Xie, J., Wan, J., Shah, S., Ghodsi, A., Wendell, P., Zaharia, M., and Xin, R. Free dolly: Introducing the world’s first truly open instruction-tuned llm, 2023. URL https: find
  20. Kim, Y. and Rush, A. M. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 1317– 1327, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/D16-1139. URL https://aclanthology.org/D… ↗ our card
  21. Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature, 518(7540): 529–533, 2015. find
  22. Narayan, S., Cohen, S. B., and Lapata, M. Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1797–1807, 2018. find
  23. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. find
  24. Kumar, A., Fu, J., Soh, M., Tucker, G., and Levine, S. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019. find
  25. OpenAI. Gpt-4 technical report, 2023. find our card
  26. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730–27744, 2022. find
  27. Lee, H., Cho, H., Kim, H., Gwak, D., Kim, J., Choo, J., Yun, S.-Y., and Yun, C. Plastic: Improving input and label plasticity for sample efficient reinforcement learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. find
  28. Papineni, K., Roukos, S., Ward, T., and Zhu, W.-J. Bleu: a method for automatic evaluation of machine translation. In Isabelle, P., Charniak, E., and Lin, D. (eds.), Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pp. 311–318, Philadelphia, Pennsylvania, USA,… find our card
  29. Lee, K. and Shin, J. Rényicl: Contrastive representation learning with skew rényi divergence. Advances in Neural Information Processing Systems, 35:6463–6477, 2022. find
  30. Lee, L. On the effectiveness of the skew divergence for statistical language analysis. In International Workshop on Artificial Intelligence and Statistics, pp. 176–183. PMLR, 2001. find
  31. Lin, A., Wohlwend, J., Chen, H., and Lei, T. Autoregressive knowledge distillation through imitation learning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6121– 6133, Online, November 2020. Association for Computational Linguistics. doi: 10.… find
  32. Lin, C.-Y. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pp. 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https: //aclanthology.org/W04-1013. find
  33. Pedro, D. A unified bias-variance decomposition and its applications. In 17th International Conference on Machine Learning, pp. 231–238, 2000. find
  34. Peng, B., Li, C., He, P., Galley, M., and Gao, J. Instruction tuning with gpt-4, 2023. find our card
  35. Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. OpenAI blog, 2018. find
  36. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  37. Liu, L., Pillutla, K., Welleck, S., Oh, S., Choi, Y., and Harchaoui, Z. Divergence frontiers for generative models: Sample complexity, quantization effects, and frontier integrals. Advances in Neural Information Processing Systems, 34:12930–12942, 2021. find
  38. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020. find
  39. Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. find
  40. Rubenstein, P., Bousquet, O., Djolonga, J., Riquelme, C., and Tolstikhin, I. O. Practical and consistent estimation of f-divergences. Advances in Neural Information Processing Systems, 32, 2019. find
  41. Mirzadeh, S. I., Farajtabar, M., Li, A., Levine, N., Matsukawa, A., and Ghasemzadeh, H. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5191–5198, 2020. find our card
  42. Sanh, V., Debut, L., Chaumond, J., and Wolf, T. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. ArXiv, abs/1910.01108, 2019. URL https://api.semanticscholar. org/CorpusID:203626972. ↗
  43. See, A., Liu, P. J., and Manning, C. D. Get to the point: Summarization with pointer-generator networks. In Barzilay, R. and Kan, M.-Y. (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1073–1083, Vancouver, Canada, July 2017… find
  44. Wang, Y., Kordi, Y., Mishra, S., Liu, A., Smith, N. A., Khashabi, D., and Hajishirzi, H. Self-instruct: Aligning language models with self-generated instructions. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 13484–13508, Toro… find
  45. Wen, Y., Li, Z., Du, W., and Mou, L. f-divergence minimization for sequence-level knowledge distillation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10817–10834, Toronto, Canada, July 2023. Association for Computational Lin… find
  46. Sun, S., Cheng, Y., Gan, Z., and Liu, J. Patient knowledge distillation for BERT model compression. In Inui, K., Jiang, J., Ng, V., and Wan, X. (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Languag… find our card
  47. Xue, L., Constant, N., Roberts, A., Kale, M., Al-Rfou, R., Siddhant, A., Barua, A., and Raffel, C. mT5: A massively multilingual pre-trained text-to-text transformer. In Toutanova, K., Rumshisky, A., Zettlemoyer, L., Hakkani-Tur, D., Beltagy, I., Bethard, S., Cotterell, R., Chakraborty, T., and Zhou… find
  48. Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/ stanford\_alpaca, 2023. ↗
  49. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. find
  50. Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022. find
  51. Wang, P., Wang, Z., Li, Z., Gao, Y., Yin, B., and Ren, X. Scott: Self-consistent chain-of-thought distillation. arXiv preprint arXiv:2305.01879, 2023a. find our card
  52. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging LLM-as-a-judge with MT-bench and chatbot arena. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Tra… find
  53. Wang, Y., Mishra, S., Alipoormolabashi, P., Kordi, Y., Mirzaei, A., Naik, A., Ashok, A., Dhanasekaran, A. S., Arunkumar, A., Stap, D., Pathak, E., Karamanolakis, G., Lai, H., Purohit, I., Mondal, I., Anderson, J., Kuznia, K., Doshi, K., Pal, K. K., Patel, M., Moradshahi, M., Parmar, M., Purohit, M.,… find