SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
f (y) \geq f (x + \delta) + \langle \nabla f (x + \delta), y - x - \delta \rangle + \frac {\mu}{2} \| y - x - \delta \| ^ {2}.
$$
Applying expectation and (30), we get
$$
f (y) \geq \mathbb {E} f (x + \delta) + \mathbb {E} \langle \nabla f (x + \delta), y - x \rangle - \mathbb {E} \langle \nabla f (x + \delta), \delta \rangle + \frac {\mu}{2} \| y - x \| ^ {2} + \frac {\mu}{2} \mathbb {E} \| \delta \| ^ {2}.
$$
It remains to bound the term  $-\mathbb{E}\langle \nabla f(x + \delta),\delta \rangle$  , which can be done using  $L$  -smoothness and applying expectation as follows:
$$
- \mathbb {E} \langle \nabla f (x + \delta), \delta \rangle \stackrel {{(2 5)}} {{\geq}} \mathbb {E} \left[ f (x) - f (x + \delta) - \frac {L}{2} \| \delta \| ^ {2} \right] = f (x) - \mathbb {E} f (x + \delta) - \frac {L}{2} \mathbb {E} \| \delta \| ^ {2}.
$$
Lemma 6. Denote  $A = 2\mathcal{L} + \left(\mathcal{L}^2 + \frac{1}{\gamma^2}\right)\alpha$  and  $B = \left(\mathcal{L} + \frac{1}{\gamma^2}\right)\nu$ , where  $\alpha, \nu$  are defined in Lemma 2. Then
$$
\begin{array}{l} \mathbb {E} \left\| \frac {\delta (x)}{\gamma} - \nabla f _ {\xi} (x + \delta (x)) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ \leq 4 A \left(f (x) - f \left(x _ {*}\right)\right) + 2 B + 4 \mathbb {E} \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2}, \tag {45} \\ \end{array}
$$
Proof. Using (43) with  $y = x$ , we get
$$
\mathbb {E} \left\| \frac {\delta (x)}{\gamma} - \nabla f _ {\xi} (x + \delta (x)) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2}
$$
$$
\stackrel {(4 3)} {\leq} \quad 2 \mathbb {E} \| \nabla f _ {\xi} (x) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} + 2 \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) \mathbb {E} \| \delta (x) \| ^ {2}
$$
$$
\begin{array}{l} \stackrel {(2 0) + (2 6) + (4 0)} {\leq} \quad 8 \mathcal {L} (f (x) - f (x _ {*})) + 4 \mathbb {E} \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} + 2 \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) (2 \alpha (f (x) - f (x _ {*})) + \nu) \\ = \qquad 4 \left(2 \mathcal {L} + \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) \alpha\right) (f (x) - f (x _ {*})) + 2 \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) \nu + 4 \mathbb {E} \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2}. \\ \end{array}
$$
E.2 Proof of Theorem 4
Denoting  $\delta^t = \delta (y^t)$  we have  $\mathcal{C}(y^t) = y^t +\delta^t$  . Then
$$
\begin{array}{l} \left\| y ^ {t + 1} - x ^ {*} \right\| ^ {2} \\ = \left\| \mathcal {C} \left(y ^ {t}\right) - \gamma \nabla f _ {\xi} \left(\mathcal {C} \left(y _ {t}\right)\right) + \gamma \lambda \nabla f _ {\xi} (\theta) - x ^ {*} \right\| ^ {2} \\ = \left\| y ^ {t} - x ^ {*} + \delta^ {t} - \gamma \nabla f _ {\xi} \left(y ^ {t} + \delta^ {t}\right) + \gamma \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ = \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + 2 \langle \delta^ {t} - \gamma \nabla f _ {\xi} (y ^ {t} + \delta^ {t}) + \gamma \lambda \nabla f _ {\xi} (\theta), y ^ {t} - x ^ {*} \rangle + \left\| \delta^ {t} - \gamma \nabla f _ {\xi} (y ^ {t} + \delta^ {t}) + \gamma \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2}. \\ \end{array}
$$
Taking conditional expectation  $\mathbb{E}_t\coloneqq \mathbb{E}\left[\cdot \mid y^t\right]$ , we get
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ \left| \left| y ^ {t + 1} - x ^ {*} \right| \right| ^ {2} \right] \\ = \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + 2 \gamma \langle \mathbb {E} _ {t} [ \nabla f (y ^ {t} + \delta^ {t}) ] - \lambda \nabla f (\theta), x ^ {*} - y ^ {t} \rangle + \mathbb {E} _ {t} [ \| \delta^ {t} - \gamma \nabla f _ {\xi} (y ^ {t} + \delta^ {t}) + \gamma \lambda \nabla f _ {\xi} (\theta) \| ^ {2} ] \\ \stackrel {(4 4)} {\leq} \quad \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + 2 \gamma \left[ f (x ^ {*}) - f (y ^ {t}) - \frac {\mu}{2} \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + \frac {L - \mu}{2} \mathbb {E} _ {t} \left\| \delta^ {t} \right\| ^ {2} \right] + 2 \gamma \lambda \langle \nabla f (\theta), y ^ {t} - x ^ {*} \rangle \\ + \gamma^ {2} \mathbb {E} _ {t} \left\| \frac {\delta^ {t}}{\gamma} - \nabla f _ {\xi} \left(y ^ {t} + \delta^ {t}\right) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ = \quad (1 - \gamma \mu) \| y ^ {t} - x ^ {*} \| ^ {2} - 2 \gamma (f (y ^ {t}) - f (x ^ {*})) + \gamma (L - \mu) \mathbb {E} _ {t} \| \delta^ {t} \| ^ {2} + 2 \gamma \lambda \langle \nabla f (\theta), y ^ {t} - x ^ {*} \rangle \\ + \gamma^ {2} \mathbb {E} _ {t} \left\| \frac {\delta^ {t}}{\gamma} - \nabla f _ {\xi} \left(y ^ {t} + \delta^ {t}\right) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(2 7) + (2 3)} {\leq} \quad (1 - \gamma \mu) \left\| y ^ {t} - x ^ {*} \right\| ^ {2} - \gamma (f (y ^ {t}) - f (x ^ {*})) - \frac {\gamma \mu}{2} \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + \gamma (L - \mu) \mathbb {E} _ {t} \left\| \delta^ {t} \right\| ^ {2} \\ + \frac {\gamma \mu}{2} \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + \frac {2 \gamma \lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} + \gamma^ {2} \mathbb {E} _ {t} \left\| \frac {\delta^ {t}}{\gamma} - \nabla f _ {\xi} (y ^ {t} + \delta^ {t}) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ \stackrel {(4 5)} {\leq} \quad (1 - \gamma \mu) \left\| y ^ {t} - x ^ {*} \right\| ^ {2} - \gamma (f (y ^ {t}) - f (x ^ {*})) + \gamma (L - \mu) \mathbb {E} _ {t} \left\| \delta^ {t} \right\| ^ {2} + \frac {2 \gamma \lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} \\ + 4 \gamma^ {2} A \left(f \left(y ^ {t}\right) - f \left(x ^ {*}\right)\right) + 2 \gamma^ {2} B + 4 \gamma^ {2} \mathbb {E} \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \\ = \quad (1 - \gamma \mu) \| y ^ {t} - x ^ {*} \| ^ {2} + \gamma (4 \gamma A - 1) (f (y ^ {t}) - f (x ^ {*})) + 2 \gamma^ {2} B + \gamma (L - \mu) \mathbb {E} _ {t} \| \delta^ {t} \| ^ {2} \\ + \frac {2 \gamma \lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} + 4 \gamma^ {2} \mathbb {E} \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \\ \stackrel {(4 0)} {\leq} \quad (1 - \gamma \mu) \| y ^ {t} - x ^ {*} \| ^ {2} + \gamma (4 \gamma A - 1) (f (y ^ {t}) - f (x ^ {*})) + 2 \gamma^ {2} B \\ + \gamma (L - \mu) (2 \alpha (f (x _ {k}) - f (x _ {*})) + \nu) \\ + \frac {2 \gamma \lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} + 4 \gamma^ {2} \mathbb {E} \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \\ = (1 - \gamma \mu) \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + \gamma (4 \gamma A + 2 \alpha (L - \mu) - 1) (f (y ^ {t}) - f (x ^ {*})) + 2 \gamma^ {2} B + \gamma (L - \mu) \nu \\ + \frac {2 \gamma \lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} + 4 \gamma^ {2} \mathbb {E} \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2}, \\ \end{array}
$$
where  $\alpha$  and  $\nu$  are as in Lemma 2 and  $A$  and  $B$  are defined Lemma 6. Next, we show that the bounds on  $\gamma$  and  $\omega$  lead to  $2\gamma A + \alpha (L - \mu)\leq 1 / 2$ . Plugging the expression of  $A$ , we the former inequality becomes
$$
2 \gamma \left(2 \mathcal {L} + \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) \alpha\right) + \alpha (L - \mu) \leq \frac {1}{2}.
$$
Rearranging terms, we get
$$
\frac {2 \omega}{\mu} = \alpha \leq \frac {1 / 2 - 4 \gamma \mathcal {L}}{2 \gamma \mathcal {L} ^ {2} + \frac {2}{\gamma} + L - \mu}.
$$
For  $\gamma \leq \frac{1}{16\mathcal{L}}$ , the above inequality holds if  $\omega = \mathcal{O}(\gamma \mu)$ . If  $\gamma = \frac{1}{16\mathcal{L}}$ , the condition on variance parameter  $\omega$  becomes  $\omega = \mathcal{O}(\mu/\mathcal{L})$ . Thus, by assumption on  $\omega$  and  $\gamma$ , we have  $2\gamma A + \alpha (L - \mu) \leq 1/2$ , and hence
$$
\mathbb {E} _ {t} \left[ \left\| y ^ {t + 1} - x ^ {*} \right\| ^ {2} \right] \leq (1 - \gamma \mu) \left\| y ^ {t} - x ^ {*} \right\| ^ {2} + D,
$$
$$
\begin{array}{l} { D } { = } { 2 \gamma ^ { 2 } B + \gamma ( L - \mu ) \nu + \frac { 2 \gamma \lambda ^ { 2 } } { \mu } \| \nabla f ( \theta ) \| ^ { 2 } + 4 \gamma ^ { 2 } \mathbb { E } \| \nabla f _ { \xi } ( x ^ { * } ) - \lambda \nabla f _ { \xi } ( \theta ) \| ^ { 2 } } \\ = 2 \gamma^ {2} B + \gamma (L - \mu) \nu + \frac {2 \gamma}{\mu} N _ {3} (\lambda), \\ \end{array}
$$
Figure 4: Validation loss statistics for the same setup as in Figure 2a.
with  $N_{3}(\lambda) = \lambda^{2}\| \nabla f(\theta)\|^{2} + 2\gamma \mu \mathbb{E}\| \nabla f_{\xi}(x^{*}) - \lambda \nabla f_{\xi}(\theta)\|^{2}$ . Applying Lemma 1 with  $c = 2\mu$  we get  $N_{3}(\lambda^{*}) = N_{3}(0)\min \left(1,\frac{1}{\gamma}\mathcal{O}(f(\theta) - f^{*})\right) = 2\mu \sigma_{*}^{2}\min \left(\gamma ,\mathcal{O}(f(\theta) - f^{*})\right)$ . Taking expectation, unrolling the recurrence, and applying the tower property, we get
$$
\mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] \leq (1 - \gamma \mu) ^ {t} \left\| y ^ {0} - x ^ {*} \right\| ^ {2} + \frac {D}{\gamma \mu}
$$
where the neighborhood is given by
$$
\begin{array}{l} \frac {D}{\gamma \mu} = \frac {1}{\gamma \mu} (2 \gamma^ {2} B + \gamma (L - \mu) \nu + \frac {2 \gamma}{\mu} N (\lambda^ {*})) = \frac {1}{\gamma \mu} (2 \gamma^ {2} (\mathcal {L} ^ {2} + 1 / \gamma^ {2}) \nu + \gamma (L - \mu) \nu + \frac {2 \gamma}{\mu} N (\lambda^ {*})) \\ { = } { \frac { \nu } { \mu } ( 2 \gamma \mathcal { L } ^ { 2 } + ^ { 2 } / \gamma + L - \mu ) + \frac { 2 } { \mu ^ { 2 } } N ( \lambda ^ { * } ) = \frac { \nu } { \mu } ( 2 \gamma \mathcal { L } ^ { 2 } + ^ { 2 } / \gamma + L - \mu ) + \frac { 4 \sigma _ { * } ^ { 2 } } { \mu } \min ( \gamma , \mathcal { O } ( f ( \theta ) - f ^ { * } ) ) . } \\ \end{array}
$$
Ignoring the absolute constants and using bounds  $\gamma \leq \frac{1}{16\mathcal{L}}$  and  $\omega = \mathcal{O}(\gamma \mu)$ , we have
$$
\mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] \leq (1 - \gamma \mu) ^ {t} \left\| x ^ {0} - x ^ {*} \right\| ^ {2} + \mathcal {O} \left(\frac {\omega}{\gamma \mu} \| x ^ {*} \| ^ {2}\right) + \frac {4 \sigma_ {*} ^ {2}}{\mu} \min  (\gamma , \mathcal {O} (f (\theta) - f ^ {*})).
$$
Applying this inequality in (38) we conclude the proof.2
F Additional Experimental Validation
In this section we provide additional experimental validation for our theory.
F.1 The impact of the learning hyperparameters
We begin by measuring the impact that the optimization parameters, in particular the step size, have on the convergence of SGD and KD. For this, we perform experiments on linear models trained on the MNIST dataset, without momentum and regularization, using a mini-batch size of 10, for a total of 100 epochs. We compute the cross entropy train loss between the student and true labels, measured as a running average over all iterations within an epoch, similar to Figure 2. We also compute the minimum train loss, as well as the average and standard deviation across the last 10 epochs. The results in Figure 5a show the impact that different learning rates have on the overall training dynamics of self-distillation. In all cases, the teacher was trained using the same hyperparameters as the self-distilled models ( $\lambda = 0$  in the plot). We can see that using a higher learning rate introduces more variance in the SGD update, and KD would have a more pronounced variance reduction effect. In all cases, however, we can find an optimum  $\lambda$  achieving a lower train loss compared to SGD.
2The rate in Theorem 4 uses  $\gamma = \frac{1}{16\mathcal{L}}$  value for the learning rate.
Figure 5: Ablation study on the training loss of self-distillation and SGD, when taking into account the training hyperparameters (learning rate) and quality of the teacher.
Figure 6: The impact of self-distillation on the train loss of compressed models. Here a random mask is computed at initialization, and kept fixed throughout training.
F.2 The impact of the teacher's quality
Next, we quantify the impact that a better trained teacher could have on self-distillation. Using the same setup of convex MNIST as above, we perform self-distillation using a better teacher, i.e. one that achieves  $93.7\%$  train accuracy and  $92.5\%$  test accuracy. (In comparison, the teacher trained using a step size of 0.05 achieved  $92\%$  train accuracy and  $90.9\%$  test accuracy) We can see in Figure 5b that this better-trained teacher has a more substantial impact on the models which inherently have higher variance, i.e. those trained with a higher learning rate; in this case, the optimal value of  $\lambda$  is closer to 1, which is also suggested by the theory (see Equation 14). We note that a similar behavior was also observed on the CIFAR-10 features linear classification task presented in Figure 2b.
F.3 The impact of knowledge distillation on compression
Now we turn our attention towards validating the theoretical results developed in the context of knowledge distillation for compressed models, presented in Section 6. We consider again the convex MNIST setting, as described in the previous section, and we perform self-distillation from the better trained teacher. We prune the weights at initialization, using a random mask at a chosen sparsity level, and we apply this fixed mask after each parameter update. The results presented in Figure 6 show that self-distillation can indeed reduce the train loss, compared to SGD, even for compressed updates. Moreover, we observe that with increased sparsity the impact of self-distillation is less pronounced, as also suggested by the theory.
Lastly, following our discussion from Section 7, we discuss some limitations of our work.
- As a theoretical paper, we used several assumptions to make our claims rigorous. However, one can always question each assumption and extend the theory under certain relaxations.
Our theoretical claims are based on strong (quasi) convexity or Polyak-Łojasiewicz condition, which are standard assumptions in the optimization literature.
- Another limitation concerning the "distillation for compression" part of our theory is the unbiasedness condition  $\mathbb{E}\left[\mathcal{C}(x)\right] = x$  in Assumption 4. Ideally, we would utilize any "biased" compression operator, such as TopK, with similar convergence properties. However, it is known that biased estimators (e.g., biased compression operators or biased stochastic gradients) are harder to analyze in theory.