SlowGuess/NeurIPS_2023_OCR huggingface dataset:
As we mentioned, our Proposition 1 does not hold precisely for arbitrary deep non-linear neural networks. However, we showed that this simple model (5) of distillation gradient approximates the true distillation gradient reasonably well both empirically (see Figure 1) and analytically (see Appendix B.3). There is much more to investigate for the case of non-convex deep networks where exact tracking of teacher's impact across multiple layers of non-linearities becomes harder. We see our results as a promising first step towards a more complete understanding of the effectiveness of distillation. One interesting direction of future work would be to construct more complex models for distillation gradient and to investigate further connections with more complex variance-reduction methods, e.g. [4], which may yield even better-performing variants of KD.
MS has received funding from the European Union's Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement No 101034413.
[1] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communication-efficient SGD via gradient quantization and encoding. In Advances in Neural Information Processing Systems 30, pages 1709-1720, 2017.
[2] Lei Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? CoRR, abs/1312.6184, 2013.
[3] C. Bucilua, R. Caruana, and A. Niculescu-Mizil. Model compression. Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD '06, pages 535-541, New York, NY, USA, 2006.
[4] Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex SGD. Advances in Neural Information Processing Systems, 32, 2019.
[5] Wojciech Czarnecki, Siddhant Jayakumar, Max Jaderberg, Leonard Hasenclever, Yee Whye Teh, Nicolas Heess, Simon Osindero, and Razvan Pascanu. Mix & match agent curricula for reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, 2018.
[6] Tri Dao, Govinda M Kamath, Vasilis Syrgkanis, and Lester Mackey. Knowledge Distillation as Semiparametric Inference. International Conference on Learning Representations (ICLR), 2021.
[7] Alexandre Galashov, Siddhant Jayakumar, Leonard Hasenclever, Dhruva Tirumala, Jonathan Schwarz, Guillaume Desjardins, Wojtek M. Czarnecki, Yee Whye Teh, Razvan Pascanu, and Nicolas Heess. Information asymmetry in KL-regularized RL. In International Conference on Learning Representations, 2019.
[8] Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235, 2023.
[9] Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent. 23rdInternational Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
[10] R.M. Gower, P. Richtárik, and F. Bach. Stochastic quasi-gradient methods: variance reduction via Jacobian sketching. Math. Program. 188, pages 135–192, 2021.
[11] Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General Analysis and Improved Rates. Proceedings of the 36th International Conference on Machine Learning, 2019.
[12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Conference on Computer Vision and Pattern Recognition, 2016.
[13] G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network. Deep Learning Workshop at NIPS, 2014.
[14] G. Hinton, O. Vinyals, and J. Dean. Distilling the Knowledge in a Neural Network. *ArXiv e-prints*, March 2015.
[15] Eugenia Iofinova, Alexandra Peste, Mark Kurtz, and Dan Alistarh. How Well Do Sparse Imagenet Models Transfer? Conference on Computer Vision and Pattern Recognition, 2022.
[16] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in Neural Information Processing Systems, 26, 2013.
[54] Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33:5776–5788, 2020.
[55] Jianqiao Wangni, Jialei Wang, Ji Liu, and Tong Zhang. Gradient sparsification for communication-efficient distributed optimization. In Advances in Neural Information Processing Systems, pages 1306-1316, 2018.
[56] Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Terngrad: Ternary gradients to reduce communication in distributed deep learning. In Advances in Neural Information Processing Systems, page 1509-1519, 2017.
[57] Lin Xiao and Tong Zhang. A Proximal Stochastic Gradient Method with Progressive Variance Reduction. SIAM Journal on Optimization, 2014.
[58] Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng. Revisiting Knowledge Distillation via Label Smoothing Regularization. Conference on Computer Vision and Pattern Recognition, 2020.
[59] Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be your own teacher: Improve the performance of convolutional neural networks via self distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 3713-3722, 2019.
[60] Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation. ICCV, 2019.
[61] Xinshao Wang Ziyun Li, Di Hu, Neil M. Robertson, David A. Clifton, Christoph Meinel, and Haojin Yang. Not All Knowledge Is Created Equal: Mutual Distillation of Confident Knowledge. NeurIPS 2022 Workshop(Trustworthy and Socially Responsible Machine Learning), 2022.
1 Introduction 1
2 Related Work 2
3 Knowledge Distillation 3
4 Convergence Theory for Self-Distillation 5
5 Removing Bias and Improving Variance Reduction 8
6 Convergence for Distillation of Compressed Models 9
7 Discussion and Future Work 10
A Basic Facts and Inequalities 16
B Proofs for Section 3 17
B.1 Binary Logistic Regression 17
B.2 Multi-class Classification with Soft-max 18
B.3 Generic Non-linear Classification 18
C Proofs for Section 4 20
C.1 Key lemma 20
C.2 Proof of Theorem 1 21
C.3 Proof of Theorem 2 23
D Proofs for Section 5 23
D.1 Proof of Theorem 3 23
E Proofs for Section 6 24
E.1 Five Lemmas 24
E.2 Proof of Theorem 4 26
F Additional Experimental Validation 28
F.1 The impact of the learning hyperparameters 28
F.2 The impact of the teacher's quality 29
F.3 The impact of knowledge distillation on compression 29
G Limitations 29
A Basic Facts and Inequalities
To facilitate the reading of technical part of the work, here we present several standard inequalities and basic facts that we are going to use in the proofs.
- Usually we need to bound the sum of two error terms by individual errors using the following simple bound
$$
\left\| a + b \right\| ^ {2} \leq 2 \left\| a \right\| ^ {2} + 2 \left\| b \right\| ^ {2}. \tag {20}
$$
A direct generalization of this inequality for arbitrary number of summation terms is the following one:
$$
\left\| \sum_ {i = 1} ^ {n} a _ {i} \right\| ^ {2} \leq n \sum_ {i = 1} ^ {n} \| a _ {i} \| ^ {2}. \tag {21}
$$
This inequality can be seen as a special case of Jensen's inequality for convex functions  $h\colon \mathbb{R}^d\to \mathbb{R}$
$$
h \left(\sum_ {i = 1} ^ {n} \alpha_ {i} x _ {i}\right) \leq \sum_ {i = 1} ^ {n} \alpha_ {i} h (x _ {i}),
$$
where  $x_{i} \in \mathbb{R}^{d}$  are any vectors and  $\alpha_{i} \in [0,1]$  with  $\sum_{i=1}^{n} \alpha_{i} = 1$ . Then, (20) follows from Jensen's inequality when  $h(x) = \|x\|^{2}$  and  $\alpha_{1} = \dots = \alpha_{n} = \frac{1}{n}$ . A more general version of the Jensen's inequality, from the perspective of probability theory, is
$$
h (\mathbb {E} z) \leq \mathbb {E} h (z) \tag {22}
$$
for any convex function  $h$  and random vector  $z\in \mathbb{R}^d$
Another extension of (20), which we actually apply in that form, is Peter-Paul inequality given by
$$
\langle a, b \rangle \leq \frac {1}{2 s} \| a \| ^ {2} + \frac {s}{2} \| b \| ^ {2}, \tag {23}
$$
for any positive  $s > 0$ . Then, (20) is the special case of (23) with  $s = 1$ .
- Typically, a function  $f$  is called  $L$ -smooth if its gradient is Lipschitz continuous with Lipschitz constant  $L \geq 0$ , namely
$$
\left\| \nabla f (x) - \nabla f (y) \right\| \leq L \| x - y \|, \tag {24}
$$
In particular, smoothness inequality (24) implies the following quadratic upper bound [35]:
$$
f (y) \leq f (x) + \left\langle \nabla f (x), y - x \right\rangle + \frac {L}{2} \| y - x \| ^ {2}, \tag {25}
$$
for all points  $x, y \in \mathbb{R}^d$ . If the function  $f$  is additionally convex, then the following lower bound holds too:
$$
f (y) \geq f (x) + \left\langle \nabla f (x), y - x \right\rangle + \frac {1}{2 L} \| \nabla f (y) - \nabla f (x) \| ^ {2}. \tag {26}
$$
- As we mentioned in the main part of the paper, function  $f$  is called  $\mu$ -strongly convex if the following inequality holds
$$
f (y) \geq f (x) + \left\langle \nabla f (x), y - x \right\rangle + \frac {\mu}{2} \| y - x \| ^ {2}, \tag {27}
$$
for all points  $x, y \in \mathbb{R}^d$ . Recall that strong quasi-convexity assumption (1) is the special case of (27) when  $y = x^{*}$ . In other words, strong convexity (27) implies strong quasi-convexity (1). Continuing this chain of conditions, strong quasi-convexity (1) implies PL condition (9), which, in turn, implies the so-called Quadratic Growth condition given below
$$
f (x) - f ^ {*} \geq \frac {\mu}{8} \| x - x ^ {*} \| ^ {2}, \tag {28}
$$
for all  $x\in \mathbb{R}^d$ . Derivations of these implications and relationship with other conditions can be found in [17].
- As many analyses with linear or exponential convergence speed, our analysis also uses a very standard transformation from a single-step recurrence relation to convergence inequality. Specifically, assume we have the following recursion for  $e^t$ ,  $t = 0,1,2,\ldots$ :
$$
e ^ {t + 1} \leq (1 - \eta) e ^ {t} + N,
$$
with constants  $\eta \in (0,1]$  and  $N\geq 0$ . Then, repeated application of the above recursion gives
$$
\begin{array}{l} e ^ {t} \leq (1 - \eta) ^ {t} e ^ {0} + (1 - \eta) ^ {t - 1} N + (1 - \eta) ^ {t - 2} N + \dots + (1 - \eta) ^ {0} N \\ \leq (1 - \eta) ^ {t} e ^ {0} + N \sum_ {j = 0} ^ {\infty} (1 - \eta) ^ {j} = (1 - \eta) ^ {t} e ^ {0} + \frac {N}{\eta}. \tag {29} \\ \end{array}
$$
- If  $z \in \mathbb{R}^d$  is a random vector and  $x \in \mathbb{R}^d$  is fixed (or has randomness independent of  $z$ ), then the following decomposition holds:
$$
\mathbb {E} \left[ \| x + z \| ^ {2} \right] = \| x \| ^ {2} + \mathbb {E} \left[ \| z \| ^ {2} \right] \tag {30}
$$
B Proofs for Section 3
For the sake of presentation, we first consider binary logistic regression problem as a special case of multi-class classification.
B.1 Binary Logistic Regression
In this case we have  $d$ -dimensional input vectors  $a_{n} \in \mathbb{R}^{d}$  with their binary true labels  $b_{n} \in \mathcal{B} = [0,1]$ , predictor  $\phi_{x}(a) = \sigma (x^{\top}\bar{a}) \in (0,1)$  with parameters  $x \in \mathcal{P} = \mathbb{R}^{d + 1}$  and lifted input vector  $\bar{a} = [a1]^{\top} \in \mathbb{R}^{d + 1}$  (to avoid additional notation for the bias terms), where  $\sigma (t) = \frac{1}{1 + e^{-t}}$ ,  $t \in \mathbb{R}$  is the sigmoid function. Besides, the loss is given by the cross entropy loss below
$$
\ell (p, q) = H \left(\left[ \begin{array}{c} q \\ 1 - q \end{array} \right], \left[ \begin{array}{c} p \\ 1 - p \end{array} \right]\right) = - q \log p - (1 - q) \log (1 - p), \quad p, q \in [ 0, 1 ].
$$
Based on (2) and (3), we have
$$
\begin{array}{l} f _ {n} (x \mid \theta , \lambda) = (1 - \lambda) f _ {n} (x) + \lambda f _ {n} (x \mid \theta) \\ = (1 - \lambda) \ell \left(\phi_ {x} \left(a _ {n}\right), b _ {n}\right) + \lambda \ell \left(\phi_ {x} \left(a _ {n}\right), \phi_ {\theta} \left(a _ {n}\right)\right) \\ = \ell \left(\phi_ {x} \left(a _ {n}\right), (1 - \lambda) b _ {n} + \lambda \phi_ {\theta} \left(a _ {n}\right)\right) \\ = \ell (\sigma (x ^ {\top} \bar {a} _ {n}), (1 - \lambda) b _ {n} + \lambda \sigma (\theta^ {\top} \bar {a} _ {n})) = \ell (\sigma (x ^ {\top} \bar {a} _ {n}), s _ {n}) \\ = s _ {n} \log \left(1 + e ^ {- x ^ {\top} \bar {a} _ {n}}\right) + (1 - s _ {n}) \log \left(1 + e ^ {x ^ {\top} \bar {a} _ {n}}\right), \\ \end{array}
$$
where  $s_n = (1 - \lambda) b_n + \lambda \sigma (\theta^\top \bar{a}_n)$  are the soft labels. Notice that  $f_n(x \mid \theta, \lambda = 0) = f_n(x)$ . Next, we derive an expression for the stochastic gradient for the above objective, namely the gradient  $\nabla_x f_n(x \mid \theta, \lambda)$  of the loss associated with  $n^{th}$  data point  $(a_n, b_n)$ .
$$
\begin{array}{l} \nabla_ {x} \left[ s _ {n} \log \left(1 + e ^ {- x ^ {\top} \bar {a} _ {n}}\right) \right] = - s _ {n} \frac {e ^ {- x ^ {\top} \bar {a} _ {n}}}{1 + e ^ {- x ^ {\top} \bar {a} _ {n}}} \bar {a} _ {n} = - s _ {n} \sigma (- x ^ {\top} \bar {a} _ {n}) \bar {a} _ {n}, \\ \nabla_ {x} \left[ (1 - s _ {n}) \log \left(1 + e ^ {x ^ {\top} \bar {a} _ {n}}\right) \right] = (1 - s _ {n}) \frac {e ^ {x ^ {\top} \bar {a} _ {n}}}{1 + e ^ {x ^ {\top} \bar {a} _ {n}}} \bar {a} _ {n} = (1 - s _ {n}) \sigma (x ^ {\top} \bar {a} _ {n}) \bar {a} _ {n}. \\ \end{array}
$$
Hence, using the identity  $\sigma(t) + \sigma(-t) = 1$  for the sigmoid function, we get
$$
\begin{array}{l} \nabla_ {x} f _ {n} (x \mid \theta , \lambda) = \left[ - s _ {n} \left(1 - \sigma \left(x ^ {\top} \bar {a} _ {n}\right)\right) + \left(1 - s _ {n}\right) \sigma \left(x ^ {\top} \bar {a} _ {n}\right) \right] \bar {a} _ {n} \\ = \left[ \sigma \left(x ^ {\top} \bar {a} _ {n}\right) - s _ {n} \right] \bar {a} _ {n} \\ = \left[ \sigma \left(x ^ {\top} \bar {a} _ {n}\right) - (1 - \lambda) b _ {n} - \lambda \sigma \left(\theta^ {\top} \bar {a} _ {n}\right) \right] \bar {a} _ {n} \\ = \left(\sigma \left(x ^ {\top} \bar {a} _ {n}\right) - b _ {n}\right) \bar {a} _ {n} - \lambda \left(\sigma \left(\theta^ {\top} \bar {a} _ {n}\right) - b _ {n}\right) \bar {a} _ {n} = \nabla f _ {n} (x) - \lambda \nabla f _ {n} (\theta). \\ \end{array}
$$
Thus, the distillation gradient for binary logistic regression tasks the same form as (5).
B.2 Multi-class Classification with Soft-max
Now we extend the above steps for multi-class problem. Here we consider a  $K$ -classification model with one hidden layer that has soft-max as the last layer, i.e., the forward pass has the following steps:
$$
a _ {n} \to X ^ {\top} a _ {n} \to \phi_ {X} (a _ {n}) := \sigma (X ^ {\top} a _ {n}) \in \mathbb {R} ^ {K},
$$
where  $X = \left[x_{1} x_{2} \ldots x_{K}\right] \in \mathbb{R}^{d \times K}$  are the student's model parameters,  $a_{n} \in \mathcal{A} = \mathbb{R}^{d}$  is the input data and  $\sigma$  is the soft-max function (as a generalization of sigmoid function). Then we simplify the loss
$$
\begin{array}{l} f _ {n} (X \mid \Theta , \lambda) = (1 - \lambda) f _ {n} (X) + \lambda f _ {n} (X \mid \Theta) \\ = (1 - \lambda) \ell \left(\phi_ {X} \left(a _ {n}\right), b _ {n}\right) + \lambda \ell \left(\phi_ {X} \left(a _ {n}\right), \phi_ {\Theta} \left(a _ {n}\right)\right) \\ = \ell \left(\phi_ {X} \left(a _ {n}\right), (1 - \lambda) b _ {n} + \lambda \phi_ {\Theta} \left(a _ {n}\right)\right) \\ = \ell (\sigma (X ^ {\top} a _ {n}), s _ {n}) \\ = - \sum_ {k = 1} ^ {K} s _ {n, k} \log \sigma (X ^ {\top} a _ {n}) _ {k} = - \sum_ {k = 1} ^ {K} s _ {n, k} \log \frac {e ^ {x _ {k} ^ {\top} a _ {n}}}{\sum_ {j = 1} ^ {K} e ^ {x _ {j} ^ {\top} a _ {n}}} \\ = \sum_ {k = 1} ^ {K} s _ {n, k} \left(\log \sum_ {j = 1} ^ {K} e ^ {x _ {j} ^ {\top} a _ {n}} - \log e ^ {x _ {k} ^ {\top} a _ {n}}\right) = \sum_ {k = 1} ^ {K} s _ {n, k} \log \sum_ {j = 1} ^ {K} e ^ {x _ {j} ^ {\top} a _ {n}} - \sum_ {k = 1} ^ {K} s _ {n, k} \log e ^ {x _ {k} ^ {\top} a _ {n}} \\ = \log \sum_ {k = 1} ^ {K} e ^ {x _ {k} ^ {\top} a _ {n}} - \sum_ {k = 1} ^ {K} s _ {n, k} \log e ^ {x _ {k} ^ {\top} a _ {n}} = \log \sum_ {k = 1} ^ {K} e ^ {x _ {k} ^ {\top} a _ {n}} - \sum_ {k = 1} ^ {K} s _ {n, k} x _ {k} ^ {\top} a _ {n}, \\ \end{array}
$$
where  $s_n = (1 - \lambda) b_n + \lambda \phi_\Theta(a_n) \in \mathbb{R}^K$  are the soft labels. Next, we derive an expression for the stochastic gradient for the above objective, namely the gradient  $\nabla_{x_k} f_n(X \mid \Theta, \lambda)$  of the loss associated with  $n^{th}$  data point  $(a_n, b_n)$ .
$$
\begin{array}{l} \nabla_ {x _ {k}} f _ {n} (X \mid \Theta , \lambda) = \nabla_ {x _ {k}} \left[ \log \sum_ {k = 1} ^ {K} e ^ {x _ {k} ^ {\top} a _ {n}} - \sum_ {k = 1} ^ {K} s _ {n, k} x _ {k} ^ {\top} a _ {n} \right] = \frac {e ^ {x _ {k} ^ {\top} a _ {n}}}{\sum_ {i = 1} ^ {K} e ^ {x _ {i} ^ {\top} a _ {n}}} a _ {n} - s _ {n, k} a _ {n} \\ = \left(\sigma (X ^ {\top} a _ {n}) - s _ {n}\right) _ {k} a _ {n} = \left(\sigma (X ^ {\top} a _ {n}) - b _ {n}\right) _ {k} a _ {n} - \lambda \left(\sigma (\Theta^ {\top} a _ {n}) - b _ {n}\right) a _ {n} \\ = \nabla_ {x _ {k}} f _ {n} (X) - \lambda \nabla_ {\theta_ {k}} f _ {n} (\Theta). \\ \end{array}
$$
Again, we get the same expression (5) for the distillation gradient
$$
\nabla_ {X} f _ {n} (X \mid \Theta , \lambda) = \nabla_ {X} f _ {n} (X) - \lambda \nabla_ {\Theta} f _ {n} (\Theta).
$$
B.3 Generic Non-linear Classification
Finally, consider arbitrary classification model that ends with linear layer and soft-max as the last layer, i.e., the forward pass has the following steps
$$
a _ {n} \rightarrow \psi_ {n} (x) \rightarrow \phi_ {x} (a _ {n}) := \sigma (\psi_ {n} (x)) \in \mathbb {R} ^ {K},
$$
where  $a_{n}\in \mathcal{A} = \mathbb{R}^{d}$  is the input data and  $\psi_n(x)\in \mathbb{R}^K$  are the logits with respect to the model parameters  $x$ . Denote  $\varphi_{n}(z)\coloneqq \ell (\sigma (z),b_{n})$  the loss associated with logits  $z$  and the true label  $b_{n}$ . In words,  $\psi_{n}$  gives the logits from the input data, while  $\varphi_{n}$  gives the loss from given logits. Then, clearly we have the following representation for the loss function  $f_{n}(x) = \varphi_{n}(\psi_{n}(x))$ . Next, let us
simplify the distillation loss as
$$
\begin{array}{l} f _ {n} (x \mid \theta , \lambda) = (1 - \lambda) f _ {n} (x) + \lambda f _ {n} (x \mid \theta) \\ = (1 - \lambda) \ell \left(\phi_ {x} \left(a _ {n}\right), b _ {n}\right) + \lambda \ell \left(\phi_ {x} \left(a _ {n}\right), \phi_ {\theta} \left(a _ {n}\right)\right) \\ = \ell \left(\phi_ {x} \left(a _ {n}\right), (1 - \lambda) b _ {n} + \lambda \phi_ {\theta} \left(a _ {n}\right)\right) \\ = \ell (\sigma (\psi_ {n} (x), (1 - \lambda) b _ {n} + \lambda \sigma (\psi_ {n} (\theta))) \\ = \ell (\sigma (\psi_ {n} (x), s _ {n}) \\ = - \sum_ {k = 1} ^ {K} s _ {n, k} \log \sigma (\psi_ {n} (x)) _ {k} = - \sum_ {k = 1} ^ {K} s _ {n, k} \log \frac {e ^ {\psi_ {n , k} (x)}}{\sum_ {j = 1} ^ {K} e ^ {\psi_ {n , j} (x)}} \\ = \sum_ {k = 1} ^ {K} s _ {n, k} \left(\log \sum_ {j = 1} ^ {K} e ^ {\psi_ {n, j} (x)} - \psi_ {n, k} (x)\right) \\ = \sum_ {k = 1} ^ {K} s _ {n, k} \log \sum_ {j = 1} ^ {K} e ^ {\psi_ {n, j} (x)} - \sum_ {k = 1} ^ {K} s _ {n, k} \psi_ {n, k} (x) \\ = \log \sum_ {k = 1} ^ {K} e ^ {\psi_ {n, k} (x)} - \sum_ {k = 1} ^ {K} s _ {n, k} \psi_ {n, k} (x), \\ \end{array}
$$
where  $s_n = (1 - \lambda) b_n + \lambda \phi_\theta(a_n)$ . Now we need to differentiate obtained expression and derive an expression for the stochastic gradient for the above objective, namely the gradient  $\nabla_x f_n(x \mid \theta, \lambda)$  of the loss associated with  $n^{th}$  data point  $(a_n, b_n)$ . Applying the gradient operator, we get
$$
\begin{array}{l} \nabla_ {x} f _ {n} (x \mid \theta , \lambda) = \nabla_ {x} \left[ \log \sum_ {k = 1} ^ {K} e ^ {\psi_ {n, k} (x)} - \sum_ {k = 1} ^ {K} s _ {n, k} \psi_ {n, k} (x) \right] \\ = \sum_ {k = 1} ^ {K} \frac {e ^ {\psi_ {n , k} (x)}}{\sum_ {j = 1} ^ {K} e ^ {\psi_ {n , j} (x)}} \nabla_ {x} \psi_ {n, k} (x) - \sum_ {k = 1} ^ {K} s _ {n, k} \nabla_ {x} \psi_ {n, k} (x) \\ = \sum_ {k = 1} ^ {K} \left(\sigma \left(\psi_ {n} (x)\right) - s _ {n}\right) _ {k} \nabla_ {x} \psi_ {n} (x) \\ = J \psi_ {n} (x) (\sigma (\psi_ {n} (x)) - s _ {n}), \\ \end{array}
$$
$$
J \psi_ {n} (x) := \frac {\partial \psi_ {n} (x)}{\partial x} = [ \nabla \psi_ {n, 1} (x) \nabla \psi_ {n, 2} (x) \dots \nabla \psi_ {n, K} (x) ] \in \mathbb {R} ^ {d \times K}
$$
is the Jacobian of vector-valued function  $\psi_n\colon \mathbb{R}^d\to \mathbb{R}^K$ . From the derivation so far we imply that
$$
\sigma \left(\psi_ {n} (x)\right) - s _ {n} = \left(\sigma \left(\psi_ {n} (x)\right) - b _ {n}\right) - \lambda \left(\sigma \left(\psi_ {n} (\theta)\right) - b _ {n}\right) = \nabla \varphi_ {n} \left(\psi_ {n} (x)\right) - \lambda \nabla \varphi_ {n} \left(\psi_ {n} (\theta)\right).
$$
Taking into account that  $f_{n}(x) = \varphi_{n}(\psi_{n}(x))$ , we show the following form for the distilled gradient
$$
\begin{array}{l} \nabla_ {x} f _ {n} (x \mid \theta , \lambda) = J \psi_ {n} (x) (\nabla \varphi_ {n} (\psi_ {n} (x)) - \lambda \nabla \varphi_ {n} (\psi_ {n} (\theta))) \\ = \frac {\partial \psi_ {n} (x)}{\partial x} \frac {\partial f _ {n} (x)}{\partial \psi_ {n} (x)} - \lambda \frac {\partial \psi_ {n} (x)}{\partial x} \frac {\partial f _ {n} (\theta)}{\partial \psi_ {n} (\theta)}. \\ \end{array}
$$
Before we proceed to the proofs of Theorems 1 and 2, we prove a key lemma that will be useful in both proofs. The lemma we are about to present covers Part 2 (Optimal distillation weight) and Part 3 (Impact of the teacher) of the proof overview discussed in the main content.
C.1 Key lemma
To simplify the expressions in our proofs, let us introduce some notation describing stochastic gradients. Denote the signal-to-noise ratio with respect to parameters  $\theta$  by
$$
\beta (\theta) := \frac {\| \nabla f (\theta) \| ^ {2}}{\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right]} \in [ 0, 1 ], \tag {31}
$$
and correlation coefficient between stochastic gradients  $\nabla f_{\xi}(x)$  and  $\nabla f_{\xi}(y)$  by
$$
\rho (x, y) := \frac {\operatorname {C o v} (x , y)}{\operatorname {V a r} (x) \operatorname {V a r} (y)} \in [ - 1, 1 ], \tag {32}
$$
$$
\operatorname {C o v} (x, y) := \mathbb {E} \left[ \langle \nabla f _ {\xi} (x) - \nabla f (x), \nabla f _ {\xi} (y) - \nabla f (y) \rangle \right], \qquad \operatorname {V a r} (x) := \sqrt {\operatorname {C o v} (x , x)},
$$
are the covariance and variance respectively.
Lemma 1. Let  $N(\lambda) = \lambda^2\| \nabla f(\theta)\|^2 +c\gamma \mathbb{E}\left[||\nabla f_{\xi}(x^{*}) - \lambda \nabla f_{\xi}(\theta)\| ^2\right]$  for some constant  $c\geq 0$  Then, the optimal  $\lambda$  that minimizes  $N(\lambda)$  is given by
$$
\lambda_ {*} = \frac {\mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}), \nabla f _ {\xi} (\theta) \rangle \right]}{\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] + \frac {1}{c \gamma} \| \nabla f (\theta) \| ^ {2}}. \tag {33}
$$
$$
\frac {N (\lambda_ {*})}{N (0)} = 1 - \rho^ {2} (x ^ {*}, \theta) \frac {1 - \beta (\theta)}{1 + \frac {1}{c \gamma} \beta (\theta)} \leq \min  \left(1, \mathcal {O} \left(\frac {1}{\gamma} (f (\theta) - f ^ {*})\right)\right).
$$
Proof. Notice that  $N(\lambda)$  is quadratic in  $\lambda$  and using the first-order optimality condition, we conclude
$$
\frac {d}{d \lambda} N (\lambda) = 2 \lambda \| \nabla f (\theta) \| ^ {2} + c \gamma \left(- 2 \mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}), \nabla f _ {\xi} (\theta) \rangle \right] + 2 \lambda \mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right]\right) = 0,
$$
we get (33). Furthermore, plugging the expression of  $\lambda_{*}$  into  $N(\lambda)$ , we get