Precise Knowledge Transfer via Flow Matching

For ease of understanding, we show the pseudo code of FM-KT in the Offline KD scenario. The implementation of FM-KT^Θ and OFM-KT only needs to modify the optimization objective as described in our main paper.

FM-KT proposes a novel serial training paradigm in order to avoid “cheating” by accessing $X^{T}$ during training:

where $Z_{1}=\alpha_{1}X^{S}$. Here we assume the loss function is $\ell_{2}$-norm. Broadly speaking, the loss function used by FM-KT only needs to ensure that it can achieve the effect of minimizing the difference in distributions similar to Kullback-Leibler Divergence. When $i\geq 1$, $Z_{1-i/N}=Z_{1-(i-1)/N}-g_{v_{\theta}}(Z_{1-(i-1)/N},1-(i-1)/N)/N$. $\mathcal{T}(\cdot)$ is used for shape alignment to ensure the calculation of $\ell_{2}$-norm. Let us define $q(\cdot|\cdot)$ as the predefined conditional probability, and $p_{v_{\theta}}(\cdot|\cdot)$ as the predicted conditional probability. We know the training objective of the classical diffusion probability model can be performed by minimizing the upper bound on negative log-likelihood:

We can rewritten it as

For $i\!\geq\!1$, if $\textrm{Law}(Z_{i/N})\stackrel{{\scriptstyle\sim}}{{=}}\textrm{Law}(\hat{Z}_{i/N})$ is guaranteed, then $\textrm{Law}(Z_{(i-1)/N})\stackrel{{\scriptstyle\sim}}{{=}}\textrm{Law}(\hat{Z}_{(i-1)/N})$ can also be guaranteed by optimizing $\mathbb{E}_{\hat{Z}_{i/N},Z_{1},Z_{0}}D_{\mathrm{KL}}(q(Z_{(i-1)/N}|Z_{i/N},Z_{0})||p_{v_{\theta}}(Z_{(i-1)/N}|\hat{Z}_{i/N}))$ in Eq. 8. Based on the prior condition $\textrm{Law}(Z_{1})\stackrel{{\scriptstyle\sim}}{{=}}\textrm{Law}(\hat{Z}_{1})$, we can deduce $\{\textrm{Law}(Z_{i/N})\stackrel{{\scriptstyle\sim}}{{=}}\textrm{Law}(\hat{Z}_{i/N})\}_{i=0}^{N-1}$ sequentially by recursive method.

Note that this derivation via Bayes’ Theorem satisfies almost all noise schedules, such as Rectified flow (Liu et al., 2022; Lipman et al., 2022), VP ODE (Song et al., 2023c; Albergo & Vanden-Eijnden, 2022) and VE ODE (Song et al., 2023c; Albergo & Vanden-Eijnden, 2022). More details can be found in Appendix O. In fact, the upper bound in Eq. 8 is precisely the optimization objective of FM-KT. The main difference between Eq. 7 and Eq. 8 is that $\hat{Z}_{i/N}$ replaces ${Z}_{i/N}$, and $\hat{Z}_{i/N}$ is obtained from the reverse sampling process. In this manner, the trajectory $\{\hat{Z}_{i/N}\}_{i}$ is obtained through the estimator $g_{v_{\theta}}(\cdot)$ and no longer contains the teacher knowledge $X^{T}$, thereby enabling the distillation process to proceed normally.

Ensemble is a method that trains multiple models, aggregates their outputs through voting, and produces a final prediction. In this section, we prove theoretically that FM-KT is essentially a unique implicit ensemble method under the assumption that its noise schedule is set as Rectified flow.

First, we define ODE in FM-KT as $X^{S}-X^{T}=\frac{dX_{t}}{dt}$ (for the convenience of derivation, this definition is slightly different from that in the main paper), so we need to fit $||\frac{dX_{t}}{dt}-g_{v_{\theta}}(X_{t},t)||_{2}^{2}$, where $X_{t}=tX^{S}+(1-t)X^{T}$, $t\sim\mathcal{U}[0,1]$. In inference, this ODE solver defaults to Euler’s method in FM-KT, and the sampling must be discrete with $N$ steps because fitting continuous time steps $t$ consumes extensive computational costs. When the meta-encoder $g_{v_{\theta}}(\cdot)$ is at the optimal solution, we assume that its error from the true value can be expressed as a function of $x_{t}$ and $t$, and that this function is at 1-Lipschitz.

Thus, we can define $\mathcal{H}(t)=\operatorname*{arg\,sup}_{X_{t}}\{||\frac{dX_{t}}{dt}-g_{v_{\theta}}(X_{t},t)||_{2}^{2}\}$, then the truncation error $\mathcal{K}(t)$ can be defined as $\left[\frac{dX_{t}}{dt}-g_{v_{\theta}}(\mathcal{H}(t),t)\right]$, which is also at 1-Lipschitz under the assumption $\mathcal{H}(t)$ is at 1-Lipschitz. After that, we also need to define the step number of sampling in inference. We set it as $K$, so $dt$ is $1/K$. Based on the aforementioned notations, we can analyse the truncation error by the recursive method.

For the sake of derivation convenience, we define $\{Z_{t}\}_{t}$ as the sampled trajectory in inference to distinguish it from $\{X_{t}\}_{t}$ in training. Thus, a step in sampling can be described as $Z_{1-(i+1)/K}\!=\!Z_{1-i/K}-g_{v_{\theta}}(Z_{1-{i/K}},1-i/K)dt$, and $Z_{1-i/K}=X_{1-i/K}+\mathcal{E}(Z_{1-i/K})$, where $\mathcal{E}(Z_{1-i/K})$ refers to the truncation error accumulated to a intermediate sample $Z_{1-i/K}$ in the sampling process. Note that $\mathcal{K}(t)$ and $\mathcal{E}(Z_{1-i/K})$ are not results of the norm, and therefore $\forall t$ and $\forall Z_{1-i/K}$, this derivation does not need to satisfy that $\mathcal{K}(t)\geq 0$ and $\mathcal{E}(Z_{1-i/K})\geq 0$. This approach avoids the accumulation of the truncation error due to $\ell_{2}$-norm $\geq$ 0. We can derive the sample $Z_{1-(i-1)/K}$ in the next step by the derivation:

where $\psi(t)=\nabla_{X_{t}}g_{v_{\theta}}(X_{t},t)$. Then, Eq. 9 can continue to be derived as

Thus, $\mathcal{E}(Z_{1-(i+1)/K})=\mathcal{E}(Z_{1-i/K})[1-(1/K)\psi(1-i/K)]+(1/K)\mathcal{K}(1-i/K)$. After that, the recursive method leads us to the following conclusions:

Looking at the first term, we can see that the truncation error comes from summing $\mathcal{K}(\cdot)$ over all time points. When treating the error sampling as Monte Carlo sampling, with a sufficient number of samplings $K$, it becomes possible for FM-KT to approximate ensemble methods and thus estimate the ground truth effectively.

In this section, we present pair decoupling (PD), a straightforward yet effective technique for enhancing performance in the feature-based distillation scenario of image classification using FM-KT. This method involves shuffling a subset of samples in a batch to achieve regularization, thereby preventing overfitting of the teacher’s refined low-level hierarchical features. Let $B$, $C$, $H$ and $W$ denote the batch size, the number of channels, the height of the feature map, and the width of the feature map, respectively. Given the teacher’s feature map $X^{T}\in\mathbb{R}^{B\times C\times H\times W}$ from a specific layer, PD is applied prior to all FM-KT related calculations. Implementing PD involves defining a hyperparameter, the dirac ratio $\beta_{d}$, and perturbing $B\!-\!\left\lfloor{\beta_{d}}B\right\rfloor$ samples in the batch. The Pytorch code for this is provided in Appendix A. Specifically, PD selects the random $B\!-\!\left\lfloor{\beta_{d}}B\right\rfloor$ samples $X^{T}\left[0:B\!-\!\left\lfloor{\beta_{d}}B\right\rfloor\right]$ in a batch and then shuffles them:

We refer to the hyperparameter $\beta_{d}$ as “dirac ratio” because, following the PD operation, $\left\lfloor{\beta_{d}}B\right\rfloor$ samples are used to compute $\mathcal{L}_{\textrm{FM-KT}}$. Here, $X^{T}$ and $X^{S}$ are treated as Dirac distributions with the objective of achieving one-to-one matching. Conversely, the remaining $B-\left\lfloor{\beta_{d}}B\right\rfloor$ samples are utilized in the computation of $\mathcal{L}_{\textrm{FM-KT}}$, where $X^{T}$ and $X^{S}$ are considered as non-Dirac distributions, targeting many-to-many matching.

Due to the specificity of the feature-based distillation scenario for image classification, PD is specifically designed to avoid over-matching the refined low-level feature thus improving the final performance of the student. Meanwhile, our experiments in Appendix E empirically demonstrate that PD is effective only in the feature-based distillation scenario of image classification, whereas in other scenarios it rather degrades the performance. This is because matching the teacher’s feature/logit at a fine-grained level is closely related to the final performance of the student in the logit-based distillation scenario for image classification as well as the feature-based distillation scenario for object detection. In other words, in the feature-based distillation scenario for image classification, it does not imply that improving similarity between the student’s low-level feature and the teacher’s low-level feature will result in the greater classification accuracy of the student.

Here, we experimentally substantiate some empirical findings on the topics of normalization layer selection in the meta-encoder, where stages used for distillation in the feature-based scenario, and ideal configuration of dirac ratio $\beta_{d}$ in different scenarios.

Ablation experiments on various normalization operations reveal instability in the FM-KT training paradigm when using BatchNorm. As observed in Table 5, the accuracy achieved with BatchNorm as the normalization layer is approximately 1%, even when the training of the student remains stable (i.e., the loss is not NAN). This indicates that using BatchNorm in FM-KT introduces instability by computing the mean and variance of inputs at different time points during inference. It is essential to note that although BatchNorm is applied in DiffKD, this choice is justified as the student converges effectively with a sufficiently high number of layers in the Diffusion Model (i.e. meta-encoder) mentioned in their work. Similar results are obtained in our studies by replacing the meta-encoder in FM-KT with Diffusion Model in DiffKD.

In Fig. 8, we investigated the optimal stages for distillation in the feature-based scenario and the ideal configuration for the dirac ratio $\beta_{d}$. As there are no specific distillation stages for the logit-based distillation, we designated it as “[0, 0, 0]” for clarity. Our observations indicate that in the feature-based distillation scenario, the distillation stage does not significantly affect the final outcomes. Meanwhile, the configuration “[1, 1, 1]” often underperforms compared with “[0, 1, 1]” and “[0, 0, 1]”. This observation aligns with the conclusions drawn from most of the prior feature-based distillation studies (Tung & Mori, 2019; Zagoruyko & Komodakis, 2016a; Zong et al., 2023). Moreover, for different values of $\beta_{d}$, the settings $\beta_{d}$=0.25 typically yields the best result in feature-based distillation, while $\beta_{d}$=1.0 excels in the logit-based distillation. This implies that the PD technique is more effective in the feature-based distillation context for image classification, and less so in the logit-based distillation.

It is worth noting that our experiments on PD in object detection revealed that $\beta_{d}$=1.0 and $\beta_{d}$=0.75 yield comparable performance, whereas a decrease in $\beta_{d}$ results in diminished performance. Based on these findings, we recommend using $\beta_{d}$=0.25 as the default in the feature-based distillation scenario for image classification and $\beta_{d}$=1.0 in other contexts.

With the development of deep learning, there are a number of architectures belonging to continuous network representation, such as RNN (Williams & Zipser, 1989), LSTM (S & J, 1997), Neural ODE (Chen et al., 2018), GflowNet family (Bengio et al., 2021; Zhang et al., 2022), diffusion model family (Song et al., 2023c; Karras et al., 2022; Ho et al., 2020; Song et al., 2023a), INN (Solodskikh et al., 2023), DiffKD (Huang et al., 2023) and KDiffusion (Yao et al., 2024). Here, we mainly emphasize the similarities and differences between our proposed FM-KT and these methods. In this way, we show the novelty of the FM-KT design and its advantages in application:

• •

  The application scenario of FM-KT is different from RNN, LSTM, Neural ODE, GflowNet family, diffusion model family and INN. Of these, only FM-KT, DiffKD and KDiffusion are applied to knowledge distillation in the form of continuous network representations.

• •

  RNN, LSTM, Neural ODE, GflowNet family, diffusion model family, INN, and FM-KT have a meta-encoder shared parameters. However, the difference is that the forward process (meaning the backward process in diffusion model family and FM-KT) in RNN, LSTM, and GflowNet family is unknown. Unlike Neural ODE, diffusion model family and FM-KT, there exists a human-designed sampling process (a.k.a predefined forward processes), which makes it impossible to use numerical integration to trade-off performance and efficiency.

• •

  INN primarily enables the continuous representation of convolutional operations (i.e. convolutional kernels), not the entire network. In contrast, the Neural ODE, diffusion model family, continuously represents the entire network.

• •

  The primary distinction between Neural ODE/FM-KT and the diffusion model family lies in their training paradigms. The diffusion model family is trained using unpaired samples, aiming to capture the entire data distribution. In contrast, Neural ODE/FM-KT utilizes paired samples, focusing on learning the Dirac distribution of the output.

• •

  The biggest difference between FM-KT and Neural ODE is that FM-KT has a deterministic a priori forward process to model the optimization objective of the intermediate points (i.e. $\{Z_{1-i/N}\}_{i}$), which ensures the stability of training. Neural ODE has no such a priori forward process, and simply expects the network itself to learn a continuous representation from input to output.