Enhancing Knowledge Distillation of Large Language Models through Efficient Multi-Modal Distribution Alignment
Paper · 2025 · Tianyu Peng and Jiajun Zhang
Summary
The paper addresses a weakness in white-box knowledge distillation (KD) of large language models. Teacher LLMs often produce multi-modal next-token distributions with several plausible candidates, and student models struggle to reproduce them. Standard objectives each fail in a different way. Forward KL tends to over-smooth the student's distribution, while reverse KL (RKL) tends to drop some peaks. Symmetric divergences such as JSD, TVD and AKL have not been shown to fix peak learning in realistic settings. The authors first measure peak agreement using two top-k metrics, CR and MOR. Comparing Llama-2 and OpenELM models against a 70B reference, they find that agreement with the stronger model tracks model quality. They then propose RLKD, which adds a word-level ranking loss based on Spearman's rank correlation coefficient (SRCC). This loss pushes the student to order its top-k candidates the same way the teacher does. Because it is an add-on term, it can be combined with existing distillation objectives. Experiments show better alignment of multi-modal predictions and consistent gains across several downstream tasks, at almost no extra compute cost.
Key points
- Metrics: Two measures of peak-prediction agreement are introduced.
- CR is the share of exact top-k matches, counting both the tokens and their order.
- MOR is the mean token overlap within the top-k, ignoring order.
- Motivation study: Llama-2 and OpenELM models from 270M to 13B were compared against Llama-2 70B, used as a stand-in for ground truth. The test set was 5,000 slices sampled from SlimPajama.
- Downstream setup: The ranking loss was added to baselines with word-level objectives (KL, RKL, JSD, TVD, AKL), aligning the top-5 predictions. A ranking-only variant, "Rank-k", was also tested.
- Case study: With TVD, the student scored 0 accuracy on GSM8K because it failed to learn the expected answer format. The authors attribute this to weak peak learning.
- Cost and scope: The extra loss adds about 1% training time. Distillation was tested only on Llama-architecture models. Code is available at github.com/Pty72/RLKD.
Fields, topics and concepts
📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works
🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,320 worksVision & language · 657 works
Cite
Cite as Tianyu Peng and Jiajun Zhang. 2025. Enhancing Knowledge Distillation of Large Language Models through Efficient Multi-Modal Distribution Alignment. Proceedings of the 31st International Conference on Computational Linguistics.
@inproceedings{peng2025enhancing,
title = {{Enhancing Knowledge Distillation of Large Language Models through Efficient Multi-Modal Distribution Alignment}},
author = {Tianyu Peng and Jiajun Zhang},
booktitle = {Proceedings of the 31st International Conference on Computational Linguistics},
year = {2025},
url = {https://aclanthology.org/2025.coling-main.169/}
}References
- Mathieu Blondel, Olivier Teboul, Quentin Berthet, and Josip Djolonga. 2020. Fast differentiable sorting and ranking. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 950–9… find
- Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.
- Hongzhan Chen, Xiaojun Quan, Hehong Chen, Ming Yan, and Ji Zhang. 2024. Knowledge distillation for closed-source language models. CoRR, abs/2401.07013.
- Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. CoRR, abs/2110.14168.
- Mike Conover, Matt Hayes, Ankit Mathur, Jianwei Xie, Jun Wan, Sam Shah, Ali Ghodsi, Patrick Wendell, Matei Zaharia, and Reynold Xin. 2023. Free dolly: Introducing the world’s first truly open instruction-tuned llm. find
- Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2023. Knowledge distillation of large language models. CoRR, abs/2306.08543.
- Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015. Distilling the knowledge in a neural network. CoRR, abs/1503.02531.
- Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. 2020. The curious case of neural text degeneration. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
- Tao Huang, Zekang Li, Hua Lu, Yong Shan, Shusheng Yang, Yang Feng, Fei Wang, Shan You, and Chang Xu. 2022. Relational surrogate loss learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net. find
- Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scalin…
- Yoon Kim and Alexander M. Rush. 2016. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 1317–1327. The Association for Computational Linguistics.
- Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun. 2024. Distillm: Towards streamlined distillation for large language models. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net.
- Chin-Yew Lin. 2004. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pages 74–81, Barcelona, Spain. Association for Computational Linguistics.
- Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net. find
- Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, and Mohammad Rastegari. 2024. Openelm: An efficient language model family with open training and inference framework. CoRR, abs/2404.14619. — [arXiv:2404.14619](https://arxiv.org/abs/2404.14619) find
- Shashi Narayan, Shay B. Cohen, and Mirella Lapata. 2018. Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. ArXiv, abs/1808.08745.
- OpenAI. 2023. GPT-4 technical report. abs/2303.08774.
- CoRR, Elias Ramzi, Nicolas Audebert, Clément Rambour, André Araujo, Xavier Bitot, and Nicolas Thome. 2023. Optimization of rank losses for image retrieval. CoRR, abs/2309.08250. — [arXiv:2309.08250](https://arxiv.org/abs/2309.08250) find
- Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. 2020. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In KDD ’20: The 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Virtual Event, CA, USA, August 23-27, 2020… find
- Ethan M. Rudd, David Krisiloff, Daniel Olszewski, Edward Raff, and James Holt. 2022. Efficient malware analysis using metric embeddings. In Proceedings of the Conference on Applied Machine Learning in Information Security, CAMLIS 2022, Arlington, Virginia, USA, October 20-21, 2022, volume 3391 of CE… find
- Sam Shleifer and Alexander M. Rush. 2020. Pre-trained summarization distillation. CoRR, abs/2010.13002. — [arXiv:2010.13002](https://arxiv.org/abs/2010.13002) find
- Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://www.cerebras.net/blog/slimpajama-a-627b-token-cleane d-and-deduplicated-version-of-redpajama. ↗
- Shicheng Tan, Weng Lam Tam, Yuanchun Wang, Wenwen Gong, Shu Zhao, Peng Zhang, and Jie Tang. 2023. GKD: A general knowledge distillation framework for large-scale pre-trained language model. In Proceedings of the The 61st Annual Meeting of the Association for Computational Linguistics: Industry Track… find
- Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller,… — Llama 2: Open Foundation and Fine-Tuned Chat Models
- Sen Wang and Jin Zheng. 2023. Monoskd: General distillation framework for monocular 3d object detection via spearman correlation coefficient. In ECAI 2023 - 26th European Conference on Artificial Intelligence, September 30 - October 4, 2023, Kraków, Poland - Including 12th Conference on Prestigious… find
- Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. 2022. Emergent abilities of large language models.
- Trans. Mach. Learn. Res., 2022. find
- Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. 2023. 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), ACL 2023, Toronto, Canada, July 9-14, 2023, pages 10817–10834.…
- Junhong Wu, Yang Zhao, Yangyifan Xu, Bing Liu, and Chengqing Zong. 2024a. Boosting LLM translation skills without general ability loss via rationale distillation. CoRR, abs/2410.13944. find
- Taiqiang Wu, Chaofan Tao, Jiahao Wang, Zhe Zhao, and Ngai Wong. 2024b. Rethinking kullback-leibler divergence in knowledge distillation for large language models. CoRR, abs/2404.02657.
- Aiyuan Yang, Bin Xiao, Bingning Wang, Borong Zhang, Ce Bian, Chao Yin, Chenxu Lv, Da Pan, Dian Wang, Dong Yan, Fan Yang, Fei Deng, Feng Wang, Feng Liu, Guangwei Ai, Guosheng Dong, Haizhou Zhao, Hang Xu, Haoze Sun, Hongda Zhang, Hui Liu, Jiaming Ji, Jian Xie, Juntao Dai, Kun Fang, Lei Su, Liang Song,… — Baichuan 2: Open Large-scale Language Models
- Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. 2023. Scaling relationship on learning mathematical reasoning with large language models. CoRR, abs/2308.01825.
- Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. 2023. GLM-130B: an open bilingual pre-trained model. In The Eleventh…
- Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. 2024. Tinyllama: An open-source small language model. CoRR, abs/2401.02385.
- Qinhong Zhou, Zonghan Yang, Peng Li, and Yang Liu. 2023. Bridging the gap between decision and logits in decision-based knowledge distillation for pre-trained language models. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023… find
- Xunyu Zhu, Jian Li, Yong Liu, Can Ma, and Weiping Wang. 2023. A survey on model compression for large language models. CoRR, abs/2308.07633. 37.83%
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/peng-2025-enhancing-knowledge-distillation-large-language/