SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
S = \gamma \int d t d \boldsymbol {x} \hat {\Delta} (t, \boldsymbol {x}) (\Delta (t, \boldsymbol {x}) - y (\boldsymbol {x})) + \ln \mathbb {E} _ {a, \boldsymbol {w}} \exp \left(\int d t d \boldsymbol {x} \hat {\Delta} (t, \boldsymbol {x}) a (t) \phi (\boldsymbol {w} (t) \cdot \boldsymbol {x})\right) \tag {114}
$$
The Hessian over  $\pmb{q} = \{\Delta_{\mu}(t),\hat{\Delta}_{\mu}(t)\}$  is
$$
\nabla^ {2} S = \left[ \begin{array}{c c} 0 & \mathbf {I} + D _ {\Delta} \\ \mathbf {I} + D _ {\Delta} & \kappa \end{array} \right]. \tag {115}
$$
where  $D_{\Delta}(t,\pmb {x};s,\pmb{x}^{\prime}) = \left\langle \frac{\partial}{\partial\Delta(s,\pmb{x}^{\prime})} a(t)\phi (\pmb {w}(t)\cdot \pmb {x})\right\rangle$  We can use the following implicit rule
$$
\begin{array}{l} \frac {\partial a (t)}{\partial \Delta (s , \boldsymbol {x})} = \gamma \Theta (t - s) p (\boldsymbol {x}) \phi (\boldsymbol {w} (s) \cdot \boldsymbol {x}) + \gamma \mathbb {E} _ {\boldsymbol {x} ^ {\prime}} \int_ {0} ^ {t} d t ^ {\prime} \Delta \left(t ^ {\prime}, \boldsymbol {x} ^ {\prime}\right) \dot {\phi} (\boldsymbol {w} \cdot \boldsymbol {x} ^ {\prime}) \boldsymbol {x} ^ {\prime} \cdot \frac {\partial \boldsymbol {w} (t)}{\partial \Delta (s , \boldsymbol {x})} \\ \frac {\partial \boldsymbol {w} (t)}{\partial \Delta (s , \boldsymbol {x})} = \gamma \Theta (t - s) p (\boldsymbol {x}) a (s) \dot {\phi} (\boldsymbol {w} (s) \cdot \boldsymbol {x}) \boldsymbol {x} \\ + \gamma \mathbb {E} _ {\boldsymbol {x} ^ {\prime}} \int_ {0} ^ {t} d t ^ {\prime} \Delta \left(t ^ {\prime}, \boldsymbol {x} ^ {\prime}\right) \left[ \frac {\partial a \left(t ^ {\prime}\right)}{\partial \Delta (s , \boldsymbol {x})} \dot {\phi} (\boldsymbol {w} \cdot \boldsymbol {x} ^ {\prime}) + a \left(t ^ {\prime}\right) \ddot {\phi} (\boldsymbol {w} \cdot \boldsymbol {x} ^ {\prime}) \frac {\partial \boldsymbol {w} \left(t ^ {\prime}\right)}{\partial \Delta (s , \boldsymbol {x})} \cdot \boldsymbol {x} ^ {\prime} \right] \tag {116} \\ \end{array}
$$
The above equations could be solved and then used to compute  $D_{\Delta}(t, \boldsymbol{x}; s, \boldsymbol{x}^{\prime})$  which must then be inverted to get the observed prediction variance.
K.2 Linear Activations
Using the ideas in the preceding sections, we can make more progress in the case of a two layer linear network in the online learning setting. The key idea is to track the kernel and prediction error projections onto the space of linear functions. In this case we get the following DMFT over the order parameter  $\beta(t) = \frac{1}{N} \mathbf{W}^\top \mathbf{a} \in \mathbb{R}^D$ .
$$
\frac {d}{d t} a (t) = \gamma (\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)) \cdot \boldsymbol {w} (t)
$$
$$
\frac {d}{d t} \boldsymbol {w} (t) = \gamma a (t) \left(\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)\right)
$$
$$
\boldsymbol {\beta} (t) = \frac {1}{\gamma} \langle a (t) \boldsymbol {w} (t) \rangle \tag {117}
$$
At infinite width, we see that the dynamics can be reduced to tracking the projection of the weights  $\boldsymbol{w}$  and  $\beta$  on the  $\beta_{\star}$  direction. The  $D - 1$  off-target dimensions vanish  $\beta_{\perp}(t) = 0$ . At infinite width, we arrive at the alignment dynamics studied in prior work [64, 9]
$$
\begin{array}{l} \frac {d}{d t} \boldsymbol {\beta} (t) = \boldsymbol {M} (t) \left(\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)\right) \\ \frac {d}{d t} \boldsymbol {M} (t) = \gamma^ {2} \boldsymbol {\beta} (t) (\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)) ^ {\top} + \gamma^ {2} \boldsymbol {\beta} (t) (\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)) ^ {\top} \\ + 2 \gamma^ {2} \left(\beta_ {\star} - \beta (t)\right) \cdot \beta (t) \mathbf {I} \tag {118} \\ \end{array}
$$
We note that  $\beta(t) = \beta(t)\beta_{\star}$  and that  $M$  has only one special eigenvector  $\beta_{\star}$  with eigenvalue  $m_{\star}(t)$ . It thus suffices to track evolution in this single direction
$$
\frac {d}{d t} \beta (t) = m _ {\star} (t) \left(\beta_ {\star} - \beta (t)\right), \frac {d}{d t} m _ {\star} (t) = 4 \gamma^ {2} \beta (t) \left(\beta_ {\star} - \beta (t)\right) \tag {119}
$$
We note that this equation is identical to the differential equation for a single training example in Appendix J. Here  $\beta_{\star} - \beta (t)$  plays the role of  $\Delta_y(t)$  and  $m_{\star}(t)$  plays the role of the kernel  $K_{y}(t)$ . A key observation is the conservation law  $4\gamma^{2}\frac{d}{dt}\beta (t)^{2} = \frac{d}{dt} m_{\star}(t)^{2}$ , from which it follows that  $m_{\star}(t)^{2} - 4 = 4\gamma^{2}\beta (t)$  [9]
$$
\frac {d}{d t} \beta (t) = 2 \sqrt {1 + \gamma^ {2} \beta (t) ^ {2}} (\beta_ {\star} - \beta (t)) \tag {120}
$$
This is identical to the differential equations for a single sample (producing prediction  $f(t)$  and kernel  $K(t)$ ) if the following substitutions are made
$$
f (t) \leftrightarrow \beta (t), K (t) \leftrightarrow m _ {\star} (t) \tag {121}
$$
We now proceed to compute finite size corrections starting from the action
$$
S = \gamma \int d t \hat {\boldsymbol {\beta}} (t) \cdot \boldsymbol {\beta} (t) + \ln \mathbb {E} \exp \left(- \int d t \hat {\boldsymbol {\beta}} (t) \cdot \boldsymbol {w} (t) a (t)\right) \tag {122}
$$
The necessary ingredients are
$$
\begin{array}{l} \boldsymbol {\kappa} (t, s) = \left\langle a (t) a (s) \boldsymbol {w} (t) \boldsymbol {w} (s) ^ {\top} \right\rangle - \gamma^ {2} \boldsymbol {\beta} (t) \boldsymbol {\beta} (s) \\ = \left\langle a (t) a (s) \right\rangle \left\langle \boldsymbol {w} (t) \boldsymbol {w} (s) ^ {\top} \right\rangle + \left\langle a (s) \boldsymbol {w} (t) \right\rangle \left\langle a (t) \boldsymbol {w} (s) ^ {\top} \right\rangle \in \mathbb {R} ^ {D \times D} \tag {123} \\ \end{array}
$$
Similarly we have to compute the sensitivity tensor
$$
\boldsymbol {D} (t, s) = \left\langle \frac {\partial}{\partial \boldsymbol {\beta} (s) ^ {\top}} a (t) \boldsymbol {w} (t) \right\rangle \in \mathbb {R} ^ {D \times D} \tag {124}
$$
We start from the dynamics
$$
\frac {d}{d t} \boldsymbol {w} (t) = \gamma a (t) \left(\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)\right), \frac {d}{d t} a (t) = \gamma \left(\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t)\right) \cdot \boldsymbol {w} (t) \tag {125}
$$
Next, we have to calculate causal derivatives for fields
$$
\begin{array}{l} \frac {\partial}{\partial \boldsymbol {\beta} (s) ^ {\top}} \boldsymbol {w} (t) = - \gamma \Theta (t - s) a (s) \mathbf {I} + \gamma \int_ {0} ^ {t} d t ^ {\prime} (\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} (t ^ {\prime})) \frac {\partial a (t ^ {\prime})}{\partial \boldsymbol {\beta} (s) ^ {\top}} \\ \frac {\partial}{\partial \boldsymbol {\beta} (s)} a (t) = - \gamma \Theta (t - s) \boldsymbol {w} (s) + \gamma \int_ {0} ^ {t} d t ^ {\prime} \left(\boldsymbol {\beta} _ {\star} - \boldsymbol {\beta} \left(t ^ {\prime}\right)\right) \cdot \frac {\partial \boldsymbol {w} \left(t ^ {\prime}\right)}{\partial \boldsymbol {\beta} (s)} \tag {126} \\ \end{array}
$$
Following an identical argument as in J, we see that  $D$  has block diagonal structure with  $D_{\beta_{\star}}(t,s)$  on the  $\beta_{\star}\beta_{\star}^{\top}$  direction and  $D_{\perp}(t,s)$  in any of the  $D - 1$  remaining directions
$$
D _ {\beta_ {\star}} (t, s) = \left\langle \frac {\partial}{\partial \beta (s)} a (t) w _ {\beta_ {\star}} (t) \right\rangle , D _ {\perp} (t, s) = \left\langle \frac {\partial}{\partial \beta_ {\perp} (s)} a (t) w _ {\perp} (t) \right\rangle \tag {127}
$$
Similarly,  $\kappa(t,s)$  has a similar decomposition
$$
\begin{array}{l} \kappa_ {\beta_ {*}} (t, s) = \langle a (t) a (s) \rangle \langle w _ {\beta_ {*}} (t) w _ {\beta_ {*}} (s) \rangle + \langle a (s) w _ {\beta_ {*}} (t) \rangle \langle a (t) w _ {\beta_ {*}} (s) \rangle \\ \kappa_ {\perp} (t, s) = \langle a (t) a (s) \rangle \langle w _ {\perp} (t) w _ {\perp} (s) \rangle + \langle a (s) w _ {\perp} (t) \rangle \langle a (t) w _ {\perp} (s) \rangle \tag {128} \\ \end{array}
$$
The processes have the following equations at infinite width
$$
\frac {d}{d t} w _ {\beta_ {\star}} (t) = \gamma a (t) (\beta_ {\star} - \beta (t)), \frac {d}{d t} a (t) = \gamma w _ {\beta_ {\star}} (t) (\beta_ {\star} - \beta (t)), \frac {d}{d t} w _ {\perp} (t) = 0 \tag {129}
$$
As a consequence we note that  $\langle w_{\perp}(t)a(s)\rangle = 0$  so that  $\kappa_{\perp}(t,s) = \langle a(t)a(s)\rangle$ . Letting  $v_{+}(t) = \frac{1}{\sqrt{2}} (w_{\beta_{\star}}(t) + a(t))$  and  $v_{-}(t) = \frac{1}{\sqrt{2}} (w_{\beta_{\star}}(t) + a(t))$ , we find the same decoupled stochastic processes as in Appendix I.3.
$$
\frac {d}{d t} v _ {+} (t) = \gamma \left(\beta_ {\star} - \beta (t)\right) v _ {+} (t), \frac {d}{d t} v _ {-} (t) = - \gamma \left(\beta_ {\star} - \beta (t)\right) v _ {-} (t) \tag {130}
$$
We can use these equations to perform the necessary averages for  $\kappa_{\beta_{\star}}$  and  $D_{\beta_{\star}}$ . Lastly, we use
$$
\frac {\partial}{\partial \beta_ {\perp} (s)} w _ {\perp} (t) = - \gamma \Theta (t - s) a (s) \tag {131}
$$
to evaluate  $D_{\perp}(t,s)$ . The observed covariances are just
$$
\boldsymbol {\Sigma} _ {\beta_ {\star}} = \left(\gamma \mathbf {I} - \boldsymbol {D} _ {\beta_ {\star}}\right) ^ {- 1} \boldsymbol {\kappa} _ {\beta_ {\star}} \left(\gamma \mathbf {I} - \boldsymbol {D} _ {\beta_ {\star}}\right) ^ {- 1 \top}, \boldsymbol {\Sigma} _ {\perp} = \left(\gamma \mathbf {I} - \boldsymbol {D} _ {\perp}\right) ^ {- 1} \boldsymbol {\kappa} _ {\perp} \left(\gamma \mathbf {I} - \boldsymbol {D} _ {\perp}\right) ^ {- 1 \top} \tag {132}
$$
We note that these expressions are identical to those in Appendix J under the substitution  $\beta_{\star} - \beta (t)\rightarrow$ $\Delta (t)$  and  $D\to P$  . Thus the expected test risk is
$$
\left\langle \left| \boldsymbol {\beta} (t) - \boldsymbol {\beta} _ {\star} \right| ^ {2} \right\rangle \sim (\beta (t) - \beta_ {\star}) ^ {2} + \frac {1}{N} \Sigma_ {\beta_ {\star}} (t, t) + \frac {(D - 1)}{N} \Sigma_ {\beta_ {\perp}} (t, t) + \mathcal {O} (N ^ {- 2}) \tag {133}
$$
This recovers the variance we obtained in the multiple-sample whitened data case J.
K.3 Connections to Offline Learning in Linear Model
Remark 1 The finite size variance of generalization error in an online learning setting with linear target function  $y = \beta^{*} \cdot x$  has an identical form as the model described above. In this setting, we sample infinitely many fresh data points  $\pmb{x} \sim \mathcal{N}(0, \mathbf{I})$  at each step leading to the flow  $\frac{d}{dt} \pmb{w}_i(t) = \gamma a_i(t) \mathbb{E}_{\pmb{x}} \Delta(\pmb{x}) \pmb{x}$  and  $\frac{d}{dt} a_i(t) = \gamma \pmb{w}_i(t) \cdot \mathbb{E}_{\pmb{x}} \Delta(\pmb{x}) \pmb{x}$ . The order parameter of interest in this setting is  $\beta(t) = \frac{1}{\gamma N} \sum_{i=1}^{N} \pmb{w}_i(t) a_i(t)$ . The precise correspondence between this setting and the offline setting is summarized in Table 2. We note that this argument could be extended to higher degree monomial activations as well, at the cost of tracking higher degree tensors (eg for quadratic activations  $M = \frac{1}{N} \sum_{i=1}^{N} a_i \pmb{w}_i \pmb{w}_i^\top \in \mathbb{R}^{D \times D}$  is sufficient).
Setting	Order Params.	Target	Off-target Dims.	Loss	Variance	Infinite Quantity
Offline	Δ = y - f	y	P - 1	Train	O(P/N)	D
Online	β* - β	β*	D - 1	Test	O(D/N)	P
Table 2: Summary of the equivalence between the leading  $1 / N$  correction in the offline setting and the online setting for two layer linear networks. In the offline training setting, the order parameters are the errors  $\pmb{\Delta} = \pmb{y} - \pmb{f} \in \mathbb{R}^{P}$  while in the online case they are  $\beta_{\star} - \beta \in \mathbb{R}^{D}$ .
As in the offline case, in Fig. 4 (c) and (d) we see that the variance contribution to test loss  $|\beta - \beta_{\star}|^2$  increases with input dimension  $D$ . We note that this perturbative effect to the loss dynamics is reminiscent of the deviations from mean field behavior studied in SGD [43, 44], though this present work concerns fluctuations driven by initialization variance rather than stochastic sampling of data. In Fig. 4 (e) we show that richer networks have lower variance at fixed  $N$ . Similarly, leading order theory for richer networks more accurately captures their dynamics as  $D / N$  increases (Fig. 4 (f)).
L Deep Linear Networks
For deep linear networks, the fields  $h_{\mu}^{\ell}(t), g_{\mu}^{\ell}(t)$  are Gaussian and have the following self-consistent equations
$$
h _ {\mu} ^ {\ell} (t) = u _ {\mu} ^ {\ell} (t) + \gamma \int_ {0} ^ {t} d s \sum_ {\nu} \left[ A _ {\mu \nu} ^ {\ell - 1} (t, s) + \Delta_ {\nu} (s) H _ {\mu \nu} ^ {\ell - 1} (t, s) \right] g _ {\nu} ^ {\ell} (s), u _ {\mu} ^ {\ell} (t) \sim \mathcal {G P} (0, H ^ {\ell - 1})
$$
$$
g _ {\mu} ^ {\ell} (t) = r _ {\mu} ^ {\ell} (t) + \gamma \int_ {0} ^ {t} d s \sum_ {\nu} \left[ B _ {\mu \nu} ^ {\ell} (t, s) + \Delta_ {\nu} (s) G _ {\mu \nu} ^ {\ell + 1} (t, s) \right] h _ {\nu} ^ {\ell} (s), r _ {\mu} ^ {\ell} (t) \sim \mathcal {G P} (0, \boldsymbol {G} ^ {\ell + 1}). \tag {134}
$$
where  $H_{\mu \nu}^{\ell}(t,s) = \left\langle h_{\mu}^{\ell}(t)h_{\nu}^{\ell}(s)\right\rangle$  and  $G_{\mu \nu}^{\ell}(t,s) = \left\langle g_{\mu}^{\ell}(t)g_{\nu}^{\ell}(s)\right\rangle$  and  $A_{\mu \nu}^{\ell}(t,s) = \left\langle \frac{\partial h_{\mu}^{\ell}(t)}{\partial r_{\nu}(s)}\right\rangle$  and  $B_{\mu \nu}^{\ell}(t,s) = \left\langle \frac{\partial h_{\mu}^{\ell}(t)}{\partial r_{\nu}(s)}\right\rangle$  [9]. Therefore, we express the action as a differentiable function of the order parameters by integrating over the Gaussian field distribution. For concreteness, we vectorize our fields over time and samples  $\pmb{h}^{\ell} = \mathrm{Vec}\{h_{\mu}^{\ell}(t)\}_{\{\mu \in [P], t \in \mathbb{R}_{+}\}}$ ,  $\pmb{g}^{\ell} = \mathrm{Vec}\{g_{\mu}^{\ell}(t)\}_{\{\mu \in [P], t \in \mathbb{R}_{+}\}}$  we consider the contribution of a single hidden layer.
$$
\mathcal {Z} _ {\ell} = \int d \hat {\boldsymbol {h}} ^ {\ell} d \hat {\boldsymbol {g}} ^ {\ell} d \boldsymbol {h} ^ {\ell} d \boldsymbol {g} ^ {\ell} \exp \left(- \frac {1}{2} \hat {\boldsymbol {h}} ^ {\ell} \boldsymbol {\Sigma} _ {u} \hat {\boldsymbol {h}} ^ {\ell} + i \hat {\boldsymbol {h}} ^ {\ell} \cdot (\boldsymbol {h} ^ {\ell} - \boldsymbol {C} ^ {\ell} \boldsymbol {g} ^ {\ell}) - \frac {1}{2} \boldsymbol {h} ^ {\ell \top} \hat {\boldsymbol {H}} ^ {\ell} \boldsymbol {h} ^ {\ell}\right)
$$
$$
\exp \left(- \frac {1}{2} \hat {\boldsymbol {g}} ^ {\ell} \boldsymbol {\Sigma} _ {r} ^ {\ell} \hat {\boldsymbol {g}} ^ {\ell} + i \hat {\boldsymbol {g}} ^ {\ell} \cdot (\boldsymbol {g} ^ {\ell} - \boldsymbol {D} ^ {\ell} \boldsymbol {h} ^ {\ell}) - \frac {1}{2} \boldsymbol {g} ^ {\ell \top} \hat {\boldsymbol {G}} ^ {\ell} \boldsymbol {g} ^ {\ell}\right)
$$
where  $C_{\mu \nu}^{\ell}(t,s) = \gamma \Theta (t - s)\left[A_{\mu \nu}^{\ell -1}(t,s) + H_{\mu \nu}^{\ell -1}(t,s)\Delta_{\nu}(s)\right]$  and  $D_{\mu \nu}^{\ell}(t,s) = \gamma \Theta (t - s)\left[B_{\mu \nu}^{\ell}(t,s) + G_{\mu \nu}^{\ell +1}(t,s)\Delta_{\nu}(s)\right]$ . Performing the joint Gaussian integrals over  $(\pmb{h}^{\ell},\pmb{g}^{\ell},\hat{\pmb{h}}^{\ell},\hat{\pmb{g}}^{\ell})$  we find
$$
\ln \mathcal {Z} _ {\ell} = - \frac {1}{2} \ln \det  \left[ \begin{array}{c c c c} - \hat {\boldsymbol {H}} ^ {\ell} & 0 & \mathbf {I} & - \boldsymbol {D} ^ {\ell \top} \\ 0 & - \hat {\boldsymbol {G}} ^ {\ell} & - \boldsymbol {C} ^ {\ell \top} & \mathbf {I} \\ \mathbf {I} & - \boldsymbol {C} ^ {\ell} & \boldsymbol {\Sigma} _ {u} & 0 \\ - \boldsymbol {D} ^ {\ell} & \mathbf {I} & 0 & \boldsymbol {\Sigma} _ {r} \end{array} \right] \tag {135}
$$
We can then automatically differentiate the DMFT action to get the propagator. For example, for a three layer linear network, the full DMFT action has the form
$$
\begin{array}{l} S = \frac {1}{2} \operatorname {T r} \left[ \hat {\boldsymbol {H}} ^ {1} \boldsymbol {H} ^ {1} + \hat {\boldsymbol {H}} ^ {2} \boldsymbol {H} ^ {2} + \hat {\boldsymbol {G}} ^ {1} \boldsymbol {G} ^ {1} + \hat {\boldsymbol {G}} ^ {2} \boldsymbol {G} ^ {2} \right] - \gamma^ {2} \operatorname {T r} \boldsymbol {A B} \\ - \frac {1}{2} \ln \det  \left[ \begin{array}{c c c c} - \hat {\boldsymbol {H}} ^ {1} & 0 & \mathbf {I} & - \boldsymbol {D} ^ {1 ^ {\top}} \\ 0 & - \hat {\boldsymbol {G}} ^ {1} & - \boldsymbol {C} ^ {1 ^ {\top}} & \mathbf {I} \\ \mathbf {I} & - \boldsymbol {C} ^ {1} & \mathbf {1 1} ^ {\top} & 0 \\ - \boldsymbol {D} ^ {1} & \mathbf {I} & 0 & \boldsymbol {G} ^ {2} \end{array} \right] \\ - \frac {1}{2} \ln \det  \left[ \begin{array}{c c c c} - \hat {\boldsymbol {H}} ^ {2} & 0 & \mathbf {I} & - \boldsymbol {D} ^ {2 ^ {\top}} \\ 0 & - \hat {\boldsymbol {G}} ^ {2} & - \boldsymbol {C} ^ {2 ^ {\top}} & \mathbf {I} \\ \mathbf {I} & - \boldsymbol {C} ^ {2} & \boldsymbol {H} ^ {1} & 0 \\ - \boldsymbol {D} ^ {2} & \mathbf {I} & 0 & \mathbf {1 1} ^ {\top} \end{array} \right] \tag {136} \\ \end{array}
$$
where  $C^1 = \gamma \Theta_{\Delta}$  and  $C^2 = \gamma \Theta_{\Delta} \odot H^1 + \gamma A$  and  $D^1 = \gamma \Theta_{\Delta} \odot G^2 + \gamma B$  and  $D^2 = \gamma \Theta_{\Delta}$ . This above example can be extended to deeper networks. The total size of the block matrices which we compute determinants over is  $4PT \times 4PT$  for a dataset of size  $P$  trained for  $T$  steps.
M Discrete Time Dynamics and Edge of Stability Effects
Large step size effects can induce qualitatively different dynamics in neural network training. For instance, if the step size exceeds that required for linear stability with the initial kernel, the kernel can decrease in order to stabilize the dynamics [57]. Alternatively, during training the kernel may exhibit a "progressive sharpening" phase where its top eigenvalue grows before reaching a stability bound set by the learning rate [19]. It is therefore well motivated to study how dynamics in this regime alter finite size effects in neural networks. We will first solve a special model which was considered in prior work [57]: a two layer linear network trained on a single training point. We will then provide the full DMFT equations for the discrete time case and provide an outline for how one could obtain finite size effects in that picture.
M.1 Two Layer Linear Equations
In a two layer linear network, the DMFT equations are
$$
h (t + 1) = h (t) + \eta \gamma \Delta (t) z (t), z (t + 1) = z (t) + \eta \gamma \Delta (t) h (t)
$$
$$
f (t) = \frac {1}{\gamma} \langle z (t) h (t) \rangle \tag {137}
$$
The NTK has the form  $K(t) = \langle h(t)^2 + z(t)^2 \rangle$ . We can easily show that the kernel and error have coupled dynamics
$$
\begin{array}{l} f (t + 1) = f (t) + \eta \left\langle h (t) ^ {2} + z (t) ^ {2} \right\rangle \Delta (t) + \eta^ {2} \gamma \Delta (t) ^ {2} \left\langle h (t) z (t) \right\rangle \\ = f (t) + \eta K (t) \Delta (t) + \eta^ {2} \gamma^ {2} \Delta (t) ^ {2} f (t) \tag {138} \\ \end{array}
$$
$$
\begin{array}{l} K (t + 1) = K (t) + 4 \eta \gamma \Delta (t) \left\langle h (t) z (t) \right\rangle + \eta^ {2} \gamma^ {2} \Delta (t) ^ {2} \left\langle h (t) ^ {2} + z (t) ^ {2} \right\rangle \\ = K (t) + 4 \eta \gamma^ {2} \Delta (t) f (t) + \eta^ {2} \gamma^ {2} \Delta (t) ^ {2} K (t) \tag {139} \\ \end{array}
$$
These equations define the infinite width evolution of  $\Delta(t)$  and  $K(t)$ . Already at this level of analysis, we can reason about the evolution of  $K(t)$ . In the small  $\eta$  limit, we could disregard terms of order  $\mathcal{O}(\eta^2)$  and arrive at the following gradient flow approximation for  $K(t) \sim 2\sqrt{1 + \gamma^2 f(t)^2}$  [9]. This evolution will not reach the edge of stability provided that  $\eta < \frac{1}{\sqrt{1 + \gamma^2 y^2}}$ . For large  $\gamma$  and  $y = 1$ , this leads to the constraint  $\eta \gamma < 1$ . However, if  $\eta$  exceeds this bound, the gradient flow approximation is no longer reasonable and the system reaches an edge of stability effect as shown in Figure 6.
To calculate the finite size effects, we need to compute  $\kappa$  and  $D(t,s) = \frac{\partial}{\partial\Delta(s)}\left\langle h(t)^2 +z(t)^2\right\rangle$ . To evaluate these quantities we utilize the same change of variables employed in Appendix I.3. In discrete time, these decoupled equations are
$$
v _ {+} (t + 1) = v _ {+} (t) + \eta \gamma \Delta (t) v _ {+} (t), v _ {-} (t + 1) = v _ {-} (t) - \eta \gamma \Delta (t) v _ {-} (t). \tag {140}
$$
Given  $\Delta(t)$ , these can be expressed as linear systems of equations. Now, we can easily compute the uncoupled kernel variance
$$
\begin{array}{l} \kappa (t, s) = 2 \left\langle h (t) h (s) \right\rangle^ {2} + 2 \left\langle z (t) z (s) \right\rangle^ {2} + 2 \left\langle h (t) z (s) \right\rangle^ {2} + 2 \left\langle z (t) h (s) \right\rangle^ {2} \\ = \left\langle v _ {+} (t) v _ {+} (s) + v _ {-} (t) v _ {-} (s) \right\rangle^ {2} + \left\langle v _ {+} (t) v _ {+} (s) - v _ {-} (t) v _ {-} (s) \right\rangle^ {2}. \tag {141} \\ \end{array}
$$
Similarly, we can calculate  $D(t,s)$  by using the fact  $\langle h(t)^2 + z(t)^2 \rangle = \langle v_+(t)^2 + v_-(t)^2 \rangle$ .
$$
D (t, s) = 2 \left\langle v _ {+} (t) \frac {\partial v _ {+} (t)}{\partial \Delta (s)} \right\rangle + 2 \left\langle v _ {-} (t) \frac {\partial v _ {-} (t)}{\partial \Delta (s)} \right\rangle
$$
$$
\frac {\partial v _ {+} (t)}{\partial \Delta (s)} = \gamma \Theta (t - s) v _ {+} (s) + \sum_ {t ^ {\prime} <   t} \Delta (t ^ {\prime}) \frac {\partial v _ {+} (t ^ {\prime})}{\partial \Delta (s)}
$$
$$
\frac {\partial v _ {-} (t)}{\partial \Delta (s)} = - \gamma \Theta (t - s) v _ {-} (s) - \sum_ {t ^ {\prime} <   t} \Delta \left(t ^ {\prime}\right) \frac {\partial v _ {-} \left(t ^ {\prime}\right)}{\partial \Delta (s)} \tag {142}
$$
These can be directly solved as a linear system of equations.
N Computing Details
Experiments for Figures 3, 6 and 2 were conducted on a Google Colab GPU with JAX. Experiments for Figures 5, A.3, 7 were performed on a NVIDIA SMX4-A100-80GB GPU. The total compute required for all Figures in the paper took around 4 hours. Jupyter Notebooks to reproduce plots can be found at https://github.com/Pehlevan-Group/dmft_fluctuations.