Beyond Logits: Aligning Feature Dynamics for Effective Knowledge Distillation
Authors: Guoqiang Gong, Jiaxing Wang, Jin Xu, Deping Xiang, Zicheng Zhang, Leqi Shen, Yifeng Zhang, JunhuaShu, ZhaolongXing, Zhen Chen, pengzhang liu, Ke Zhang
Accepted at: acl25

Beyond Logits: Aligning Feature Dynamics for Effective Knowledge Distillation

Guoqiang Gong1,*, Jiaxing Wang1,*, Jin Xu2, Deping Xiang1, Zicheng Zhang1, Leqi Shen3, Yifeng Zhang1, Junhua Shu1, Zhaolong Xing1, Zhen Chen1, Pengzhang Liu1,†, Ke Zhang1
1 JD.com   2 University of Oxford   3 Tsinghua University
{gongguoqiang1, wangjiaxing41,liupengzhang}@jd.com
*Equal contribution. †Corresponding author.

Abstract

Knowledge distillation (KD) compresses large language models (LLMs), known as teacher models, into lightweight versions called student models, enabling efficient inference and downstream applications. However, prevailing approaches accomplish this by predominantly focusing on matching the final output distributions of student/teacher models. Drawing on the perspective that transformers can be viewed as discretizing ordinary differential equation (ODEs) on integer time steps (corresponding to layer indices), where intermediate features evolve across layers, we argue that effective KD requires aligning the entire feature dynamics between teacher and student models, which we call feature dynamics distillation (FDD). This alignment involves matching both the feature trajectory and its first-order derivative, rather than just the final states. Our approach extends the original KD objective with two additional loss terms: layer-wise feature KD, which matches discretized feature trajectory, and layer feature delta KD, which matches first-order changes in features across adjacent layers. Extensive experiments on various tasks validate the effectiveness of our distillation method.

1 Introduction

Large language models (LLMs) have demonstrated remarkable capabilities in various natural language processing tasks (OpenAI, 2022, 2024; Qwen, 2024; DeepSeek-AI-Group, 2024). The success is, however, largely driven by scaling up the model parameters. (e.g., 175B parameters for GPT-3 (Brown et al., 2020), 70B parameters for Qwen2.5 (Qwen, 2024) and 671B for Deepseek-V3 (DeepSeek-AI-Group, 2024)). The large model size requires substantial computation and poses significant challenges in practical deployment, especially in resource-constrained environments. By transferring knowledge from a large teacher model to a smaller student model, knowledge distillation (KD) emerges as a promising solution. The rich information learned by the teacher model helps optimize the student to achieve performance levels that would be unattainable through training the student model alone.

The backbone of LLMs, the transformer architecture (Vaswani et al.), with layer-wise skip connections, can be viewed as discretized ODEs where integer time steps correspond to layer indices (Chen et al., 2018; Lu et al., 2019). The ODE states represent intermediate features that evolve across layers from initial input embeddings to final layer features, which are then projected to logits for next-token predictions. Through this lens, student models can be understood as a coarse discretization of the underlying feature dynamics. The teacher models’ more accurate representation of feature dynamics can thus guide student models toward better discretizations within their representation capacity. From this perspective, we argue that effective and comprehensive knowledge distillation should align the entire feature dynamics between student and teacher models, extending beyond the conventional approach of solely matching output distributions through various divergence losses, including Kullback-Leibler divergence (Hinton et al., 2015; Muralidharan et al., 2024; Liu et al., 2024) or reverse Kullback-Leibler divergence (Gu et al., 2024; Ko et al., 2024). To this end, we propose feature dynamics distillation (FDD), which optimizes student models to match both the discretized feature trajectory and its first-order derivative (estimated through finite differences) with the teacher models. This comes down to two additional KD loss terms: the layer-wise KD loss for trajectory matching and the layer-delta KD loss for derivative matching. We demonstrate that these additional KD loss terms consistently improve distillation effectiveness, outperforming previous state-of-the-art approaches, validating the importance of feature
dynamics distillation.

Our ODE perspective on KD is further supported by recent empirical observations about LLMs, including findings that intermediate layers can encode richer representations (Skean et al., 2025; Geva et al., 2022; Dar et al., 2023; Elhoushi et al., 2024) and that decoding through layer contrasting can better surface factual knowledge in LLMs (Chuang et al., 2024). The former demonstrates the necessity of conducting layer-wise KD, while the latter corroborates the importance of layer-delta KD. In addition, there is a growing body of literature exploring layer-wise distillation schemes (Jiao et al., 2020; Wang et al., 2020) by matching various aspects of teacher and student models such as intermediate features and attention scores. These studies, though heuristically established and primarily focus on BERT-based models (Devlin et al., 2019), show the promise of developing KD scheme that can more comprehensively mimic the teacher.

Extensive experiments are conducted on various tasks. The results show that our proposed FDD ensures sufficient exploitation of teacher knowledge, leading to more effective and nuanced student models that demonstrate state-of-the-art performance.

2 Preliminary

2.1 Knowledge Distillation of LLMs

A transformer-based large language model (LLM) typically consists of three components: An embedding layer \( f_{\text{embed}}(\cdot) \), a stack of \( L \) transformer blocks, and a language modeling (LM) head \( f_{\text{head}}(\cdot) \). Given a tokenized one-hot encoded input sequence \( \mathbf{x}_{1:N} := [\mathbf{x}_1, \ldots, \mathbf{x}_N] \), the embedding function \( f_{\text{embed}}(\cdot) \) maps it into dense feature representations \( \mathbf{h}_{1:N}(0) \). These features are then processed sequentially through \( L \) transformer layers, producing intermediate features \( \{ \mathbf{h}_{1:N}(1), \ldots, \mathbf{h}_{1:N}(L) \} \). Finally, the LM head \( f_{\text{head}}(\cdot) \), typically a linear soft-max layer, projects the final features \( \mathbf{h}_{1:N}(L) \) to vocabulary distribution for next-token predictions.

Larger LLMs typically achieve better performance but are computationally expensive during inference. Knowledge distillation (KD) is a model compression technique where the “knowledge” learned by a large, complex model (teacher \( T \)) is transferred to a smaller, more efficient model (student \( S \)) (Hinton et al., 2015). This can be achieved by minimizing the Kullback-Leibler divergence (KLD) between their output distributions:

\[
\mathcal{L}_{\text{KD}} = \frac{1}{N} \sum_{i=1}^{N} \mathcal{D}_{\text{KL}}(p^T(\mathbf{x}_{i+1}|\mathbf{x}_{\leq i}) \parallel p^S_w(\mathbf{x}_{i+1}|\mathbf{x}_{\leq i}))
\]
\[
= \frac{1}{N} \sum_{i=1}^{N} \mathcal{D}_{\text{KL}}(f^T_{\text{head}}(\mathbf{h}_i^T(L)) \parallel f^S_{\text{head}}(\mathbf{h}_i^S(L)))
\]

where \( w \) denotes parameters of the student model to be optimized.

2.2 Transformers as ODEs

Neural network with layer-wise residual connections (He et al., 2016) can be interpreted as discretizing an ordinary differential equation (ODE) at integer time steps (Chen et al., 2018), where the time variable corresponds to layer indices. Transformer (Vaswani et al.) also incorporates residual connections and stacks layers consisting of one multi-head attention (MHA) module and a multi-layer perceptron (MLP) module, mapping from sequences to sequences:

\[
\mathbf{h}_{1:N} \leftarrow \mathbf{h}_{1:N} + \text{MHA}(\text{LN}(\mathbf{h}_{1:N}))
\]
\[
\mathbf{h}_{1:N} \leftarrow \mathbf{h}_{1:N} + \text{MLP}(\text{LN}(\mathbf{h}_{1:N})),
\]

where LN denotes layer normalization (Ba et al., 2016). Therefore, it can be seen as an ODE with intermediate feature sequences as states:

\[
\frac{d\mathbf{h}_{1:N}(t)}{dt} = f_{\text{ode}}(\mathbf{h}_{1:N}(t); t, \mathbf{w}(t))
\]
\[
= \text{MHA}^{(t)}(\text{LN}^{(t)}(\mathbf{h}(t)), \mathbf{w}_{\text{MHA}}(t))
\]
\[
+ \text{MLP}^{(t)}(\text{LN}^{(t)}(\mathbf{h}(t)), \mathbf{w}_{\text{MLP}}(t))
\]

where \( \mathbf{w}(t) = [\mathbf{w}_{\text{MHA}}(t), \mathbf{w}_{\text{MLP}}(t)] \) denotes time (layer) dependent parameters. To recover the standard transformer architecture and accommodate for alternating MHAs and MLPs, one must follow the Lie-Trotter splitting scheme (Trotter, 1959) for discretization. This ODE fully characterizes the dynamics of feature evolution in the model forward pass (Lu et al., 2019).

3 Method

Existing knowledge distillation approaches predominantly focus on output-layer distillation. Here we argue that KD can be made more effective by aligning the entire feature dynamics between the teacher and student models, incorporating both the feature trajectory and its first-order derivative. The approach, derived from the ODE perspective, is
Figure 1: An illustration of feature dynamics distillation (FDD) from an ODE perspective. The student model (in red), which represents a coarser discretization of the underlying feature dynamics ODE, is optimized to match the teacher model (in green), which is more accurate. As the feature dimensions are distinct, we map features to intermediate predictive distributions through the LM heads (green and red small rectangles) and conduct matching directly in this space. Both feature trajectory (dots) and its derivative (arrows) are matched, which comes down to a layer-wise KD loss term and a layer-delta KD loss term. The final output distribution KD loss is also incorporated.

concretized in Section 3.1. It then translates into incorporating layer-wise KD (Section 3.2) and layer-delta KD (Section 3.3) after discretization. The framework of feature dynamics distillation (FDD) is illustrated in Figure 1.

3.1 Distillation under the ODE Perspective

Given the underlying corresponding ODEs (Equation 2) for the teacher \( \mathcal{T} \) and the student \( \mathcal{S} \), our goal is to align their feature dynamics. Nevertheless, features in the teacher and student typically reside in different spaces with different dimensions. To bridge the gap and thus facilitate the alignment, we project both teacher and student features into the shared vocabulary prediction distribution space by multiplexing the language modeling head \( f_{\text{head}}(\cdot) \) of their own. The overall dynamic is then 1:

\[
\frac{d h^{\mathcal{S} \text{ or } \mathcal{T}}_{i:N}(t)}{dt} = f_{\text{ode}} \circ \mathcal{T}(h^{\mathcal{S} \text{ or } \mathcal{T}}_{i:N}(t); t) \\
y^{\mathcal{T} \text{ or } \mathcal{S}}_{i:N}(t) = \log f^{\mathcal{T} \text{ or } \mathcal{S}}_{\text{head}}(h^{\mathcal{T} \text{ or } \mathcal{S}}_{i:N}(t))
\]

The prediction dynamics \( y(t) \) effectively capture knowledge contained in the feature dynamics \( h(t) \). In the following, we abuse "feature dynamics" as this "prediction dynamics". Actually, discretization of its exponential \( e^{y(t)} \) exactly corresponds to the Logits Lens (nostalgiaerist, 2020), which has been widely applied to interpret the internals of LLM.

With the dynamics established, we align teacher-student prediction dynamics \( y^{\mathcal{T}}(t) \) and \( y^{\mathcal{S}}(t) \) from \( t = 0 \) to \( T \) with the trajectory loss:

\[
\mathcal{L}_{\text{KD}}^{\text{Traj}} = \frac{1}{TN} \int \sum_{i=1}^{N} \mathcal{D}_{\text{KL}}(e^{y^{\mathcal{T}}_i(t)} \| e^{y^{\mathcal{S}}_i(t)}) dt
\]

Matching feature trajectories alone can be insufficient due to discretization and matching errors, making it advantageous to also align the temporal evolution of features. We thus propose incorporating an additional loss term to align the first derivative of the feature trajectory:

\[
\mathcal{L}_{\text{KD}}^{\text{Der}} = \frac{1}{TN} \int \sum_{i=1}^{N} \mathcal{D}_{\text{Cos}}(\frac{dy^{\mathcal{T}}_i(t)}{dt} \| \frac{dy^{\mathcal{S}}_i(t)}{dt}) dt
\]

where \( \mathcal{D}_{\text{Cos}}(\mathbf{a}, \mathbf{b}) = 1 - \frac{\mathbf{a}^T \mathbf{b}}{\|\mathbf{a}\| \|\mathbf{b}\|} \). The feature dynamics become very uncertain at times far from the discretization sampling points. The trajectory derivative thus provides necessary supplementary information characterizing the underline feature dynamics.

These additional knowledge distillation loss terms, \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) and \( \mathcal{L}_{\text{KD}}^{\text{Der}} \) derived from the ODE perspective, provide much richer training signals compared to the conventional loss based solely on output distributions.

3.2 Layer-wise Knowledge Distillation

For practical use, we discretize the ODEs Equation 3, where layer indices correspond to integer time steps. This turns out to be a layer-wise distillation scheme. As the layer number of teacher \( L^{\mathcal{T}} \) and student \( L^{\mathcal{S}} \) are different, we choose a distillation schedule. Teacher intermediate layers indexed by \( \mathcal{I}^{\mathcal{T}} = \{ l^{\mathcal{T}}_{I_1}, \ldots, l^{\mathcal{T}}_{I_{L_D}} \} \) are distilled to the student layers indexed by \( \mathcal{I}^{\mathcal{S}} = \{ l^{\mathcal{S}}_{I_1}, \ldots, l^{\mathcal{S}}_{I_{L_D}} \} \) correspondingly, where \( L_D \) is the number of intermediate layers to be distilled. Typically \( l^{\mathcal{T}}_{I_j} = K \cdot l^{\mathcal{S}}_{I_j} \),

1Here we omit the dependence of \( f_{\text{ode}}(\cdot) \) on time variant parameters \( \mathbf{w}(t) \) and keep only dependence on time \( t \) for brevity.
where \( K = \frac{L_T}{L_S} \). The feature trajectory loss Equation 4 then becomes:

\[
\mathcal{L}_{\text{KD}}^{\text{Traj}} = \frac{1}{LN} \sum_{j=1}^{L_D} \sum_{i=1}^{N} \mathcal{D}_{\text{KL}}(e^{v_i^T(l_{T_j}^T)} \| e^{v_i^S(l_{S_j}^S)})
\]

(6)

As intermediate layers encode equal or even richer representations than those of the final layer (Skean et al., 2025), such a layer-wise distillation scheme facilitates capturing rich semantic knowledge contained in the distinct layers so that the teacher model is better exploited.

Previous layer-wise distillation methods match the teacher and student features by minimizing their L2 distance: \( \frac{1}{L_D N} \sum_{j=1}^{L_D} \sum_{i=1}^{N} \| h_i^T(l_{T_j}^T) - \mathbf{P} h_i^S(l_{S_j}^S) \|_2 \). Generally the feature dimension \( d_h^T \) and \( d_h^S \) are different so the projection matrix \( \mathbf{P} \) is incorporated to project the teacher and student features into the same space. In FDD, we utilize the already trained language model heads \( f_{\text{head}}(\cdot) \) instead. This treatment bears a strong resemblance to the task-aware filter proposed in (Liang et al., 2023). Projection matrix trained with L2 loss indiscriminately matches the features of teacher and student. While LM head, pretrained for predicting the next token, is much more relevant to the downstream tasks(e.g. instruction tuning). As a result, it serves as a better filter for selecting knowledge that is useful for the target tasks from the hidden representations (features).

3.3 Layer-delta Knowledge Distillation

While matching the feature trajectory (layer-wise distillation) extracts rich information encoded in each layer, it does not capture the temporal evolution of features across layers due to discretization and matching errors. The feature derivative loss Equation 5 is thus incorporated. Again, with proper discretization of the ODEs, the feature derivative loss turns into a layer-delta knowledge distillation scheme:

\[
\mathcal{L}_{\text{KD}}^{\text{Der}} = \frac{1}{NL_\Delta} \sum_{j=1}^{L_\Delta} \sum_{i=1}^{N} \mathcal{D}_{\text{Cos}}(\Delta_i^T(l_{T_j}^T) \| \Delta_i^S(l_{S_j}^S))
\]

(7)

Where \( L_\Delta = L_D - 1 \). The feature changes (\( \Delta \)) between adjacent selected layers are:

\[
\begin{align*}
\Delta_i^T(l_{T_j}^T) &= y_i^T(l_{T_j}^T) - y_i^T(l_{T_{j-1}}^T) \\
\Delta_i^S(l_{S_j}^S) &= y_i^S(l_{S_j}^S) - y_i^S(l_{S_{j-1}}^S)
\end{align*}
\]

This layer-delta distillation loss measures the similarity of feature transitions between the teacher and the student, ensuring that the student model learns not only the knowledge encoded in each layer but also the pattern of representation refinement across layers.

3.4 Overall Optimization

The final distillation objective combines the feature trajectory loss \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \), the feature derivative loss \( \mathcal{L}_{\text{KD}}^{\text{Der}} \) as well as the conventional KL divergence aligning the output distribution of teacher and student.

\[
\mathcal{L}_{\text{overall}} = \mathcal{L}_{\text{KD}} + \alpha \mathcal{L}_{\text{KD}}^{\text{Traj}} + \beta \mathcal{L}_{\text{KD}}^{\text{Der}}
\]

(8)

where \( \alpha, \beta \) are hyperparameters controlling the relative contributions of each loss. This comprehensive objective ensures the alignment of the teacher and the student on both semantic representations and their evolutionary patterns across layers. Ensuring a more effective student model. Note that although we utilize the conventional Kullback-Leibler divergence to establish FDD, it is compatible to other recently developed KD losses like Jensen-Shannon divergence (Agarwal et al., 2024), the reverse Kullback-Leibler divergence (Gu et al., 2024) and the Skew KLD (Ko et al., 2024).

An overall procedure of FDD is given in Algorithm 1.

Algorithm 1 Feature Dynamic Distillation (FDD)

Require:
Teacher model \( \mathcal{T} \); Student model \( \mathcal{S} \) parameterized by \( w \).
Learning rate \( \eta \); Hyper-parameters \( \alpha, \beta \).
Selected teacher layer indexes \( \mathcal{I}^T \) and Student layer indexes \( \mathcal{I}^S \) for distillation.

Ensure:
Distilled student model.

1: while not done do
2:   Forwarding both \( \mathcal{T} \) and \( \mathcal{S} \), compute the KD loss following Equation 1.
3:   Apply LM heads on the selected layers. Compute the feature trajectory loss \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) following Equation 6.
4:   Compute the feature derivative loss \( \mathcal{L}_{\text{KD}}^{\text{Der}} \) following Equation 7.
5:   Update \( w \) via stochastic gradient descent:
      \[ w \leftarrow w - \eta \nabla_w \left\{ \mathcal{L}_{\text{KD}} + \alpha \mathcal{L}_{\text{KD}}^{\text{Traj}} + \beta \mathcal{L}_{\text{KD}}^{\text{Der}} \right\} \]
6: end while
4 Experiments

In this section, we first compare FDD to state-of-the-art algorithms in Section 4.2. Then we analyze the effectiveness of each design ingredient in Section 4.3.

4.1 Experimental Setup

Datasets. To evaluate the effectiveness of our method, we conduct experiments on seven instruction-following benchmarks:

• Dolly Evaluation is a sampled test set of the databricks-dolly-15k dataset3, consisting of 500 samples.

• Self-Instruct (Wang et al., 2023) is a user-oriented instruction-following dataset containing 252 samples.

• Vicuna Evaluation (Chiang et al., 2023) comprises 80 challenging questions designed to evaluate the Vicuna model.

• Wizard Evaluation (Xu et al., 2024) contains 218 real-world human instructions of varying difficulty across different domains from diverse sources.

• Koala (Geng et al., 2023) consists of 180 real user queries that span various topics and represent real-world chat system use cases.

• Super-Natural Instructions (S-NI) (Wang et al., 2022) includes 1,616 diverse NLP tasks with expert-written instructions, covering 76 distinct task types. The test set includes 9,000 samples spanning 119 tasks.

• Unnatural Instructions (UnNI) (Honovich et al., 2023) includes 240K AI-generated instructions with minimal human involvement. Following previous works (Gu et al., 2024; Ko et al., 2024), we randomly sampled 10,000 samples from the core set for evaluation.

Evaluation Metrics. Following previous methods (Gu et al., 2024; Ko et al., 2024), we utilize ROUGE-L (Lin, 2004) and GPT-4o feedback (Zheng et al., 2023) to assess the quality of model-generated outputs. ROUGE-L is a widely used evaluation metric that measures the longest common subsequence between the generated and reference texts. For the other metric, GPT-4o is tasked with comparing the model’s responses to the reference answers and assigning scores ranging from 1 to 10 for each response. We calculate and report the ratio of the total scores between the model’s responses and the reference answers. For the evaluation, we sample the responses using 5 random seeds and report the average scores.

Base Models. Our FDD is evaluated on three types of models with various sizes: LLaMA2 (Touvron et al., 2023) (13B teacher, 7B student), OpenLLaMA2 (Geng and Liu, 2023) (7B teacher, 3B student) and GPT2 (Radford et al., 2019) (1.5B teacher, 0.1B student).

Baselines. We compare our FDD against various state-of-the-art approaches:

• SFT directly fine-tunes the student model on the fixed datasets.

• KD (Hinton et al., 2015) applies KLD on fixed datasets.

• SeqKD (Kim and Rush, 2016) fine-tunes the student model using teacher-generated outputs.

• ImitKD (Lin et al., 2020) employs KLD on student-generated outputs (SGOs).

• GKD (Agarwal et al., 2024) uses Jensen-Shannon Divergence (JSD) on a combination of SGOs and fixed datasets.

• MiniLLM (Gu et al., 2024) utilizes reverse KLD and a policy gradient method on SGOs.

• DistiLLM (Ko et al., 2024) applies Skew KLD and an adaptive off-policy approach to enhance efficiency in utilizing SGOs.

Implementation Details. Our implementation is based on the experimental setup outlined by (Gu et al., 2024; Ko et al., 2024). We randomly select 14,000 samples in databricks-dolly-15k for training while allocating 500 samples each for validation and testing. Following previous works (Gu et al., 2024; Ko et al., 2024), the student model is first fine-tuned for 3 epochs on the training set. To ensure a fair comparison with other methods, such as ImitKD (Lin et al., 2020), GKD (Agarwal et al., 2024), MiniLLM (Gu et al., 2024), and
<table>
  <tr>
    <th rowspan="2">Methods</th>
    <th colspan="2">Dolly</th>
    <th colspan="2">Self-Instruct</th>
    <th colspan="2">Vicuna</th>
    <th colspan="2">Wizard</th>
    <th colspan="2">Koala</th>
    <th colspan="2">S-NI</th>
    <th colspan="2">UnNI</th>
    <th rowspan="2">Average</th>
  </tr>
  <tr>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>GPT-4o</th>
    <th>R-L</th><th>R-L</th>
  </tr>
  <tr>
    <th colspan="18"><i>LLAMA2-13B → LLAMA2-7B</i></th>
  </tr>
  <tr>
    <td>Teacher (SFT)</td><td>29.27</td><td>73.46</td><td>21.68</td><td>75.73</td><td>19.77</td><td>57.98</td><td>47.87</td><td>48.50</td><td>34.10</td><td>35.49</td><td>28.26</td>
  </tr>
  <tr>
    <td>SFT</td><td>28.38</td><td>65.59</td><td>20.32</td><td>65.68</td><td>17.72</td><td>51.93</td><td>44.56</td><td>45.00</td><td>34.37</td><td>33.25</td><td>26.81</td>
  </tr>
  <tr>
    <td>KD (<b>Hinton et al., 2015</b>)</td><td>28.56</td><td>67.93</td><td>19.92</td><td>63.67</td><td>18.20</td><td>53.17</td><td>43.95</td><td>45.82</td><td>31.85</td><td>31.12</td><td>25.93</td>
  </tr>
  <tr>
    <td>SeqKD (<b>Kim and Rush, 2016</b>)</td><td>28.29</td><td>72.30</td><td>19.82</td><td>64.14</td><td>17.70</td><td>51.62</td><td>44.72</td><td>43.67</td><td>30.80</td><td>31.82</td><td>25.69</td>
  </tr>
  <tr>
    <td>ImitKD (<b>Lin et al., 2020</b>)</td><td>27.03</td><td>65.59</td><td>21.76</td><td>65.99</td><td>17.90</td><td>43.87</td><td>41.90</td><td>42.41</td><td>31.18</td><td>31.07</td><td>25.79</td>
  </tr>
  <tr>
    <td>GKD (<b>Agarwal et al., 2024</b>)</td><td>29.27</td><td>75.07</td><td>21.25</td><td>72.79</td><td>18.61</td><td>54.10</td><td>46.49</td><td>47.31</td><td>36.94</td><td>34.69</td><td>28.15</td>
  </tr>
  <tr>
    <td>MiniLLM (<b>Gu et al., 2024</b>)</td><td>30.75</td><td>76.14</td><td>23.61</td><td>73.65</td><td>20.80</td><td>60.24</td><td>47.10</td><td>47.87</td><td>36.17</td><td>36.82</td><td>29.63</td>
  </tr>
  <tr>
    <td>DistiLLM (<b>Ko et al., 2024</b>)</td><td>31.06</td><td>76.67</td><td>23.07</td><td>70.47</td><td>20.60</td><td>58.75</td><td>47.43</td><td>47.94</td><td>37.11</td><td>37.28</td><td>29.82</td>
  </tr>
  <tr>
    <td>Ours</td><td><b>32.57</b></td><td><b>80.17</b></td><td><b>24.56</b></td><td><b>75.11</b></td><td><b>19.88</b></td><td><b>62.48</b></td><td><b>48.42</b></td><td><b>48.85</b></td><td><b>40.43</b></td><td><b>43.01</b></td><td><b>32.09</b></td>
  </tr>
  <tr>
    <th colspan="18"><i>OpenLLAMA2-7B → OpenLLAMA2-3B</i></th>
  </tr>
  <tr>
    <td>Teacher (SFT)</td><td>27.09</td><td>59.32</td><td>18.24</td><td>57.49</td><td>17.89</td><td>42.01</td><td>39.86</td><td>40.59</td><td>31.44</td><td>33.20</td><td>25.57</td>
  </tr>
  <tr>
    <td>SFT</td><td>24.95</td><td>54.22</td><td>17.38</td><td>50.23</td><td>14.58</td><td>38.75</td><td>33.28</td><td>35.06</td><td>28.65</td><td>27.69</td><td>22.65</td>
  </tr>
  <tr>
    <td>KD (<b>Hinton et al., 2015</b>)</td><td>25.10</td><td>52.47</td><td>17.11</td><td>51.15</td><td>16.64</td><td>38.29</td><td>35.60</td><td>37.23</td><td>28.93</td><td>28.32</td><td>23.22</td>
  </tr>
  <tr>
    <td>SeqKD (<b>Kim and Rush, 2016</b>)</td><td>25.22</td><td>55.68</td><td>17.30</td><td>52.39</td><td>16.06</td><td>42.48</td><td>32.84</td><td>37.44</td><td>29.26</td><td>27.90</td><td>23.14</td>
  </tr>
  <tr>
    <td>ImitKD (<b>Lin et al., 2020</b>)</td><td>23.81</td><td>53.64</td><td>17.55</td><td>51.15</td><td>16.62</td><td>40.77</td><td>37.04</td><td>37.09</td><td>29.83</td><td>28.97</td><td>23.35</td>
  </tr>
  <tr>
    <td>GKD (<b>Agarwal et al., 2024</b>)</td><td>26.54</td><td>60.34</td><td>20.21</td><td>57.80</td><td>19.14</td><td>47.28</td><td>39.53</td><td>39.75</td><td>35.44</td><td>31.65</td><td>26.60</td>
  </tr>
  <tr>
    <td>MiniLLM (<b>Gu et al., 2024</b>)</td><td>28.69</td><td>64.13</td><td>20.42</td><td>55.95</td><td>18.93</td><td>42.01</td><td>39.41</td><td>40.10</td><td>35.46</td><td>35.11</td><td>27.72</td>
  </tr>
  <tr>
    <td>DistiLLM (<b>Ko et al., 2024</b>)</td><td>28.99</td><td>63.41</td><td>19.99</td><td>53.94</td><td>20.41</td><td>47.13</td><td>39.75</td><td>40.61</td><td>36.60</td><td>35.27</td><td>28.25</td>
  </tr>
  <tr>
    <td>Ours</td><td><b>30.15</b></td><td><b>66.61</b></td><td><b>20.93</b></td><td><b>59.96</b></td><td><b>19.63</b></td><td><b>48.06</b></td><td><b>40.85</b></td><td><b>42.55</b></td><td><b>39.88</b></td><td><b>38.51</b></td><td><b>29.82</b></td>
  </tr>
  <tr>
    <th colspan="18"><i>GPT-2 XL (1.5B) → GPT-2 (0.1B)</i></th>
  </tr>
  <tr>
    <td>Teacher (SFT)</td><td>26.34</td><td>47.95</td><td>15.27</td><td>38.63</td><td>16.96</td><td>34.88</td><td>29.13</td><td>32.05</td><td>26.71</td><td>30.18</td><td>23.09</td>
  </tr>
  <tr>
    <td>SFT</td><td>23.53</td><td>35.42</td><td>10.49</td><td>26.58</td><td>14.98</td><td>22.79</td><td>20.56</td><td>22.18</td><td>16.78</td><td>19.95</td><td>17.15</td>
  </tr>
  <tr>
    <td>KD (<b>Hinton et al., 2015</b>)</td><td>22.31</td><td>35.42</td><td>10.37</td><td>25.96</td><td>15.50</td><td>23.10</td><td>20.07</td><td>22.67</td><td>16.64</td><td>19.18</td><td>16.80</td>
  </tr>
  <tr>
    <td>SeqKD (<b>Kim and Rush, 2016</b>)</td><td>24.11</td><td>36.58</td><td>11.60</td><td>28.90</td><td>14.99</td><td>21.86</td><td>19.79</td><td>22.39</td><td>19.58</td><td>22.12</td><td>18.48</td>
  </tr>
  <tr>
    <td>ImitKD (<b>Lin et al., 2020</b>)</td><td>22.04</td><td>36.37</td><td>10.19</td><td>27.04</td><td>14.83</td><td>22.17</td><td>20.62</td><td>22.74</td><td>18.11</td><td>21.03</td><td>17.24</td>
  </tr>
  <tr>
    <td>GKD (<b>Agarwal et al., 2024</b>)</td><td>24.36</td><td>34.83</td><td>10.45</td><td>27.20</td><td>15.73</td><td>24.65</td><td>21.01</td><td>23.23</td><td>17.16</td><td>19.86</td><td>17.51</td>
  </tr>
  <tr>
    <td>MiniLLM (<b>Gu et al., 2024</b>)</td><td>24.35</td><td>36.88</td><td>10.38</td><td>27.66</td><td><b>16.15</b></td><td>24.96</td><td>21.12</td><td>23.65</td><td>24.76</td><td>25.11</td><td>20.15</td>
  </tr>
  <tr>
    <td>DistiLLM (<b>Ko et al., 2024</b>)</td><td>25.20</td><td>37.31</td><td>12.32</td><td>28.59</td><td>15.60</td><td>24.34</td><td>21.39</td><td>23.79</td><td>23.48</td><td>26.38</td><td>20.60</td>
  </tr>
  <tr>
    <td>Ours</td><td><b>25.70</b></td><td><b>38.19</b></td><td><b>12.48</b></td><td><b>30.75</b></td><td>15.94</td><td><b>26.04</b></td><td><b>22.06</b></td><td><b>24.77</b></td><td><b>27.62</b></td><td><b>29.04</b></td><td><b>22.16</b></td>
  </tr>
</table>

Table 1: Comparison of state-of-the-art knowledge distillation methods. ROUGE-L (R-L) metric (Lin, 2004) and GPT-4o feedback scores are reported. ‘Average’ represents the mean ROUGE-L score across the five benchmarks. The best performance is highlighted.

DistiLLM (<b>Ko et al., 2024</b>), we initialize the student models using the same fine-tuned checkpoint. The hyperparameters \( \alpha \) and \( \beta \) are both set to 1, and a constant learning rate of 5e-4 is applied across all experiments. For GPT-2 models, we train all parameters for 20 epochs. For LLaMA2 and OpenLLaMA2 models, we adopt the LoRA (<b>Hu et al., 2022</b>) technique, with training conducted over 10 epochs. Following DistiLLM (<b>Ko et al., 2024</b>), LoRA is applied to the query and value weights with a rank of 16. After the knowledge distillation process, the student models are selected based on their ROUGE-L scores on the validation set, and these selected checkpoints are subsequently evaluated on the test set.

For feature trajectory and its first-order derivative matching, without loss of generality, we select intermediate layers uniformly from both the student and the teacher models. Taking the student as an example, let \( L^S \) denote the total number of intermediate layers and \( L_D \) the desired number of sampled layers. The selection interval is then \( Q = \lceil L^S / (L_D + 1) \rceil \). The indices of the sampled layers are then \( \mathcal{I} = \{ l \times Q \mid l \in [1, L_D] \} \). FDD enables direct decoding of intermediate features into the vocabulary space using the model’s pre-trained LM head (<b>nostalgebraist, 2020</b>). However, the LM head may be less effective in processing the intermediate features because the LM head is specifically tuned to handle the final layer output representations during pretraining. Following the methodology of Tuned Lens (<b>Belrose et al., 2023</b>), we optionally train lightweight adapters for each selected layer before feature dynamic distillation.
<table>
  <tr>
    <th>Methods</th>
    <th>Dolly</th>
    <th>Self-Instruct</th>
    <th>Vicuna</th>
    <th>S-NI</th>
    <th>UnNI</th>
    <th>Average</th>
  </tr>
  <tr>
    <th colspan="7">LLaMA2-13B → TinyLLaMA-1B</th>
  </tr>
  <tr>
    <td>Teacher (SFT)</td>
    <td>29.27</td>
    <td>21.68</td>
    <td>19.77</td>
    <td>34.10</td>
    <td>35.49</td>
    <td>28.26</td>
  </tr>
  <tr>
    <td>SFT</td>
    <td>22.60</td>
    <td>15.80</td>
    <td>16.03</td>
    <td>27.37</td>
    <td>26.72</td>
    <td>21.70</td>
  </tr>
  <tr>
    <td>KD</td>
    <td>22.99</td>
    <td>16.34</td>
    <td>16.00</td>
    <td>28.31</td>
    <td>26.75</td>
    <td>22.08</td>
  </tr>
  <tr>
    <td>SeqKD</td>
    <td>22.43</td>
    <td>15.89</td>
    <td>15.73</td>
    <td>29.17</td>
    <td>26.74</td>
    <td>21.99</td>
  </tr>
  <tr>
    <td>ImitKD</td>
    <td>20.65</td>
    <td>16.04</td>
    <td>15.92</td>
    <td>26.46</td>
    <td>25.15</td>
    <td>20.84</td>
  </tr>
  <tr>
    <td>GKD</td>
    <td>23.61</td>
    <td>18.89</td>
    <td>17.03</td>
    <td>32.51</td>
    <td>29.89</td>
    <td>24.38</td>
  </tr>
  <tr>
    <td>MiniLLM</td>
    <td>26.53</td>
    <td>18.94</td>
    <td>17.34</td>
    <td>34.70</td>
    <td>32.74</td>
    <td>26.05</td>
  </tr>
  <tr>
    <td>DistiLLM</td>
    <td>26.72</td>
    <td>19.05</td>
    <td><b>18.57</b></td>
    <td>35.08</td>
    <td>32.83</td>
    <td>26.45</td>
  </tr>
  <tr>
    <td>Ours</td>
    <td><b>28.97</b></td>
    <td><b>19.19</b></td>
    <td>18.18</td>
    <td><b>37.09</b></td>
    <td><b>36.30</b></td>
    <td><b>27.95</b></td>
  </tr>
</table>

Table 2: Comparison of state-of-the-art knowledge distillation methods. ROUGE-L metric is reported. ‘Average’ represents the mean ROUGE-L score across the five benchmarks. The best performance is highlighted.

4.2 Comparation with State of the Arts

Table 1 presents the comparison results between FDD and previous state of the arts. We make the following observations:

• First, FDD consistently outperforms previous KD methods, achieving the highest average ROUGE-L scores and generally the best GPT-4o scores. For instance, in the LLaMA2-13B to LLaMA2-7B distillation setting, FDD improves the average ROUGE-L score from the previous best result (29.82) to 32.09. This demonstrates FDD’s superior ability to transfer knowledge from the teacher to the student model. Unlike baseline methods that focus solely on the final output distribution, FDD provides a more comprehensive way for the student model to leverage the rich information encoded in the teacher.

• Second, the results validate the scalability of FDD across different model architectures and parameter scales. When teacher model sizes range from 1.5B to 13B, FDD consistently achieves the best or near-best performance, demonstrating its adaptability to varying model sizes. For example, in the GPT-2 XL to GPT-2 distillation setting, FDD achieves an average ROUGE-L score of 22.16, outperforming DistiLLM (20.60) and MiniLLM (20.15). Notably, in some cases, the distilled student model even surpasses the performance of the teacher model, consistent with findings from prior studies (Gu et al., 2024; Ko et al., 2024).

Table 2 presents the results when there is a large gap between models, with LLaMA2-13B (Touvron et al., 2023) serving as the teacher model and TinyLLaMA-1B (Zhang et al., 2024a) as the student model. From Table 2, we see that in this large model size gap setting, FDD still consistently outperforms baselines in most cases, showing the effectiveness of matching the whole trajectory in mitigating the distribution mismatch problem.

<table>
  <tr>
    <th>\( \mathcal{L}_{\text{KD}} \)</th>
    <th>\( \mathcal{L}_{\text{KD}}^{\text{Traj}} \)</th>
    <th>\( \mathcal{L}_{\text{KD}}^{\text{Der}} \)</th>
    <th>Average</th>
  </tr>
  <tr>
    <td>✓</td>
    <td>-</td>
    <td>-</td>
    <td>19.28</td>
  </tr>
  <tr>
    <td>✓</td>
    <td>✓</td>
    <td>-</td>
    <td>21.35</td>
  </tr>
  <tr>
    <td>✓</td>
    <td>-</td>
    <td>✓</td>
    <td>20.52</td>
  </tr>
  <tr>
    <td>✓</td>
    <td>✓</td>
    <td>✓</td>
    <td>22.16</td>
  </tr>
</table>

Table 3: The impact of different loss functions. ‘Average’ denotes the average ROUGE-L score on the five test datasets.

4.3 Ablation Studies

To evaluate the effectiveness of each design component, we conduct ablation experiments under the GPT-2 (1.5B \( \rightarrow \) 0.1B) setting and assess the results with the ROUGE-L metric.

Effect of the Number of Distilled Layers  The results in Figure 2 demonstrate how incorporating intermediate-layer information affects the performance of the student model. When no intermediate-layer information is used (Number of Distilled Layers = 0), the average ROUGE-L score is 19.28, serving as the baseline. Performance improves as we increase the number of distilled intermediate layers, reaching a peak ROUGE-L score of 22.16 with four sampled layers. However, further increasing the number of layers leads to slight decline in performance. We conjecture this is due to over-constraining the student model, while the student struggles to mimic the teacher behavior due to their capacity gap. Additionally, our current approach uses evenly spaced integer step sizes (layer indices) to match student and teacher models, rather than employing adaptive step sizes as suggested by ODE theory.

Effect of \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) and \( \mathcal{L}_{\text{KD}}^{\text{Der}} \)  The experimental results in Table 3 demonstrate the effectiveness of \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) and \( \mathcal{L}_{\text{KD}}^{\text{Der}} \) in enhancing the performance of the student model. Specifically, incorporating \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) with \( \mathcal{L}_{\text{KD}} \) improves the average performance from 19.28 to 21.35, highlighting the importance of aligning semantic representations between corre-
Figure 2: Effect of selecting different number of intermediate layers. Averaged ROUGE-L score on the five benchmarks are reported.

Figure 3: Ablation results of employing the LM head on the Dolly Evaluation (Dolly), Super-Natural Instructions (S-NI), and Unnatural Instructions (UnNI) datasets.

sponding intermediate layers of the teacher and student models. Similarly, adding \( \mathcal{L}_{\text{KD}}^{\text{Det}} \) to \( \mathcal{L}_{\text{KD}} \) achieves an average performance of 20.52, demonstrating the value of capturing the evolution of representations across layers. Notably, the combination of \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) and \( \mathcal{L}_{\text{KD}}^{\text{Det}} \) with \( \mathcal{L}_{\text{KD}} \) achieves the highest performance of 22.16, underscoring the complementary nature of these two objectives. \( \mathcal{L}_{\text{KD}}^{\text{Traj}} \) ensures layer-wise alignment of semantic distributions, \( \mathcal{L}_{\text{KD}}^{\text{Det}} \) models the evolution of representations. The integration of these complementary objectives facilitates comprehensive knowledge transfer from the teacher to the student.

Effect of Multiplexing the LM Head We evaluated the effectiveness of using the LM head for intermediate layer dimension alignment between teacher and student models. As illustrated in Figure 3, we compared three alignment strategies: (1) Projection: This strategy randomly initializes a projection matrix as proposed in (Jiao et al., 2020). The projection matrix and student model are then jointly optimized. (2) Projection-PT: This approach consists of two steps. First, we freeze both teacher and student models, only the projection matrix is trained. Then, we jointly optimize the projection matrix and student model. (3) LM Head: This strategy directly maps intermediate layer features to the vocabulary space through the Language Model head. While Projection-PT demonstrates modest improvements over the basic Projection approach, the LM Head strategy consistently outperforms both alternative approaches across all evaluation metrics. This superior performance can be attributed to the LM Head’s ability to provide stronger task-relevant signals through direct mapping of intermediate features to the vocabulary space.

5 Related Work

Knowledge Distillation of LLMs Knowledge distillation (KD) has long been an effective way to compress neural networks (Hinton et al., 2015). Standard KD minimizes the Kullback-Leibler divergence (KLD) between the student’s and the teacher’s output distributions (Sanh et al., 2019; Wen et al., 2023; Liu et al., 2024; Muralidharan et al., 2024; Zhong et al., 2024; Zhang et al., 2024b; Wu et al., 2025) on a fixed dataset. Recently, (Lin et al., 2020; Agarwal et al., 2024; Gu et al., 2024) argues that applying KLD incurs train-inference mismatch, thus turn to reverse Kullback-Leibler divergence (RKLD) and student-generated outputs, casting KD into a reinforcement learning problem. Orthogonal to the above works, builds upon an ODE perspective, our FDD focuses mainly on matching different quantities to fully exploit the teacher. Along this line of research, hidden states (features) (Jiao et al., 2020) and attention scores (Wang et al., 2020) have been considered. However, these approaches are only successful on BERT (Devlin et al., 2019) models. TED (Liang et al., 2023) employs task-aware filters to align hidden representations between student and teacher models at each layer. While task-aware filters function similarly to LM head, they differ in that the LM head utilizes pre-trained parameters. Besides, our FDD goes further beyond layerwise distillation. From an ODE perspective, FDD aligns the entire feature dynamics between student and teacher models, including the discretized feature trajectories (layer-wise features) as well as its first-order derivative. The effectiveness of the proposed method is validated on modern GPT-like LLMs.

Another line of research is black-box KD, also
known as data KD. In practice, cases are that teacher output distribution is not accessible e.g., ChatGPT (OpenAI, 2022) APIs so the teacher is utilized as data generators (Taori et al., 2023; Peng et al., 2023; Hsieh et al., 2023; Xu et al., 2024), the resulting generated data are used for supervised fine-tuning smaller LMs. Black-box KD compromises the effectiveness of the student for teacher utility. In this paper, we consider only the white-box distillation where the teacher model is fully exploitable.

Neural ODEs Neural networks with layer-wise residual connections (He et al., 2016) have been found corresponding to Euler discretization of solutions to ordinary differential equations (ODEs) (Chen et al., 2018), with layer indices being integer time steps. Specifically, following the Lie-Trotter splitting scheme (Trotter, 1959) to accommodate for alternating MHA and MLP blocks, (Lu et al., 2019) cast the transformer (Vaswani et al.) into a convection-diffusion equation. Various novel deep architectures are proposed then from an ODE perspective, including Macaron Net (Lu et al., 2018), TransEvolve (Dutta et al., 2021), ODETransformer (Li et al., 2022) etc. Different from the above works aiming at improving the model architecture, we develop a new knowledge distillation approach from an ODE perspective to fully exploit the knowledge encoded in the teacher.

6 Conclusion

In this paper, we proposed feature dynamics distillation, a general knowledge distillation (KD) method for LLMs, inspired by viewing transformers through an ODE perspective. FDD extends beyond traditional KD methods that only match output logits by considering the entire feature dynamics, including both the discretized feature trajectory and its first-order derivative estimated through finite differences. This approach enables student models to leverage the rich information encoded in intermediate layers while mimicking the pattern of representation refinement across layers. Through extensive experiments and analysis, we demonstrated the effectiveness of this strategy.

7 Limitations

The limitations of this paper are as follows: (1) Our proposed FDD method utilizes the LM head to project intermediate layer representations of both the teacher and student models into a shared vocabulary space. Consequently, its applicability may be restricted when applied across different model architectures or vocabulary configurations. (2) Due to limited computational resources, our experiments were conducted with a teacher model of up to 13 billion parameters and a student model of up to 7 billion parameters. While the experimental results demonstrate that our method significantly outperforms existing methods at these model scales, the limitation in model size prevents us from verifying the generalizability of our findings to extremely large models, such as those with 405 billion parameters.

References

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. Layer normalization. In 1607.06450.

Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. 2023. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08172.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Mateusz Litwin Eric Sigler, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. arXiv preprint arXiv:2412.19437.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. 2018. Neural ordinary differential equations. Advances in neural information processing systems, 31.

Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. 2023. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality.

Yung-Sung Chuang, Yujia Xie, Hongyin Luo, Yoon Kim, James Glass, and Pengcheng He. 2024. Dola: Decoding by contrasting layers improves factuality in large language models. In International Conference on Learning Representations.
Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. 2023. Analyzing transformers in embedding space. In Association for Computational Linguistics.

DeepSeek-AI-Group. 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Chapter of the Association for Computational Linguistics.

Subhabrata Dutta, Tanya Gautam, Soumen Chakrabarti, and Tanmoy Chakraborty. 2021. Redesigning the transformer architecture with insights from multi-particle dynamical systems. In Advances in Neural Information Processing Systems.

Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, and Carole-Jean Wu. 2024. Layerskip: Enabling early exit inference and self-speculative decoding. In Association for Computational Linguistics.

Xinyang Geng, Arnav Gudibande, Hao Liu, Eric Wallace, Pieter Abbeel, Sergey Levine, and Dawn Song. 2023. Koala: A dialogue model for academic research.

Xinyang Geng and Hao Liu. 2023. Openllama: An open reproduction of llama.

Mor Geva, Avi Caciularu, Kevin Wang, and Yoav Goldberg. 2022. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Conference on Empirical Methods in Natural Language Processing.

Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. Minilm: Knowledge distillation of large language models. In International Conference on Learning Representations.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.

Or Honovich, Thomas Scialom, Omer Levy, and Timo Schick. 2023. Unnatural instructions: Tuning language models with (almost) no human labor. In Annual Meeting of the Association for Computational Linguistics.

Cheng-Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alex Ratner, Ranjay Krishna, Chen-Yu Lee, and Tomas Pfister. 2023. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. In Findings of the Association for Computational Linguistics.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.

Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. Tiny{bert}: Distilling {bert} for natural language understanding. In Conference on Empirical Methods in Natural Language Processing.

Yoon Kim and Alexander M. Rush. 2016. Sequence-level knowledge distillation. In Conference on Empirical Methods in Natural Language Processing.

Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun. 2024. Distillm: Towards streamlined distillation for large language models. In International Conference on Machine Learning.

Bei Li, Quan Du, Tao Zhou, Shuhan Zhou, Xin Zeng, 2 Tong Xiao1, and Jingbo Zhu. 2022. Ode transformer: An ordinary differential equation-inspired model for neural machine translation. In Association for Computational Linguistics.

Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao. 2023. Less is more: Task-aware layer-wise distillation for language model compression. In International Conference on Machine Learning.

Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. 2020. Autoregressive knowledge distillation through imitation learning. In Conference on Empirical Methods in Natural Language Processing.

Chin-Yew Lin. 2004. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out.

Jiaheng Liu, Chenchen Zhang, Jinyang Guo, Yuanxing Zhang, Haoran Que, Ken Deng, Zhiqi Bai, Jie Liu, Ge Zhang, Jiakai Wang, Yanan Wu, Congnan Liu, Wenbo Su, Jiamang Wang, Lin Qu, and Bo Zheng. 2024. Ddk: Distilling domain knowledge for efficient large language models. In Advances in Neural Information Processing Systems.

Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie-Yan Liu. 2019. Understanding and improving transformer from a multi-particle dynamic system point of view. arXiv preprint arXiv:1906.02762.

Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. 2018. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference on Machine Learning.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochoowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. 2024. Compact language models via pruning and knowledge distillation. Advances in Neural Information Processing Systems.

nostalgebraist. 2020. Interpreting gpt: the logit lens.

OpenAI. 2022. Introducing chatgpt.

OpenAI. 2024. Gpt-4 technical report. arXiv preprint arXiv:2412.19437.

Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, and Jianfeng Gao. 2023. Instruction tuning with gpt-4. arXiv preprint arXiv:2304.03277.

Qwen. 2024. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. In Workshop on Energy Efficient Machine Learning and Cognitive Computing - Advances in Neural Information Processing Systems.

Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Patel, Jalal Naghibyev, Yann LeCun, and Raviv Shwartz-Ziv. 2025. Layer by layer: Uncovering hidden representations in language models. arXiv preprint arXiv:2502.02013.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.

Hale F Trotter. 1959. On the product of semi-groups of operators. Proceedings of the American Mathematical Society, 10(4):545–551.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems.

Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. 2020. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. In Advances in Neural Information Processing Systems.

Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A. Smith, Daniel Khashabi, and Hannaneh Hajishirzi. 2023. Self-instruct: Aligning language models with self-generated instructions. In Annual Meeting of the Association for Computational Linguistics.

Yizhong Wang, Swaroop Mishra, Pegah Alipoormalabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap, Eshaan Pathak, Giannis Karamanolakis, Haizhi Gary Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby Kuznia, Krima Doshi, Kuntal Kumar Pal, Maitreyaa Patel, Mehrad Moradshahi, Mihr Parmar, Mirali Purohit, Neeraj Varshtey, Phani Rohitha Kaza, Pulkit Verma, Ravechaj Singh Puri, Rushang Karia, Savan Doshi, Shaileja Keyur Sampat, Siddhartha Mishra, Sujan Reddy A, Sumanta Patro, Tanay Dixit, and Xudong Shen. 2022. Super-naturalinstructions: Generalization via declarative instructions on 1600+ NLP tasks. In Conference on Empirical Methods in Natural Language Processing.

Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. 2023. f-divergence minimization for sequence-level knowledge distillation. In Annual Meeting of the Association for Computational Linguistics.

Taiqiang Wu, Chaofan Tao, Jiahao Wang, Runming Yang, Zhe Zhao, and Ngai Wong. 2025. Rethinking kullback-leibler divergence in knowledge distillation for large language models. In International Conference on Computational Linguistics.

Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, Qingwei Lin, and Daxin Jiang. 2024. WizardLM: Empowering large pre-trained language models to follow complex instructions. In The Twelfth International Conference on Learning Representations.

Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. 2024a. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.02385.

Songming Zhang, Xue Zhang, Zengkui Sun, Yufeng Chen, and Jinan Xu. 2024b. Dual-space knowledge distillation for large language models. In Conference on Empirical Methods in Natural Language Processing, pages 18164–18181.

Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuoban Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. 2023. Judging llm-as-a-judge with mt-bench and chatbot arena. In Advances in Neural Information Processing Systems.

Qihuang Zhong, Liang Ding, Li Shen, Juhua Liu, Bo Du, and Dacheng Tao. 2024. Revisiting knowledge distillation for autoregressive language models. In Association for Computational Linguistics.

Reviewer paper summaries: This paper proposes a novel knowledge distillation method, named Feature Dynamics Distillation (FDD). From the perspective of regarding the Transformer as a discretized Ordinary Differential Equation (ODE), FDD emphasizes aligning the entire feature dynamics of the teacher and student models, including the feature trajectories and their first-order derivatives. FDD extends the original KD objective by adding two additional loss terms: layer-wise feature KD and layer feature delta KD. And the effectiveness of FDD has been verified on multiple tasks.

This paper introduces Feature Dynamics Distillation (FDD), a novel knowledge distillation approach for large language models. Based on viewing transformers as discretized ordinary differential equations (ODEs), FDD aligns not just the final output distributions between teacher and student models, but also the entire feature dynamics across layers. The method extends traditional knowledge distillation with two additional components: layer-wise feature distillation to match feature trajectories, and layer-delta distillation to align feature derivatives across adjacent layers. Experiments across various model architectures (LLAMA2, OpenLLAMA2, GPT2) and instruction-following benchmarks demonstrate that FDD consistently outperforms state-of-the-art distillation methods.
Summary of strengths: 1. FDD presents a new perspective by treating the Transformer as an ODE, emphasizing the alignment of the entire feature dynamics, rather than just the final output distribution. It provides new ideas for knowledge distillation in LLMs.
2. The experiments are explicit. For example, ablation experiments and experiments with multiple different model structure scales illustrate the effectiveness of the method.

1. The paper provides a well-motivated theoretical framework by connecting transformers to ODEs, giving clear mathematical justification for the proposed approach.
2. Experiments across multiple model architectures (from 0.1B to 13B parameters) and five diverse benchmarks show consistent improvements over existing methods. For example, in LLAMA2-13B to LLAMA2-7B distillation, FDD improved average ROUGE-L scores from 29.82 to 32.09.
3. The paper offers clear implementation guidance, including hyperparameter settings, optimization techniques, and ablation studies that evaluate different components, allowing for practical reproducibility.
4. The paper resolves the challenge of aligning different feature dimensions between teacher and student models by using the language model head to project features into a shared vocabulary space, showing better performance than alternative projection methods.
Summary of weaknesses: 1. The cost analysis of FDD compared with the baseline method during the distillation process has not been presented. Currently, there is a lack of explanations regarding resource consumption such as training time.

2. FDD introduces additional loss terms and alignment strategies, which may increase the complexity of implementation and parameter tuning. However, there is a lack of relevant parameter analysis experiments.

3. The description of the author's motivation in the current version is not clear enough. Can it be understood as a different way to solve the problem of distribution mismatch, that is, the distribution of the final output is difficult to learn, while the intermediate features are more fine-grained and can transfer knowledge more comprehensively? Whether this situation can be reflected in that the distillation results can still be achieved well when the distribution differences are too large, for example, the author can try to show the distillation results when comparing with existing methods in the case of a large gap between models such as from a 13B model to a 1B model. 
4. The current discussion of related work lacks insights. There have been relevant discussions on the impact of distribution mismatch on KD, such as [1], [2], [3] and [4]. Clarifying the various potential connections between these papers and the KD work mentioned in this paper and reorganizing the related work will help the reader better understand the motivation for this paper.

[1] Shadow Knowledge Distillation: Bridging Offline and Online Knowledge Transfer

[2] Parameter-Efficient and Student-Friendly Knowledge Distillation

[3] Tailoring Instructions to Student’s Learning Levels Boosts Knowledge Distillation

[4] TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models

1. While the paper demonstrates improved performance, it doesn't thoroughly analyze the additional computational cost of implementing layer-wise and layer-delta distillation compared to traditional KD, which would be important for practical adoption. Like do we need additional saving or storing for the distilling and the speed comparison?
2. The experiments only go up to 13B parameters for teacher models, leaving uncertainty about effectiveness for very large models
Comments, suggestions and typos: I would be happy to discuss and address the above with the author.

1. It seems the distilled models can even surpass the original teacher model in most of the benchmarks. Are there any explanations for it?
2. What's the cost of using additional hidden states for the distilling? Are you first store all the hidden states on the disk and then load them during distilling, or you server the teacher and student model at the same time and get the teacher's hidden states in real-time?
Overall assessment scores: 3.5

4