SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} N (\lambda_ {*}) = c \gamma \left(\frac {\lambda_ {*} ^ {2}}{c \gamma} \| \nabla f (\theta) \| ^ {2} + \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] - 2 \lambda_ {*} \mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}), \nabla f _ {\xi} (\theta) \rangle \right] + \lambda_ {*} ^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right]\right) \\ = c \gamma \left(\mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) \| ^ {2} \right] - 2 \lambda_ {*} \mathbb {E} \left[ \langle \nabla f _ {\xi} \left(x ^ {*}\right), \nabla f _ {\xi} (\theta) \rangle \right] + \lambda_ {*} ^ {2} \left(\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] + \frac {1}{c \gamma} \| \nabla f (\theta) \| ^ {2}\right)\right) \\ = c \gamma \left(\mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) \| ^ {2} \right] - \frac {\left(\mathbb {E} \left[ \langle \nabla f _ {\xi} \left(x ^ {*}\right) , \nabla f _ {\xi} (\theta) \rangle \right]\right) ^ {2}}{\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] + \frac {1}{c \gamma} \| \nabla f (\theta) \| ^ {2}}\right) \\ = c \gamma \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] \left(1 - \frac {(\mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}) , \nabla f _ {\xi} (\theta) \rangle \right]) ^ {2}}{\left(\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] + \frac {1}{c \gamma} \| \nabla f (\theta) \| ^ {2}\right) (\mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right])}\right). \\ \end{array}
$$
Note that  $N(0) = c\gamma \sigma_{*}^{2}$ . From  $\mathbb{E}\left[\nabla f_{\xi}(x^{*})\right] = \nabla f(x^{*}) = 0$ , we imply that
$$
\operatorname {C o v} \left(x ^ {*}, \theta\right) = \mathbb {E} \left[ \left\langle \nabla f _ {\xi} \left(x ^ {*}\right), \nabla f _ {\xi} (\theta) - \nabla f (\theta) \right\rangle \right] = \mathbb {E} \left[ \left\langle \nabla f _ {\xi} \left(x ^ {*}\right), \nabla f _ {\xi} (\theta) \right\rangle \right],
$$
$$
\rho (x ^ {*}, \theta) = \frac {\mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}) , \nabla f _ {\xi} (\theta) \rangle \right]}{\sqrt {\mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right]} \sqrt {\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] - \| \nabla f (\theta) \| ^ {2}}}.
$$
Now we can simplify the expression for  $N(\lambda_{*})$  as follows
$$
\begin{array}{l} \frac {N (\lambda_ {*})}{N (0)} = 1 - \frac {\left(\mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}) , \nabla f _ {\xi} (\theta) \rangle \right]\right) ^ {2}}{\left(\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] + \frac {1}{c \gamma} \| \nabla f (\theta) \| ^ {2}\right) \left(\mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right]\right)} \\ { = } { 1 - \frac { ( \mathbb { E } [ \langle \nabla f _ { \xi } ( x ^ { * } ) , \nabla f _ { \xi } ( \theta ) \rangle ] ) ^ { 2 } } { ( \mathbb { E } [ \| \nabla f _ { \xi } ( \theta ) \| ^ { 2 } ] - \| \nabla f ( \theta ) \| ^ { 2 } ) ( \mathbb { E } [ \| \nabla f _ { \xi } ( x ^ { * } ) \| ^ { 2 } ] ) } \frac { \mathbb { E } [ \| \nabla f _ { \xi } ( \theta ) \| ^ { 2 } ] - \| \nabla f ( \theta ) \| ^ { 2 } } { \mathbb { E } [ \| \nabla f _ { \xi } ( \theta ) \| ^ { 2 } ] + \frac { 1 } { c \gamma } \| \nabla f ( \theta ) \| ^ { 2 } } } \\ = 1 - \rho^ {2} (x ^ {*}, \theta) \frac {1 - \beta (\theta)}{1 + \frac {1}{c \gamma} \beta (\theta)}. \\ \end{array}
$$
Here,  $\rho(x^{*},\theta)$  is the correlation coefficient between stochastic gradients at  $x^{*}$  and  $\theta$ . Hence, we showed with tuned distillation weight the neighborhood can shrink by some factor depending on the teacher's parameters. In the extreme case when the teacher  $\theta = x^{*}$  is optimal, we have  $\rho(x^{*},\theta) = 1$ ,  $\beta(\theta) = 0$  and, thus, no neighborhood  $N(\lambda_{*}) = 0$ . This hints us on the fact that the reduction factor  $N(\lambda_{*}) / N(0)$  of the neighborhood is controlled by the "quality" of the teacher.
To make this argument rigorous, consider the teacher's model to be away from the optimal solution  $x^{*}$  within the limit described by the following inequality
$$
f (\theta) - f ^ {*} \leq \frac {\sigma \left(x ^ {*}\right) \sigma (\theta)}{\mathcal {L}}, \tag {34}
$$
where  $\sigma^2 (x)\coloneqq \mathbb{E}\left[\| \nabla f_\xi (x)\| ^2\right]$  is the second moment of the stochastic gradients. Without loss of generality, we assume that  $\sigma^2 (x) > 0$  for all parameter choices  $x\in \mathbb{R}^d$ : otherwise we have  $\sigma_{*}^{2} = 0$  and even plain SGD ensures full variance reduction. Then, we can simplify the reduction factor as
$$
\begin{array}{l} 1 - \rho^ {2} (x ^ {*}, \theta) \frac {1 - \beta (\theta)}{1 + \frac {1}{c \gamma} \beta (\theta)} = 1 - \frac {\left(\mathbb {E} \left[ \langle \nabla f _ {\xi} (x ^ {*}) , \nabla f _ {\xi} (\theta) \rangle \right]\right) ^ {2}}{\mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) \| ^ {2} \right] \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right]} \frac {1}{1 + \frac {1}{c \gamma} \beta (\theta)} \\ = \quad 1 - \left(\frac {\sigma^ {2} (\theta) + \sigma^ {2} (x ^ {*}) - \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \nabla f _ {\xi} (\theta) \| ^ {2} \right]}{2 \sigma (\theta) \sigma (x ^ {*})}\right) ^ {2} \frac {1}{1 + \frac {1}{c \gamma} \beta (\theta)} \\ = \quad 1 - \left(\frac {\sigma^ {2} (\theta) + \sigma^ {2} \left(x ^ {*}\right)}{2 \sigma (\theta) \sigma \left(x ^ {*}\right)} - \frac {\mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) - \nabla f _ {\xi} (\theta) \| ^ {2} \right]}{2 \sigma (\theta) \sigma \left(x ^ {*}\right)}\right) ^ {2} \frac {1}{1 + \frac {1}{c \gamma} \beta (\theta)} \\ \stackrel {(1 0) + (3 4)} {\leq} \quad 1 - \left(1 - \frac {\mathcal {L} (f (\theta) - f ^ {*})}{\sigma (\theta) \sigma (x ^ {*})}\right) ^ {2} \frac {1}{1 + \frac {2 L}{c \gamma} \frac {f (\theta) - f ^ {*}}{\sigma^ {2} (\theta)}} \\ \leq \quad \left(2 + \frac {2 L}{c \gamma \mathcal {L}} \frac {\sigma \left(x ^ {*}\right)}{\sigma (\theta)}\right) \frac {\mathcal {L} (f (\theta) - f ^ {*})}{\sigma (\theta) \sigma \left(x ^ {*}\right)} \\ = \left(\frac {2 \mathcal {L}}{\sigma \left(x ^ {*}\right) \sigma (\theta)} + \frac {1}{c \gamma} \frac {2 L}{\sigma^ {2} (\theta)}\right) (f (\theta) - f ^ {*}) = \mathcal {O} \left(\frac {1}{\gamma} (f (\theta) - f ^ {*})\right), \\ \end{array}
$$
where the last inequality used  $1 - \frac{(1 - u)^2}{1 + uv} \leq (2 + v)u$  for all  $u, v \geq 0$ .
Denote by  $\mathbb{E}_t[\cdot ]\coloneqq \mathbb{E}[\cdot \mid x^t ]$  the conditional expectation with respect to  $x^{t}$ . Then, we start bounding the error using the update rule (7).
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ \left\| x ^ {t + 1} - x ^ {*} \right\| ^ {2} \right] \\ = \| x ^ {t} - x ^ {*} \| ^ {2} - 2 \gamma \left\langle x ^ {t} - x ^ {*}, \nabla f (x ^ {t}) - \lambda \nabla f (\theta) \right\rangle + \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ = \| x ^ {t} - x ^ {*} \| ^ {2} - 2 \gamma \left\langle x ^ {t} - x ^ {*}, \nabla f (x ^ {t}) \right\rangle + 2 \gamma \lambda \left\langle x ^ {t} - x ^ {*}, \nabla f (\theta) \right\rangle + \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(1) + (2 0)} {\leq} \quad \left(1 - \gamma \mu\right) \| x ^ {t} - x ^ {*} \| ^ {2} - 2 \gamma \left(f \left(x ^ {t}\right) - f \left(x ^ {*}\right)\right) + 2 \gamma \lambda \left\langle x ^ {t} - x ^ {*}, \nabla f (\theta) \right\rangle \\ + 2 \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(2 8) + (2 3)} {\leq} \quad (1 - \gamma \mu) \| x ^ {t} - x ^ {*} \| ^ {2} - \gamma (f (x ^ {t}) - f (x ^ {*})) - \frac {\gamma \mu}{8} \| x ^ {t} - x ^ {*} \| ^ {2} + \frac {\gamma \mu}{8} \left\| x ^ {t} - x ^ {*} \right\| ^ {2} + \frac {8 \gamma}{\mu} \lambda^ {2} \| \nabla f (\theta) \| ^ {2} \\ + 2 \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(3)} {\leq} \quad (1 - \gamma \mu) \| x ^ {t} - x ^ {*} \| ^ {2} - \gamma (f (x ^ {t}) - f (x ^ {*})) + \frac {8 \gamma}{\mu} \lambda^ {2} \| \nabla f (\theta) \| ^ {2} \\ + 4 \gamma^ {2} \mathcal {L} \left(f \left(x ^ {t}\right) - f \left(x ^ {*}\right)\right) + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} = \quad (1 - \gamma \mu) \| x ^ {t} - x ^ {*} \| ^ {2} - \gamma (1 - 4 \gamma \mathcal {L}) (f (x ^ {t}) - f (x ^ {*})) \\ + \frac {8 \gamma}{\mu} \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \end{array}
$$
$$
\leq (1 - \gamma \mu) \| x ^ {t} - x ^ {*} \| ^ {2} + \frac {8 \gamma}{\mu} \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right],
$$
where we used Peter-Paul inequality (23) with parameter  $s = \frac{8}{\mu}$  and the step-size bound  $\gamma \leq \frac{1}{4\mathcal{L}}$  in the last inequality. Applying full expectation and unrolling the recursion, we get
$$
\begin{array}{l} \mathbb {E} \left[ \| x ^ {t} - x ^ {*} \| ^ {2} \right] \leq (1 - \gamma \mu) \mathbb {E} \left[ \| x ^ {t - 1} - x ^ {*} \| ^ {2} \right] + \frac {8 \gamma}{\mu} \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + 2 \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \stackrel {(2 9)} {\leq} \quad (1 - \gamma \mu) ^ {t} \| x ^ {0} - x ^ {*} \| ^ {2} + \frac {8 \lambda^ {2}}{\mu^ {2}} \| \nabla f (\theta) \| ^ {2} + \frac {2 \gamma}{\mu} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ = (1 - \gamma \mu) ^ {t} \| x ^ {0} - x ^ {*} \| ^ {2} + \frac {8}{\mu^ {2}} N _ {1} (\lambda), \tag {35} \\ \end{array}
$$
where  $N_{1}(\lambda)\coloneqq \lambda^{2}\| \nabla f(\theta)\|^{2} + \frac{\gamma\mu}{4}\mathbb{E}\left[\| \nabla f_{\xi}(x^{*}) - \lambda \nabla f_{\xi}(\theta)\|^{2}\right]$ . Applying Lemma 1 with  $c = \mu /4$  we imply that for some  $\lambda = \lambda_{*}$  the neighborhood size is
$$
N _ {1} (\lambda_ {*}) \stackrel {\mathrm {L e m m a 1}} {\leq} N _ {1} (0) \cdot \min \left(1, \mathcal {O} \left(\frac {1}{\gamma} (f (\theta) - f ^ {*})\right)\right) = \frac {\mu \sigma_ {*} ^ {2}}{4} \cdot \min \left(\gamma , \mathcal {O} (f (\theta) - f ^ {*})\right).
$$
Plugging the above bound of  $N_{1}$  into (35) completes the proof.
We start the recursion from the  $L$ -smoothness condition of  $f$ . As before,  $\mathbb{E}_t$  denotes conditional expectation with respect to  $x^t$ .
$$
\begin{array}{l} \mathbb {E} _ {t} \left[ f (x ^ {t + 1}) - f ^ {*} \right] \\ \stackrel {(2 5)} {\leq} \quad \left(f (x ^ {t}) - f ^ {*}\right) - \gamma \left\langle \nabla f (x ^ {t}), \nabla f (x ^ {t}) - \lambda \nabla f (\theta) \right\rangle + \frac {L \gamma^ {2}}{2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \stackrel {(2 0)} {\leq} \quad \left(f (x ^ {t}) - f ^ {*}\right) - \gamma \| \nabla f (x ^ {t}) \| ^ {2} + \gamma \lambda \left\langle \nabla f (x ^ {t}), \nabla f (\theta) \right\rangle \\ + L \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + L \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \stackrel {(9) + (2 3)} {\leq} \left(1 - \gamma \mu\right) \left(f \left(x ^ {t}\right) - f ^ {*}\right) - \frac {\gamma}{4} \| \nabla f \left(x ^ {t}\right) \| ^ {2} - \frac {\gamma \mu}{2} \left(f \left(x ^ {t}\right) - f ^ {*}\right) + \frac {\gamma}{4} \| \nabla f \left(x ^ {t}\right) \| ^ {2} + \gamma \lambda^ {2} \| \nabla f (\theta) \| ^ {2} \\ + L \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + L \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \stackrel {(1 0)} {\leq} \quad (1 - \gamma \mu) \left(f (x ^ {t}) - f ^ {*}\right) - \frac {\gamma \mu}{2} \left(f (x ^ {t}) - f ^ {*}\right) + \gamma \lambda^ {2} \| \nabla f (\theta) \| ^ {2} \\ + 2 L \mathcal {L} \gamma^ {2} (f (x ^ {t}) - f ^ {*}) + L \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ = \quad (1 - \gamma \mu) \left(f \left(x ^ {t}\right) - f ^ {*}\right) - \frac {\gamma \mu}{2} \left(1 - \gamma \frac {4 L \mathcal {L}}{\mu}\right) \left(f \left(x ^ {t}\right) - f ^ {*}\right) \\ + \gamma \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + L \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \leq \quad (1 - \gamma \mu) \left(f \left(x ^ {t}\right) - f ^ {*}\right) + \gamma \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + L \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} \left(x ^ {*}\right) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right], \\ \end{array}
$$
where in the last inequality we used step-size bound  $\gamma \leq \frac{1}{4\mathcal{L}}\frac{\mu}{L}$ . Applying full expectation and unrolling the recursion, we get
$$
\begin{array}{l} \mathbb {E} \left[ f (x ^ {t}) - f ^ {*} \right] \leq (1 - \gamma \mu) \mathbb {E} \left[ f (x ^ {t - 1}) - f ^ {*} \| ^ {2} \right] + \gamma \lambda^ {2} \| \nabla f (\theta) \| ^ {2} + L \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ \stackrel {(2 9)} {\leq} \quad (1 - \gamma \mu) ^ {t} \left(f (x ^ {0}) - f ^ {*}\right) + \frac {\lambda^ {2}}{\mu} \| \nabla f (\theta) \| ^ {2} + \frac {L \gamma}{\mu} \mathbb {E} \left[ \| \nabla f _ {\xi} (x ^ {*}) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} \right] \\ = (1 - \gamma \mu) ^ {t} (f \left(x ^ {0}\right) - f ^ {*}) + \frac {1}{\mu} N _ {2} (\lambda), \tag {36} \\ \end{array}
$$
where  $N_{2}(\lambda)\coloneqq \lambda^{2}\| \nabla f(\theta)\|^{2} + L\gamma \mathbb{E}\left[\| \nabla f_{\xi}(x^{*}) - \lambda \nabla f_{\xi}(\theta)\|^{2}\right]$ . Similar to the previous case, we applying Lemma 1 with  $c = L$  and conclude that for some  $\lambda = \lambda_{*}$  the neighborhood size is
$$
N _ {2} \left(\lambda_ {*}\right) \stackrel {\text {L e m m a 1}} {\leq} N _ {2} (0) \cdot \min  \left(1, \mathcal {O} \left(\frac {1}{\gamma} (f (\theta) - f ^ {*})\right)\right) = L \sigma_ {*} ^ {2} \cdot \min  \left(\gamma , \mathcal {O} (f (\theta) - f ^ {*})\right).
$$
Plugging the above bound of  $N_{2}$  into (36) completes the proof.
Again, we start the recursion from the smoothness condition of  $f$ .
$$
\mathbb {E} _ {t} \left[ f (x ^ {t + 1}) - f ^ {*} \right]
$$
$$
\begin{array}{l} \stackrel {(2 5)} {\leq} \quad \left(f (x ^ {t}) - f ^ {*}\right) - \gamma \left\langle \nabla f (x ^ {t}), \nabla f (x ^ {t}) \right\rangle + \frac {L \gamma^ {2}}{2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (\theta) + \nabla f (\theta) \| ^ {2} \right] \\ \stackrel {(2 1)} {\leq} \quad \left(f (x ^ {t}) - f ^ {*}\right) - \gamma \| \nabla f (x ^ {t}) \| ^ {2} \\ + \frac {3}{2} L \gamma^ {2} \mathbb {E} _ {t} \left[ \| \nabla f _ {\xi} (x ^ {t}) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + \frac {3}{2} L \gamma^ {2} \mathbb {E} \left[ \| \nabla f _ {\xi} (\theta) - \nabla f _ {\xi} (x ^ {*}) \| ^ {2} \right] + \frac {3}{2} L \gamma^ {2} \| \nabla f (\theta) \| ^ {2} \\ \stackrel {(9) + (1 0)} {\leq} \left(1 - \gamma \mu\right) \left(f \left(x ^ {t}\right) - f ^ {*}\right) - \gamma \mu \left(f \left(x ^ {t}\right) - f ^ {*}\right) \\ + 3 L \mathcal {L} \gamma^ {2} (f (x ^ {t}) - f ^ {*}) + 3 L \mathcal {L} \gamma^ {2} (f (\theta) - f ^ {*}) + 3 L ^ {2} \gamma^ {2} (f (\theta) - f ^ {*}) \\ \leq \quad (1 - \gamma \mu) \left(f \left(x ^ {t}\right) - f ^ {*}\right) + 3 L (L + \mathcal {L}) \gamma^ {2} \left(f (\theta) - f ^ {*}\right), \\ \end{array}
$$
where we used step-size bound  $\gamma \leq \frac{\mu}{3L\mathcal{L}}$  in the last step. Therefore,
$$
\begin{array}{l} \mathbb {E} \left[ f (x ^ {t}) - f ^ {*} \right] \leq (1 - \gamma \mu) \mathbb {E} \left[ f (x ^ {t}) - f ^ {*} \right] + 3 L (L + \mathcal {L}) \gamma^ {2} (f (\theta) - f ^ {*}) \\ \stackrel {(2 9)} {\leq} \left(1 - \gamma \mu\right) ^ {t} (f (x ^ {0}) - f ^ {*}) + \frac {3 L (L + \mathcal {L})}{\mu} \cdot \gamma \left(f (\theta) - f ^ {*}\right), \\ \end{array}
$$
E Proofs for Section 6
First, we break the update rule (18) into two parts by introducing an auxiliary model parameters  $y^{t} \in \mathbb{R}^{d}$ :
$$
\begin{array}{l} {y ^ {t + 1}} = {x ^ {t} - \gamma (\nabla f _ {\xi} (x ^ {t}) - \lambda \nabla f _ {\xi} (\theta)),} \\ x ^ {t + 1} = \mathcal {C} (y ^ {t + 1}). \\ \end{array}
$$
Without loss of generality, we assume that initialization satisfies  $x^0 = y^0 = \mathcal{C}(y^0)$ . Then, for all  $t \geq 0$  we have  $x^{t} = \mathcal{C}(y^{t})$  and the update rule of  $y^{t + 1}$  can be written recursively without  $x^{t}$  via
$$
y ^ {t + 1} = \mathcal {C} \left(y ^ {t}\right) - \gamma \left(\nabla f _ {\xi} \left(\mathcal {C} \left(y ^ {t}\right)\right) - \lambda \nabla f _ {\xi} (\theta)\right). \tag {37}
$$
Using unbiasedness of the compression operator  $\mathcal{C}$ , we decompose the error
$$
\begin{array}{l} \mathbb {E} \left[ \left\| x ^ {t} - x ^ {*} \right\| ^ {2} \right] \stackrel {(3 0)} {=} \mathbb {E} \left[ \left\| x ^ {t} - y ^ {t} \right\| ^ {2} \right] + \mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] \\ = \mathbb {E} \left[ \left\| \mathcal {C} (y ^ {t}) - y ^ {t} \right\| ^ {2} \right] + \mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(1 9)} {\leq} \omega \mathbb {E} \left[ \left\| y ^ {t} \right\| ^ {2} \right] + \mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] \\ \stackrel {(2 0)} {\leq} \quad (2 \omega + 1) \mathbb {E} \left[ \left\| y ^ {t} - x ^ {*} \right\| ^ {2} \right] + 2 \omega \| x ^ {*} \| ^ {2}. \tag {38} \\ \end{array}
$$
Thus, our goal would be to analyze iterates (37) and derive the rate for  $x^t$ . In fact, the special case of (37) was analyzed by [18] in the non-stochastic case and without distillation ( $\lambda = 0$ ). The analysis we provide here for (37) is based on [18], and from this perspective, our analysis can be seen as an extension of their analysis.
To avoid another notation for the expected smoothness constant (see Assumption 3), analogous to (24) we assume that  $\mathcal{L}$  also satisfies the following smoothness inequality:
$$
\mathbb {E} \left[ \| \nabla f _ {\xi} (x) - \nabla f _ {\xi} (y) \| ^ {2} \right] \leq \mathcal {L} ^ {2} \| x - y \| ^ {2}. \tag {39}
$$
E.1 Five Lemmas
First, we need to upper bound the compression error by function suboptimality. Denote  $\delta(x) \coloneqq \mathcal{C}(x) - x$ . Let  $\mathbb{E}_{\delta}$  and  $\mathbb{E}_{\xi}$  be the expectations with respect to the compression operator  $\mathcal{C}$  and sampling  $\xi$  respectively.
Lemma 2 (Lemma 1 in [18]). Let  $\alpha = \frac{2\omega}{\mu}$  and  $\nu = 2\omega \| x^{*}\|^{2}$ . For all  $x\in \mathbb{R}^d$  we have
$$
\mathbb {E} _ {\delta} \| \delta (x) \| ^ {2} \leq 2 \alpha \left(f (x) - f \left(x ^ {*}\right)\right) + \nu , \tag {40}
$$
Proof. From (20) we imply  $\| x\|^2 \leq 2\| x - x^*\|^2 + 2\| x_*\|^2$ , and from  $\mu$ -strong convexity condition (27) we get  $\| x - x^*\|^2 \leq \frac{2}{\mu} (f(x) - f(x^*))$ . Putting these inequalities together, we arrive at
$$
\mathbb {E} _ {\delta} \| \delta (x) \| ^ {2} \leq \omega \| x \| ^ {2} \leq 2 \omega \| x - x ^ {*} \| ^ {2} + 2 \omega \| x ^ {*} \| ^ {2} \leq \frac {4 \omega}{\mu} (f (x) - f (x ^ {*}) + 2 \omega \| x ^ {*} \| ^ {2}.
$$
Lemma 3. For all  $x,y\in \mathbb{R}^d$  we have
$$
\mathbb {E} \| \nabla f _ {\xi} (x + \delta (x)) - \nabla f _ {\xi} (y) \| ^ {2} \leq \mathcal {L} ^ {2} \left(\| x - y \| ^ {2} + \mathbb {E} \| \delta (x) \| ^ {2}\right), \tag {41}
$$
$$
f (x) \leq \mathbb {E} f (x + \delta (x)) \leq f (y) + \langle \nabla f (y), x - y \rangle + \frac {L}{2} \| x - y \| ^ {2} + \frac {L}{2} \mathbb {E} \| \delta (x) \| ^ {2}. \tag {42}
$$
Proof. Fix  $x$  and let  $\delta = \delta(x)$ . Inequality (41) follows from Lipschitz continuity of the gradient, applying expectation and using (19):
$$
\mathbb {E} \| \nabla f _ {\xi} (x + \delta) - \nabla f _ {\xi} (y) \| ^ {2} \overset {(3 9)} {\leq} \mathcal {L} ^ {2} \mathbb {E} _ {\delta} \| x + \delta - y \| ^ {2} \overset {(1 9) + (3 0)} {=} \mathcal {L} ^ {2} \left(\| x - y \| ^ {2} + \mathbb {E} _ {\delta} \| \delta \| ^ {2}\right).
$$
The first inequality in (42) follows by applying Jensen's inequality (22) and using (19). Since  $f$  is  $L$ -smooth, we have
$$
\begin{array}{l} \mathbb {E} f (x + \delta) \stackrel {(2 5)} {\leq} \mathbb {E} f (y) + \langle \nabla f (y), x + \delta - y \rangle + \frac {L}{2} \| x + \delta - y \| ^ {2} \\ { } ^ { ( 1 9 ) \pm ( 3 0 ) } f ( y ) + \langle \nabla f ( y ) , x - y \rangle + \frac { L } { 2 } \| x - y \| ^ { 2 } + \frac { L } { 2 } \mathbb { E } \| \delta \| ^ { 2 } . \\ \end{array}
$$
Lemma 4. For all  $x, y \in \mathbb{R}^d$  it holds
$$
\begin{array}{l} \mathbb {E} _ {\delta} \left\| \frac {\delta (x)}{\gamma} - \nabla f _ {\xi} (x + \delta (x)) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ \leq 2 \left\| \nabla f _ {\xi} (y) - \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} + 2 \mathcal {L} ^ {2} \| x - y \| ^ {2} + 2 \left(\mathcal {L} ^ {2} + \frac {1}{\gamma^ {2}}\right) \mathbb {E} _ {\delta} \| \delta (x) \| ^ {2}. \tag {43} \\ \end{array}
$$
Proof. Fix  $x$ , and let  $\delta = \delta(x)$ . Then for every  $y \in \mathbb{R}^d$  we can write
$$
\begin{array}{l} \mathbb {E} _ {\delta} \left\| \frac {\delta}{\gamma} - \nabla f _ {\xi} (x + \delta) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} \\ { = } { \mathbb { E } _ { \delta } \left\| \frac { \delta } { \gamma } \pm \nabla f _ { \xi } ( y ) - \nabla f _ { \xi } ( x + \delta ) + \lambda \nabla f _ { \xi } ( \theta ) \right\| ^ { 2 } } \\ \stackrel {(2 0)} {\leq} \quad 2 \mathbb {E} _ {\delta} \left\| \frac {\delta}{\gamma} - \nabla f _ {\xi} (y) + \lambda \nabla f _ {\xi} (\theta) \right\| ^ {2} + 2 \mathbb {E} _ {\delta} \| \nabla f _ {\xi} (y) - \nabla f _ {\xi} (x + \delta) \| ^ {2} \\ \stackrel {(4 1)} {\leq} \frac {2}{\gamma^ {2}} \mathbb {E} _ {\delta} \| \delta \| ^ {2} - \frac {2}{\gamma} \mathbb {E} _ {\delta} \langle \delta , \nabla f _ {\xi} (y) - \lambda \nabla f _ {\xi} (\theta) \rangle + \| \nabla f _ {\xi} (y) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} + 2 \mathcal {L} ^ {2} \left(\| x - y \| ^ {2} + \mathbb {E} _ {\delta} \| \delta \| ^ {2}\right) \\ \stackrel {(1 9)} {=} \frac {2}{\gamma^ {2}} \mathbb {E} _ {\delta} \| \delta \| ^ {2} + 2 \| \nabla f _ {\xi} (y) - \lambda \nabla f _ {\xi} (\theta) \| ^ {2} + 2 \mathcal {L} ^ {2} \left(\| x - y \| ^ {2} + \mathbb {E} _ {\delta} \| \delta \| ^ {2}\right). \\ \end{array}
$$
The next lemma generalizes the strong convexity inequality (27) (special case of  $\delta(x) \equiv 0$ ).
Lemma 5. If  $f$  is  $L$ -smooth and  $\mu$ -strongly convex, then for all  $x, y \in \mathbb{R}^d$ , it holds
$$
f (y) \geq f (x) + \left\langle \mathbb {E} \left[ \nabla f (x + \delta) \right], y - x \right\rangle + \frac {\mu}{2} \| y - x \| ^ {2} - \frac {L - \mu}{2} \mathbb {E} \| \delta (x) \| ^ {2}. \tag {44}
$$
Proof. Fix  $x$  and let  $\delta = \delta(x)$ . Using (27) with  $x \gets x + \delta$ , we get