ToDi: Token-wise Distillation via Fine-Grained Divergence Control
ToDi adaptively combines FKL and RKL per token via probability ratio, significantly improving distillation accuracy.
Key Findings
Methodology
This work employs gradient analysis to reveal that FKL promotes under-estimated tokens, while RKL suppresses over-estimated ones, highlighting their complementary roles. Based on this, ToDi introduces a token-wise adaptive weighting mechanism, using the teacher-student probability log-ratio as input to a sigmoid function. This dynamic weighting adjusts the contribution of FKL and RKL for each token, enabling fine-grained distribution alignment. Experiments on datasets like Dolly-15k, GPT2-120M, and TinyLLaMA-1.1B compare ToDi against baselines such as FKL, RKL, JS, SKL, SRKL, and AKL, demonstrating superior performance in instruction-following benchmarks. Ablation studies confirm the stability and efficiency of the method, showing that token-wise adjustment outperforms uniform weighting strategies.
Key Results
- Across five instruction datasets, ToDi achieves the highest average ROUGE-L scores, surpassing all baselines by 2-3 points; for instance, on GPT2-120M, it reaches 18.66 versus 18.12 (FKL) and 18.38 (RKL).
- In diverse model configurations, including GPT2-1.5B and LLaMA2-7B, ToDi maintains robust improvements, especially on complex tasks, with consistent gains in accuracy and stability.
- Ablation results show that token-wise dynamic weighting yields 2-3 points higher scores than fixed global weights, with faster convergence and more stable training curves.
Significance
This research advances knowledge distillation by introducing a fine-grained, token-level divergence control mechanism, addressing the limitations of uniform loss application. It enhances the ability of smaller models to closely mimic large teachers, especially in complex, resource-constrained scenarios. The method bridges a gap between theoretical understanding of divergence roles and practical efficiency, offering a scalable solution for deploying high-performance LLMs in real-world applications, thus impacting both academia and industry.
Technical Contribution
The paper develops a theoretically grounded framework where the gradient behaviors of FKL and RKL are exploited to design a token-wise adaptive weighting function based on the teacher-student probability ratio. The sigmoid-based weight ensures smooth, monotonic adjustment, maintaining linear complexity with respect to vocabulary size. Empirical results demonstrate that this approach significantly outperforms existing static or globally adaptive methods, providing a new paradigm for fine-grained distribution matching in knowledge distillation.
Novelty
This work is the first to incorporate the complementary roles of FKL and RKL into a token-wise dynamic weighting scheme driven by the teacher-student probability ratio. Unlike prior methods that uniformly apply divergence across the vocabulary, ToDi enables fine-grained, per-token control, leading to more precise distribution alignment. The theoretical analysis and extensive experiments establish its novelty and effectiveness, filling a critical gap in the literature.
Limitations
- The method relies on accurate estimation of teacher-student probability ratios; in scenarios with extreme prediction biases, the weighting may be less effective.
- While maintaining linear complexity, the per-token computation still incurs additional overhead compared to simpler uniform strategies, especially with very large vocabularies.
- Generalization to multi-modal or multi-task settings remains to be validated, and the approach may require adaptation for such scenarios.
Future Work
Future directions include extending ToDi to multi-modal distillation, integrating reinforcement learning for adaptive weight tuning, and exploring hybrid models combining sparse and dense representations. Additionally, optimizing computational efficiency for extremely large vocabularies and testing in real-world deployment environments are promising avenues.
AI Executive Summary
Large language models (LLMs) have revolutionized natural language processing, yet their deployment faces challenges due to high computational costs and resource demands. Knowledge distillation (KD) offers a promising solution by transferring knowledge from a large teacher model to a smaller student. However, traditional KD methods like Forward KL (FKL) and Reverse KL (RKL) apply uniform divergence losses across the entire vocabulary, neglecting token-specific prediction discrepancies. This coarse approach limits the student’s ability to accurately mimic the teacher’s distribution, especially at a fine-grained level.
In this context, the paper introduces ToDi, a novel token-wise distillation method that adaptively combines FKL and RKL for each token based on the teacher-student probability ratio. By analyzing the gradient behaviors, the authors reveal that FKL encourages under-estimated tokens, while RKL suppresses over-estimated ones, highlighting their complementary roles. ToDi leverages this insight by computing a dynamic weight via a sigmoid function, which smoothly adjusts the contribution of each divergence per token. This fine-grained control enables the student model to better align with the teacher, capturing subtle distribution nuances.
Extensive experiments across instruction-following benchmarks demonstrate that ToDi consistently outperforms existing methods, achieving higher ROUGE-L scores and better human evaluation metrics. For example, on the Dolly-15k dataset, ToDi improves average scores by 2-3 points over baselines like AKL and fixed-weight approaches. Ablation studies confirm that token-wise dynamic weighting significantly enhances performance and training stability, with linear complexity relative to vocabulary size. The method’s robustness across different models and tasks underscores its practical value.
This work marks a significant step forward in knowledge distillation, providing a theoretically grounded, efficient, and highly effective approach for fine-grained distribution matching. It opens new avenues for deploying high-quality LLMs in resource-constrained environments and sets a foundation for future multi-modal and multi-task extensions, promising broader impacts in AI model compression and efficiency.
Deep Dive
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.