TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models
Authors: Makoto Shing
Kou Misaki
Han Bao
Sho Yokoi
Takuya Akiba
Venue: ICLR 2025 Spotlight (2025)
Keywords: Lanauge Models
Knowledge Distillation
Primary area: generative models
TL;DR: We propose TAID, a novel knowledge distillation method for language models that uses a time-dependent intermediate distribution to dynamically bridge student-teacher gaps, addressing common challenges in distilling large language models.

Abstract: Causal language models have demonstrated remarkable capabilities, but their size poses significant challenges for deployment in resource-constrained environments. Knowledge distillation, a widely-used technique for transferring knowledge from a large teacher model to a small student model, presents a promising approach for model compression.
A significant remaining issue lies in the major differences between teacher and student models, namely the substantial capacity gap, mode averaging, and mode collapse, which pose barriers during distillation.
To address these issues, we introduce $\textit{Temporally Adaptive Interpolated Distillation (TAID)}$, a novel knowledge distillation approach that dynamically interpolates student and teacher distributions through an adaptive intermediate distribution, gradually shifting from the student's initial distribution towards the teacher's distribution. We provide a theoretical analysis demonstrating TAID's ability to prevent mode collapse and empirically show its effectiveness in addressing the capacity gap while balancing mode averaging and mode collapse.
Our comprehensive experiments demonstrate TAID's superior performance across various model sizes and architectures in both instruction tuning and pre-training scenarios. Furthermore, we showcase TAID's practical impact by developing two state-of-the-art compact foundation models: $\texttt{TAID-LLM-1.5B}$ for language tasks and $\texttt{TAID-VLM-2B}$ for vision-language tasks.
These results demonstrate TAID's effectiveness in creating high-performing and efficient models, advancing the development of more accessible AI technologies.

Published as a conference paper at ICLR 2025
TAID: TEMPORALLY ADAPTIVE INTERPOLATED DIS- TILLATION FOR EFFICIENT KNOWLEDGE TRANSFER IN
LANGUAGE MODELS
Makoto Shing', Kou Misaki', Han Bao", Sho Yokoi Takuya Akiba' 'Sakana AI, "Kyoto University, 'NINIAL, "Tohoku University, 'RIKEN [mkshing, kou. miiuakiiski takaai,takiaaaaaaaaaaaaaaaaaaaaa aaaaaaaaaaaaa.... ai, baoli u. ac. yokoi@ninjal.ac.
ABSTRACT
Causal language models have demonstrated remarkable capabilities, but their size poses significant challenges for deployment in resource-constrained environ- ments. Knowledge distillation, a widely-used technique for transferring knowl- edge from large teacher model to small student model, presents a promising approach for model compression. A significant remaining issue lies in the major differences between teacher and student models, namely the substantial capacity gap, mode averaging, and mode collapse, which pose barriers during distillation.s To address these issues, we introduce Temporally Adaptive Interpolated Distil- lation (TAID), a novel knowledge distillation approach that dynamically interpo- lates student and teacher distributions through an adaptive intermediate distribu- tion, gradually shifting from the student's initial distribution towards the teacher's s distribution. We provide a theoretical analysis demonstrating TAID's ability to prevent mode collapse and empirically show its effectiveness in addressing the capacity gap while balancing mode averaging and mode collapse. Our com- prehensive experiments demonstrate TAID's superior performance across various model sizes and architectures in both instruction tuning and pre-training scenar- ios. Furthermore, showcase TAID's practical impact by developing two state- of-the-art compact foundation models: FAID-LLM- 5B for language tasks and TAID-VIM-2 2B for rision-languaee tasks. These results demonstrate TAID's ef- fectiveness in creating opment of more accessible AI technologies.
and efficient models, advancing the devel-
INTRODUCTION
Large language models are too large. Causal language models (LMMs) are increasingly becoming
essential tools across various sectors (Malinka et al., (2023) Wu et al. 2023 Zhang et al. 2023a He
al. 2024). Scaling data size, model size, and training steps has been the primary approach to
improve LM performance (Kaplan et al. 2020 Hoffmann et al. 2022 OpenAl et al. 2024). ) leading to rapid advancements both proprietary and open-source L.Ms (Touvron et al. [023] Abdin al.)
2024 Yang al.) 2024). However, the success of large LMs creates challenges: they are too large for edge devices (Qu et al., 2024| Thawakar et al. 2024, Liu et al. 2024), have decoding times too long for real-time applications (Wan et al. 2023 Leviathan et al. [023] Miao et al.) 2024), and
consume significant energy resources (Luccioni et al.
2023) Faiz et al. 22024). This paradox of scale
hinders the widespread deployment and use of LMs despite their potential and high demand.
Knowledge distillation offers promising prescription. One promising approach to developing compact yet high perfhheeforming models is knowledge distillation (KD) (Hinton et al. 2015). KD aims to transfer the knowledge, specifically the predicted distributions, from well-trained, high-capacity teacher model to a more compact student model, often achieving better performance than small models trained solely (Bucillundefined al. 2006 Ba & Caruana, 2014: Hinton et al.) 2015). In the context of compressing large LMs, KD is becoming a mainstream approach, with many specialized
KD methods actively being developed (Xu al. 2024 Team et al. 2024 Muralidharan al. 2024)

---

Published as a conference paper at ICLR 2025
Standard Knowledge Distillation
Ours: Temporally Adaptive Interpolated Distillation (TAID)
Direct optimization of student towards fixed teacher distribution
Adaptive gradual optimization student
via time-dependent intermediate teacher distributions
student
teacher
student
teacher
Time (
Figure I: Comparison of standard KD and TAID. (Left) Standard KD methods typically employ direct optimization towards fixed teacher distribution. (Right) TAID creates dynamic bridge through adaptive, time-dependent intermediate teacher distributions (green dashed lines), enabling gradual optimization of the student. This approach facilitates flexible transition from the student's initial distribution towards the teacher's distribution over time, effectively addressing the capacity gap and balancing knowledge transfer across varying model sizes.
The formidable, unresolved challenge of teacher-student differences. Nevertheless, KD is a flawless method, and two significant issues remain, both stemming from the differences between teacher models and the student models.
(i) Capacity gap - the substantial capacity gap between large teacher model and compact student model makes effective knowledge transfer more difficult (Mirzadeh et al., 2020 Cho & Hariha- ran 2019: Zhang et al. (023b). As LMs continue to grow in size and complexity, this capacity gap becomes increasingly pronounced, making even more challenging to distill knowledge ef- fectively. (ii) Mode averaging and mode collapse due to the disparity in model capacity, KD methods often struggle with mode-averaging and mode-collapse issues, where student models either fail to oversmooth rich output distributions of teacher model or become overly focused on specific modes (Wen et al., 2023:
Gu et al. 2024 Agarwal et al. 2024).
A new method to overcome the teacher-student difference. To overcome the fundamental issue of differences between teacher and student models, we introduce Temporally Adaptive Interpolated Distillation (TAID), new approach KD for LMs. TAID reduces the gap between teacher and student model throughout the training process by dynamically introducing an intermediate teacher that interpolates teacher and student models provide target distribution with a modest capabil- ity (see Figure l). This simple technique allows for learning higher-quality student model than with existing KD methods (Section 6). scales student's performance with teacher's size even un- der large capacity gaps (Section 6.3. 2). and suppresses mode-averagíng and mode-collapse issues theoretically and empirically (Section and (63.3).
Our main contributions to this paper are as follows:
We introduce TAID (Section 3). a new knowledge distillation method that reimagines the distilla- tion process as a dynamic, adaptive knowledge transfer from student to teacher distributions. This approach addresses common challenges in distilling large language models.
We provide theoretical analysis of TAID (Section 4 with a regression model as proxy to the language modeling objective, demonstrating its ability to prevent mode collapse in the distillation process. This theoretical guarantee sets TAID apart from traditional self-distillation methods, which can suffer from mode collapse.
We conduct extensive experiments (Section across various model sizes and architectures, demonstrating TAID's 's superiority in both instruction tuning and pre- we experimentally reveal TAID's robustness to capacity gaps (Section 6.3.2. and its ability bal- ance between mode averaging and mode collapse, unlike existing KD methods (Section 6.3.3
p pre-training scenarios. Moreover,
We demonstrate TAID's practical impact by developing two ofate-of-the compact models (Sec- tion 7 TAID- -ILD--LMM-1 1 5B achieves the best performance for language models under 2B param- eters, while TAID VLM- 23 outperforms vision- Languulanguugge models up to 4B parameters, showcas- ing TAID's effectiveness across different domains.
2

---

Published as a conference paper at ICLR 2025
2 PRELIMINARIES
Problem setting for language model distillation. A language model is defined probability distribution p over token sequences = (81, .2..... ys) E yus. where Y is the vocabulary and is the sequence length. The distribution is obtained by applying the softmax function logit values: p(y= 5²*)= = soflmax(logit, (V+ y y²*)) = exp(logit, (v.lvv "omlv>"))//>> "D/// exp(logit, ())). The model
satisfies the autoregressive property: ply) II ply= I y yo where y : (y), U2,.... Vy-1). and p(ys =)* = p(v1) for = 1. In KD for language models, we aim to transfer knowledge from a well-trained teacher model p to a parametric student model qg. The objective is to find parameters 0 that minimize distance measure J between their distributions.
Traditional knowledge distillation approaches. |Hinton al.) (2015) introduced KD using the Kullback-Leibhe (KI.) divergence, which is formulated for language models as: JKL 4o) == p(zal="*))
$/ EE= Ew.ey p(ys y y log qo(y= yy However, KD based on the standard KL divergence often suffers from the mode-averaging problem, where a student model attempts to aggressively cover all modes of a teacher distribution despite being incapable, potentially resulting a over- smoothed and less accurate distribution (Wen et al.] [0233 Gu et al. 2024). To address this, Wen al. (2023) proposed using the Reverse KL (RKL) divergence: JRKL (p. qo == (ge. p). While this approach mitigates the mode-averaging problem, it can lead to mode collapse, where the student model focuses only on the dominant modes of the teacher distribution.
Curse of capacity gap. Mirzadeh et al. (2020), Cho & Hariharan (2019). and ZZhang et al. (2023b) reported a curse of capacity gap, where an excessively large model can negatively impact the per- formance of the student model. This phenomenon poses a significant challenge in KD, particularly for language models. As state-of-the-ar language models continue to grow in size and complexity, the capacity gap becomes increasingly critical in developing high-performing and compact student models. Addressing the capacity gap is crucial for effectively transferring knowledge from large- scale language models to more portable ones without sacrificing performance. Our experiments (Section 6.3.2) provide empirical evidence of the capacity gap and demonstrate how our proposed method addresses this challenge.
3 PROPOSED METHOD: TAID
We introduce Temporally Adaptive Interpolated Distillation (TAID), novel knowledge distillation method for large language models. TAID uses a dynamic, time-dependent intermediate teacher bridge the gap between student and teacher models (see Figure I This approach facilitates smoother knowledge transfer, addressing the capacity gap and balancing mode-averaging and mode- collapse issues. We show how TAID mitigates these issues in Sections 6.3.2 and 6.3.3 respectively.
3.1 TEMPORALLY INTERPOLATED DISTRIBUTION
The key idea behind TAID is employ time-dependent intermediate teacher to bridge the gap between student and teacher models. We formally define the intermediate distribution as follows:
Definition 3.1 (TAID Interpolated Distribution). For any input sequence y E y and any output token ys E V, the TAID interpolated distribution pt is defined as:
p == softmax (11 ((l-t)-loge (v. ( + t. logit (v.lbt))
where E [0, 1] is time-dependent interpolation parameter, logit represents a detached version of the student logits (i.e., treated as a constant without being backpropagated), and logit, represents the teacher logits.
The interpolation is performed at the logit level to preserve relative confidence between predictions. The TAID objective function with the interpolation parameter is defined as the KL divergence between the intermediate distribution pt and the student distribution go:
3

---

Published as a conference paper at ICLR 2025
Definition 3.2 (TAID Objective). The TAID objective function time t is defined
JAD (p. qo) = (KL 90)
Σ 5 s=1 M.EY
Σ
Pe(ualy
log go (v(ll
(2)
We gradually increase the interpolation parameter from to I during training so that the intermedi- ate distribution pa adaptively transitions from the student's initial distribution towards the teacher's s distribution. Refer to Section 3.2 for the scheduling of the interpolation parameter. The detached % in pr ensures that we only optimize the student model go in the denominator of the KL divergence, effectively treating the intermediate distribution as a target.
The intermediate distribution provides a crucial advantage in addressing the capacity gap and mode- averaging/collapse issues. By smoothly transitioning from the student's initial distribution to the teacher's distribution, TAID facilitates gradual transfer of knowledge. This approach effectively mitigates issues associated with significant capacity gaps between teacher and student models. This be understood as follows: When t is small, the student model is encouraged focus its own modes, reinforcing its unique characteristics. In this phase, TAID behaves similarly to self- distillation (using the student model the teacher), which amplifies generalization by sparsifying the model (Mobahi et al.] (020). Thus, the student model tends to capture dominant features of the student's distribution. As increases, the student gradually incorporates the teacher's knowledge, capturing more nuanced and rich signals from the teacher distribution. This balanced approach re- sults in a student model that not only captures the essential knowledge from the teacher but also maintains its ability to generalize effectively. Despite TAID's relevance to self-distillation the in- terpolation parameter is essential to avoid mode collapse, which self-distillation cannot escape. We will theoretically demonstrate it in Section 4
3.2 ADAPTIVE INTERPOLATION PARAMETER UPDATE
While TAID demonstrates effectiveness even with simple linear increase of the interpolation pa- rameter t, we propose an adaptive update mechanism to achieve more efficient learning and im- proved accuracy. The key motivation to dynamically adjust based on the student's s learning progress. The adaptive update strategy is designed to aggressively increase in the early stages when the interpolated distribution Pr is close to the student model go. as the model fitting is not chal- lenging in this phase. As the student model approaches the teacher model, the increase in t becomes more gradual, allowing for careful fitting to the more complex teacher distribution.
Our adaptive update strategy is based the relative change in the objective function: 8n == (J (/J(Emm + c). where J((2) is the value of the TAID objective function inter- polation parameter ln, in is the interpolation parameter at step n. and is a small constant to prevent division by zero. We update in using a momentum-based approach to smooth out short-term flue- tuations: mn fm,-1 (1 3)6n, where 3 is the momentum coefficient. The interpolation parameter is then updated as: ta + mi., max(funsar fn-1 + a sigmoid(m,.) (1 fn),))), where is the step size for and linear is linear increase schedule
lower bound for t. To allow for flexible initialization, is set to a start value sunn which is a hyperparameter. The complete TAID
training procedure is summarized in Algorithm in Appendix A
This update mechanism allows for more aggressive increases in during the early stages of train- ing when the student is learning rapidly (high &)), and more gradual increases as the student model approaches the teacher's complexity (low 8,). The sigmoid function bounds the update, ensuring stable learning, while the max and min operations guarantee monotonic increase within the pre- defined range. A detailed analysis of how different a values affect the behavior of and the learning dynamics presented in Section 6.3.1
4 THEORETICAL ANALYSIS
TAID distills from the intermediate distribution Pt- partially containing the student model qo the mixture component. This may apparently cause the collapse because student's modes are amplified repeatedly during the fitting recursion. Such collapse phenomenon has been theoretically observed
4

---

Published as a conference paper at ICLR 2025
for self-distillate where the teacher and student models are identical (Mobahi et al. 2020). We aim to demonstrate that TAID avoids mode collapse, unlike self-distillation.
We borrow the analysis framework of Mobahi et al. (2020) to study least-square regression proxy to language modeling. In each training step. the student model is updated by fitting to the interpolated label (1 f)ye fyeacher where Y/ and yleacher the labels of the current student and teacher models, respectively, and is the interpolation parameter (being linearly increased) at the current step. Here, we suppose the student model achieves e-interpolation of the training signals so that the regression loss is minimized near-perfectly in each time step.
Theorem 4.1 (Non-collapse Nature (Informally)). Suppose we run distillation for T steps in total.
If the teacher model has sufficiently large signals so that the label is at least as large as 91 De),)), then the student model does not collapse for any time t.
Notably, self-distillation inevitably collapses for sufficiently large steps (Mobahi et al. 2020 Propo- sition 4), corroborating the benefit of the intermediate distribution and its adaptive update. The formal statement and more discussions can be found in Appendix B
5 RELATED WORKS
Improving objective functions. To address the mode-averaging and mode-collapse issues that the traditional KL divergence-based methods (Section 2 face, various alternative objective func- tions have been applied to knowledge distillation. Wen et al. (2023) applied the Total Varia- tion Distance, formulated at the sequence level similar to Kim & Rush (0166) J/vD(p. qo) == Ey Ip(y) - qo (-)] Agarwal et al. (2024) utilized the Generalized Jensen Shannon (JS) Diver- gence: JGUSD qe) >_KD(p.r)
(RKD(PP7) r). where r(y) Ap(y) + (1 X)ge (y) and E [0. Additionally, Ko et al. (0244) employed the Skew KL Divergence: JSKD qa) := (p.r.. They also defined the Skew Reverse KL. Divergence as (PRKD qo) := JKL (go, r). These approaches aim to balance preserving teacher knowledge and allowing student generaliza- tion. However, they typically use fixed teacher distribution throughout distillation, potentially hindering knowledge transfer when there is a significant capacity gap between teacher and student. In contrast, our TAID method introduces a time-dependent intermediate distribution, gradually tran- sitioning from the student's initial distribution to the teacher's, mitigating the capacity gap issue and enabling more stable learning. While Skew KL divergence also adopts an intermediate distri- bution, its approach differs significantly from TAID. Skew KL divergence uses fixed intermediate distribution and transfers the teacher's 's knowledge to it, whereas TAID employs a time-dependent intermediate distribution and transfers it to the student. This distinction, particularly the dynamic nature of TAID's s intermediate distribution, makes TAID more suitable for adaptive updates of the student model the interpolation parameter changes over time (see Appendix C for detailed com- parison).
Utilizing student-generated outputs (SGOs). Recent research in KD for language models has explored utilizing on-policy data sampled from teacher and student models during training (Gu jet al., 2024. Zhang et al., 20246). Within this approach, some studies have specifically focused on leveraging student-generated outputs (SGOs) (Agarwal et al. 2024 Ko et al. 2024). While these methods show promise in improving distillation performance and addressing the distribution mismatch between training and inference due to the autoregressive nature of LMs when training on a fixed dataset (Pomerleau 1991 Ross & Bagnell 2010). they are computationally expensive for large-scale models. TAID achieves superior performance without relying on on-policy data or SGOs, offering improved computational efficiency for large-scale datasets and models (see Section 6.1 Future work could explore combining TAID with on- on-policy approaches to potentially achieve even better performance.
KD methods from image classification. KD has been extensively studied in image classification tasks, with some logit-based methods being applicable to language model distillation. Notable ex- amples include CTKD (Li et al. 2023b) and DKD (Zhao et al. 2022). which have shown remarkable performance using standard KL divergence. CTKD shares similar curriculum learning approach with TAID, gradually increasing task difficulty. CTKD achieves this through a learnable temperature parameter that modifies both student and teacher distributions. In contrast, TAID modifies only the
5

---

Published as a conference paper at ICLR 2025
Table Evaluating distillation methods for LLM instruction tuning. The MT-Bench scores after training are listed, where higher scores indicate better conversational performance. For each of the three reacher-student pairs, different distillation algorithms, including the proposed TAID method, are compared. The highest score in each column is highlighted in bold.
Teacher Student
Phi-3-mini (3.8B)
Llama-2 (6.7B) TinyLlama (1.1B)
StableLM Zephyr (2.8B)
Method
TinyLlama (1.1B)
Pythia (0.4B)
SFT
2.00
3.94
2.57
KL Hinton al.
2015
2.71
3,99
2.74
RKL
Wen ct
2023
Gu
2024
3,48
3,92
2.53
TVD
(Wen et al.
2023
3.27
3.64
2.57
Adaptive KL Wu ct a.
2024
3.27
3.77
2.64
GKD (Agarwal
2024
2.24
3.82
2.59
DistiLLM (Ko
2024
3.23
3.97
2.97
CTKD
2023b
1.78
2.84
1.39
DKD (Zhao et al. 2022
2.70
4.14
2.90
(Ours) TATD w/o adaptive update
3,44
4.18
2.88
(Ours) TAID
4.05
4.27
3.05
teacher distribution through interpolation, potentially preserving more of the student's learned infor- mation. DKD decomposes KL divergence into target and non-target class components, allowing for better weight adjustment in tasks of varying difficulty. However, these image classification-base methods are not sufficiently effective language modeling due to the unique characteristics of the language domain. We experimentally verified in Section 6.3.4] TAID addresses these challenges through its adaptive interpolation, while remaining flexible enough to be combined with methods like DKD for simpler tasks.
6 EMPIRICAL ANALYSIS
We evaluate TAID across instruction tuning and pre-training scenarios, using various model sizes and architectures. Our experiments compare TAID against state-of-the-art methods, demonstrating its superior performance and efficiency, while providing insights into its behavior across different capacity gaps and its ability to balance mode-averaging and mode-collapse issues.
6.1 INSTRUCTION TUNING
Experimental setup. For the instruction- instruction-followinng task, we used the UltraChat 200k dataset (Ding et al.) 2023) for training. Performance was assessed using MT-Bench (Zheng et al. 2023]. a benchmark designed to evaluate model's instruction-following ability, with scor- ing conducted by GPT-4. For our experiments, we utilized three teacher-student pairs: 1es-mmiii4-iiine
(Abdin et al.] 2024) as teacher with TinyLlama (Zhang et al.) 2024a) student, 1-ama-2-7b-chhe student, and StableLM Zephyr 3B (Team) 2023) as teacher with Pythia CM (Biderman et al. 2023) student. To evaluate the pure effectiveness of our distillation method, focused solely on distillation using instruction data, unlike previous studies (Gu et al. 2024 Agarwal et al., 2024 Ko al. 2024) that often perform supervised fine-tuning (SFT) before distillation include additional cross-entropy loss on pre-training corpora. Furthermore, to simulate a more practical sce- nario, used powerful teacher models trained on in-house data with open weights for distillation to smaller student models. We compared TAID against prior works, including KL divergence (Hinton
(Touvron et al. 2023) teacher with TinyLlan ama as
al. 2015), RKL (Wen et al., 2023), Total Variation Distance (TVD) (Wen et al., 2023). Adaptive KL ( Wu et al. 2024), as well as methods utilizing SGOs such as Generalized KD (GKD) (Agarwal
ct al. 2024) and DistiLLM (Ko et al., 2024). Additionally, we included two methods originally
proposed for image classification tasks: CTKD (Li et al.] 2023b) and DKD ZZaao al. 2022). to
assess their effectiveness in language model distillation. We also included a supervised fine-tuning (SFT) baseline to demonstrate the benefits of knowledge distillation. To isolate the impact of our adaptive update mechanism, we evaluated TAID both with and without this feature, where TAID without adaptive update uses linear increase of the interpolation parameter with respect to the

---

Published as a conference paper at ICLR 2025
Table 2: Evaluating distillation methods for LLM continued pre-training. The Open LLM Leaderboard scores after training are listed, with higher scores indicating better performance. The average score across the tasks (Average column) commonly used as an indicator of overall language proficiency. The highest score in each column is highlighted in bold.
Method
ARC
HellaSwag
MMLU
TrustfulQA
Winogrande
GSM8K
Average
SFT
41.38
63.66
25.89
35.64
61.25
1.21
38.17
KT. (Hinton
2015
44.97
65.43
25.11
37.95
63.22
2.80
39.91
TVDJ Wen
43.52
64,50
25.95
36.38
63.14
2.96
39.41
Adaptive KL Wü a 2024
43.77
63.09
26.04
36.42
63.22
2.12
39.11
GIS (Agarwal ci al. 2024
44.71
65.67
25.27
37.76
62.12
3.34
39.81
Skew KT. eK A124
44.62
65.25
25.79
37.45
62.51
3.41
39.84
Skew RKL al.
2024
44.11
64.80
36.76
62.83
3.03
39.60
(Ours) TAID
45.48
65.43
25.43
37.92
63.38
2.96
40.10
training steps. Detailed
and implementation specifics for TAID and all baseline
methods are provided in Appendix D.I
presents Results. Presentss the MTT-Bench scores for all methods across the three different teacher- student pairs. Our proposed TAID method consistently outperforms all baseline methods, including those proposed for image classification (CTKD and DKD) and methods utilizing SGOs such as GKD and DistiLLM. Notably, TAID achieves superior performance without relying on expensive SGO sampling strategies, resulting in significantly faster training times- approximaaa times faster than DistiLLM and 10 times faster than GKD. This combination of superior performance and computational efficiency, achieved without SGOs, makes TAID particularly attractive for real- wooldd applications where both model quality and training speed are crucial. An ablation study comparing TAID with and without adaptive updates shows improvements ranging from 2.2% to 17.7% across different teacher-student pairs, underlining the importance of our proposed adaptive mechanism.
6.2 PRE-TNAINING
Experimental setup. Due the limited resources, we performed continued pre-trate trainingnn initializ- ing the student model with a pre-trained model and further refining it through additional pre-training using distillation. We used the first 10% of the SmolLM-Corpus (Ben Allal et 2024) dataset, amounting to approximately 20 billion tokens. We used al. 2024) as the teacher model and TinyLlama as the student model. Similar to our instruc- tion tuning experiments, we focused solely on distillation without additional supervised fine-tuning or pre-training losses. Due to the computational cost associated with sampling from the student model in large-scale pre-training and the absence of prompts as in instruction-following tasks, we adapted the baseline methods to use only their objective functions without SGOs. We compared TAID against these modified baselines, including KL divergence, TVD, Adaptive KL, GJS (used in GKD), and Skew KL/RKL (used DistiLLM). To evaluate the pre-trained models, we followed the Open LLM Leaderboard (Beeching al. 2023) methodology, which is commonly used to assess the underlying capabilities of models through few-shot evaluation. This methodology includes six di- verse tasks, with evaluation settings and metrics adhering the Open LLM Leaderboard standards. Detailed hyperparameters and implementation specifics are provided in Appendix D.2
ruct (Abdin
Results. Table 2 presents the results of our pre- pre-training experiments. Following the standard prac- tice in the L.LM community, we reported the average scores across diverse tasks. TAID achieves the highest average score across all six tasks, outperforming all baseline methods. This superior average performance demonstrates TAID's effectiveness in transferring knowledge from the teacher to the student model across diverse range of tasks. While TAID shows the best overall performance, it worth noting that it achieves the highest scores on two individual tasks (ARC and Winogrande) and competitive performance pled with the highest average score, underscores TAID's robustness and effectiveness in knowledge distillation for large language models.
the others. The consistently strong performance across tasks,

---

Published as a conference paper at ICLR 2025
67
1.0
Interpolation Parameter
1.5
KL
Objective ve Value
TAID
Test Accuracy (%)
52
1.4
0.8
a=
1.3
21
0.6
1e-3
1.2
Teacher
5=
TAID
linear
1.1
0.4
4000
8000
12000
4000
8000 Steps
12000
1310M
6.96
Steps
Teacher Size
Figure 2: Analysis of TAID's behavior and performance. (Left) Interpolation parameter be- havior: Higher a values lead to faster initial growth compared to linear increase, allowing for more aggressive knowledge transfer in early stages when the capacity gap is small. (Middle) Objective value comparison: TAID exhibits a more stable objective value with lower variance compared to standard KL divergence throughout training, indicating consistent learning difficulty that aligns with the student's 's evolving capabilities. (Right) Performance across different teacher sizes: TAID shows monotonic improvement and outperforms other methods as teacher size increases, demon- strating its effectiveness in addressing the curse of capacity gap.
6.3 ANALYSIS
6.3.1 ANALYSIS OF INTERPOLATION PARAMETER AND TRAINING STABILITY
We analyzed TAID's interpolation parameter and learning dynamics to validate design. Fig- ure 2 (Left) shows how different learning rates affect I's behavior over time under the setting of Section 6. with stat set to 0.4. We can confirm that is smoothly increasing thanks to our adap- tive update mechanism. Higher values lead to faster initial growth of t, enabling more aggressive early knowledge transfer, which is particularly beneficial when the capacity gap between student and teacher models is small.
Figure 2 (Middle) compares the objective value of TAID (using the intermediate distribution) with the standard KL divergence between the teacher and student during training. TAID demonstrates a constant value with low variance throughout the training process, in contrast to the higher and more variable loss of standard KI.. This stability in loss indicates that TAID's adaptive interpolation mechanism keeps the learning task at a consistent level of difficulty, aligning with the student's s current capabilities. This controlled learning environment potentially leads to more efficient and stable knowledge transfer throughout the training process.
6.3.2 PERFORMANCE ACROSS VARIOUS CAPACITY GAPS
TAID's design, which gradually transfers knowledge from the teacher model, is expected to address the curse of capacity gap described in Section [ To evaluate this, we conducted an experiment using a fixed-size student model (70m) trained with teachers of varying capacities (410M to 6.9B) from the Pythia Suite (Biderman et al. 2023). Models were trained on a random IB token subset of the SmolLM-Corpus for epoch, due to computational cost constraints. We chose the LAMBADA dataset (Paperno et al. (016)) for evaluation, as it tests a model's 's ability to predict the final word of a passage, directly assessing language modeling capability without relying on specific knowledge, making it suitable for comparing models with small-scale training.
Figure (Right) shows that TAID consistently outperforms both KI. and RKL divergence methods across all teacher model sizes. Notably, TAID exhibits consistent upward trend in performance as the teacher model size increases while KL and RKL methods show inconsistent performance trends. This inconsistency in KL and RKL methods aligns with the curse of capacity gap, where larger teacher models do not always lead to better student performance, described Section TAID's consistent improvement with larger teachers indicates its robustness in handling varying capacity gaps, making it particularly suitable for distilling knowledge from state-of-th of-te----hh-ar large language models into more compact and deployable student models.
8

---

Published as a conference paper at ICLR 2025
6.3.3 BALANCING MODE AVERAGING AND MODE COLLAPSE
To demonstrate TAID's effectiveness balancing mode-averagien and mode-collapse issues, we an- alyzed the distributions of student models trained using KL divergence, RKL divergence, and TAID. We used the trained models of the with distributions calculated from the UltraChat 200k train set.
mini-44-inse
(teacher) and TinyLlame (stu-
dent) pair in Section 6.1
Table 3 presents a summary of our analysis, showing the probability mass distribution for the head and tail of the vocabulary as ranked by the teacher model. We observe that TAID consistently main- tains probability masses between those of KL and RKL for both the head and tail of the distribution. In the head, TAID captures dominant vocabulary in the teacher's distribution more than KL, effectively avoiding the mode-averaging issue. While RKL captures the dominant vocabulary more than TAID, it significantly fails to capture low-frequent vocabulary in the tail of the teacher distribu- tion, which TAID captures reasonably, preventing the mode- collapse issue. These results indicate that TAID successfully navigates the trade-off between mode averaging and mode col- lapse, achieving more balanced and faithful representation of the teacher's 's distribution across both common and rare tokens. This balanced approach contributes to TAID's superior perfor- mance in knowledge distillation tasks, it more effectively captures the full spectrum of the teacher's ' knowledge while maintaining focused distribution.
Table 3: Probability mass distribu- tion analysis. Head: sum of proba- bilities for top-10 tokens. Tail: sum of probabilities for tokens in the 80- 100th percentile
Method
Head
Tail
KL
0.216
40.2 x x10-7
RKL
0.227
8.1 x10-7
TAID
0.218
39.0 x10-7
6.3.4 COMPARISON WITH IMAGE CLASSIFICATION TASKS
Our experiments revealed that KD methods developed for image classification, such CTKD (Li et al., 2023b) and DKD (hhao et al.] 2022), underperform in language model distillation. We hypothesize that this is due to fundamen- tal differences in the distributions between language model- ing tasks and image classification tasks. Figure B illustrates the entropy of the distribution and the probabilities of ground- truth classes (target-class probabilities) for two representative models: ResNet-56 (He et al. 2016) for image classification and GPT-2 (Radford et al. 2019) for language modeling 2 Im- age classification typically involves predicting a one-hot dis- tribution with high target-class probability and low entropy. In contrast, language modeling predicts a more diverse prob- ability distribution, resulting in lower target-class probabili- ties and higher entropy. These characteristics lead to two key challenges in language model distillation. First, there is an increased susceptibility
Entropy
Target-class Probabilly
1.0
0.8
0.6
0.4
0.2
0.0
Res 56
GPT-2
AesNet 56 GPT-2
Figure 3: Comparison between im- age classification and language mod- eling tasks. Language modeling (GPT-2) exhibits significantly higher entropy and lower target-class proba- bilities compared to image classifica- tion (ResNet-56). These fundamental differences highlight the unique chal- lenges in language model distillation.
mode collapse, the model can easily be pulled toward non-target modes. Second, language modeling poses a significant challenge for smaller models with limited capacity: predicting extremely low-frequency classes. This difficulty is compounded by power law distribution of word frequencies (Zipf's law), resulting in a large number of extremely low-frequency classes in the long tail of the distribution. To test this hypothesis and to assess TAID's flexibility, we evaluated TAID on multiple image classification tasks (results in Ap- pendix De3). While gains were modest on CIFAR- -000 TAID consistently outperformed CTKD and DKD on the more complex ImageNet task. This aligns with our observation that ImageNet (en- tropy: 6.67, target-class probability: 0.00130) presents a more challenging distribution compared to CIFAR- 100 (entropy: 0.485, target-class probability: 0.613). These findings highlight the need for distillation methods tailored to language modeling's unique challenges. TAID's 's strong performance
TTpically, probabilities range from 10 to 10-2 for Head tokens and from 10 10 10 for Tail tokens. For this analysis, used the CIFAR-100 (Krizhevsky 2009) dataset for ResNet-56 and the Open Web- Text (Gokaslan & Cohen 2019) dataset for GPT-2.

---

Published as a conference paper at ICLR 2025
Table 4: Performance TAID-LIM-1. 5B, our new otate-of-the-art LLM for models under 2B parameters. See Table for task breakdown.
Table 5: Performance of TAID-VIM-22 our new state-of-the-art VLM for models to 4B parameters. See Table 10 for task breakdown.
Model
vightEval (t)
Model
Open-VI.M-LB (1)
Qwen2-1. 5B
Yang 2024
46.19
PaliGemma (Beyer
2024
46.56
Phi-1. 1.58 Lie
50.39
Mini CPmiC Van 2024
47.93
StablelM-22
68 Rellagente et al 2024
51.24
Uni-3-vision
Abdin et al
2024
53.60
Smol LM- 7B Allal al.
2024
51.31
InternVi2-22
Chen et
2024
53.96
MAID-LIM-1.55
52.27
TAID-VIM-2B
56.43
across domains, particularly in complex tasks, demonstrates its potential as versatile approach to knowledge distillation. Future work could explore its application other tasks involving long-tail distributions or complex probability predictions beyond language modeling.
7 APPLICATION TO STATE-OF-THE-ARE MODEL DEVELOPMENT
Building upon systematic evaluation of TAID, further demonstrate its effectiveness in devel- oping state-of-the models. We introduce two models: TAID-LIM-1- 5B and TAID-VIM-22 which have achieved state-of-the-art performance in their respective size categories for large lan- guage models (LLMs) and langugge models (VLMs).
TAID-LLM-1.5B. We developed TAID-LLM-1 5B, a new 1.5B 1.5B- --B-paraeeeer language model, us- ing our TAID method. Following recent conventions in evaluating language models of this size (AI-) lal et al. 2024), we evaluated using LightEval a comprehensive benchmark suite for small language models. Table Sable4]shon shows that TAID- - LLM 1144-.14 1.55..5BBB 5B achieves the highest score, setting new state-of-ther torte-of-hhearttt for models with fewer than billion parameters. Detailed settings and results can be found in Appendix E.I
TLMD-B To showcase TAID's versatility, we developed TAID-VIM-2B, new 2B- parameter vision-language model. We evaluated following the Open VLM Leaderboard proto- col (OpenCompass Contributors 2023 9 As shown in Table 5 TAID-VIM-28 achieves the high- est score among state-of-the-art Vision- langugg models up to 4B parameters, even surpassing the performance of larger models like Phi (0s-3-vision (4.2B parameters). This success highlights TAID's capability in transferring multimodal knowledge across significant capacity gaps. Detailed settings and results can be found in Appendix E.2
CONCLUSION
We introduced Temporally Adaptive Interpolated Distillation (TAID), novel knowledge distilla- tion approach that effectively addresses the challenges of compressing large language models. Our experiments demonstrated TAID's superior performance across various model sizes and architec- tures, consistently outperforming state-of-the-ard methods. The development of TAID-I.IMM] .B and TAID-VIM-2B, achieving state-of-the-an performance their categories, underscores TAID's practical impact. TAID's dynamic bridge mechanism effectively mitigates mode-averaging and mode-collapse problems, leading to more stable and efficient training. These advantages contribute to more accessible deployment of advanced language technologies in resource-constrained environ- ments. Future research could extend TAID to other distance metrics, explore non-linear interpo- lations, adapt it for multi-teacher distillation (Wan et al. 2024), and investigate its application in other modalities and tasks beyond classification. In conclusion, TAID represents a significant ad- vancement in knowledge distillation, offering both theoretical insights and practical benefits. As AI evolves, techniques like TAID will be crucial in making these advancements more accessible and deployable real-world applications.
httro:://hugiingfcccccc///
1m
https::/hugginggfaas co/ spaces http://ccccc.c.ccc/ooooaaaaaaaan 1eade__eeddrbrrr
10

---

Published as a conference paper at ICLR 2025
AUTHOR CONTRIBUTIONS
Makoto Shing and Takuya Akiba initiated this project. Makoto Shing is the main contributor who conceptualized and proposed the TAID method, designed and conducted all experiments, performed theoretical analysis, implemented the main code, wrote the initial draft of the manuscript, and was responsible for data analysis and interpretation of results. Consistently led and executed all pects of the project from inception to completion. Kou Misaki contributed to data processing for the TAID-LLM-1.5B model. Han Bao provided crucial feedback on theoretical interpretations and anal- ysis. Sho Yokoi offered valuable insights and feedback, especially based on his expertise in Natural Language Processing. Takuya Akiba served as the primary advisor throughout the project, offer- ing guidance, technical insight, advice, and supervision from inception to completion. All authors reviewed and edited the final manuscript.
ACKNOWLEDGEMENTS
The authors would like to thank Masanori Suganuma and Tianyu Zhao for providing valuable dis- cussions and feedback while drafting the text. This work is based on results obtained from a project, JPNP20017, subsidized by the New Energy and Industrial Technology Development Organization (NEDO). This work was supported by JSPS KAKENHI (Grant No. 22H05106), JST FOREST (Grant No. JPMJFR2331), and JST PRESTO (Grant No. JPMJPR24K6).
REFERENCES
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, et al. Phi-3 technical re- port: A highly capable language model locally on your phone. arXiv preprint arXiv:2404. 14219, 2024.
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, and Matthieu Geist. On-policy distillation of language models: Learning from self self-generated mistakes. In International Conference on Learning Representations, 2024.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. Smollm blazingly fast and remarkably powerful, 2024.
Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (eds.), Advances Neural Information Processing Systems. Curran Associates, Inc., 2014.
Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf. Open 11m
leaderboard.
Lttp:/p/p.eeeppee-----eeeeeeeeee eaderboardd 2023.
open 1m
Marco Bellagente, Jonathan Tow, Dakota Mahan, Duy Phung, Maksym Zhuravinskyi, Reshinth Adithyan, James Baicoianu, Ben Brooks, Nathan Cooper, Ashish Datta, et al. Stable Im 2 1.6 technical report. arXiv preprint arXiv:2402.17833 2024.
Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von URI. htaps::/hhuggngghen cot.ceeeetc/ttaattttt.s.
Werra. Smollm-corpus, 2024.
HucgingFaceTB/seen stt.///smlllle 1m corpus
Lucas Beyer, Andreas Steiner, André Susano Pinto, Alexander Kolesnikov, Xiao Wang, Daniel Salz, Maxim Neumann, Ibrahim Alabdulmohsin, Michael Tschannen, Emanuele Bugliarello, et al. Paligemma: A versatile 3b vlm for transfer. arXiv preprint Xii::240..07726, 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, Usven Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In An- dreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt. Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.

---

Published as a conference paper at ICLR 2025
Cristian Buciluundefined, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 535-541. Association for Computing Machinery, 2006.
Zhe Chen, Jiannan Wu, Wenhai Wang, Weijie Su, Guo Chen, Sen Xing, Muyan Zhong, Qinglong Zhang, Xizhou Zhu, Lewei Lu, et al. Internvl: Scaling up vision foundation models and aligning for generic visual-linguistic tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11185-24198, June 2024.
Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation, In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou Enhancing chat language models by scaling high-quality instructional conversa- tions. In Empirical Methods in Natural Language Processing, 2023.
Ahmad Faiz, Sotaro Kaneda, Ruhan Wang, Rita Osi, Prateek Sharma, and Fan Chen. LLMCarbon: Modeling the end-to-end carbon footprint of large language models. In International Conference on Learning Representations, 2024.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. httt:://kkyLion007.pii... /o/ Texet Corus 2019.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. MiniLLM: Knowledge distillation of large
language models. In International Conference on Learning Representations, 2024.
Kai He, Rui Mao, Qika Lin, Yucheng Ruan, Xiang Lan, Mengling Feng, and Erik Cambria. A survey of large language models for healthcare: from data, technology, and applications to accountability and ethics. arXiv preprint arXiv:23/0.05662 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. In NIPS Deep Learning and Representation Learning Workshop, 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, et al. Train- ing compute-optimal large language models. S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.). Advances in Neural Information Processing Systems, pp. 30016-30030. Curran Associates, Inc., 2022.
Dongfu Jiang, Xuan Hc, Huaye Zeng, Cong Wei, Max Ku, and Qian Liu. Mantis: Interleaved multi-image instruction tuning. arXiv preprint 240V::2405.01483, 2024.
Ying Jin, Jiaqi Wang, and Dahua Lin. Multi-level logit distillation. In Proceedings the IEEE/CVF Conference on Computer Vision and Puttern Recognitio, pp. 24276-24285, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
Yoon Kim and Alexander M. Rush. Sequence-level knowledge distillation. In Jian Su, Kevin Duh, and Xavier Carreras (eds.), Empirical Methods in Natural Language Processing, pp. 1317-1327. Association for Computational Linguistics, 2016
Jongwoo Ko. Sungnyun Kim, Tianyi Chen, and Se- Young Yun. DistiLLM: Towards streamlined distillation for large language models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), International Con- ference Machine Learning, pp. 24872-24895. PMLR, 2024.
12

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, Depart- ment of Computer Science, University of Toronto, 2009.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference Machine Learning, pp. 19274-19288 PMLR, 2023.
Yuanzhi Li, Séhastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint (tti:223090054663 2023a.
Zheng Li, Xiang Li, Lingfeng Yang. Borui Zhao, Renjie Song, Lei Luo, Jun Li, and Jian Yang. Curriculum temperature for knowledge distillation. Proceedings of the AAAI Conference on Artificial Intelligence, pp. 1504-1512, 20233.
Zechun Liu, Changsheng Zhao, Forrest landola, Chen Lai. Yuandong Tian, Igor Fedorov, Yunyang Xiong. Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, MobileLLM: Optimizing sub-billion parameter language models for on-device use cases. In International Conference on Machine Learning, 2024.
Alexandra Sasha Luccioni, Sylvain Viguier, and Anne-Laure Ligozat. Estimating the carbon foot- print of bloom, 176b parameter language model. Journal of Machine Learning Research, pp- 1-15, 2023.
Kamil Malinka, Martin Peresíni, Anton Fire, Ondrej Hujnák, and Filip Janus. On the educational impact of chatgpt: Is artificial intelligence ready obtain university degree? In Innovation and Technology in Computer Science Education V. 1, pp. 47-53, 2023.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang. Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, al. Specinfer: Accelerating large lan- guage model serving with tree-based speculative inference and verification. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and 2024.
Operating Systems, Volume 3, pp.
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings the the AAAI Conference on Arlificial Intelligence, pp. 5191-5198, 2020.
Hossein Mobahi, Mehrdad Farajtabar, and Peter Bartlett. Self-distillation amplifies regularization in Hilbert space. Advances in Neural Information Processing Systems, pp. 3351-3361, 2020.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint mmXiv:2207.146673 2024.
OpenAl, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2024.
OpenCompass Contributors. Opencompass: A universal evaluation platform for foundation models. http:///hhhhabb.cooooo/eooooooee
httpaaaassspppccpppsssssssss 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Association for Computational Linguistics, pp. 1525-1534. Association for Computational Lin- guistics, 2016.
Dean A. Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation, 1991.
Guanqiao Qu, Qiyuan Chen, Wei Wei, Zheng Lin, and Xianhao Chen. Mobile edge intelligence for large language models: contemporary survey. arXiv preprint arXiv:2407.189211 2024.
Qwen Team. Qwen2.5: A party of foundation models, September 2024, URL ht qwenlm.
github. lo/blog//////o/ooo tww.2...//
13

---

Published as a conference paper at ICLR 2025
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Stephane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Yee Whye Teh and Mike Titterington (eds.), Proceedings of the Thirteenth International Conference Artificial Intelligence and Statistics, pp. 661-668. PMLR, 2010.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, al. Gemma 2: Improving open language models at a practical size. arXiv preprint arXiv:2408.00112 2024.
Stability AI Language Team. Stablelm zephyr 3b, 2023. URL
stabilttyaal/saaallelm-zzeii
Omkar Thawakar, Ashmal Vayani, Salman Khan, Hisham Cholakal, Rao M. Anwer, Michael Fels- berg, Tim Baldwin, and Eric P. Xing. Mobillama: Towards accurate and lightweight fully trans- parent gpt. arXiv preprint arXiv:2402.1684 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2300.0922 200.009888 2023.
Fanqi Wan, Xinting Huang, Deng Cai, Xiaojun Quan, Wei Bi, and Shuming Shi. Knowledge fusion of large language models. In International Conference Learning Representations, 2024.
Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, et Efficient large language models: A survey. arXiv preprint arXiv:23 12.0322..86.. 2023.
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. In Anna Rogers, Jordan Boyd-Grabes and Naoaki Okazaki (eds.), Associ- ation for Computational Linguistics, pp. 10817-10834, Association for Computational Linguis- tics, 2023.
Shijie Wu, Ozan Irsoy, Steven Lu, Vadim Dabravolski, Mark Dredze, Sebastian Gehrmann, Prab- hanjan Kambadur, and David Rosenberg. Bloomberggpt: A large language model for finance. arXiv preprint arXiv:2303. /7564, 2023.
Taiqiang Wu, Chaofan Tao, Jiahao Wang, and Zhe Zhao. Rethinking kullback-leibler divergence in knowledge distillation for large language models. arXiv preprint arXiv:2404.026377 2024.
Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can Xu, and Dacheng Tao. A survey on knowledge distillation of large language models. arXiv preprint arXiv:2402. 13166, 2024.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Yuan Yao. Tianyu Yu, Ao Zhang, Chongyi Wang. Junbo Cui, Hongji Zhu, Tianchi Cai, Haoyu Li, Weilin Zhao, Zhihui He, et al. Minicpm-v: A gpt-4v level milm on your phone. arXiv preprint a4Xiv:2408. 01 800, 2024.
Chaoning Zhang, Chenshuang Zhang. Chenghao Li, Yu Qiao, Sheng Zheng, Sumit Kumar Dam, Mengchun Zhang, Jung Uk Kim, Seong Tae Kim, Jinwoo Choi, et al. One small step for gen- erative ai, one giant leap for agi: complete survey on chatgpt in aige era. arXiv preprint arXiv:2304.0644 064888 2023a.
Chen Zhang, Yang Yang, Jiahao Liu, Jingang Wang, Yunsen Xian, Benyou Wang, and Dawei Song. Lifting the curse of capacity gap in distilling language models. Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki (eds.), Association for Computational Linguistics. Association for Computational Linguistics, 2023b.
Peiyuan Zhang, Guangtao Zeng, and Tianduo Wang. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.0m 002385, 2024a.
14

---

Published as a conference paper at ICLR 2025
Rongzhi Zhang, Jiaming Shen, Tianqi Liu, Haorui Wang, Zhen Qin, Feng Han, Jialu Liu, Simon Baumgartner, Michael Bendersky, and Chao Zhang. PLaD: Preference-basee based large language model distillation with pseudo-preferenee pairs. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Association for Computational Linguistics, pp. 15623-15636. Association for Computational Linguistics, 2024b.
Borui Zhao, Quan Cui, Renjie Song, Yiyu Qiu, and Jiajun Liang. Decoupled knowledge distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 11953- 11962, 2022.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging LLM-as-a-judge with MT-bench and chatbot arena. In seirth-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023.
A TAID TRAINING ALGORITHM
Algorithm | provides a detailed description of the TAID training procedure, including the adap- tive update mechanism for the interpolation parameter t. The TAID algorithm utilizes several key
Algorithm TAID training algorithm
Input: Learning rate 7), learning rate of the interpolation parameter a, momentum coefficient B. total iterations N, start value ixtats end value lend
2: Initialize student model parameters 0
3: Initialize 1 = start Me 0, TAID OC
(fo)
4: for each training iteration n 1 N do
5: Compute linear increase value: 4inear = Istart + (Lend (stun) n/N
6: Sample batch {1155 _x y3) } from dataset D
7:
Compute Pt,, (y=læ* (*) using Eq.
(I)
8:
Compute TAID using Eq. 2
9:
Update 0: 0 4- - V TAID
8n (J(A JTAID / (f(t-1)) +
11:
my 3mn-1 + 3)0,
AL a sigmoid .1 (1- (n) mn)
13:
tn+1 min(tend, max(tiinaas tn At))
14: end for
hyperparameters that control the behavior of the interpolation parameter i. and the adaptive update mechanism. We discuss the effects of these parameters below:
a (learning rate of t): This parameter controls the speed of the adaptive update for t. Figure 2 (Left) shows the behavior of t for different values of a, including linear increase for comparison. As a increases, we observe that t grows more rapidly in the early stages when the student model is close the initial interpolation distribution. This allows for more efficient learning when the task is relatively easy for the student.
3 (momentum coefficient): This parameter controls the smoothness the adaptive update. A higher value of 3 results in more stable updates by reducing the impact of short-term fluctuations in the objective function. In our experiments, we found that 3 value around 0.99 worked well across different scenarios.
laa (initial value of f): This parameter determines the starting point of the interpolation. It is particularly useful for skipping the initial stages of learning when the task is very easy for the student. The choice of Eat should be based on the intuitive gap between the initial student and teacher models. In our experiments, we found that values between 0.2 and 0.4 often yield good results, depending the initial similarity between the student and teacher models.
tend (maximum value of f): This parameter sets the upper limit for f, typically set to 1.0 to ensure that the final distribution matches the teacher model.
15

---

Published as a conference paper at ICLR 2025
The algorithm uses linear increase schedule (Finear) lower bound for t. ensuring that increases at least linearly over the course of training. This approach maintains the adaptive nature of TAID while guaranteeing minimum rate of progression towards the teacher distribution.
In our experiments, TAID demonstrated robust performance across various tasks with minimal hy- perparameter tuning. We usually used / = 0.99 and 5c-4, with faat typically ranging between 0.2 and 0.4, depending on the initial student-teacher similarity. While these default values often yield good results, practitioners may achieve further improvements by fine-tuning these parameters for their specific tasks and model architectures, particularly in cases that differ significantly from our experimental settings.
B THEORETICAL ANALYSIS OF MODE COLLAPSE
In this section, we formally study the mode-collapse behavior of TAID.
B.1 ANALYSIS MODEL
To study the collapse phenomenon, we leverage the analysis framework used by Mobahi et al.) 2020). We study the regression problem in the interpolation regime
1 (f(x)) - (()) y1, Sc f arg min R(f) s.t. feF
(3)
where D = {(x,, Bi) is a finite training set with d-dimensional covariates x; E X' C Rd desired loss tolerance parameter, R(f)
and one-dimensional outcome yi E R. c regularization functional, and F C R is hypothesis space. Since we are interested in a large model regime, is reasonably assumed to encompassing all measurable functions. The mean- squared loss is used in (3) instead of the KL divergence, which is convenient obtain analytical solutions later. The regularizer in the following form is considered:
f u(x, x (x)((x))
R(f)
where u symmetric kernel inducing R(J) with equality only when 0. The interpolation problem (3) may collapse depending on the teacher signals. Let us stack labels into a vector:
y [u1 1/2 yN] E R.M.
When ||y|| < Ne holds, the problem 3 has trivial solution f 0. Such collapse may happen particularly in the self-distillation paradigm because the teacher signals are (partially) given by our hypothesis itself. Thus, it is crucial investigate when and whether the non-collapse condition || > Ne satisfied to ensure that our hypothesis learns meaningful signals.
Variational problem. The Lagrangian variational problem of 3 given as follows:
N
K = arg min
(f(x)) y.)² + X (xxx,)))((x))
JEF
N
1=1
(5)
N
1
where
Σ N
())
and X-1 > 0 the Lagrange multiplier. The solution to the variational problem (5) can be analytically written down. Let be the Green function of the linear operator [x] [u(x x' fx(')xx such that
f [ ('xxx)))(xx xo)dx' = 6(x xx);
(6)
Thee interpolation regime must be distinguished from the time interpolation used in the proposed TAID.
16

---

Published as a conference paper at ICLR 2025
where 8(x) is the Dirac delta. Let G R N and 8x RVV be
G = Nx(x;))
xj)
and gx,i
N xx,xx
for all i,j E [N].
Then, the analytical solution to is given as follows (Mobahi al. 2020 Proposition 1):
fx(x) = - (AI + G)
If diagonalize G (which is positive definite) as G V DV, the prediction vector over the training inputs , is given as
f (x)) ... (xx)) - +
(8)
The solution (8) is essentially a nonlinear extension of the ridge estimator. Note that V RN is an orthogonal matrix and D = ti. dN) has positive eigenvalues solely.
Importantly, is the solution the variational problem (5), which is parametrized by À satisfving y:)² - O. Solving this in À is hard because of its non non-linearity, but Mobahi al. 2020 Eq. (24)) evaluate its upper and lower bound:
X
av Ne
A
for some [diins dmaa);
|
VNee
where dmax = max; de and dmin = min, d, Thus, the analytical solution 7 with this range of A a solution to the original interpolation problem (3) too.
Remark on connection language modeling. The interpolation formulation based on the standard (one-dimensional) regression problem, which obviously deviates from the language model- ing problem introduced in (2). Nonetheless, we believe that this formulation not only beneficial for our transparent understanding owing to its simplicity but also has connection to multi-categorical distributions. In distributional modeling, a student model go outputs probability distribution over D, and falls into mode collapse when qe has only few numbers of non-zero probabilities, that is, {o E y ( = c) > 0} < [ To deal with the multi cateircate outputs, we can extend the one-dimensional problem as follows:
N (fd(()) - vi,c)² Sc,
1
VCED, y, f = argmin R(fe) s.t.
f.eF
N
where teacher signal Vi.c given in the one-hot format such that Ecey Vi.c and Vi.e (0, for all c E y. We can follow the subsequent analysis straightforwardly. In this multi-categorical problem, model (fe)eey regarded falling into mode collapse if fe for many This is measured by the teacher signal condition ||ya|| < Ne for each c, where y. E 10 is the stacked labels for class Thus, studying (3) is directly relevant mode collapse language modeling.
B.2 FORMAL THEORETICAL STATEMENT
To study TAID in fashion of the interpolation problem 33. procedure listed in Algorithm 2 Here, the input signals yo are deemed as the well-trained teacher- can deem y1 the well-trained teacher, but the resulting distillation dynamics would not change much.
consider the following learning
Theorem B.I. Let k := dmax/dmin 1) be the condition number of G. The prediction vector
Y does not collapse, namely Yt+1 0 cannot be solution to the interpolation problem 3 if for some E [0, 11, either of the following holds:
1 1 (ro (100-7)- o(1), ro To
2T) or - kb
min
(10)
where ro ||yoll/ /olll//Nee > and o(1) is asymptotic term in the large ro limit.
17

---

Published as a conference paper at ICLR 2025
Algorithm 2 TAID learning procedure for least-square regression
Input: T number of iterations, Yo E IR- input signals
1: 1-00 0
2: while do
3:
Yt (1 +)y+ $oo
Compose intermediate teacher
4: A - 0y Ne/ (15. - Ne
Choose an appropriate X, by ( D Solve the variational problem with teacher y, and X
5: Y1+1 V D -V-
6:
++++1
7:
end while
To make the asymptotics in To work well, we need to ensure sufficienlly strong initial signals ||yol
and/or near-interpolation (small e). The first bound in (T0) is non-vacuous when 7 = (rro). Though is a rather strong requirement, the asymptotic term becomes negligible numerically with a moder- ate magnitude of ro (like 5 to 10).
To see how TAID benefits from the intermediate teacher, compare the non-collapse condition (10)
with that of self-distillation (Mobahi al. 2020 Proposition 4):
5 T0 1
(11)
We have two observations. First, TAID is beneficial in the latter phase of recursion (namely, step t closer to T), where self-distillation never escape from collapse eventually. This is intuitive feature of TAID because the intermediate teacher partly consists strong signals yo that does not de- pend on leameed student predictors. Second, TAID is worse in the early phase of recursion (namely, step è closer to 1) than self-distillation by a constant factor. Specifically, TAID and self-distillation have critical steps of collapse
O(ro/(7 + =)) and O(ro/=) respectively. To ensure that TAID learns meaningful features in the early phase, should be reasonably bounded away from O, leading to a worse critical point than self-distillation. This is a price that TAID has to pay for the stabilization in the latter phase.
By setting 7 Theorem
10P. we get more interpretable corollary, which is the formal version of
Corollary B.1.1. If initialization |yool
satisfies
1
47 KK)
yoll = Ω
/Ne
2
the prediction vector Y141 does not collapse for any t.
B.3 PROOF
Vy.. where the norm is preserved ||z,|| ly=ll We also write 2, Vý, and || //e for convenience. At each
Proof of Theorem
B.1
Subsequently,
the change-of-variable
time t, the non-collapse criterion is given by l l Ne( < ri > 1): if holds, the next update in Line would not collapse. Let A, == D D I D) We first show the second case, namely, the
18

---

Published as a conference paper at ICLR 2025
prediction avoids collapse when TT < t. Then, ž, is recursively expanded.
To
t +
2,-
Zt
T 20
T
1
A 2 + T
(12)
T
1
(1 11-1-1 1-1
1-1
t
1
A--1
T 20
+
Zo
T
T
T
t
(1
1
t 1-11
A, 12-12
1
t
t -
1
A
+
t -
T
T
T
T
T
1 t t-r T
t-8 8
T
A
Zo +
A
Zo
+
20
T
T
T
T
T=
LT=
8==0
T!
II
A
T) (T - + - 1)!
II A
=1
t
T1+1. (T -T-
T1
(T-
T
===
T=1
A,zo.
To evaluate Ar, we first look at A, for T E [0, 1]. Since A, is a diagonal matrix, its k-th element of A, can be expressed as follows:
1/ic -1 [iel/Ve-11 +1)
dx
ar /d
S
(A+)s
(13)
X7 + dx
12..l//Ne-11
|| 1/ve-1 +1)
2
where a, given in P9. The last inequalities can be formally shown by induction in 5 E Thus, the minimum singular value of A, is evaluated as follows:
[0, 1).
Omin At)
=Omin
T!
A
+
(t- (t-r)
- -T (T
1)!
A
+
7
t
T141 (T - 1)!
T1+1. (T -t--1 1)!
LT=0
7!
(1 (1-T) (T - + 1)!
A - ))
A
+ T'inn
Σ
Tt+l. (T -1-1-10 1)!
T1+1. (T 1- 1]!
t
+ amin
T
I
Omin
I
T
where the second identity holds because all matrices evaluated are diagonal. This implies
t |
| min (A,)
2
T ||ool|--
T
The last equality uses Zu Zo- Thus, the non-collapse criterion ||ž.||
/Ne holds long as
t (VNe/|@oll)T (VNe/lyyol)>)
Next, supposing & small enough such that t < Pe with 7 E (0, 1), we show that the prediction avoids collapse when f < (3 + o(1))(ro 7). To the non-collapse criterion 7t 1, first
19

---

Published as a conference paper at ICLR 2025
derive lower bound of ry:
2
2o
A (-1
+
T
NN
T VNe
IV
2 1
Zo
VNee
T
VNe
IV
min A1-1)
T
T
7
Miin (A1-1)n2-1 17
TO
13
Y
Tt-1
To
8
+1
rt-1-1
Iv
Y
(Borc-1 - B1) -/--
TO
where (a) due to the "reverse" triangle inequality and (b) is due to Mobahi (2020
Eq. (137))
(which is essentially a linear lower bound of a convex function in ro) with
30 (ro 1)2 + x(2ro - 1) (ro 1 x)²
r8xx
and B1 = (ro - 1 +32
By recursively lower bounding rt, we obtain the following bound:
-/ 31
[
2
2 [11- 3 3o ro
ro
31 36 1
=: Biroo
7. To
Bo
3u
where 30 (1 x))
go and B1 ==
1
B1. To derive the non-collapse condition, we solve
= 1 to derive the critical f, which is equivalent to
log (1t71(1-8)44
(_+ro(1-Po)
t
log 8o
By simple algebra,
711 + (xc 2)ro 11] (NT + +x(x 1)ro)
log
+ x-1 1 7(x 1)(x
t)) 1+r
t
log x
+ log
xiK-11
-
7 fro 1)- a 1 1 1)(N 2-Po m
))(())66100 14
1
+x 2)ro (N 1]] + (NFg +N( 1)ro]
Rla
(n 11a7
s(x- ((x-1)(7o 1)+x(r + (2 3)ro 1)x+3) m-1 (ro 2 -
710 + (x 2)ro
(x 1)
R( 1)ro
goÔ [2x -e
(xt*)r - h( x 2)+
K7o (xx
yr (2 1)
1x+1) l)7
(ro 7) 2(x
1)ro (1c 1)
where the inequality is due (in large ro) expressed as follows:
1. < logg S x 1. The last lower bound can be asymptotically
7+o(1)
-+*+0(1)
1
(ro 7) o())
7+oo1
(ro ([+o))
20

---

Published as a conference paper at ICLR 2025
Table 6: Performance comparison between TAID and Skew KL across different teacher sizes.
TAID shows consistent improvement with larger teachers, while Skew KL's performance degrades.
Method
410M
1B
2.8B
6.9B
TAID
20.82
21.17
21.70
22.01
SKL
18.65
18.50
18.28
18.20
Thus, the non-collapse condition in the second case < y+v (ro - 7) o())
Proof of Corollary B.I.7 By the non-collapse criterion TT) with y
1
1 1+ (ro o(1) T
suffices for Ye not being collapsed for any t. By solving this quadratic inequality, we can verify the statement.
C DETAILED COMPARISON WITH SKEW KL
We provide detailed comparison between TAID and Skew KL to highlight their fundamental differ- ences, focusing on two key aspects: the direction of knowledge flow and the nature of interpolation design.
The first key difference lies in the direction of knowledge flow, which can be understood through their objective functions. The TAID objective is formulated as JTAID (p. 2o JKL. (Pt, go). while the Skew KL objective takes the form (PK 9o) = (p, r). where r(y) >p(y) + (1 X) (y) and À [0, 1]. In TAID, the interpolated distribution pt teaches the student model ga. creating direct path for knowledge transfer from the interpolated distribution to the student. Conversely, in Skew KL, the teacher teaches the interpolated distribution r, establishing indirect path where the student's 's knowledge is mixed into the target distribution.
The second fundamental difference is in the design of the interpolation mechanism. TAID employs a time-dependent parameter t that gradually changes during training, enabling adaptive knowledge transfer that evolves with the student's learning progress. In contrast, Skew KL uses fixed inter- polation parameter throughout the training process, maintaining a constant mixing ratio between teacher and student distributions.
Our empirical study validates the benefits of these design choices, particularly in handling the ca- pacity gap between teacher and student models. 6 the performance comparison across different teacher sizes, demonstrating that TAID achieves consistent improvement as teacher size increases from 410M 6.9B parameters, while Skew KL's performance degrades with larger teach- ers.
D EXPERIMENTAL DETAILS
D. INSTRUCTION TUNING EXPERIMENTS
For our instruction tuning experiments, we utilized the UltraChat 200k dataset. We preprocessed the dataset by removing samples exceeding a maximum length 2048 tokens, resulting in approx- imately 150k training samples and 2k validation samples.
All models were trained for epochs using batch size of 64. We employed the AdamW optimizer with learning rate of 1e-4 and cosine learning rate scheduler. To select the best checkpoint for evaluation, we calculated the ROUGE-L score on the validation set after each epoch and chose the checkpoint with the highest score.
For our proposed TAID method, we used a momentum coefficient (5) of 0.99 across all experiments. The learning rate of t. (a) was 5e-4, The initial value of & (Iamn) was set 0.4 for the
21

---

Published as a conference paper at ICLR 2025
Table 7: Top-1 accuracies (%) on the CIFAR-100 dataset. Results for different teacher-student
pairs are shown.
Teacher Student
ResNet ResNet20
ResNet -10 ResNet
ResNet x ResNet8Xe
PPRN-0--2 WRN-
WRN-40
VGG13
Method
WRN-40
VGG8
Hinton 2015
70.66
73.08
73.33
74.92
73.51
72.93
CTKD Li 20216
71.19
73.52
73.39
75.45
73.93
73.52
DKD
71.97
74.11
76.32
76.24
74.81
74.68
MLKD Jin
72.19
74.11
77.08
76.63
75.35
75.18
(Ours) TAID
72.25
73.51
74.85
75.81
74.51
74.38
Phi thi-3 33-miii-44 mmini44-innstree pair and 0.2 for the other two pairs. The final value of (Hend) was set to 1.0 for all experiments.
Regarding baseline methods, we implemented GKD using Generalized Jensen- Shannon Divergence (GJSD) with À 0.1 as the objective function and student data fraction of 0.5. For DistiLLM, we used Skew KI. divergence with 0.1 and an initial student data fraction of 0.0. We selected the better performing skew divergence between Skew Forward KI. and Skew Reverse KI. based on the best ROUGE-L score. Following the original DistiLLM paper, we calculated the validation loss twice per epoch, totaling 10 times, leverage the Adaptive SGO scheduler. For Adaptive KL, our implementation was used since no official implementation was available. For CTKD and DKD, we followed their settings used in the training ImageNet (Deng et al. 2009).
Jensen-Shannon
In terms of computational efficiency, we observed significant differences in training times among the different methods. TAID completed its training approximately 0.7 hours per epoch hardware setup using 8 NVIDIA H100 GPUs. In comparison, DistiLL.M required about 2 hours per epoch, while GKD took approximately 9.8 hours per epoch under the same conditions. These differ- ences in training time are primarily attributed to the computational complexity of methods utilizing SGOs. TAID's ability achieve competitive performance without relying SGOs contributes to its faster training times.
D.2 PRE-INAINING EXPERIMENTS
For our pre-training experiments, we used the first 10% of the SmolLM-Corpus (Ben Allal al. 2024) dataset, which amounted to approximately 20 billion tokens.
The pre-training was conducted for epoch using a distributed setup with 80 NVIDIA H100 GPUs, each processing batch size of 8, resulting in an effective batch size of 640. We used the Adam optimizer with a learning rate of 1e-4 and a cosine learning rate scheduler.
The TAID-specific parameters for the pre-training experiments were kept consistent with those used in the Phi-3- mi mini-4k-inn 1irk-iissrruct pair in the instruction tuning experiments. Also, the base- line methods in the pre-training experiments were implemented similarly to the instruction tuning experiments, with adjustments made to exclude SGOs due to the computational constraints of large- scale pre-training. Specifically, for methods like DistiLLM, we only used the core divergence com- ponents without the SGO-based additions.
D.3 IMAGE CLASSIHICATION RESULTS
To explore TAID's 's applicability beyond language models, we conducted experiments image clas- sification tasks using the CIFAR-100 and ImageNet datasets.
D.4 CIFAR-100 RESULTS
We evaluated TAID the CIFAR-100 dataset, which consists of 100 classes. Table 7 presents the top-1 accuracies achieved by TAID and other knowledge distillation methods on various teacher- student model pairs.
22

---

Published as a conference paper at ICLR 2025
Table 8: Top-1 accuracies (%) on the ImageNet validation set. Results for different teacher- student pairs are shown.
Teacher Student
ResNet 34
ResNet
Method
ResNet18
MN-V1
KD Hinton et al
2015
71.03
70.50
CTKD (Li
2023b
71.38
71.16
DKD Zhap al.
2022
71.70
72.05
MLKD (Tn
71.90
73.01
(Ours) TAID
72.10
72.71
As shown in Table 7 TAID performs competitively on CIFAR-I00, consistently outperforming KL divergence across all model pairs. However, the gains are modest compared to state-of-the-art methods specifically designed for image classification, such MLKD.
Interestingly, based on the analysis of DKD, we can interpret that for simpler tasks like CIFAR-100 where the teacher's target class probabilities are close to the weight of the NCKD component in DKD becomes small. This suggests that combining TAID with DKD could potentially lead to further performance improvements, leveraging the strengths of both approaches handling different aspects of the distillation process.
D.5 IMAGENET RESULTS
To assess TAID's performance on larger-scale image classification task, we conducted experi- ments on the ImageNet dataset, which contains 1000 classes. Table S presents the top-1 accuracies achieved by TAID and other methods on ImageNet.
On ImageNet, TAID shows more pronounced improvements, consistently outperforming CTKD and DKD across both teacher-student pairs. For the Mee.eet3--ReNeelll pair, TAID achieves the highest accuracy among all methods. For the outperforming CTKD and DKD, and achieving results close to MIKDD
pair, TAID performs competitively,
These results on ImageNet demonstrate that TAID's performance improves relative to other methods as the task complexity increases. With its larger number of classes and more diverse images, Ima- geNet presents a more challenging scenario where TAID's adaptive interpolation mechanism shows more significant gains. This aligns with our observations in the main text that TAID's strengths are particularly evident in tasks with higher complexity and entropy.
E
MODEL DETAILS
E. TAID-LLM-L.5B
For the development of TATD-LIM-1 55, we utilized the full M--oLps dataset. The train- ing process consisted of epochs, employing the AdamW optimizer with cosine learning rate scheduler. We set the initial learning rate to le-5.
this experiment, we used Qwen2- Besen2-72BIInttree Instruct as the teacher model and Qwen2-1.50 100-11.5-Intttuu Instruce used a momentum coefficient (3) 0.99 and learning rate of t (a) of 5e-5. The initial value of (tstan) was set to 0.4, and the final value (teoa) was set to 1.0.
the student model. For the TAID-specific parameters, we
To enhance training efficiency, we pre-computed the probabilities from the teacher model. Further- more, to manage storage costs effectively, only utilized the top 50 probabilities. This approach allowed us to balance computational resources and model performance, enabling efficient knowl- edge transfer from the large teacher model to the smaller student model.
Table 9 presents the detailed results for LEDD-LM-11.5 and other state-of ofee---ftthhttt the-hh--t art small language models across various tasks evaluated using the LightEval benchmark suite (Allal et al. 2024).
23

---

Published as a conference paper at ICLR 2025
Table 9: Performance of TAID-LIM-1. 5B, our new state-of-the-art LLM for models under 2B
parameters.
Model
MMLU
TriviaQA
ARC
PIQA
Hellaswag
OBQA
Winogrande
Average
Owen2-1 1onn2-1.00 1.56
arg nd. N04
37.91
1.38
48.12
75.30
63.87
36.80
59.98
16.19
Qwen2.
41.15
58.41
76.01
66.40
40.00
59.35
48.86
Phi-1 5. Lia 3125
35.92
6.06
60.53
75.62
60.72
46.00
67.88
50.39
StabletMM-2
37.20
58.72
8mo TM-1 18 Alblea 304
39.97
22.56
59.95
76.06
62.91
42.80
54.91
51.31
LLM--1.5
Table 10: Performance of TAID-VIM-2B, our new state-of-the-art VLM for models up to 4B parameters.
Model
MMBBenc.V11 MMStar
MMMU_YAL
MathVVsta
OCRBench
Al2D
HallusionBench
MMVet
Germsw 3,5
65.6
46.3
34.9
25.7
61.4
68.3
52.2
53.1
46.6
Mini Cani 2
65.8
39.1
36.2
39.8
60.5
62.9
36.1
41.3
47.9
2hi-3-10 ------ieee mi
65.2
47.3
46.1
44.6
63.7
78.4
39,0
44.1
53.6
2
696
49.3
36,3
46.0
78.1
74.1
38.0
39.7
54.0
2AID-VEN-28
70.7
19.5
25.1
$1.6
78.6
74,0
56.8
35.1
56.4
LightEval is designed comprehensively assess the capabilities of small language models through series of seven zero-shot tasks. Note that the scores in denotss the average scores in Table
As shown in Table 9 TAID-LIM-1 5B achieves competitive superior performance across all tasks, with particularly strong results in PIQA and Hellaswag, This demonstrates the effectiveness of our distillation approach in creating compact model that maintains high performance across diverse range of language tasks.
E.2 TAID- VLM-2B
we trained on the Mantis-Instruce dataset (Jiang et al. 2024). The training process spanned 3 epochs, using the AdamW optimizer with cosine learning rate scheduler. The initial learning rate was set to 1c-6.
For
In this vision-languaa model distillation task, employed InternVL2-82 (Chen et al. 2024) as the teacher model and InternVL2-28 as the student model. The TAID-specific parameters remained largely consistent with those used for TAID-IILM11 5B, with momentum coefficient (3) of 0.99 and start of 0.4. However, we adjusted the learning rate of / to 5e-4 to accommodate the characteristics of vision-language model training. The fend value was maintained at 1.0.
Table 10 presents the detailed results for TAID PAID-VLM- 2B and other state-of-the-art small vision- language models across various tasks. Note that the scores in Table denotes the average scores in Table 10
TAID-VLM-29 achieves competitive or superior performance across most tasks, with particularly strong results in MMStar, and HallusionBench. This demonstrates the ef- fectiveness of our distillation approach creating compact language model that maintains high performance across diverse range of multimodal tasks.
As shown Table 10
24


Reviews:
[{"review_id": "ahh6Pj8CzS", "reviewer": "Reviewer_RYft", "summary": "This paper proposes TAID, introducing a temporally adaptive interpolation between the student and teacher distributions. The interpolation is upon relative change in the objective function. Numerical experiments over instruction tuning and pretraining stage well demonstrated the efficacy of TAID.", "questions": "See the weakness.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well motivated, written well and easy to follow.\n\n- The numerical experiments are comprehensive to well show the effectiveness of TAID.", "weaknesses": "- The technical contribution feels somewhat limited, primarily focusing on the difference between TAID and Skew KLD through the proposed temporal adaptive interpolation mechanism. This design choice, however, appears heuristic in nature. Given that multiple adaptive mechanisms could potentially meet the provided criteria, it would be helpful for the authors to elaborate on why this particular format was chosen for the adaptive mechanism. Providing additional rationale or theoretical support would strengthen the justification for this approach and enhance the perceived novelty of the contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HBIvvItz7p", "reviewer": "Reviewer_jQRU", "summary": "The paper focuses on improving the distillation procedure for LLMs and VLMs. It discusses the student-teacher distribution mismatch in standard distillation and identifies two crucial problems: the capacity gap and mode averaging/collapse. A novel approach based on time-dependent corrections to intermediate teacher distributions is proposed.\nThe method is compared to other techniques in the LLM setting, showing superior results. Finally, the paper presents results for distillation in both LLMs and VLMs.", "questions": "1. It is known that top-k logit distillation aids knowledge distillation (KD). Did the authors conduct experiments involving this? By top-k, I mean applying KD only on the top 100/1000 neurons to reduce the capacity gap.\n2. The approach potentially bears similarities to linearly increasing the loss weight during distillation via t. Any thoughts on this?\n3. What does m represent in line 199?\n4. Were the results in Table 1 computed by the authors using the same data and settings?\n5. How would this approach compare to adjusting the learning rate during the distillation process? Could it have a similar effect?\n6. Will the code be open-sourced?\n7. Is it beneficial to combine this with cross-entropy (CE) loss? If so, what should the ratio be?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem formulation is well done. The paper discusses issues with current KD, identifies two specific problems, and addresses them effectively, creating a compelling narrative.\n- The paper addresses both LLM and VLM distillation.\n- The approach is straightforward, which is a strength. It requires adjusting only a single hyperparameter.\n- The resulting models demonstrate state-of-the-art (SOTA) performance compared to models in the literature, as shown in Tables 4 and 5.", "weaknesses": "1. The results in Table 2 show that TAID performs very similarly to standard KL. They also indicate that other techniques underperform compared to KL. It is unclear if this is due to implementation.\n2. The authors discuss the capacity gap and mode collapse in standard distillation. Is there any experimental evidence of this for LLM and VLM cases?\n3. The paper claims to outperform all other distillation techniques. An analysis comparing the differences in loss functions between methods would be important to support this claim and to explain why the proposed approach is superior.\n4. For Tables 4 and 5, it is important to assess improvements by comparing them to standard cross-entropy (CE) loss training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q4oFsDTdea", "reviewer": "Reviewer_qQNq", "summary": "This method addresses key challenges such as the capacity gap between teacher and student models and issues like mode averaging and mode collapse during distillation. The core idea behind TAID is the introduction of an intermediate, dynamic distribution that gradually shifts from the student model’s distribution to the teacher’s, enabling a smoother and more effective transfer of knowledge. The authors provide both theoretical and empirical analysis, showing that TAID prevents mode collapse, enhances performance across different model sizes, and results in more efficient and high-performing LLMs.", "questions": "1. Does TAID also hold for reverse KL divergence? Are there any results regarding this?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. **Adaptive Approach**: The use of a dynamic, time-dependent intermediate teacher distribution addresses the capacity gap effectively, making it easier to transfer knowledge between models of different sizes.\n\n2. **Versatility**: TAID is effective across a variety of model sizes and tasks, including language-only and vision-language models, making it broadly applicable.\n\n3. **Theoretical Perspective**: TAID’s approach is supported by solid theoretical analysis, proving its ability to prevent mode collapse and balance between mode averaging and collapse, which traditional KD methods often struggle with. However, I didn't check the detailed proof for the theorem.", "weaknesses": "1. Despite the basic concept of the objective function being very similar to Skew KL in DistiLLM, there is limited comparison between Skew KL and TAID. I am wondering if there is a reasonable analysis of the main differences between the Skew KL and TAID functions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)