Elamur Research⌕
☰
home / papers / 2025 / ToDi: Token-wise Distillation via Fine-Grained Divergence Control

ToDi: Token-wise Distillation via Fine-Grained Divergence Control

Paper · 2025 · Seongryong Jung, Suwan Yoon, DongGeon Kim et al.

Summary

Knowledge distillation for LLMs usually applies one divergence, Forward KL (FKL) or Reverse KL (RKL), with the same weight to every vocabulary entry. Some methods mix the two per sample or per time step, such as Adaptive KL (AKL), but none adjusts the mix for individual tokens. The authors studied the gradients of both losses. FKL gives a strong upward signal for tokens the student assigns too little probability, while RKL gives a strong downward signal for tokens it over-weights.

Their method, Token-wise Distillation (ToDi), uses this finding. For each vocabulary token, it sets an FKL/RKL mixing weight with a sigmoid of the log-ratio between teacher and student probabilities, scaled by a parameter β. The cost grows linearly with vocabulary size.

Students based on GPT-2 and TinyLLaMA were evaluated on instruction-following benchmarks. ToDi beat recent distillation baselines on ROUGE-L and won more often in GPT-4 pairwise judgments. This shows that fine-grained divergence control yields smaller models that track their teachers more closely.

Key points

  • Setup: students were GPT2-120M and TinyLLaMA-1.1B. Evaluation used Dolly, S-NI, UnNI, Self and Vicuna. Code is on GitHub (jungseongryong/ToDi).

  • Effect of β (GPT-2 student, average ROUGE-L):

    • β = 1.0 scored best at 18.66.
    • β = 0.6 dropped to 18.40, because the weights flatten toward a fixed 50/50 split.
    • β → ∞ fell to 18.24, because the hard switch between FKL and RKL destabilizes training.

  • Token-wise vs. fixed mixing (TinyLLaMA, average ROUGE-L): ToDi scored about 24.83. Fixed FKL–RKL blends of 0.25/0.75, 0.5/0.5 and 0.75/0.25 stayed around 24.5. Reversing the weighting direction hurt every dataset.

  • Teacher–student Pearson similarity (computed on databricks/dolly-15k):

    StudentAKLToDi
    GPT-20.4770.482
    TinyLLaMA0.5990.610

  • Limitations:

    • Teacher and student must share a vocabulary.
    • The method needs the teacher's full logits, so it only works with open models.
    • Very large models were not tested.

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

 Institutions

🎓Chung-Ang UniversitySeoul · South Korea3 works · 12 citations

 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

Cite

Cite as Seongryong Jung, Suwan Yoon, DongGeon Kim and Hwanhee Lee. 2025. ToDi: Token-wise Distillation via Fine-Grained Divergence Control. Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2025.emnlp-main.409

@inproceedings{jung2025todi,
  title = {{ToDi: Token-wise Distillation via Fine-Grained Divergence Control}},
  author = {Seongryong Jung, Suwan Yoon, DongGeon Kim and Hwanhee Lee},
  booktitle = {Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing},
  year = {2025},
  doi = {10.18653/v1/2025.emnlp-main.409},
  url = {https://doi.org/10.18653/v1/2025.emnlp-main.409}
}

References

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

  1. Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In The Twelfth International Conference on Learning Representations.
  2. Ibtihel Amara, Nazanin Sepahvand, Brett H Meyer, Warren J Gross, and James J Clark. 2022. Bd-kd: balancing the divergences for online knowledge distillation. arXiv preprint arXiv:2212.12965. — [arXiv:2212.12965](https://arxiv.org/abs/2212.12965) find
  3. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, and… find
  4. Alan Chan, Hugo Silva, Sungsu Lim, Tadashi Kozuno, A Rupam Mahmood, and Martha White. 2022. Greedification operators for policy optimization: Investigating forward and reverse kl divergences. Journal of Machine Learning Research, 23(253):1–79. find
  5. Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, and 1 others. 2024. Scaling instruction-finetuned language models. Journal of Machine Learning Research, 25(70):1–53. find
  6. Tao Feng, Yicheng Li, Li Chenglin, Hao Chen, Fei Yu, and Yin Zhang. 2024. Teaching small language models reasoning through counterfactual distillation. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 5831–5842, Miami, Florida, USA. Association for Com…
  7. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. MiniLLM: Knowledge distillation of large language models. In The Twelfth International Conference on Learning Representations. find
  8. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.
  9. Or Honovich, Thomas Scialom, Omer Levy, and Timo Schick. 2023. 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), pages 14409–14428, Toronto, Canada. Associat…
  10. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, and 1 others. 2022. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3.
  11. Tao Huang, Shan You, Fei Wang, Chen Qian, and Chang Xu. 2022. Knowledge distillation from a stronger teacher. In Advances in Neural Information Processing Systems, volume 35, pages 33716–33727. Curran Associates, Inc. find
  12. Harold Jeffreys. 1946. An invariant form for the prior probability in estimation problems. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 186(1007):453–461. find
  13. Jongwoo Ko, Tianyi Chen, Sungnyun Kim, Tianyu Ding, Luming Liang, Ilya Zharkov, and Se-Young Yun. 2025. DistiLLM-2: A contrastive approach boosts the distillation of LLMs. In Forty-second International Conference on Machine Learning.
  14. Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Seyoung Yun. 2024. Distillm: Towards streamlined distillation for large language models. In The Forty-first International Conference on Machine Learning. ICML.
  15. Daphne Koller and Nir Friedman. 2009. Probabilistic graphical models: principles and techniques. MIT press. find
  16. Hyoje Lee, Yeachan Park, Hyun Seo, and Myungjoo Kang. 2023. Self-knowledge distillation via dropout. Computer Vision and Image Understanding, 233:103720. find
  17. 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.
  18. Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, and 21 others. 2025. 2 ol… find
  19. Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap, Eshaan Pathak, Giannis Karamanolakis, Haizhi Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby Kuznia, Krima Doshi, Kuntal Kumar Pal, and 16 others. 2022. Super-NaturalInstructions: Generalization via declarative instr… find
  20. Qwen, :, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, and 25 others. 2025. Qwen2.5 technical report. Preprint, arXiv:2412.15115. find
  21. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, and 1 others. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9. find
  22. Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szczechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, Debajyoti Datta, and 21 others… — Multitask Prompted Training Enables Zero-Shot Task Generalization
  23. Makoto Shing, Kou Misaki, Han Bao, Sho Yokoi, and Takuya Akiba. 2025. TAID: Temporally adaptive interpolated distillation for efficient knowledge transfer in language models. In The Thirteenth International Conference on Learning Representations. find
  24. Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, Louis Rouillard, Thomas Mesnard, Geoffrey Cideron, Jean bastien Grill, Sabela Ramos, Edouard Yvinec, Michelle Casbon, Etienne Pot, Ivo Penche… find
  25. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, and 1 others. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
  26. Chaoqi Wang, Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. 2024. Beyond reverse KL: Generalizing direct preference optimization with diverse divergence constraints. In The Twelfth International Conference on Learning Representations. find
  27. Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. 2022. Finetuned language models are zero-shot learners. In International Conference on Learning Representations.
  28. 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), pages 10817– 10834, Toronto, Canada. Association for Computati…
  29. Taiqiang Wu, Chaofan Tao, Jiahao Wang, Runming Yang, Zhe Zhao, and Ngai Wong. 2025. Rethinking Kullback-Leibler divergence in knowledge distillation for large language models. In Proceedings of the 31st International Conference on Computational Linguistics, pages 5737–5755, Abu Dhabi, UAE. Associati…
  30. Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. 2024a. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.02385.
  31. Songming Zhang, Xue Zhang, Zengkui Sun, Yufeng Chen, and Jinan Xu. 2024b. Dual-space knowledge distillation for large language models. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 18164–18181, Miami, Florida, USA. Association for Computational Ling…
  32. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, and 1 others. 2023. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36:46595–46623.
  33. Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A. Smith, Daniel Khashabi, and Hannaneh Hajishirzi. 2023. 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… find
  34. Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/jung-2025-todi-token-wise-distillation-fine/