文章地址

在以往的知识蒸馏中,我们要求 teacher 和 student 的 feature 或者 logits 直接数值上的靠近。但是由于 teacher 和 student 网络的差异,直接使用数值上的靠近可能对 student 比较困难。

Method

本文提出的方法是,蒸馏 teacher 和 student 之间 feature 的分布,而忽略数值上的一一对应,从而具有更好的灵活性。
Fig 2.

RKD 公式:

$$\mathcal{L}_{\mathrm{RKD}}=\sum_{\left(x_{1}, \ldots, x_{n}\right) \in \mathcal{X}^{N}} l\left(\psi\left(t_{1}, \ldots, t_{n}\right), \psi\left(s_{1}, \ldots, s_{n}\right)\right)$$

其中 $(x_1, x_2,...,x_n)$ 是互不相同的n元组, $\psi$ 是 relational potential 函数,度量 $(x_1, x_2,…,x_n)$ 之间的relation structure。

Distance-wise distillation loss

relation potential 函数定义为:

$$\psi_{\mathrm{D}}\left(t_{i}, t_{j}\right)=\frac{1}{\mu}\left\|t_{i}-t_{j}\right\|_{2}$$

其中 $\mu$ 为归一化系数,消除 student 和 teacher 在 scale 上的差异:

$$\mu=\frac{1}{\left|\mathcal{X}^{2}\right|} \sum_{\left(x_{i}, x_{j}\right) \in \mathcal{X}^{2}}\left\|t_{i}-t_{j}\right\|_{2}$$

基于该relation potential函数,distance-wise损失定义为:

$$\mathcal{L}_{\mathrm{RKD-D}}=\sum_{\left(x_{i}, x_{j}\right) \in \mathcal{X}^{2}} l_{\delta}\left(\psi_{\mathrm{D}}\left(t_{i}, t_{j}\right), \psi_{\mathrm{D}}\left(s_{i}, s_{j}\right)\right)$$

其中 $l_{\delta}$ 是 Huber loss:

$$l_{\delta}(x, y)=\left\{\begin{array}{ll}{\frac{1}{2}(x-y)^{2}} & {\text { for }|x-y| \leq 1} \\ {|x-y|-\frac{1}{2},} & {\text { otherwise }}\end{array}\right.$$

Angle-wise distillation loss

作者给了针对三元组的损失,这里 $\psi$ 定义为:

$$\begin{array}{l}{\psi_{\mathrm{A}}\left(t_{i}, t_{j}, t_{k}\right)=\cos \angle t_{i} t_{j} t_{k}=\left\langle\mathbf{e}^{i j}, \mathbf{e}^{k j}\right\rangle} \\ {\text { where } \quad \mathbf{e}^{i j}=\frac{t_{i}-t_{j}}{\left\|t_{i}-t_{j}\right\|_{2}}, \mathbf{e}^{k j}=\frac{t_{k}-t_{j}}{\left\|t_{k}-t_{j}\right\|_{2}}}\end{array}$$

Angle-wise蒸馏损失定义为:

$$\mathcal{L}_{\mathrm{RKD-A}}=\sum_{\left(x_{i}, x_{j}, x_{k}\right) \in \mathcal{X}^{3}} l_{\delta}\left(\psi_{\mathrm{A}}\left(t_{i}, t_{j}, t_{k}\right), \psi_{\mathrm{A}}\left(s_{i}, s_{j}, s_{k}\right)\right)$$

Experiment

Tab 1.

可以发现一些有趣的事情:

  1. 可以自己蒸馏自己(自蒸馏)
  2. student 可能超过 teacher
  3. 可以用来做domain adaptation
  4. 不加 L2 可能效果更好
文章目录