Theoretical analysis: link between curvature and gradient approximation error
We discussed how generalization can be linked to curvature, and we now examine the link between curvature and gradient alignment during learning dynamics. We represent approximate gradients in a rule-agnostic manner, where an arbitrary approximation is represented in terms of its component along the gradient direction plus an arbitrary orthogonal vector (Figure 4 A):

\vec{\hat{g}}=\rho\vec{g}+\vec{e}, (6)

where \vec{g}:=\nabla L(W_{h})\in\mathbb{R}^{N^{2}} and \vec{\hat{g}}\in\mathbb{R}^{N^{2}} are the exact and approximate gradients, respectively (we reshaped \nabla L(W_{h}) into a vector here); \vec{e}\in\mathbb{R}^{N^{2}} is an arbitrary orthogonal vector to \vec{g} . Here, scalar \rho\in\mathbb{R} represents the relative step length along the gradient direction that the approximate rule is making. As we will see in Theorem 1 , \rho is an important quantity in our analysis. One can easily compute \rho from \vec{g} and \vec{\hat{g}} by \rho=\frac{\vec{\hat{g}}^{T}\vec{g}}{\vec{g}^{T}\vec{g}} .

We express weight updates as discrete dynamical systems (with weights W_{h} as the state variables):

 W^{+}_{h} \leftarrow W^{-}_{h}+F(W^{-}_{h})=W^{-}_{h}-\eta abla L(W^{-}_{h}), for BPTT (7) W^{+}_{h} \leftarrow W^{-}_{h}+\hat{F}(W^{-}_{h})=W^{-}_{h}-\eta\tilde{ abla}L(W^{-}_{h}), for an approximate rule, (8)

where \eta is the learning rate, \tilde{\nabla} is an approximate gradient, and F:\mathbb{R}^{N^{2}}\rightarrow\mathbb{R}^{N^{2}} (resp. \hat{F}:\mathbb{R}^{N^{2}}\rightarrow\mathbb{R}^{N^{2}} ) denotes the map defined by BPTT (resp. an approximate) weight update rule. Notation W^{+} and W^{-} denotes W at the next and current step, respectively. We note in passing that dynamical systems view of weight updates have been used previously [ 141 , 142 ] .

We introduce additional notations before presenting Theorem 1 . J\in\mathbb{R}^{N^{2}\times N^{2}} (resp. \widehat{J}\in\mathbb{R}^{N^{2}\times N^{2}} ) is the Jacobian of the dynamical system of BPTT (resp. an approximate rule) in Eq. 7 (resp. Eq. 8 ). \lambda^{J}_{1}\in\mathbb{R} (resp. \widehat{\lambda_{1}^{J}}\in\mathbb{R} ) is the leading eigenvalue for BPTT (resp. an approximate rule) Jacobian. \lambda^{H}_{1}\in\mathbb{R} is the leading eigenvalue of the loss’ Hessian matrix. W^{*}_{B}\in\mathbb{R}^{N\times N} (resp. W^{*}_{e}\in\mathbb{R}^{N\times N} ) is the final fixed point for BPTT (resp. an approximate rule). We now present Theorem 1 .

Consider an RNN defined in Eq. 1 with a single scalar output \hat{y} and least squares loss as in Eq. 3 presented only at the last time step T , and weights are updated according to the difference equation for BPTT 7 (resp. an approximate rule 8 ) on a single example (stochastic gradient descent) using learning rate \eta_{B} (resp. \eta_{e} ). In the limit of stable fixed point convergence with zero training error, the dominant loss’ Hessian eigenvalue attained by BPTT (resp. approximate rule) is bounded by |\lambda^{H}_{1}(W^{*}_{B})|<\frac{1}{\eta_{B}} (resp. |\lambda^{H}_{1}(W^{*}_{e})|<\frac{1}{|\rho|\eta_{e}}) .

Full proof is in Appendix B . Here is a summary of the main steps involved:

The Jacobian of BPTT dynamical system (Eq. 7 ) is the loss’ Hessian scaled by a constant;

Using the above relationship, we can bound |\lambda^{H}_{1}| from |\lambda^{J}_{1}|<1 , which is the condition for a discrete-time dynamical system to converge to a fixed point

However, the link between the Jacobian of an approximate update rule and the loss’ Hessian is less obvious. Thus, we derive a link between |\widehat{\lambda_{1}^{J}}| and |\lambda^{J}_{1}| , and then apply the step above

∎

The consequence of the upper bound derived in Theorem 1 is that truncated gradient rules can converge to minima with a higher dominant Hessian eigenvalue than BPTT, with the leading eigenvalue bound inversely proportional to |\rho| ( |\rho|<1 usually). In practice, \rho can vary depending on task settings and in our setup, we observed it to be somewhere between 0.02 and 0.3 .
We remark that this higher upper bound is consistent with the increased spread of curvature and generalization observed for bio-plausible rules in experiments.
Theorem 1 highlights scalar \rho (Eq. 6 ), relative step length along the gradient direction, as an important factor in the curvature bound. To test that, we eliminate the factor of \rho by reducing the learning rule of BPTT such that its step length is matched to that of the three-factor rule. This resulted in the blue curves in Figure 4 , which is still trained using BPTT but with the update scaled by a factor of \rho . By matching the step length of BPTT and a three-factor rule along the gradient direction, similar convergence behaviors were observed. Similar observations were also made when the matching step experiment was repeated at three times the learning rate for all rules (Appendix Figure 12 C). This result then attributes the curvature preference behavior to relative step length along the gradient direction, and thereby indicating a link between curvature and gradient alignment under certain conditions . Consistent with earlier results, when the step length of BPTT is matched to that of a three-factor rule, its generalization performance also worsened (Appendix Figure 6 ).

We make a few more remarks regarding Theorem 1 . \rho is usually less than 1 because otherwise the additional orthogonal component \vec{e} would imply a larger update size for the approximate rule compared to BPTT. This would make the approximate rule more prone to numerical instabilities (Eq. 9 ) and would require a change in solver hyperparameters such as overall learning rate which in turn, influence the update size. In our experiments, when the three-factor rule was scaled up to match the along-gradient update sizes of BPTT, we quickly ran into numerical overflow (values of NaNs in the network), which is expected for a very large learning rate. The exact point for when this numerical instability is reached depends on many factors such as the model, task, numerical precision as well as the consistency of update direction. We have also equated numerical instabilities in simulations as a proxy for situations problematic for the brain in Discussion. Curiously, \vec{e} did not seem to play a helpful role in finding flatter minima; this could be due to that approximation error \vec{e} is not well aligned with the sharpest directions (Appendix Figure 10 ). In fact, \vec{e} being orthogonal to the leading Hessian eigenvectors is a consequence of the assumptions behind Theorem 1 . Despite making assumptions including scalar output, MSE loss and loss available at the last step, our empirical results indicate that our conclusions also extend to other setups: vector output, cross-entropy loss and loss that accumulates over time steps (Figure 4 ). In Appendix B.3 , we discuss the generality of Theorem 1 by using quadratic expansion of the loss function and assuming that \vec{e} is orthogonal to the leading Hessian eigenvectors.

One might assume increasing the learning rate of the approximate rules could compensate for the reduced step length due to \rho . However, this can raise other issues. Suppose \Delta W=-\eta\vec{g} for the exact gradient and we increase the learning rate of the approximate update until \widehat{\Delta W}=-\eta\vec{g}-\eta\vec{e} (magnitude of \vec{e} is scaled accordingly), and because \vec{g}\perp\vec{e} :

\|\widehat{\Delta W}\|^{2}=\eta^{2}\|\vec{g}\|^{2}+\eta^{2}\|\vec{e}\|^{2}>\eta^{2}\|g\|^{2}=\|\Delta W\|^{2} (9)

In other words, if approximate updates \widehat{\Delta W} make the same amount of progress as BPTT along the gradient direction, \widehat{\Delta W} would have a larger magnitude due to the orthogonal component \vec{e} , and large update magnitude can be very problematic for numerical stability [ 143 ] . Thus, the magnitude of \vec{e} limits the learning rate that can be used. Because of the numerical issues associated with increasing the learning rate for approximate rules (due to \vec{e} ), the differences in generalization and curvature convergence between rules cannot be reduced by increasing the learning rate for approximate rules. To balance between this numerical stability issue and the potential benefit of large learning rates, one can consider using a large learning rate early in training to prevent premature stabilization in sharp minima followed by gradual decay to mitigate the stability issue (see Appendix Figure 5 ).

Taken together, Theorem 1 and Eq. 9 connect gradient alignment with the curvature of the converged solution under certain conditions. For an approximation that is aligned poorly with the exact gradient, large orthogonal approximation error vector \vec{e} limit the step length for numerical stability reasons (Eq. 9 ) and small relative step length \rho correspond to a larger curvature bound (Theorem 1 ).