SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \nabla_ {W _ {j}} y _ {d} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} \left(y _ {j - 1} (\theta , s _ {a})\right) \nabla_ {y _ {j}} y _ {d} (\theta , s _ {a}) \\ \nabla_ {W _ {j}} y _ {d} ^ {(i)} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} ^ {\prime} (y _ {j - 1} (\theta , s _ {a})) \odot y _ {j - 1} ^ {(i)} (\theta , s _ {a}) \nabla_ {y _ {j}} y _ {d} (\theta , s _ {a}) \quad \forall j > i \\ \nabla_ {W _ {j}} y _ {d} ^ {(i)} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} (y _ {j - 1} (\theta , s _ {a})) \nabla_ {y _ {j}} y _ {d} (\theta , s _ {a}) \quad \forall j \leq i \\ \nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} ^ {\prime} \left(y _ {j - 1} \left(\theta , s _ {a}\right)\right) \odot \gamma_ {j - 1} ^ {(i)} \left(\theta , s _ {a}\right) \nabla_ {y _ {j}} y _ {d} \left(\theta , s _ {a}\right) \quad \forall j > i + 1 \tag {37} \\ \nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} ^ {\prime \prime} (y _ {j - 1} (\theta , s _ {a})) \nabla_ {y _ {j}} y _ {d} (\theta , s _ {a}) \quad \forall j = i + 1 \\ \nabla_ {W _ {j}} \gamma_ {d} ^ {(i)} (\theta , s _ {a}) = c _ {j} \sigma_ {j - 1} (y _ {j - 1} (\theta , s _ {a})) \nabla_ {y _ {i + 1}} y _ {d} (\theta , s _ {a}) (c _ {i + 1} W _ {i + 1}) \times \quad \forall j <   i + 1 \\ \operatorname {d i a g} \left(\sigma_ {i} ^ {\prime \prime \prime} \left(y _ {i} (\theta , s _ {a})\right)\right) \nabla_ {y _ {j}} y _ {i} (\theta , s _ {a}) \\ \end{array}
$$
We rely on following Local Lipschitz property to prove the stability of gradients during training.
Lemma 7 (Local Lipschitz Condition). Let  $v_{j-1}(\theta)$  to denote any of following values:
$$
\sigma_ {j - 1} \left(y _ {j - 1} \left(\theta , s _ {a}\right)\right),
$$
$$
\sigma_ {j - 1} ^ {\prime} \left(y _ {j - 1} \left(\theta , s _ {a}\right)\right) \odot y _ {j - 1} ^ {(i)} \left(\theta , s _ {a}\right), \tag {38}
$$
$$
\sigma_ {j - 1} ^ {\prime} \left(y _ {j - 1} \left(\theta , s _ {a}\right)\right) \odot \gamma_ {j - 1} ^ {(i)} \left(\theta , s _ {a}\right), \tag {38}
$$
$$
\sigma_ {j - 1} ^ {\prime \prime} (y _ {j - 1} (\theta , s _ {a})),
$$
there is a constant  $K > 0$  such that for every  $R > 0$ , with probability increasing to 1 at limit of  $n \to \infty$ , or equivalently, with high probability over random initialization (w.h.p.o.r.i.) of  $\theta$ , for any  $\theta_1, \theta_2 \in B(\theta, R)$ , where  $B(\theta, R) = \{\theta' | \| \theta' - \theta \| < R\}$ , the following holds
$$
\left\| v _ {j} \left(\theta_ {1}\right) - v _ {j} \left(\theta_ {2}\right) \right\| \leq K \| \theta_ {1} - \theta_ {2} \|, \tag {39}
$$
$$
\left\| v _ {j} \left(\theta_ {1}\right) \right\| \leq K \sqrt {n}.
$$
Moreover, let  $v_{j}(\theta)$  to denote any of following values:
$$
\begin{array}{l}\nabla_ {y _ {j}} y _ {d} (\theta , s _ {a}),\\\nabla_ {u _ {j}} u _ {j} (\theta , s _ {a}) (s _ {a}, W _ {u _ {j}}) \operatorname {d i a g} \left( \right.\tau^ {\prime \prime \prime} \left(u _ {j} (\theta , s _ {a}))\right) \nabla_ {u _ {j}} (\theta , s _ {a})\end{array}\tag {40}
$$
$$
\nabla_ {y _ {i + 1}} y _ {d} (\theta , s _ {a}) (c _ {i + 1} W _ {i + 1}) \operatorname {d i a g} (\sigma_ {i} ^ {\prime \prime \prime} (y _ {i} (\theta , s _ {a}))) \nabla_ {y _ {j}} y _ {i} (\theta , s _ {a}),
$$
there is a constant  $K > 0$  such that for every  $R > 0$ , w.h.p.o.r.i. of  $\theta$ , for any  $\theta_1, \theta_2 \in B(\theta, R)$
$$
\left\| v _ {j} \left(\theta_ {1}\right) - v _ {j} \left(\theta_ {2}\right) \right\| \leq \frac {1}{\sqrt {n}} K \| \theta_ {1} - \theta_ {2} \|, \tag {41}
$$
$$
\left\| v _ {j} \left(\theta_ {1}\right) \right\| \leq K.
$$
By using the norm upper bound in Lemma 7, we have that for any  $R > 0$  and  $\theta \in B(\theta(0), R)$ , the gradient norm  $\mathrm{Tr}\big((\nabla_{W_j}y_d(\theta, s_a))^T\nabla_{W_j}y_d(\theta, s_a)\big) = c_j^2\|\sigma_{j-1}(y_{j-1}(\theta, s_a))\|^2\|\nabla_{y_j}y_d(\theta, s_a)\|^2 \leq K$  is bounded w.h.p.o.r.i. For any finite training time,  $R$  could be selected large enough to ensure  $\theta(t) \in B(\theta(0), R)$  during training. The norm bound of difference in Lemma 7 also ensures that the change of gradient norm is bounded by  $\frac{1}{\sqrt{n}}K\|\theta(t) - \theta(0)\|$  in  $B(\theta(0), R)$  w.h.p.o.r.i. which is diminishing at the limit  $n \to 0$ .
Proof of Lemma 7. All vectors  $v_{j}$  in Lemma 7 could be computed sequentially by three kinds of operations: matrix multiplication  $\zeta = W\varepsilon$ , element-wise non-linearity  $\varepsilon = \phi (\zeta)$  with  $\phi$  having bounded derivative and element-wise multiplication  $\varepsilon = \phi (\zeta)\odot \zeta^{\prime}$ , where the element-wise non-linear function  $\phi$  is bounded and Lipschitz.
For the first kind of operation, given that  $W$  follows standard normal distribution, we have the  $\zeta$  is also normally distributed and enjoy following high probability bound.
Eq. (2.19) of [71] For  $n$  dimensional standard Gaussian random vector  $\zeta \sim \mathcal{N}(0, I_n)$ , for any  $t$
$$
\Pr \left(\| \zeta \| ^ {2} \geq n + t\right) \leq 2 e ^ {- n t ^ {2} / 8}. \tag {42}
$$
Therefore we have that for  $\zeta = W\varepsilon$ ,  $\| \zeta \| \leq K\sqrt{n}\| \varepsilon \|$  for some  $K$  with high probability.
For element-wise non-linearity  $\varepsilon = \phi (\zeta)$  where  $\zeta$  is a standard Gaussian random vector with i.i.d. coordinates, we derive the Chernoff bound as follows.
$$
\begin{array}{l} \Pr \big (\| \phi (\zeta) \| \geq K \sqrt {n} \big) = \Pr \left(\sum_ {i = 1} ^ {n} \phi (\zeta_ {i}) ^ {2} \geq K ^ {2} n\right) \\ = \inf  _ {t \geq 0} \Pr \left(\exp \left(t \sum_ {i = 1} ^ {n} \phi \left(\zeta_ {i}\right) ^ {2}\right) \geq \exp \left(t K ^ {2} n\right)\right) \\ \leq \inf  _ {t \geq 0} \frac {\mathbb {E} \left[ \exp \left(t \sum_ {i = 1} ^ {n} \phi \left(\zeta_ {i}\right) ^ {2}\right)\right) ]}{e ^ {t K ^ {2} n}} \tag {43} \\ = \inf  _ {t \geq 0} e ^ {- t K ^ {2} n} \prod_ {i = 1} ^ {n} \mathbb {E} \left[ \exp \left(t \phi \left(\zeta_ {i}\right) ^ {2}\right)\right) ] \\ \leq \inf  _ {t \geq 0} e ^ {- t K ^ {2} n} \prod_ {i = 1} ^ {n} \mathbb {E} \left[ \exp \left(t \left(\mathcal {O} (1) + \mathcal {O} (1) \zeta_ {i} ^ {2}\right)\right) \right] \\ \leq e ^ {- \mathcal {O} (1) K ^ {2} n + \mathcal {O} (1) n} \\ \end{array}
$$
Therefore, we have  $\| \varepsilon \| = \| \phi (\zeta) \| \leq K\sqrt{n}$  w.h.p.o.r.i. for large enough  $K$ .
For element-wise multiplication  $\varepsilon = \phi (\zeta)\odot \zeta^{\prime}$ , since  $\phi$  is bounded,  $\| \varepsilon \| \leq \mathcal{O}(1)\| \zeta '\|$  all the time.
Recursively apply above three operations gives the bound of norm  $\| v_{j} \|$ .
In order to control the norm of difference, we note that for matrix multiplication:
$$
\begin{array}{l} \left\| W _ {1} \varepsilon_ {1} - W _ {2} \varepsilon_ {2} \right\| = \left\| W _ {1} \varepsilon_ {1} - W _ {1} \varepsilon_ {2} + W _ {1} \varepsilon_ {2} - W _ {2} \varepsilon_ {2} \right\| \\ \leq \left\| W _ {1} \left(\varepsilon_ {1} - \varepsilon_ {2}\right) \right\| + \left\| W _ {1} - W _ {2} \right\| _ {\mathrm {o p}} \| \varepsilon_ {2} \| \tag {44} \\ \leq \| W _ {1} (\varepsilon_ {1} - \varepsilon_ {2}) \| + \| W _ {1} - W _ {2} \| _ {\mathrm {F r o}} \| \varepsilon_ {2} \| \\ \end{array}
$$
The first term could be controlled with high probability as  $\| W_{1}(\varepsilon_{1} - \varepsilon_{2})\| \leq \| W_{0}(\varepsilon_{1} - \varepsilon_{2})\| + \| (W_{1} - W_{0})(\varepsilon_{1} - \varepsilon_{2})\| \leq K\sqrt{n}\| \varepsilon_{1} - \varepsilon_{2}\| + R\| \varepsilon_{1} - \varepsilon_{2}\| \leq \mathcal{O}(1)\sqrt{n}\| \varepsilon_{1} - \varepsilon_{2}\|$  for large enough  $n$ , and  $W_{0}$  is the parameter at initialization,  $R$  is the maximum distance between  $W_{1}$  and  $W_{0}$ .
For element-wise non-linearity  $\varepsilon = \phi (\zeta)$ , since  $\phi$  is Lipschitz, we have the difference as
$$
\left\| \varepsilon_ {1} - \varepsilon_ {2} \right\| = \left\| \phi \left(\zeta_ {1}\right) - \phi \left(\zeta_ {2}\right) \right\| \leq \mathcal {O} (1) \| \zeta_ {1} - \zeta_ {2} \| \tag {45}
$$
For element-wise multiplication  $\varepsilon = \phi (\zeta)\odot \zeta^{\prime}$ , since  $\phi$  is bounded and Lipschitz, we have the difference as
$$
\begin{array}{l} \left\| \varepsilon_ {1} - \varepsilon_ {2} \right\| = \left\| \phi (\zeta_ {1}) \odot \zeta_ {1} ^ {\prime} - \phi (\zeta_ {2}) \odot \zeta_ {2} ^ {\prime} \right\| \\ \leq \left\| \left(\phi \left(\zeta_ {1}\right) - \phi \left(\zeta_ {2}\right)\right) \odot \zeta_ {1} ^ {\prime} \right\| + \left\| \phi \left(\zeta_ {2}\right) \odot \left(\zeta_ {1} ^ {\prime} - \zeta_ {2} ^ {\prime}\right) \right\| \tag {46} \\ \leq \mathcal {O} (1) \| (\zeta_ {1} - \zeta_ {2}) \odot \zeta_ {1} ^ {\prime} \| + \mathcal {O} (1) \| \zeta_ {1} ^ {\prime} - \zeta_ {2} ^ {\prime} \| \\ \leq \mathcal {O} (1) \| \zeta_ {1} - \zeta_ {2} \| + \mathcal {O} (1) \| \zeta_ {1} ^ {\prime} - \zeta_ {2} ^ {\prime} \| \\ \end{array}
$$
The first term is reduced by noticing  $\zeta_1^\prime$  has  $\mathcal{O}(1)$  elements.
Recursively apply above three operations gives the bound of norm of difference.
Proof of Corollary 6. According to Theorem 5, in the infinite width limit, we have
$$
\frac {\mathrm {d}}{\mathrm {d} t} \left(\sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (t))\right) = - \left(\sum_ {i = 1} ^ {d} \Theta^ {(i)}\right) \left(\nabla_ {Y} \mathcal {L}\right) ^ {T} \tag {47}
$$
$$
\frac {\mathrm {d}}{\mathrm {d} t} \left(\sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (t))\right) = - \left(\sum_ {i = 1} ^ {d - 1} \Phi^ {(i)}\right) \left(\nabla_ {Y} \mathcal {L}\right) ^ {T}
$$
We use  $\mathcal{O}(1)$  to denote constants that is irrelevant to t. Given that  $\nabla_Y\mathcal{L}$  are bounded, we have
$$
\left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (t)) \right\| \leq \left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \right\| + \mathcal {O} (1) t \tag {48}
$$
$$
\left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (t)) \right\| \leq \left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \right\| + \mathcal {O} (1) t
$$
It follows that the Hessian is also bounded.
$$
\left. \operatorname {T r} \left\{H _ {a} (\theta (t)) \right\} \leq \mathcal {O} (1) \right\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \| + \mathcal {O} (1) t \tag {49}
$$
The time derivative and finite time change of energy and entropy is
$$
\left| \frac {\mathrm {d}}{\mathrm {d} t} V (\theta (t)) \right| \leq \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \right\| + \mathcal {O} (1) t
$$
$$
\left| \Delta_ {t} V (\theta (0)) \right| = \left| V (\theta (t)) - V (\theta (0)) \right| \leq \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) t ^ {2} \tag {50}
$$
$$
\left| \frac {\mathrm {d}}{\mathrm {d} t} \left(\Delta_ {t} S (\theta (0))\right) \right| \leq \mathcal {O} (1) + \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \right\| + \mathcal {O} (1) t
$$
$$
\left| \Delta_ {t} S (\theta (0)) \right| \leq \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) t + \mathcal {O} (1) t ^ {2}
$$
The KL divergence is
$$
\begin{array}{l} D _ {\mathrm {K L}} \left(q _ {t} \| p\right) = \mathbb {E} \left[ \Delta_ {t} V (\theta (0)) - \Delta_ {t} S (\theta (0)) \right] \\ \leq \mathbb {E} \left[ \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) t + \mathcal {O} (1) t ^ {2} \right] \tag {51} \\ \leq \mathcal {O} (1) t + \mathcal {O} (1) t ^ {2} \\ \end{array}
$$
The last inequality is because  $\sum_{i=1}^{d} Y^{(i)}(\theta(0))$  follows Gaussian distribution and the expectation of norm is finite.
The sampling efficiency could also be controlled.
$$
\begin{array}{l} \frac {1}{\operatorname {e f f} _ {\lambda}} = \frac {1}{Z _ {\lambda} ^ {2}} \mathbb {E} [ \exp (- 2 \lambda r (\theta (t)) - 2 \Delta_ {t} V (\theta (0)) + 2 \Delta_ {t} S (\theta (0))) ] \\ \leq \frac {1}{Z _ {\lambda} ^ {2}} \mathbb {E} [ \exp (- 2 \Delta_ {t} V (\theta (0)) + 2 \Delta_ {t} S (\theta (0))) ] \\ \leq \mathcal {O} (1) \mathbb {E} \left[ \exp \left(\mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) \left\| \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \right\| t + \mathcal {O} (1) t + \mathcal {O} (1) t ^ {2}\right) \right] \\ \leq \mathcal {O} (1) \exp \big (\mathcal {O} (1) t + \mathcal {O} (1) t ^ {2} \big) \mathbb {E} \left[ \exp \left(\mathcal {O} (1) \left\| \left[ \begin{array}{c} \sum_ {i = 1} ^ {d} Y ^ {(i)} (\theta (0)) \\ \sum_ {i = 1} ^ {d - 1} \Gamma^ {(i)} (\theta (0)) \end{array} \right] \right\| t\right) \right] \\ \leq \mathcal {O} (1) \exp (\mathcal {O} (1) t + \mathcal {O} (1) t ^ {2}) \tag {52} \\ \end{array}
$$
The first inequality comes from  $r(\theta) \geq 0$ . The last inequality follows from the fact that, for random vectors with tail similar to log-normal distribution, the expectation exists. More specifically, let  $\sigma_{i}(\Sigma)$  is eigenvalues of covariance matrix  $\Sigma$ , we have  $\mathbb{E}_{x \sim \mathcal{N}(0,\Sigma)}[\exp(c\|x\|)] \leq \prod_{i} \mathbb{E}_{y_{i} \sim \mathcal{N}(0,\sigma_{i}(\Sigma))}[\exp(c\|y_{i}\|)] = \prod_{i} \exp\big(c^{2}\sigma_{i}^{2}(\Sigma)/2\big)(1 + \mathrm{erf}(c\sigma_{i}(\Sigma)/2))$ .
In Figure 4, we show the correlation between KL divergence and generalization gap during training. We also show how weight distribution evolves.
Figure 4: The KL divergence, training and testing loss gap, and time dependence of weight distribution of experiment in Section 8.3.
F Additional Discussion on Algorithms
F.1 An illustrative example of TransBL
(a) Prior
(b) Loss
(c) Posterior
(a) Gradient flow
Figure 5: Bayesian learning process
(b) TransBL
Figure 6: I'llsutration of TransBL
In the presented motivating example, we explore a univariate loss function that presents two distinct global minima with zero loss. One of these is characterized as a sharp minimum, while the other represents a flat minimum. If the function were to be randomly initialized and then optimized, it might converge to either of the two localities. However, insights from PAC Bayesian indicate that the flat minimum is surrounded by a higher posterior probability density. A direct initialization from the prior, followed by training using gradient flow, often results in a significant deviation of the ensemble from the posterior. This is primarily because the optimization process fails to recognize the presence of a sharp minimum. The intuitive approach of the TransBL method is to apply a small weight to the solution found within the sharp minimum. Consequently, TransBL can adeptly recreate the posterior, as depicted in the Figure 6b.
Furthermore, we demonstrate the interpolation between the ensemble distribution obtained from optimization and Bayesian posterior through weight clipping in Section 6.1. In Figure 7a, with  $\beta = 0.2$ , the curve leans more towards the optimization result, revealing a broader spread and
less-defined peaks. This shows a scenario where optimization has a stronger influence than Bayesian learning. Yet with  $\beta = 5$ , the distribution is almost akin to what one would expect from a Bayesian posterior. In essence, the parameter  $\beta$  serves as a tuning knob, allowing us to interpolate and traverse the spectrum from pure optimization-driven results to outcomes heavily influenced by Bayesian learning. This interpolation mechanism offers a flexible approach to merge the strengths of both methodologies.
(a)  $\beta = 0.2$
(b)  $\beta = 1$
(c)  $\beta = 5$
Figure 7: Interpolation between optimization and Bayesian learning
F.2 Optimization Algorithms
In the main paper, we analyze the change of energy  $\Delta V$  and entropy  $\Delta S$  for gradient descent (GD). However, with finite step size, GD might not be reversible. Another difficulty lies in the computation of entropy change. As shown in Eq. (9), the entropy change is related to the curvature of loss  $l^s$  and Hessian trace of neural network, where the second part is hard to compute exactly.
We note that this problem could be alleviated for algorithms with enriched state space, but at a price of lower sampling efficiency.
Momentum SGD For parameter  $x$  and momentum  $v$ , the Momentum SGD with fraction  $\gamma$ , step size  $h$  and gradient oracle  $g$  can be formulated into following two steps:
$$
\left\{ \begin{array}{l} v \leftarrow \gamma v + g (x) \\ x \leftarrow x - v \end{array} . \right. \tag {53}
$$
By defining  $y = [x \quad v]$ , Momentum SGD turns into a transform  $T(y)$  composed by two sub transforms  $T(y) = L_2 \circ L_1(y)$  that correspond to two lines separately.
The determinant of Jacobian for  $L_{1}$  is  $\operatorname*{det}(\gamma I) = \gamma^d$ , where  $d$  is dimension of parameters. The determinant of Jacobian for  $L_{2}$  is 1, therefore, the determinant of Jacobian for Momentum SGD is  $\gamma^d$ .
Adagrad For parameter  $x$ , first-order momentum  $v$  and second-order momentum  $m$ , the Adagrad is:
$$
\left\{ \begin{array}{l} m \leftarrow m + g (x) ^ {2} \\ v \leftarrow \gamma v + (1 - \gamma) \frac {g (x)}{\sqrt {m}} \\ x \leftarrow x - h v \end{array} \right.. \tag {54}
$$
The square, square root and division in above formula are all element-wise.
Similar to Momentum SGD, the whole transform of Adagrad could be decomposed into three sub transform. The first and third step are both volume-preserving, and second step compress momentum by  $\gamma$  in each dimension. Therefore the determinant of Jacobian for Adagrad is  $\gamma^d$ .
Both above two algorithms enjoy simple form of Jacobian determinant. However, the sampling efficiency degenerates since  $\chi^2$  divergence increases when we consider joint probability distribution of parameter and auxiliary variables, like momentum.
F.3 Approximate DNN Hessian Trace by Forward Propagation
We first recall that due to Theorem 4, we have
$$
\operatorname {T r} \left(H _ {a} (\theta (t))\right) = \sum_ {j = 1} ^ {d - 1} \left(\nabla_ {x _ {j}} y _ {d}\right) \left(\sigma_ {j} ^ {\prime \prime} (y _ {j}) \odot \left(\sum_ {i = 1} ^ {j} \xi^ {(i, j)}\right)\right).
$$
Our motivation is based on the stability of  $\xi^{(i,j)}$ . We have shown in Theorem 5 that in the infinite width limit,  $\xi^{(i,j)}$  converge to a vector with same elements and is stable during training. For the wide but finite width network, this property is largely preserved.
If we replace  $\xi^{(i,j)}$  with its limit value, we obtain following formula:
$$
\mathrm {T r} (H _ {a} (\theta (t))) \approx \sum_ {j = 1} ^ {d - 1} (\nabla_ {x _ {j}} y _ {d} (\theta (t), s _ {a})) \sigma_ {j} ^ {\prime \prime} (y _ {j} (\theta (t), s _ {a})) \left(\sum_ {i = 1} ^ {j} (\Xi^ {(i, j)}) _ {a}\right).
$$
In order to calculate the above value efficiently, we notice that above summation can be regarded as propagation of a tangent vector.
We first change the network definiton into
$$
x _ {0} (\theta , s _ {a}) = s _ {a},
$$
$$
y _ {i} (\theta , s _ {a}) = c _ {i} W _ {i} x _ {i - 1} (\theta , s _ {a}), \quad \forall i = 1, \dots , d,
$$
$$
x _ {i} (\theta , s _ {a}) = \sigma_ {i} \left(y _ {i} (\theta , s _ {a})\right) + b _ {i}, \quad \forall i = 1, \dots , d - 1.
$$
Notice that compared to original network, we add vectors  $b_{i}$ , and when  $b_{i}$  is set as 0, the output is same as original network.
It can be easily seen that  $\nabla_{x_j}y_d = \nabla_{b_j}y_d$ , therefore we could just let  $\sigma_j''(y_j(\theta(t), s_a)) \left( \sum_{i=1}^j (\Xi^{(i,j)})_a \right)$  be the tangent vector for  $b_i$  and let it propagate along the forward pass, the result tangent vector in output space is just an estimation of  $\operatorname{Tr}(H_a(\theta(t)))$ .
F.4 Final PAC Bayes bound
$$
\mathbb {E} _ {\theta \sim q _ {t}} [ R (\theta) ] = \Phi_ {\frac {\lambda}{m}} ^ {- 1} \left(\mathbb {E} _ {\theta \sim q _ {t}} [ r (\theta) ] + \frac {D _ {K L} (q _ {t} | | p) + \log \frac {1}{\delta}}{\lambda}\right)
$$
$$
D _ {K L} (q _ {t} | | p) = \mathbb {E} _ {\{(Y _ {0} ^ {(i)}, \Gamma_ {0} ^ {(i)}) \} _ {i = 1} ^ {d}} [ V _ {t} - S _ {t} ]
$$
$$
V _ {0} = 0, S _ {0} = 0
$$
$$
Y _ {t} = Y _ {t} ^ {(d)}
$$
$$
\frac {d}{d t} V _ {t} = - \nabla \mathcal {L} \left(Y _ {t}\right) \sum_ {i = 1} ^ {d} Y _ {t} ^ {(i)} \tag {55}
$$
$$
\frac {d}{d t} S _ {t} = - T r \left(\nabla^ {2} \mathcal {L} (Y _ {t}) \Theta^ {(d)}\right) - \nabla \mathcal {L} (Y _ {t}) \sum_ {j = 1} ^ {d} \sum_ {i = 1} ^ {j} \Gamma^ {(i)} \odot \Xi^ {(i, j)}
$$
$$
\frac {d}{d t} Y _ {t} ^ {(i)} = - \Theta^ {(i)} \nabla \mathcal {L} (Y _ {t}) ^ {\top}
$$
$$
\frac {d}{d t} \Gamma_ {t} ^ {(i)} = - \Phi^ {(i)} \nabla \mathcal {L} (Y _ {t}) ^ {\top}
$$
The closed form solution of the KL divergence is hard to obtain, but the numerical solution can be computed efficiently on a computer (as shown in Figure 4) by solving the above ODE.