SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To bound lazy training distance of training linearized network, we would need the following auxiliary Lemma about high probability upper bound for the final layer output of linearized network at initialization.
Lemma D.1. Fix any data record  $\mathbf{x}$ , then with high probability  $1 - \delta$  over random initialization Eq. (5) of model weight matrices  $W^{1},\dots ,W^{L}$  for layer  $1,\dots ,L$ , i.e.,  $W_{l}\sim \mathcal{N}(0,\beta_{l}\mathbb{I})$ , it satisfies that
$$
\left\| \boldsymbol {f} _ {\boldsymbol {W}} (\boldsymbol {x}) \right\| _ {2} \leq \left\| \boldsymbol {x} \right\| _ {2} ^ {2} \tilde {\mathcal {O}} \left(\beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)\right) \tag {52}
$$
where  $\tilde{\mathcal{O}}$  ignores logarithmic terms with regard to width  $m$ , depth  $L$  and tail probability  $\delta$ .
Proof. To bound the term  $\| f_{\mathbf{W}_i^{lin}}(\mathbf{x}_i) \|_2^2$ , by definition Eq. (2), for any  $\mathbf{x}$ , we have that
$$
\left\| \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x}) \right\| _ {2} ^ {2} = \frac {\left\| h _ {L} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {L - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2}} \dots \frac {\left\| h _ {1} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {0} (\boldsymbol {x}) \right\| _ {2} ^ {2}} \cdot \frac {\left\| h _ {0} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| \boldsymbol {x} \right\| _ {2} ^ {2}} \| \boldsymbol {x} \| _ {2} ^ {2} \tag {53}
$$
We now bound the terms in the right-hand-side of Eq. (53) one by one.
Regarding the first term  $\frac{\|h_L(\pmb{x})\|_2^2}{\|h_{L - 1}(\pmb{x})\|_2^2}$  in Eq. (53), observe that by the network output definition Eq. (2), we have that
$$
\left\| h _ {L} (\boldsymbol {x}) \right\| _ {2} ^ {2} = \left\| \boldsymbol {W} ^ {L} h _ {L - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2} \overset {d} {=} \beta_ {L} \left\| h _ {L - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2} \tilde {w} ^ {2} \tag {54}
$$
where  $\tilde{w} \sim \mathcal{N}(0,1)$  and the last equality is by rotaional invariance of Gaussian distribution used for initializing  $L$ -th layer weight matrix  $\pmb{W}^{L} \in \mathbb{R}^{1 \times m}$ . Therefore, by tail probability expression for standard Gaussian random variable, we have that with high probability  $1 - \frac{\delta}{L}$  over random initialization of  $\pmb{W}^{L} \in \mathbb{R}^{1 \times m}$ , it satisfies that
$$
\frac {\left\| h _ {L} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {L - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2}} \leq 2 \beta_ {L} \log \frac {L}{\delta} \tag {55}
$$
Regarding the terms  $\frac{\|h_l(\pmb{x})\|_2^2}{\|h_{l - 1}(\pmb{x})\|_2^2}$  in Eq. (53) for layer  $l = 1,\dots ,L - 1$  , by setting  $\pmb {H}_1 = \pmb {h}_2 = h_{l - 1}$  in Eq. (32), we immediately prove that over random initialization of weight matrix  $\pmb{W}^{l}\in \mathbb{R}^{m_{l}\times m_{l - 1}}$  it satisfies that
$$
\frac {\left\| h _ {l} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {l - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2}} \stackrel {d} {=} \sum_ {i = 1} ^ {m _ {l}} \rho_ {i} \beta_ {l} \tilde {w} _ {i} ^ {2} \tag {56}
$$
where  $\rho_1, \dots, \rho_{m_l} \stackrel{i.i.d.}{\sim} \operatorname{Ber}(1, \frac{1}{2})$  and  $\tilde{w}_1, \dots, \tilde{w}_{m_l} \stackrel{i.i.d.}{\sim} \mathcal{N}(0, 1)$  and  $\rho_i$  and  $\tilde{w}_i$  are independent.
By tail probability expression for Gaussian random variable  $\tilde{w}_i$ , we have that  $P(\tilde{w}_i^2 \geq t) \leq e^{-t/2}$  for any  $t > 0$ . By union bound over  $i = 1, \dots, m_l$ , we prove that for any layer  $l = 1, \dots, L - 1$ , with high probability  $1 - \frac{\delta}{2L}$  over random initialization of weight matrix  $\mathbf{W}^l \in \mathbb{R}^{m_l \times m_{l-1}}$ , it satisfies that
$$
\max  _ {i} \tilde {w} _ {i} ^ {2} \leq 2 \log \frac {2 m _ {l} L}{\delta} \tag {57}
$$
Moreover, by applying Hoeffding's inequality to i.i.d. Bernoulli r.v.s  $\rho_{1},\dots ,\rho_{m}$ , we prove with high probability  $1 - \frac{\delta}{2L}$ , it satisfies that  $\sum_{i = 1}^{m}\rho_{i}\leq \frac{m_{l}}{2}\big(1 + \log (2L / \delta)\big)$ . By combining it with Eq. (57) via union bound, and plugging the result into Eq. (59), we prove for any  $l = 1,\dots ,L - 1$ , it satisfies with high probability  $1 - \frac{\delta}{L}$  over random initialization of weight matrix  $\pmb{W}^{l}\in \mathbb{R}^{m_{l}\times m_{l - 1}}$  that
$$
\frac {\left\| h _ {l} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {l - 1} (\boldsymbol {x}) \right\| _ {2} ^ {2}} = m _ {l} \beta_ {l} \log \frac {2 m _ {l} L}{\delta} \cdot (1 + \log (2 L / \delta)) \tag {58}
$$
By using union bound over Eq. (58) for layer  $l = 1, \dots, L - 1$  and Eq. (55), we have that with high probability  $1 - \delta$  over random initialization Eq. (5), it satisfies that
$$
\frac {\left\| h _ {L} (\boldsymbol {x}) \right\| _ {2} ^ {2}}{\left\| h _ {0} (\boldsymbol {x}) \right\| _ {2} ^ {2}} \leq \tilde {\mathcal {O}} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right) \tag {59}
$$
where  $\tilde{\mathcal{O}}$  ignores logarithmic factors with regard to  $m$ ,  $L$  and  $\delta$ .
By plugging Eq. (59) into Eq. (53), we prove that with high probability  $1 - \delta$  over initialization Eq. (5), the following bound holds.
$$
\left\| \boldsymbol {f} _ {\boldsymbol {W}} (\boldsymbol {x}) \right\| _ {2} \leq \| \boldsymbol {x} \| _ {2} ^ {2} \tilde {\mathcal {O}} \left(\beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)\right), \tag {60}
$$
Lemma 6.3. (Bounding lazy training distance via smallest eigenvalue of the NTK matrix) Under Assumption 2.4 and single-output linearized network Eq. (3) with  $o = 1$ , assume that the per-layer network widths  $m_0, \dots, m_L = \tilde{\Omega}(n)$  are large. Let  $\mathcal{L}_0^{lin}(\boldsymbol{W})$  be the empirical risk Eq. (1) for linearized network expanded at initialization vector  $\boldsymbol{W}_0^{lin}$ . Then for any  $\boldsymbol{W}_0^{lin}$ , there exists a corresponding solution  $\boldsymbol{W}_0^{\frac{1}{n^2}}$ , s.t.  $\mathcal{L}_0^{lin}(\boldsymbol{W}_0^{\frac{1}{n^2}}) - \min_{\boldsymbol{W}} \mathcal{L}_0^{lin}(\boldsymbol{W}; \mathcal{D}) \leq \frac{1}{n^2}$ ,  $\text{rank}(M_0) = n$  and
$$
R \leq \tilde {\mathcal {O}} \left(\max  \left\{\frac {1}{d \beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)}, 1 \right\} \frac {n}{\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}}\right), \tag {61}
$$
Proof. Given an arbitrary initialization parameter vector  $\mathbf{W}_0^{lin}$ , we first construct an solution  $\mathbf{W}_0^{\frac{1}{n^2}}$  that is nearly optimal for the ERM problem over  $\mathcal{L}_0^{lin}(\mathbf{W})$ , as follows.
$$
\boldsymbol {W} _ {0} ^ {\frac {1}{n ^ {2}}} - \boldsymbol {W} _ {0} ^ {l i n} = M _ {0} ^ {\dagger} \left( \begin{array}{c} 2 \ln n \cdot y _ {1} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {1}) \\ \vdots \\ 2 \ln n \cdot y _ {n} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {n}) \end{array} \right) \tag {62}
$$
where  $M_0 = \begin{pmatrix} \nabla f_{W_0^{lin}}(\boldsymbol{x}_1)^\top \\ \vdots \\ \nabla f_{W_0^{lin}}(\boldsymbol{x}_n)^\top \end{pmatrix}$  is the gradient matrix at initialization and  $\dagger$  denotes pseudo-inverse.
We now prove that the solution  $W_0^{\frac{1}{n^2}}$  is close to the initialization parameters  $W_0^{lin}$  in  $\ell_2$  distance with high probability. By applying the holder inequality in (62), we have that
$$
\begin{array}{l} R = \left\| \boldsymbol {W} _ {0} ^ {\frac {1}{n ^ {2}}} - \boldsymbol {W} _ {0} ^ {l i n} \right\| _ {2} ^ {2} \leq \left\| M _ {0} ^ {\dagger} \right\| _ {2} ^ {2} \cdot \left\| \left( \begin{array}{c} 2 \ln n \cdot y _ {1} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {1}) \\ \vdots \\ 2 \ln n \cdot y _ {n} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {n}) \end{array} \right) \right\| _ {2} ^ {2} (63) \\ \leq \frac {1}{\lambda_ {0}} \cdot \left\| \left( \begin{array}{c} 2 \ln n \cdot y _ {1} - f _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {1}) \\ \vdots \\ 2 \ln n \cdot y _ {n} - f _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {n}) \end{array} \right) \right\| _ {2} ^ {2} (64) \\ \end{array}
$$
where  $\lambda_0$  is the smallest non-zero eigenvalue of the PSD matrix  $M_0M_0^\top$ . When the data regularity assumption Assumption 2.4 holds and the per-layer width  $m_0, \dots, m_L = \tilde{\Omega}(n)$ , by applying existing bounds for the smallest eigenvalue of the NTK matrix  $M_0M_0^\top$  for single-output network in [40,
Theorem 4.1], we prove that with high probability over data sampling and random initialization Eq. (5), it satisfies that
$$
\begin{array}{l} O \left(\left(d \prod_ {l = 1} ^ {L - 1} m _ {l}\right) \cdot \left(\prod_ {l = 1} ^ {L} \beta_ {l}\right) \cdot \left(\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}\right)\right) \geq \lambda_ {0} \\ \geq \Omega \left(\left(d \prod_ {l = 1} ^ {L - 1} m _ {l}\right) \cdot \left(\prod_ {l = 1} ^ {L} \beta_ {l}\right) \cdot \left(\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}\right)\right), \tag {65} \\ \end{array}
$$
where we have set the auxiliary variables  $\xi_{l} = 1$  in [40, Theorem 4.1] because the per-layer width  $m_{l} = \widetilde{\Omega}(n)$  are large enough for  $l = 0, \dots, L$  (where  $n$  is the number of training data).
Therefore, by using Eq. (65), we have that with high probability over data sampling and random initialization Eq. (5), it satisfies that
$$
\frac {1}{\lambda_ {0}} \leq O \left(\frac {1}{\left(d \prod_ {l = 1} ^ {L - 1} m _ {l}\right) \cdot \left(\prod_ {l = 1} ^ {L} \beta_ {l}\right) \cdot \left(\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}\right)}\right), \tag {66}
$$
where  $m_l$  is the width of layer  $l$  and  $n$  is the number of training data.
For the second term in Eq. (64), by Cauchy-Schwarz inequality, we have that
$$
\left\| \binom {2 \ln n \cdot y _ {1} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {1})} {\vdots} \right\| _ {2} ^ {2} \leq 2 \cdot (2 \ln n) ^ {2} \sum_ {i = 1} ^ {n} y _ {i} ^ {2} + 2 \sum_ {i = 1} ^ {n} \| \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {i}) \| _ {2} ^ {2} \tag {67}
$$
By Lemma D.1, we have that  $\| f_{\pmb{W}_0^{lin}}(\pmb{x})\|_2 \leq \tilde{\mathcal{O}}\left(d \cdot \beta_L\left(\prod_{i=1}^{L-1}\beta_im_i\right)\right)$  with high probability over the random initialization Eq. (5). By plugging this result and  $y_i \in \{-1,1\}$  into Eq. (67), we have that with high probability over random initialization Eq. (5), it satisfies that
$$
\left\| \binom {2 \ln n \cdot y _ {1} - \boldsymbol {f} _ {\boldsymbol {W} _ {0} ^ {l i n}} (\boldsymbol {x} _ {1})} {\vdots} \right\| _ {2} ^ {2} = \tilde {\mathcal {O}} \left(n + n d \beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)\right), \tag {68}
$$
where  $\tilde{O}$  undergoes logarithmic factors with regard to  $n$ ,  $m$ ,  $L$ , and tail probability  $\delta$ . Therefore, by combining Eq. (66) and Eq. (68) with union bound, and by plugging the result into Eq. (64), we have that with high probability over data sampling and random initialization Eq. (5), it satisfies that
$$
\begin{array}{l} R \leq \tilde {\mathcal {O}} \left(\frac {n + n d \beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)}{\left(d \beta_ {L} \prod_ {l = 1} ^ {L - 1} \beta_ {l} m _ {l}\right) \cdot \left(\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}\right)}\right) \\ \leq \tilde {\mathcal {O}} \left(\max  \left\{\frac {1}{d \beta_ {L} \left(\prod_ {i = 1} ^ {L - 1} \beta_ {i} m _ {i}\right)}, 1 \right\} \frac {n}{\sum_ {l = 1} ^ {L} \beta_ {l} ^ {- 1}}\right). \\ \end{array}
$$
where  $m_l$  is the width of layer  $l$ ,  $n$  is the number of training data, and  $\tilde{\mathcal{O}}$  ignores logarithmic factors with regard to  $n$ ,  $m$ ,  $L$ , and tail probability  $\delta$ .
We finally prove that  $\mathbf{W}_0^{\frac{1}{n^2}}$  is a  $\frac{1}{n^2}$ -near-optimal solution. Note that Eq. (65) implies that with high probability  $M_0M_0^\top$  is full rank, i.e.,  $\mathrm{rank}(M_0M_0^\top) = n$ . Therefore  $M_0^\dagger = M_0^\top (M_0M_0^\top)^{-1}$
and  $\left( \begin{array}{c} \boldsymbol{f}_{\boldsymbol{W}_0^{n^2}}(\boldsymbol{x}_1) \\ \vdots \\ \boldsymbol{f}_{\boldsymbol{W}_0^{1 / 2}}(\boldsymbol{x}_n) \end{array} \right) = \left( \begin{array}{c} 2 \ln n \cdot y_1 \\ \vdots \\ 2 \ln n \cdot y_n \end{array} \right)$  with high probability over the training data sampling and
random initialization Eq. (5). By plugging it into the cross-entropy loss for the single-output network
defined below Eq. (1), we have that with high probability over the training data sampling and random initialization Eq. (5), the solution  $\mathbf{W}_0^{\frac{1}{n^2}}$  satisfies
$$
\mathcal {L} _ {0} ^ {\text {l i n}} \left(\boldsymbol {W} _ {0} ^ {\frac {1}{n ^ {2}}}\right) - \min  _ {\boldsymbol {W}} \mathcal {L} _ {0} ^ {\text {l i n}} (\boldsymbol {W}; \mathcal {D}) \leq \log (1 + \exp (- 2 \ln n)) <   \frac {1}{n ^ {2}}. \tag {69}
$$
D.3 Deferred proof for Corollary 6.4
Corollary 6.4. (Privacy utility trade-off for linearized network) Assume that all conditions in Lemma 6.3 hold. Let  $B = d\left(\prod_{i=1}^{L-1}\frac{\beta_i m_i}{2}\right)\sum_{l=1}^{L}\frac{\beta_L}{\beta_l}$  be the gradient norm constant proved in Eq. (7), and let  $R \leq \tilde{\mathcal{O}}\left(\max\left\{\frac{1}{d\beta_L\left(\prod_{i=1}^{L-1}\beta_i m_i\right)}, 1\right\}\frac{n}{\sum_{l=1}^{L}\beta_l^{-1}}\right)$  be the lazy training distance bound proved in Lemma 6.3. Then for  $\sigma^2 = \frac{2BT}{\varepsilon n^2}$  and  $T = \sqrt{\frac{\varepsilon nR}{2B}}$ , releasing all iterates of Langevin diffusion with time  $T$  satisfies  $\varepsilon$ -KL privacy, and has empirical excess risk upper bounded by
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}\right) \right] - \min  _ {\boldsymbol {W}} \mathcal {L} _ {0} ^ {l i n} (\boldsymbol {W}; \mathcal {D}) \leq \tilde {\mathcal {O}} \left(\frac {1}{n ^ {2}} + \sqrt {\frac {B R}{\varepsilon n}}\right) (70) \\ = \tilde {\mathcal {O}} \left(\frac {1}{n ^ {2}} + \sqrt {\frac {1}{2 ^ {L - 1} \varepsilon} \max  \{1 , d \beta_ {L} \prod_ {l = 1} ^ {L - 1} \beta_ {l} m _ {l} \}}\right) (71) \\ \end{array}
$$
with high probability over random initialization Eq. (5), where the expectation is over Brownian motion  $B_{T}$  in Langevin diffusion, and  $\tilde{O}$  ignores logarithmic factors with regard to width  $m$ , depth  $L$ , number of training data  $n$  and tail probability  $\delta$ . A summary of our excess empirical risk bounds under different initializations is in Table 1.
Proof. By setting  $\mathbf{W}_0^* = \mathbf{W}_0^{\frac{1}{n^2}}$  in Lemma 6.1 with  $\mathbf{W}_0^{\frac{1}{n^2}}$  constructed as Lemma 6.3, we have with high probability over random initialization Eq. (5), we have
$$
\mathbb {E} \left[ \mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}\right) \right] - \min  _ {\boldsymbol {W}} \mathcal {L} _ {0} ^ {l i n} (\boldsymbol {W}; \mathcal {D}) \leq \frac {1}{n ^ {2}} + \frac {R}{2 T} + \frac {\sigma^ {2} n}{2}. \tag {72}
$$
where  $R \leq \tilde{\mathcal{O}}\left(\max \left\{\frac{1}{d\beta_L\left(\prod_{i=1}^{L-1}\beta_im_i\right)}, 1\right\} \frac{n}{\sum_{l=1}^L\beta_l^{-1}}\right)$  by Lemma 6.3.
Meanwhile, to ensure  $\varepsilon$ -KL privacy, by Corollary 4.2, we only need to set  $\sigma^2 = \frac{2BT}{\varepsilon n^2}$  where  $B = d\left(\prod_{i=1}^{L-1} \frac{\beta_i m_i}{2}\right) \sum_{l=1}^{L} \frac{\beta_L}{\beta_l}$  by for single-output network with  $o = 1$ . By plugging  $\sigma^2 = \frac{2BT}{\varepsilon n^2}$  into (72), we prove that
$$
\mathbb {E} \left[ \mathcal {L} _ {0} ^ {\text {l i n}} \left(\bar {\boldsymbol {W}} _ {T} ^ {\text {l i n}}\right) \right] - \min  _ {\boldsymbol {W}} \mathcal {L} _ {0} ^ {\text {l i n}} (\boldsymbol {W}; \mathcal {D}) \leq \frac {1}{n ^ {2}} + \frac {R}{2 T} + \frac {B T}{\varepsilon n}. \tag {73}
$$
Setting  $T = \sqrt{\frac{\varepsilon nR}{2B}}$  in (73) and elaborating the computations suffice to prove the result.
E Discussion on extending our results to Noisy GD with constant step-size
In this section, we discuss how to extend our privacy analyses to noisy GD with constant step-size. Specifically, we only need to extend the KL composition theorem under possibly unbounded gradient difference, i.e., Theorem 3.1, to the noisy GD algorithm.
Theorem E.1 (KL composition for noisy GD under possibly unbounded gradient difference). Let the iterative update in noisy GD algorithm be defined by:  $\pmb{W}_{(k + 1)} = \pmb{W}_{(k)} - \eta \nabla \mathcal{L}(\pmb{W}_{(k)};\mathcal{D})+$
$\sqrt{2\eta\sigma^2} Z_k$ , where  $Z_k \sim \mathcal{N}(0, \mathbb{I})$ . Then the KL divergence between running noisy GD for DNN (2) on neighboring datasets  $\mathcal{D}$  and  $\mathcal{D}'$  satisfies
$$
K L \left(\boldsymbol {W} _ {(1: K)}, \boldsymbol {W} _ {(1: K)} ^ {\prime}\right) = \frac {1}{2 \sigma^ {2}} \sum_ {k = 0} ^ {K - 1} \eta \cdot \mathbb {E} \left[ \left\| \nabla \mathcal {L} \left(\boldsymbol {W} _ {(k)}; \mathcal {D}\right) - \nabla \mathcal {L} \left(\boldsymbol {W} _ {(k)}; \mathcal {D} ^ {\prime}\right) \right\| _ {2} ^ {2} \right]. \tag {74}
$$
Proof. Denote  $p_{(k)}$  as the distribution of model parameters after running noisy GD on dataset  $\mathcal{D}$  with  $k$  steps, and similarly denote  $p_{(k)}'$  as the distribution of model parameters after running noisy GD on dataset  $\mathcal{D}'$  with  $k$  steps. Similarly, denote  $p_{(1:k)}$  as the joint distribution of  $(W_{(1)},\dots ,W_{(k)})$ , and denote  $p_{(1:k)}'$  as the joint distribution of  $(W_{(1)}',\dots ,W_{(k)}')$ . Now we expand the term  $KL(p_{(1,k+1)},p_{(1,k+1}')$  by the Bayes rule as follows.
$$
\begin{array}{l} K L \left(p _ {(1: k + 1)}, p _ {(1: k + 1)} ^ {\prime}\right) (75) \\ = \mathbb {E} _ {p _ {(1: k + 1)} \left(\boldsymbol {W} _ {(1: k + 1)}\right)} \left[ \log \left(\frac {p _ {(k + 1) | (1 : k)} \left(\boldsymbol {W} _ {(k + 1)} \mid \boldsymbol {W} _ {(1 : k)}\right) p _ {(1 : k)} \left(\boldsymbol {W} _ {(1 : k)}\right)}{p _ {(k + 1) | (1 : k)} ^ {\prime} \left(\boldsymbol {W} _ {(k + 1)} \mid \boldsymbol {W} _ {(1 : k)}\right) p _ {(1 : k)} ^ {\prime} \left(\boldsymbol {W} _ {(1 : k)}\right)}\right) \right] \\ = \mathbb {E} _ {p _ {(1: k + 1)} \left(\boldsymbol {W} _ {(1: k + 1)}\right)} \left[ \log \left(\frac {p _ {(k + 1) | (1 : k)} \left(\boldsymbol {W} _ {(k + 1)} \mid \boldsymbol {W} _ {(1 : k)}\right)}{p _ {(k + 1) | (1 : k)} ^ {\prime} \left(\boldsymbol {W} _ {(k + 1)} \mid \boldsymbol {W} _ {(1 : k)}\right)}\right) \right] + \mathbb {E} _ {p _ {(1: k)} \left(\boldsymbol {W} _ {(1: k)}\right)} \left[ \log \left(\frac {p _ {(1: k)} \left(\boldsymbol {W} _ {(1: k)}\right)}{p _ {(1: k)} ^ {\prime} \left(\boldsymbol {W} _ {(1: k)}\right)}\right) \right] \\ = \mathbb {E} _ {p _ {(1: k)} \left(\boldsymbol {W} _ {(1: k)}\right)} \left[ K L \left(p _ {(k + 1) | (1: k)}, p _ {(1: k + 1) | (1: k)} ^ {\prime}\right) \right] + K L \left(p _ {(1: k)}, p _ {(1: k)} ^ {\prime}\right) (76) \\ \end{array}
$$
Observe that conditioned on fixed model parameters  $W_{(1:k)}$  at iteration  $1, \dots, k$ , the distributions  $p_{(k+1)|(1:k)}, p_{(k+1)|(1:k)}'$  are Gaussian with per-dimensional variance  $\sigma^2$ . Therefore, by computing the KL divergence between two multivariate Gaussians, we have that
$$
K L \left(p _ {(k + 1) | (1: k)}, p _ {(k + 1) | (1: k)} ^ {\prime}\right) = \frac {1}{2 \sigma^ {2}} \cdot \eta \cdot \left\| \nabla \mathcal {L} \left(\boldsymbol {W} _ {(k)}; \mathcal {D}\right) - \nabla \mathcal {L} \left(\boldsymbol {W} _ {(k)}; \mathcal {D} ^ {\prime}\right) \right\| _ {2} ^ {2} \tag {77}
$$
Therefore, by plugging Eq. (77) into Eq. (76), we have that
$$
K L \left(p _ {(1: k + 1)}, p _ {(1: k + 1)} ^ {\prime}\right) = \frac {\eta}{2 \sigma^ {2}} \mathbb {E} \left[ \left\| \nabla \mathcal {L} (\boldsymbol {W} _ {(k)}; \mathcal {D}) - \nabla \mathcal {L} (\boldsymbol {W} _ {(k)}; \mathcal {D} ^ {\prime}) \right\| _ {2} ^ {2} \right] + K L \left(p _ {(1: k)}, p _ {(1: k)} ^ {\prime}\right) \tag {78}
$$
By summing (78) over  $k = 0, \dots, K - 1$  and observing that  $KL(p_{(0)}, p_{(0)}') = 0$  (as the initialization distribution is the same between noisy GD on  $\mathcal{D}$  and  $\mathcal{D}'$ ), we finish the proof for Eq. (74).