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

• URL: http://arxiv.org/abs/2505.16297v1
• Date: Thu, 22 May 2025 06:51:16 GMT
• Title: ToDi: Token-wise Distillation via Fine-Grained Divergence Control
• Authors: Seongryong Jung, Suwan Yoon, DongGeon Kim, Hwanhee Lee,
• Abstract summary: Token-wise Distillation (ToDi) is a novel method that adaptively combines Forward KL and Reverse KL per token using a sigmoid-based weighting function.<n>ToDi consistently outperforms recent distillation baselines using uniform or less granular strategies.
• Score: 3.6152232645741025
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Large language models (LLMs) offer impressive performance but are impractical for resource-constrained deployment due to high latency and energy consumption. Knowledge distillation (KD) addresses this by transferring knowledge from a large teacher to a smaller student model. However, conventional KD, notably approaches like Forward KL (FKL) and Reverse KL (RKL), apply uniform divergence loss across the entire vocabulary, neglecting token-level prediction discrepancies. By investigating these representative divergences via gradient analysis, we reveal that FKL boosts underestimated tokens, while RKL suppresses overestimated ones, showing their complementary roles. Based on this observation, we propose Token-wise Distillation (ToDi), a novel method that adaptively combines FKL and RKL per token using a sigmoid-based weighting function derived from the teacher-student probability log-ratio. ToDi dynamically emphasizes the appropriate divergence for each token, enabling precise distribution alignment. We demonstrate that ToDi consistently outperforms recent distillation baselines using uniform or less granular strategies across instruction-following benchmarks. Extensive ablation studies and efficiency analysis further validate ToDi's effectiveness and practicality.

Related papers

• KDRL: Post-Training Reasoning LLMs via Unified Knowledge Distillation and Reinforcement Learning [72.53466291156604]
  We present textbfKDRL, a textitunified post-training framework that jointly optimize a reasoning model through teacher supervision (KD) and self-exploration (RL)<n>We first formulate a unified objective that integrates GRPO and KD, and systematically explore how different KL approximations, KL coefficients, and reward-guided KD strategies affect the overall post-training dynamics and performance.
  arXiv  Detail & Related papers  (2025-06-02T19:46:41Z)
• Kendall's $τ$ Coefficient for Logits Distillation [33.77389987117822]
  We propose a ranking loss based on Kendall's $tau$ coefficient, called Rank-Kendall Knowledge Distillation (RKKD) RKKD balances the attention to smaller-valued channels by constraining the order of channel values in student logits. Our experiments show that our RKKD can enhance the performance of various knowledge distillation baselines and offer broad improvements across multiple teacher-student architecture combinations.
  arXiv  Detail & Related papers  (2024-09-26T13:21:02Z)
• Multi-Granularity Semantic Revision for Large Language Model Distillation [66.03746866578274]
  We propose a multi-granularity semantic revision method for LLM distillation. At the sequence level, we propose a sequence correction and re-generation strategy. At the token level, we design a distribution adaptive clipping Kullback-Leibler loss as the distillation objective function. At the span level, we leverage the span priors of a sequence to compute the probability correlations within spans, and constrain the teacher and student's probability correlations to be consistent.
  arXiv  Detail & Related papers  (2024-07-14T03:51:49Z)
• Direct Preference Knowledge Distillation for Large Language Models [73.50849692633953]
  We propose Direct Preference Knowledge Distillation (DPKD) for large language models (LLMs)<n>We re-formulate KD of LLMs into two stages: first optimizing and objective consisting of implicit reward and reverse KL divergence.<n>We prove the value and effectiveness of the introduced implicit reward and output preference in KD through experiments and theoretical analysis.
  arXiv  Detail & Related papers  (2024-06-28T09:23:40Z)
• Sinkhorn Distance Minimization for Knowledge Distillation [97.64216712016571]
  Knowledge distillation (KD) has been widely adopted to compress large language models (LLMs) In this paper, we show that the aforementioned KL, RKL, and JS divergences respectively suffer from issues of mode-averaging, mode-collapsing, and mode-underestimation. We propose the Sinkhorn Knowledge Distillation (SinKD) that exploits the Sinkhorn distance to ensure a nuanced and precise assessment of the disparity between teacher and student distributions.
  arXiv  Detail & Related papers  (2024-02-27T01:13:58Z)
• Robustness-Reinforced Knowledge Distillation with Correlation Distance and Network Pruning [3.1423836318272773]
  Knowledge distillation (KD) improves the performance of efficient and lightweight models.<n>Most existing KD techniques rely on Kullback-Leibler (KL) divergence.<n>We propose a Robustness-Reinforced Knowledge Distillation (R2KD) that leverages correlation distance and network pruning.
  arXiv  Detail & Related papers  (2023-11-23T11:34:48Z)
• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes [44.97759066341107]
  Generalized Knowledge Distillation (GKD) trains the student on its self-generated output sequences by leveraging feedback from the teacher. We demonstrate the efficacy of GKD for distilling auto-regressive language models on summarization, translation, and arithmetic reasoning tasks.
  arXiv  Detail & Related papers  (2023-06-23T17:56:26Z)
• Knowledge Distillation Performs Partial Variance Reduction [93.6365393721122]
  Knowledge distillation is a popular approach for enhancing the performance of ''student'' models. The underlying mechanics behind knowledge distillation (KD) are still not fully understood. We show that KD can be interpreted as a novel type of variance reduction mechanism.
  arXiv  Detail & Related papers  (2023-05-27T21:25:55Z)
• MixKD: Towards Efficient Distillation of Large-scale Language Models [129.73786264834894]
  We propose MixKD, a data-agnostic distillation framework, to endow the resulting model with stronger generalization ability. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the error and the empirical error. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.
  arXiv  Detail & Related papers  (2020-11-01T18:47:51Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.