Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning Rules

Abstract

Learning rules—prescriptions for updating model parameters to improve performance—are typically assumed rather than derived. Why do some learning rules work better than others, and under what assumptions can a given rule be considered optimal? We propose a theoretical framework that casts learning rules as policies for navigating (partially observable) loss landscapes, and identifies optimal rules as solutions to an associated optimal control problem. A range of well-known rules emerge naturally within this framework under different assumptions: gradient descent from short-horizon optimization, momentum from longer-horizon planning, natural gradients from accounting for parameter space geometry, non-gradient rules from partial controllability, and adaptive optimizers like Adam from online Bayesian inference of loss landscape shape. We further show that continual learning strategies like weight resetting can be understood as optimal responses to task uncertainty. By unifying these phenomena under a single objective, our framework clarifies the computational structure of learning and offers a principled foundation for designing adaptive algorithms.

1. Introduction

A central concern in machine learning is identifying parameters that optimize model performance. Because directly searching for optimal parameters (e.g., via a grid search) is prohibitively costly in the high-dimensional parameter spaces characteristic of models based on artificial neural networks, optimization typically involves seeking iterative improvements in performance rather than directly searching for a global optimum. Procedures for iterative parameter improvement, or learning rules, are most commonly some variant of gradient descent [1, 2], with the backpropagation algorithm [3, 4] being a notable example. In biological neural networks, the plausibility of gradient descent is hotly debated [5, 6], and alternative rules that do not follow gradients have been proposed [7, 8].

Variants of gradient descent can largely be classified in terms of the presence or absence of three elements: momentum [9, 10], an adaptive learning rate [11-13], and loss approximation. From a geometric perspective, gradient descent corresponds to moving down the steepest part of the local loss landscape. In this view, momentum helps ensure smooth parameter changes, even when the loss landscape changes abruptly; an adaptive learning rate allows parameters to change more quickly when gradients are steady, or equivalently in regions where the loss landscape is flat; and computing the loss or its gradients approximately, for example over mini-batches, both improves efficiency and may add noise useful for generalization [14-16].

Why prefer one learning rule over another? Is a variant of gradient descent always optimal? If so, which one? If not, by what criteria do we construct or decide on an alternative? These questions are usually answered empirically, with optimizers like Adam [13] popular because they have proven performant in a wide variety of contexts [17-20]. It would be helpful to have a normative framework for answering them in a principled fashion, which, given a set of assumptions, identifies some learning rule as ‘optimal’. In this paper, our aim is to provide such a unifying framework.

Our three key insights are as follows. First, one can improve performance by optimizing not just over the next small parameter update, but over the whole sequence of future updates; this allows optimization to be less ‘myopic’ and more ‘farsighted’. Second, optimal learning dynamics ought to be sensitive to structure in parameter space. This insight is related to, but slightly more general than, the line of thought that leads to natural gradient descent [21]. Third, one can view the loss landscape as being only partially observable, which implies optimal learning ought to depend on beliefs about the loss landscape. Assuming partial observability is a useful way to model the fact that the training loss is typically only a proxy for the test loss, which is the true optimization target. We show this idea naturally yields adaptive optimizers like Adam, which update parameters using inferred loss shape.

Our framework is significant for two reasons. First, it makes it easier to generate new learning rules from a principled starting point, and to justify them without empirical guesswork. Second, it helps clarify which features of existing learning rules are essential for performance, and which are incidental. In the following sections, we discuss in more detail how different classes of well-known learning rules—including gradient descent with momentum (Sec. 3), natural gradient descent (Sec. 4), and rules with adaptive learning rates (Sec. 5)—can all be derived from our framework. Finally, Sec. 6 uses our framework to justify recently-identified rules for continual learning [22].

2. Mathematical formulation: learning rules as loss landscape navigation

Gradient descent and Newton’s method as optimal single steps through parameter space.

To motivate our framework, it is useful to observe that gradient descent and Newton’s method (a second-order analogue [23]) minimize a certain objective. Given a loss $ℒ(\theta )$ that depends on a parameter vector $\theta$, and a local first- or second-order approximation of it, we would like an update $\Delta \theta$ that decreases the loss as much as possible. Since large enough updates would invalidate our local loss approximation, we want to do this subject to the constraint that $\Vert \Delta \theta \Vert$ is not too large. We can do this by minimizing a combination of the loss and a step-size-related regularization term:

$$J(\Delta \theta )=\frac{{\Vert \Delta \theta \Vert }^2}{2\eta }+ℒ(\theta +\Delta \theta )\approx \{\begin{array}{c}\frac{{\Vert \Delta \theta \Vert }^2}{2\eta }+L(\theta )+{[{\nabla }_{\theta }ℒ(\theta )]}^T(\Delta \theta )\hfill \\ \frac{{\Vert \Delta \theta \Vert }^2}{2\eta }+ℒ(\theta )+{[{\nabla }_{\theta }ℒ(\theta )]}^T(\Delta \theta )+\frac{1}{2}{(\Delta \theta )}^T\mathit{H}(\theta )(\Delta \theta )\hfill \end{array}\phantom{\}}$$

where $\eta >0$ is the learning rate (which determines the size of the ‘trust region’), and $\mathit{H}(\theta )$ is the Hessian of $ℒ$ at $\theta$. Given the first-order loss approximation, the solution is $\Delta \theta =-\eta {\nabla }_{\theta }ℒ(\theta )$, i.e., gradient descent with a learning rate $\eta$. Geometrically, this step is down the ‘steepest’ part of the loss landscape near $\theta$. Given the second-order loss approximation, the solution is $\Delta \theta =-{[\frac{1}{\eta }\mathit{I}+\mathit{H}(\theta )]}^{-1}{\nabla }_{\theta }ℒ(\theta )$, i.e., gradient descent with a curvature-sensitive effective learning rate ${\eta }_{\text{eff}}≔{[\frac{1}{\eta }\mathit{I}+\mathit{H}(\theta )]}^{-1}$. One still takes a step down the steepest part of the loss landscape, but takes a larger step if the landscape is very flat; this avoids the slowdown gradient descent faces near a local minimum, where gradients get smaller. Newton’s method technically only corresponds to the large $\eta$ limit, where ${\eta }_{\text{eff}}=\mathit{H}{(\theta )}^{-1}$; more generally, we obtain a regularized version, which (given a specific Hessian approximation) is sometimes called the Levenberg-Marquardt algorithm [24].

Learning rules as partially observable loss landscape navigation.

The single-step view of loss optimization is arguably myopic, since the best short-term improvement of the loss may be suboptimal in the long term; famously, gradient descent finds local rather than global minima. An obvious way to address this issue is to instead optimize over multiple steps, and ask what sequence of steps should be taken in order to minimize the loss (Fig. 1a, b). In going from single- to multi-step optimization, we convert the problem of deciding on parameter updates into a navigation problem: what path through the loss landscape minimizes the loss, while simultaneously avoiding overly fast parameter changes?

Figure 1: Basic idea of our framework and a simple example.

a. Single-step approaches (top) optimize over short-term changes to the loss, while a multi-step approach (bottom) optimizes over longer-term changes. b. Gradient descent vs multi-step optimization for a double-well loss, with the optimal multi-step trajectory computed by directly minimizing the objective (see Appendix B for details). Note that the multi-step rule converges to the global rather than local minimum. c. Values of the kinetic (top) and potential (bottom) terms along the optimal trajectory from (b). The loss/potential does not decrease monotonically, since the learner must first escape a local minimum.

To formalize this, we define a continuous-time optimal control problem [25] over learning trajectories:

$$\begin{gathered} J(\{{\theta }_t\})={\mathbb{E}}_{\text{loss landscape belief}}\{[\text{parameter change cost}]+[\text{loss cost}]\} \\ ={\mathbb{E}}_{\{{\widehat{ℒ}}_t\}}\left\{{\int }_0^{\infty }\left(\frac{1}{2\eta }{[{\stackrel{.}{\theta }}_t-\mathit{f}({\theta }_t)]}^T\mathit{G}({\theta }_t)[{\stackrel{.}{\theta }}_t-\mathit{f}({\theta }_t)]+k{\widehat{ℒ}}_t({\theta }_t)\right)e^{-\gamma t}dt\right\} \end{gathered}$$ (1)

where $\eta >0$ is the learning rate, $k>0$ weights the influence of the loss, and $\gamma \ge 0$ is the temporal discounting rate. This objective generalizes the single-step objective (see Appendix A) and defines a continuous-time reinforcement learning [26, 27] and control [25] problem. It penalizes a combination of abrupt parameter changes (through the first term) and high loss (through the second), while discounting temporally distant costs through the factor $e^{-\gamma t}$. This objective also effectively turns learning into a physics problem [28], with these two terms analogous to kinetic and potential energy, and $\mathit{G}$ (the parameter space metric) and $\mathit{f}$ (the ‘drift’ or ‘bias’) encoding assumptions about the ambient (parameter space) geometry. Note: the potential/loss term may not decrease monotonically, for example if optimal dynamics involves exiting a local minimum of the loss (Fig. 1c).

Why partial observability?

In most practical settings, the learner does not have full knowledge of the loss landscape. We model this by assuming the learner possesses a structured model ${\widehat{ℒ}}_t$ of $ℒ$ that can evolve in time, which is informed by past observations (e.g., of gradients and curvature). This partial observability assumption is realistic for three reasons. First, due to the sheer size of parameter space, it is difficult to plan using more than a small part of the loss landscape at any given time. Second, the training loss is generally different from the test loss, but the two are not unrelated; it can be helpful to view the training loss as a noise-corrupted version of the test loss. Finally, the loss is usually evaluated on a batch of examples rather than the full training set. While these constraints make it difficult to plan an entire trajectory in one shot, one can instead work iteratively: take a step in the optimal direction, re-estimate your belief given new observations, and repeat.

Deriving learning rules.

Our objective (Eq. 1) takes as input four pieces of data: a loss function ${ℒ}_t$, a parameter space metric $\mathit{G}$, a drift term $\mathit{f}$, and a random variable ${\widehat{ℒ}}_t$, which models the learner’s belief about the loss landscape. This belief can evolve in time, e.g., via online Bayesian inference. Given this data, an optimal learning trajectory is a choice of ${\{{\theta }_t\}}_{t\in [0,\infty )}$ that minimizes $J$. Given an initial state ${\theta }_0$ and a desired time step $\Delta t>0$, an optimal learning rule is the difference ${\theta }_{\Delta t}-{\theta }_0$ along the optimal trajectory, which is proportional to the initial velocity ${\stackrel{.}{\theta }}_0$ when $\Delta t$ is small. This defines a navigation policy and continuous-time analogue of the optimal ‘next step’. The naive way to estimate an optimal trajectory is to consider a parameterized family of trajectories, and then directly minimize $J$ with respect to those parameters; Strang et al. [29] use this type of approach to numerically solve classical mechanics problems, and we followed their approach to generate Fig. 1b, c (see Appendix B for details). But there is a powerful alternative approach: framing learning rule optimization in terms of continuous-time learning trajectories allows us to leverage a result from the calculus of variations [28, 30] to identify optimal dynamics with solutions of the Euler-Lagrange (EL) equations (see Appendix C). These second-order ODEs are often solvable, and in our case, yield familiar learning rules like gradient descent and Adam under different assumptions. See Appendix D for discussion of technical and conceptual subtleties related to boundary conditions.

3. Momentum as a generic consequence of multi-step trajectory optimization

Unlike gradient descent, learning rules with momentum [9, 31] are more ‘inertial’: instead of following the current gradient, one follows a weighted combination of current and past gradients. Like second-order methods, this change helps speed up movement through flat regions of the loss landscape. In this section, we show that momentum is a generic consequence of multi-step optimization, and derive first- and second-order learning rules with momentum using our framework.

Momentum as a generic consequence of multi-step optimization.

A straightforward, multi-step generalization of the objective we considered to justify gradient descent is (see Appendix E)

$$J(\{{\theta }_t\})={\int }_0^{\infty }\left(\frac{1}{2\eta }{\Vert {\stackrel{.}{\theta }}_t\Vert }^2+kℒ({\theta }_t)\right)e^{-\gamma t}dt.$$ (2)

This objective is the simplest version of Eq. 1 ($\mathit{G}=\mathit{I}$, $\mathit{f}\equiv 0$, and the loss is not approximated). It penalizes only two things: abrupt parameter changes, and the loss. Optimal learning dynamics follow the EL equations, which in this case take the simple form (see Appendix E)

$${\stackrel{.}{\theta }}_t={\mathit{p}}_t{\stackrel{.}{\mathit{p}}}_t=\gamma {\mathit{p}}_t+\eta k{\nabla }_{{\theta }_t}ℒ({\theta }_t)$$ (3)

where we defined² momentum as ${\mathit{p}}_t≔{\stackrel{.}{\theta }}_t$. These equations have two interesting consequences. First, we obtain momentum essentially for free, simply from going from one-step to multi-step optimization. Second, the temporal discounting rate $\gamma$ allows one to interpolate between a momentum-based rule and standard gradient descent, since in the ‘overdamped’ (large $\gamma$) limit these equations become

$${\stackrel{.}{\theta }}_t=-\frac{\eta k}{\gamma }{\nabla }_{{\theta }_t}ℒ({\theta }_t).$$ (4)

This objective has a straightforward analogy with a mechanical system [28]: a particle with a mass $1∕\eta$ moves in a potential $kℒ$ and experiences an amount of friction proportional to $\gamma$. When ‘friction’ is sufficiently high, dynamics become non-inertial, and dominated by the shape of the potential/loss.

Deriving first- and second-order learning rules with and without momentum.

We can derive first- and second-order learning rules by tweaking Eq. 2 to locally approximate $ℒ$ near ${\theta }_0$ as

$$\widehat{ℒ}({\theta }_t)≔ℒ({\theta }_0)+{\mathit{g}}^T({\theta }_t-{\theta }_0)+\frac{1}{2}{({\theta }_t-{\theta }_0)}^T\mathit{H}({\theta }_t-{\theta }_0)$$ (5)

where $\mathit{g}$ and $\mathit{H}$ are the gradient and Hessian of $ℒ$ at ${\theta }_0$. For any $\Delta t\ge 0$, we find (Appendix E) that

$${\theta }_{\Delta t}-{\theta }_0=-(\mathit{I}-e^{\frac{\gamma }{2}\Delta t-\sqrt{\frac{{\gamma }^2}{4}\mathit{I}+\eta k\mathit{H}}\Delta t}){\mathit{H}}^{-1}\mathit{g}$$ (6)

is the optimal update. See Fig. 2a for example ${\theta }_t$ traces and Fig. 2b for example loss traces. Importantly, we can derive three learning rules from this result. When the step size $\Delta t$ is sufficiently large, we recover the Hessian-conditioned rule (i.e., Newton’s method) that ‘jumps’ straight to the minimum at $-{\mathit{H}}^{-1}\mathit{g}$. When $\Delta t$ is somewhat smaller than the other characteristic time scales, we either get gradient descent (in the overdamped large $\gamma$ limit) or a ballistic learning rule:

$${\theta }_{\Delta t}-{\theta }_0=\{\begin{array}{cc}-(\mathit{I}-e^{\frac{\gamma }{2}\Delta t-\sqrt{\frac{{\gamma }^2}{4}\mathit{I}+\eta k\mathit{H}}\Delta t}){\mathit{H}}^{-1}\mathit{g}\hfill & \text{Newton}’\text{s method};\Delta t\text{large}\hfill \\ -\frac{\eta k}{\gamma }\mathit{g}\Delta t\hfill & \text{Gradient descent};\Delta t\text{small and}\gamma \gg 1\hfill \\ -\sqrt{\eta k}{\mathit{H}}^{-1∕2}\mathit{g}\Delta t\hfill & \text{Ballistic};\Delta t\text{small and}\gamma =0\hfill \end{array}\phantom{\}}.$$ (7)

Figure 2: Effect of modulating temporal discounting rate.

a. Example optimal ${\theta }_t$ traces for a 1D quadratic loss, assuming different values of the temporal discounting rate ($\gamma =0$, 1, 10). b. Loss over time given the ${\theta }_t$ from (a), same values of $\gamma$. Note that lower values of $\gamma$ (‘longer’ planning horizon) produce loss curves that converge more quickly. c. Shape of trajectories for different $\gamma$ given a 2D anisotropic loss. In the gradient-descent-like regime ($\gamma \gg 1$), ${\theta }_2$ converges much more quickly than ${\theta }_1$ due to the anisotropy. In the ballistic ($\gamma \approx 0$) regime, the difference in convergence rates is not as extreme. d. Ratio of convergence rates $r_i≔\sqrt{{\gamma }^2∕4+\eta k{H}_{ii}}-\gamma ∕2$ assuming a diagonal Hessian. In the gradient-descent-like regime, directions with four times as much curvature converge 4 times faster; in the ballistic regime, they only converge $\sqrt{4}=2$ times faster. e. An Adam-like implementation of the ballistic ($\gamma \approx 0$) rule was used to train a small multilayer perceptron (MLP) to classify MNIST digits. Left: loss over training, right: test set accuracy over training. f. Same as in (e), but for a small convolutional neural network (CNN) trained to classify CIFAR-10 images. The ballistic rule generally performs better than SGD (black), and similarly to or worse than Adam (red).

The last learning rule is not usually considered, but has a a square-root normalization similar to Adam [13]. We call such a rule ballistic because it corresponds to the frictionless limit; it strikes a compromise between following the gradient, as gradient descent does, and jumping straight to the local minimum of the loss, as Newton’s method does. Because it depends on the square root of the Hessian, differences in how quickly different parameter directions converge are reduced somewhat relative to gradient descent (Fig. 2c, d). A heuristic implementation of this rule performs well on standard datasets (Fig. 2e, f), which suggests its behavior may be reasonable even for non-quadratic losses. See Appendix B for experiment details.

4. Parameter space geometry modulates optimal learning dynamics

Parameter space geometry modulates distances.

An insight due to Amari [21] is that learning rules ought to be sensitive to the structure of parameter space. If distances in parameter space follow a non-Euclidean metric $\mathit{G}(\theta )$ (e.g., the Fisher information matrix), the objective we used to derive gradient descent must be modified to involve a parameter change penalty $\frac{1}{2\eta }{(\Delta \theta )}^T\mathit{G}(\theta )\Delta (\theta )$, which causes the optimal first-order (one-step) rule to become $\Delta \theta =-\eta {\mathit{G}}^{-1}(\theta ){\nabla }_{\theta }ℒ(\theta )$. We can implement Amari’s insight in the continuous-time, multi-step setting by adding a metric to Eq. 2:

$$J(\{{\theta }_t\})={\int }_0^{\infty }\left(\frac{1}{2\eta }{\stackrel{.}{\theta }}_t^T\mathit{G}({\theta }_t){\stackrel{.}{\theta }}_t+kℒ({\theta }_t)\right)e^{-\gamma t}dt.$$ (8)

The corresponding EL equations that describe optimal learning dynamics read (see Appendix F)

$${\stackrel{.}{\theta }}_t={\mathit{p}}_t{\stackrel{.}{\mathit{p}}}_t=\gamma {\mathit{p}}_t+\eta k{\mathit{G}}^{-1}({\theta }_t){\nabla }_{{\theta }_t}ℒ({\theta }_t)$$ (9)

and by an argument analogous to the one in the previous section, we recover three types of rules:

$${\theta }_{\Delta t}-{\theta }_0=\{\begin{array}{cc}-(\mathit{I}-e^{\frac{\gamma }{2}\Delta t-\sqrt{\frac{{\gamma }^2}{4}\mathit{I}+\eta k{\mathit{G}}^{-1}\mathit{H}}\Delta t}){\mathit{H}}^{-1}\mathit{g}\hfill & \text{Newton}’\text{s method};\Delta t\text{large}\hfill \\ -\frac{\eta k}{\gamma }{\mathit{G}}^{-1}\mathit{g}\Delta t\hfill & \text{Nat. gradient};\Delta t\text{small and}\gamma \gg 1\hfill \\ -\sqrt{\eta k}\sqrt{{\mathit{G}}^{-1}\mathit{H}}{\mathit{H}}^{-1}\mathit{g}\Delta t\hfill & \text{Nat. ballistic};\Delta t\text{small and}\gamma =0\hfill \end{array}\phantom{\}}$$ (10)

These follow from making the replacements $\mathit{g}\to {\mathit{G}}^{-1}\mathit{g}$ and $\mathit{H}\to {\mathit{G}}^{-1}\mathit{H}$ in a model that assumes $\mathit{G}$ is locally $\theta$-independent (i.e., here, $\mathit{G}$ refers to $\mathit{G}({\theta }_0)$). In addition to recovering natural gradient descent (middle), which modifies gradient descent to account for non-Euclidean parameter space geometry [21], we also recover two other learning rules. The first is Newton’s method with a metric- and Hessian-dependent learning rate, and the last is a ballistic method which looks like Newton’s method, but involves $\sqrt{{\mathit{G}}^{-1}\mathit{H}}{\mathit{H}}^{-1}$ instead of ${\mathit{H}}^{-1}$. This method compromises between using ${\mathit{G}}^{-1}$ and ${\mathit{H}}^{-1}$ as preconditioners, or equivalently between moving based on parameter space and loss landscape geometry. While prior work has attempted to combine natural gradients and momentum heuristically [32], our framework shows how these elements arise jointly from a principled objective. As is well-known, using an anisotropic $\mathit{G}$ can add anisotropy to learning dynamics (Fig. 3a, b).

Figure 3: Parameter space geometry affects optimal learning trajectories.

a. Optimal trajectory through ${\theta }_1\text{-}{\theta }_2$ space for an isotropic quadratic loss, assuming no nontrivial $\mathit{G}$ and $\mathit{f}\equiv 0$. The heatmap and contours show the value of the loss at each (${\theta }_1$, ${\theta }_2$) value. Black line: optimal trajectory, red dot: global minimum of loss. Note that, because the loss is isotropic, the optimal trajectory is too. b. Same as (a), but given a strongly anisotropic constant metric $\mathit{G}$. Note that the optimal trajectory is no longer the same along each direction, but converges much more quickly along the ${\theta }_1$ direction. c. Same as (a), but given $\mathit{f}$ that corresponds to purely rotational dynamics. Note two differences: it spirals about the origin, and no longer converges to the global minimum of the loss, but to a different point closer to the origin (orange dot). d. Same as (a), but given $\mathit{f}$ that corresponds to weight decay. There is no anisotropy, but the trajectory does not converge to the minimum of the loss.

Natural gradient descent is not second-order optimization in disguise.

It is often argued that natural gradient descent behaves like a second-order method, with the metric $\mathit{G}$ acting as a surrogate for the Hessian $\mathit{H}$. Martens [33] explores this view in detail, and shows that natural gradient descent is sometimes equivalent to a Generalized Gauss-Newton method. But Martens also notes a variety of problems with this view, like the fact that existing theory (e.g., convergence rates) portrays the Hessian as more performant, while empirical work shows natural gradient descent is more performant. Our results suggest that the analogy between second-order methods and natural gradient descent is misguided, and that $\mathit{G}$ and $\mathit{H}$ play fundamentally different roles: $\mathit{G}$ governs the parameter velocity penalty, while $\mathit{H}$ measures loss landscape curvature. Physically, the former defines ambient geometry (as in general relativity [34]), whereas the latter pushes and pulls particles along that geometry. It is also clear when one examines the different roles of $\mathit{G}$ and $\mathit{H}$ in the rules we derived above.

Optimizing partially controllable parameters.

Parameter space geometry can also influence what kind of learning rule is optimal in a different, less well-appreciated way: suppose an external ‘force’ pushes on the parameter in a state-dependent fashion, like a ‘wind’ that determines which directions are more or less difficult to travel in. This feature can arise from partial controllability; if our control of a parameter is partial, e.g., ${\stackrel{.}{\theta }}_t=\mathit{u}+\mathit{f}({\theta }_t)$, controlling the parameter via ${\mathit{u}}_t$ involves not just moving ${\theta }_t$ in a direction that improves the loss, but also fighting against the ‘default’ dynamics $\mathit{f}(\theta )$. Here, we show that this feature can produce non-gradient rules. If we add a drift term $\mathit{f}$ to Eq. 8,

$$J(\{{\theta }_t\})={\int }_0^{\infty }\left(\frac{1}{2\eta }{[{\stackrel{.}{\theta }}_t-\mathit{f}({\theta }_t)]}^T\mathit{G}({\theta }_t)[{\stackrel{.}{\theta }}_t-\mathit{f}({\theta }_t)]+kℒ({\theta }_t)\right)e^{-\gamma t}dt.$$ (11)

The presence of $\mathit{G}$ and $\mathit{f}$ can greatly complexify the EL equations, especially if they are state-dependent. Consider the effect of the drift $\mathit{f}$ by itself (i.e., assume $\mathit{G}=\mathit{I}$; see Appendix G):

$${\ddot{\theta }}_t-[\gamma \mathit{I}+\mathit{J}({\theta }_t)-\mathit{J}{({\theta }_t)}^T]{\stackrel{.}{\theta }}_t=[{\mathit{J}({\theta }_t)}^T-\gamma \mathit{I}]f({\theta }_t)+\eta k{\nabla }_{\theta }ℒ({\theta }_t)$$

where $\mathit{J}({\theta }_t)$ is the Jacobian of $\mathit{f}$ at ${\theta }_t$. One can interpret $\mathit{J}$ as contributing to dynamics in two ways: it contributes an effective ‘potential-like’ term $[\mathit{J}{({\theta }_t)}^T-\gamma \mathit{I}]\mathit{f}({\theta }_t)$, or equivalently an effective loss; and it produces an effective state-dependent discounting factor ${\gamma }_{\text{eff}}({\theta }_t)≔\gamma \mathit{I}+\mathit{J}({\theta }_t)-\mathit{J}{({\theta }_t)}^T$. Note that this reduces to the usual discounting factor if $\mathit{J}{(\theta )}^T=\mathit{J}(\theta )$ for all $\theta$, which is true if and only if $\mathit{f}$ is the gradient of some function.

In Appendix G, we consider two examples in more detail: the case where $\mathit{f}$ corresponds to rotational dynamics (Fig. 3c), and the case where $\mathit{f}$ corresponds to weight decay (Fig. 3d). The former case is ‘nonconservative’ (e.g., $\mathit{J}(\theta )\ne \mathit{J}{(\theta )}^T$) but the second is not. Interestingly, rotational dynamics can make ‘spiraling’ trajectories optimal, and both choices of $\mathit{f}$ make optimal trajectories converge to a point different from the global minimum of the loss.

Optimal learning dynamics are generally non-gradient.

One significant consequence of including a drift term $\mathit{f}$ is that it can cause the optimal learning rule to involve non-gradient dynamics, which cannot be described as gradient descent (with or without momentum) down any objective. In our $\mathit{G}=\mathit{I}$ example, this happens if and only if $\mathit{f}$ itself is not the gradient of any function (or equivalently, if $\mathit{J}(\theta )$ is not symmetric for all $\theta$). In terms of a Helmholtz decomposition [35] $\mathit{f}(\theta )={\nabla }_{\theta }V(\theta )+\mathit{R}(\theta )$, where $V$ is some non-unique ‘potential’ function and $\mathit{R}$ is divergence-free (i.e., ${\nabla }_{\theta }\cdot \mathit{R}=0$), $\mathit{R}$ is the interesting component of $\mathit{f}$. See Appendix G for more discussion of this point. This interpretation suggests a concrete diagnostic: when learning rules involve update components that do not point downhill, such as decay terms or rotational drift, they may be optimal responses to implicit background dynamics.

5. Adaptive optimizers arise from beliefs about how gradients evolve in time

Adaptive optimizers like Adam [13] and RMSprop [12] use a learning rate that depends on the variance of recent gradients. If recent gradients were consistent, learn quickly; if not, make smaller updates. Doing this is known to work well in practice, with Adam variants outperforming most other known optimizers [17-20] in typical settings, but it is unclear why, or under what assumptions Adam is optimal. In this section, we show that an Adam-like strategy is optimal under a certain Bayesian model of how the local loss landscape shape evolves in time. More generally, we show adaptive optimizers arise from using past observations of landscape shape to estimate current and future shape.

Modeling uncertainty in the time evolution of the local loss landscape.

One strategy for navigating a partially observable loss landscape is to maintain running estimates of quantities that characterize its local shape, like the gradient and curvature. Motivated by this idea, assume that the learner maintains a local model of loss landscape shape near their current location ${\theta }_0$, i.e.,

$$\mathbb{E}[\widehat{ℒ}({\theta }_t)]=ℒ({\theta }_0)+{\mathit{m}}_t^T({\theta }_t-{\theta }_0)+\frac{\kappa }{2}{({\theta }_t-{\theta }_0)}^T{\mathit{V}}_t({\theta }_t-{\theta }_0)$$ (12)

where the gradient estimate ${\mathit{m}}_t$ and Hessian estimate $\kappa {\mathit{V}}_t$ are assumed to be time-varying. Here, $\kappa >0$ is a fixed (known) scaling factor. Since ${\mathit{m}}_t$ and ${\mathit{V}}_t$ represent average values of gradients and curvature, respectively, they are not directly observable—but the learner can estimate them via their assumed link to observable gradients ${\mathit{g}}_t$. Assume that the learner has a Bayesian model of how local landscape geometry evolves with two components: an observation model, which connects ${\mathit{g}}_t$ to ${\mathit{m}}_t$ and ${\mathit{V}}_t$; and a prior belief about how landscape geometry evolves.

Motivated by a simple model of gradient drift and diffusion in a quadratic loss (see Appendix H), assume that ${\mathit{g}}_t\sim 𝒩({\mathit{m}}_t,{\mathit{V}}_t∕\Delta t)$, which implies that $\mathbb{E}[{\mathit{g}}_t]={\mathit{m}}_t$ and $\mathrm{Cov}({\mathit{g}}_t)\Delta t={\mathit{V}}_t$. To make inference simpler, we will use a crude approximation³ associated with method-of-moment-based strategies, and assume ${\mathit{g}}_t$ and ${({\mathit{g}}_t-{\mathit{m}}_t)({\mathit{g}}_t-{\mathit{m}}_t)}^T$ provide independent observations of ${\mathit{m}}_t$ and ${\mathit{V}}_t$, i.e., ${\mathit{g}}_t\sim 𝒩({\mathit{m}}_t,({\sigma }_1^2∕\Delta t)\mathit{I})$ and ${({\mathit{g}}_t-{\mathit{m}}_t)({\mathit{g}}_t-{\mathit{m}}_t)}^T\sim 𝒩({\mathit{V}}_t,({\sigma }_2^2∕\Delta )\mathit{I})$. Here, ${\sigma }_1^2>0$ and ${\sigma }_2^2>0$ are ‘observation noise’ parameters.

One reasonable prior, which assumes that both parameters tend towards zero and become more uncertain in the absence of observations, is an Ornstein-Uhlenbeck prior. Assume

$${\stackrel{.}{\mathit{m}}}_t=-{\alpha }_1{\mathit{m}}_t+{\xi }_1{\eta }_{1t}{\stackrel{.}{\mathit{V}}}_t=-{\alpha }_2{\mathit{V}}_t+{\xi }_2{\eta }_{2t}$$ (13)

where ${\alpha }_1\ge 0$ and ${\alpha }_2\ge 0$ parameterize decay, ${\xi }_1>0$ and ${\xi }_2>0$ parameterize the uncertainty growth rate, and ${\eta }_{1t}$ and ${\eta }_{2t}$ are Gaussian white noise terms. This means ${\mathit{m}}_{t+1}\sim 𝒩({\mathit{m}}_t-{\alpha }_1{\mathit{m}}_t\Delta t,{\xi }_1^2\Delta t)$ and ${\mathit{V}}_{t+1}\sim 𝒩({\mathit{V}}_t-{\alpha }_2{\mathit{V}}_t\Delta t,{\xi }_2^2\Delta t)$.

Time-evolving beliefs yield adaptive optimal learning rule.

Assume an objective that prioritizes a mix of small loss, smooth parameter changes, and good inference (implemented via a $\log p$ term):

$$J(\{{\theta }_t\})=\underset{\Delta t\to 0}{\text{lim}}{\int }_0^{\infty }\left(\frac{{\Vert {\theta }_t\Vert }^2}{2\eta }-\frac{\log p({\mathit{m}}_{t+1},{\mathit{V}}_{t+1}\mid {\mathit{g}}_t,{\mathit{m}}_t,{\mathit{V}}_t)}{\Delta t}+k\mathbb{E}[\widehat{ℒ}({\theta }_t)]\right)e^{-\gamma t}dt$$

where $p({\mathit{m}}_{t+1},{\mathit{V}}_{t+1}\mid {\mathit{g}}_t,{\mathit{m}}_t,{\mathit{V}}_t)$ is the learner’s posterior belief about local landscape shape dynamics. This objective has a well-defined $\Delta t\to 0$ limit (see Appendix H) with EL equations

$$\begin{gathered} {\ddot{\theta }}_t-\gamma {\stackrel{.}{\theta }}_t=\eta k[{\mathit{m}}_t+\kappa {\mathit{V}}_t({\theta }_t-{\theta }_0)] \\ {\ddot{\mathit{m}}}_t-\gamma ({\stackrel{.}{\mathit{m}}}_t+{\alpha }_1{\mathit{m}}_t)={\alpha }_1^2{\mathit{m}}_t+{\xi }_1^2\left[k({\theta }_t-{\theta }_0)+\frac{1}{{\sigma }_1^2}({\mathit{m}}_t-{\mathit{g}}_t)\right] \\ {\ddot{\mathit{V}}}_t-\gamma ({\stackrel{.}{\mathit{V}}}_t+{\alpha }_2{\mathit{V}}_t)={\alpha }_2^2{\mathit{V}}_t+{\xi }_2^2\left[k\frac{\kappa }{2}({\theta }_t-{\theta }_0){({\theta }_t-{\theta }_0)}^T+\frac{1}{{\sigma }_2^2}[{\mathit{V}}_t-({\mathit{g}}_t-{\mathit{m}}_t){({\mathit{g}}_t-{\mathit{m}}_t)}^T]\right]. \end{gathered}$$

These equations are somewhat complicated, but considerably simplify if one makes assumptions about the relative sizes of parameters. If we assume that parameter changes are ballistic ($\eta \gg \gamma$) but landscape beliefs change somewhat more slowly (${\xi }_1^2$, ${\xi }_2^2\ll \gamma$),

$$\begin{gathered} {\ddot{\theta }}_t=\eta k[{\mathit{m}}_t+\kappa {\mathit{V}}_t({\theta }_t-{\theta }_0)] \\ {\stackrel{.}{\mathit{m}}}_t=-{\alpha }_1{\mathit{m}}_t-\frac{{\xi }_1^2}{\gamma }\left[k({\theta }_t-{\theta }_0)+\frac{1}{{\sigma }_1^2}({\mathit{m}}_t-{\mathit{g}}_t)\right] \\ {\stackrel{.}{\mathit{V}}}_t=-{\alpha }_2{\mathit{V}}_t-\frac{{\xi }_2^2}{\gamma }\left[k\frac{\kappa }{2}({\theta }_t-{\theta }_0){({\theta }_t-{\theta }_0)}^T+\frac{1}{{\sigma }_2^2}[{\mathit{V}}_t-({\mathit{g}}_t-{\mathit{m}}_t){({\mathit{g}}_t-{\mathit{m}}_t)}^T]\right]. \end{gathered}$$

To good approximation, this means that the optimal learning rule has

$${\theta }_{t+\Delta t}-{\theta }_t=-\sqrt{\eta k}{\mathit{V}}_t^{-1∕2}{\mathit{m}}_t\Delta t.$$ (14)

Note that the effective learning rate goes like ${\mathit{V}}_t^{-1∕2}$ rather than ${\mathit{V}}_t^{-1}$. Our framework identifies the square root as a consequence of assuming highly inertial parameter changes, but gradient-descentlike landscape shape estimate dynamics. Moreover, contrary to other ideas about Adam [36], the square root is a feature rather than a bug, and need not be ‘fixed’. Other features of Adam, like estimating ${\mathit{V}}_t$ using the uncentered averages of ${\mathit{g}}_t{\mathit{g}}_t^T$ and approximating ${\mathit{V}}_t$ as diagonal, appear to be approximations that improve efficiency and scalability, as is usually believed.

6. Noisy continual learning strategies reflect parsimony and task uncertainty

In continual learning settings, agents must balance their ability to learn new tasks with their ability to retain information about previous tasks [37, 38]. Poorly performing agents exhibit catastrophic forgetting, but even in the absence of such forgetting, many algorithms exhibit a progressive loss in plasticity as an ever larger number of tasks are learned [22]. Dohare et al. observed that periodically resetting weights that do not strongly respond to task gradients empirically seems to ameliorate this issue. Others have observed that injecting noise in different ways, like via randomly perturbing little-used weights, also helps address this issue [39]. Thus far, it has been somewhat unclear why any of these strategies ought to work, and to what extent various details (e.g., how noise is injected) matter. Our framework provides qualitative guidance here: the reason different noise injection strategies empirically work is that injecting noise at all is more important than how that noise is injected.

Modeling weight-uncertainty-sensitive learning dynamics.

Assume that the learner uses distributional estimates, rather than point estimates, of the model parameters $\theta$. For simplicity, assume each ${\theta }_i$ is associated with a normal distribution $𝒩({\mu }_i,v_i)$, where ${\mu }_i$ denotes the mean estimate of ${\theta }_i$, and $v_i$ denotes the posterior variance. Instead of penalizing abrupt weight changes, in this setting it makes more sense to penalize abrupt changes in weight distribution:

$$J(\{{\theta }_t\})=\underset{\Delta t\to 0}{\text{lim}}{\int }_0^{\infty }\left[\frac{D_{KL}(p(\theta \mid {\mu }_{t+1},{\mathit{v}}_{t+1})\parallel p(\theta \mid {\mu }_t,{\mathit{v}}_t))}{{\eta (\Delta t)}^2}-\mathscr{H}(p(\theta \mid {\mu }_t,{\mathit{v}}_t))+kℒ({\mu }_t,{\mathit{v}}_t)\right]e^{-\gamma t}dt$$

where we have also included an entropy term to explicitly penalize ‘model complexity’. When written more explicitly, this objective reduces to our standard form (Eq. 1) with a nontrivial metric $\mathit{G}$ (see Appendix I). Given a local (quadratic) loss approximation with gradient $\mathit{g}$ and Hessian $\mathit{H}$, the EL equations read

$${\ddot{\mu }}_i-\left[\frac{{\stackrel{.}{v}}_i}{v_i}+\gamma \right]{\stackrel{.}{\mu }}_i=\eta k{v}_i\left[g_i+\sum _j{H}_{ij}({\mu }_j-{\mu }_{j0})\right]$$ (15)

$${\ddot{v}}_i-\gamma {\stackrel{.}{v}}_i=2\eta \left[\frac{k}{2}{H}_{ii}{v}_i^2-\frac{1}{2}{v}_i\right]+\frac{{\stackrel{.}{v}}_i^2}{v_i}-{\stackrel{.}{\mu }}_i^2.$$ (16)

The equation for ${\mu }_i$ (Eq. 15) has an interesting feature: it involves a variance-dependent effective discounting rate ${\gamma }_{\text{eff}}≔\gamma +\frac{{\stackrel{.}{v}}_i}{v_i}$ and effective learning rate ${\eta }_{\text{eff}}≔\eta v_i$. One can interpret the former as saying that planning should become more short-term when weight variances are changing quickly (i.e., ${\stackrel{.}{v}}_i∕v_i$ is high), and that learning should speed up when uncertainty (i.e., $v_i$) is high.

Variance reflects both parsimony and task-driven instability.

The equation for $v_i$ (Eq. 16) specifies how variance ought to increase or decrease along the optimal learning trajectory. The two terms on the right-hand side tell us how this happens: variance decreases if the nearby loss landscape is very curved (intuitively, the nearest minimum is easier to find); the entropic term incentivizes increases in variance, in order to favor the ‘simplest’ model associated with a given loss value; and the final term affects variance in a more subtle way. It increases variance if the mean is changing more quickly than the variance, something that can happen in continual learning settings when there is a transition between tasks. If the variance is changing quickly but the mean is not changing, it decreases variance somewhat. See Appendix I for more discussion.

The entropic term enforces behavior analogous to the weight resetting of Dohare et al. [22]: when the loss is not particularly sensitive to a given weight’s value (i.e., $H_{ii}$ is small), uncertainty about that weight should increase, with one possible mechanism being a reset. The last term on the right-hand side, which compares the speed of mean and variance changes, appears sensitive to task uncertainty. Together, these two effects—one tied to parsimony, the other to volatility—help explain the behavior of continual learning algorithms that inject noise or reset unused weights, like Dohare et al.’s method.

7. Discussion

We proposed a unifying framework that treats learning rules as solutions to an optimal control problem. By varying assumptions about geometry, planning horizon, and uncertainty, it recovers a wide range of familiar rules—including gradient descent, momentum, natural gradient descent, Adam, and continual learning strategies—from a single objective. This framework rests on three core ideas: learning unfolds over multiple steps, not one; parameter space can have nontrivial geometry and dynamics; and the loss landscape may only be partially observable, in which case it must be inferred. Our framework not only recovers these components individually, but also shows how they can be combined, yielding principled algorithms that integrate elements like momentum and natural gradients. It also suggests a shift in how learning rules should be evaluated: rather than focusing only on convergence rates or optimization guarantees, we should consider what assumptions a rule implicitly encodes and what problem it is actually solving (e.g., one related to generalization). This perspective is especially relevant for adaptive optimizers like Adam, which our framework portrays as optimizing an inferred test loss based on a specific Bayesian model of loss landscape shape dynamics. In many deep learning settings, performant optimizers ought to get things slightly ‘wrong’ to generalize [14, 40-42].

Connection to physics.

Our optimal control formulation of learning dynamics can be precisely mapped to a classical mechanics problem, which may allow physics tools to be adapted to study learning algorithms. For example, Noether’s theorem [43] implies that (quasi-) symmetries of objectives like Eq. 1 imply conserved quantities along optimal trajectories. Analogous ideas linking the consequences of symmetry to noisy recurrent dynamics [44], learning dynamics [45], and optimal dimensionality reduction [46] have already begun to be explored.

Biological relevance.

Neural circuits in the brain are thought to implement learning rules—such as Hebbian plasticity, homeostatic mechanisms, or spike-timing-dependent plasticity—that often lack a clear gradient-descent interpretation [7, 8, 47]. These rules may include time-asymmetric or rotational dynamics, operate under metabolic or architectural constraints, or adjust weights in response to activity thresholds rather than loss functions. Our framework suggests that such rules may still be optimal under constraints imposed by biological dynamics, such as synaptic decay, intrinsic drift, or limited access to global error signals [5, 48].

By treating learning as constrained control in a partially observable environment, our framework offers a normative lens on how non-gradient rules can potentially emerge as efficient strategies in biologically realistic regimes. This perspective is analogous to and consistent with observations that complex behavioral strategies can emerge from a mix of simple rewards and naturalistic constraints (e.g., partial observability), for example in foraging tasks [49].

Related work.

Several recent efforts aimed to unify learning rules under broader frameworks. Khan and Rue [50] portray various optimizers as natural gradient descent with respect to a fairly general objective, and Shoji et al. [51] similarly observe that many learning rules can be viewed as instances of natural gradient descent. However, both works take the idea that natural gradient descent is optimal for granted, and do not incorporate multi-step planning, belief updating, or partial observability.

Our framework has parallels with earlier foundational work by Wibisono et al. [52] which relates accelerated optimization methods to a continuous-time variational objective. However, there are also important qualitative differences between their framework and ours, with the most important being that their objective cannot be interpreted as a sum of costs, unlike ours. Our objective contains two terms—a quadratic parameter velocity penalty (or ‘kinetic energy’), and a loss term (or ‘potential energy’)—which are added together. In their objective, as in classical mechanics, the potential term is subtracted from the kinetic term. While this sign difference means that their approach is more directly related to classical mechanics, it also means that it is more distantly related to optimal control. A practical benefit of our sign choice is that optimal trajectories minimize the objective, rather than make it stationary in general.

Concurrent work by Orvieto and Gower [53] also proposes a view of Adam related to loss landscape shape inference, and also formalizes it in terms of an objective with a $-\log p$ term, but our picture and theirs differ in certain details. Perhaps the most important is that we link the appearance of the square root of the Hessian to operating in the ‘ballistic’ regime, or equivalently longer-term planning.

Our work is related in spirit to recent efforts to unify training objectives in deep learning, such as the framework proposed by Alshammari et al. [54].

Limitations.

We consider a setting in which our notion of ‘optimal’ does not factor in concerns which often affect optimization in practice, like memory requirements, computational simplicity, and efficiency. Relatedly, even if a rule is identified as optimal given our framework, it is unclear how useful it may be in practice, especially if it requires potentially expensive matrix operations. On the other hand, given that our formulation is in terms of an objective function, it may be possible to penalize things like efficiency explicitly in order to partially address this issue.

Finally, we do not claim to be able to explain all phenomena related to learning dynamics or learning rules; our goal in this work is merely to propose a useful framework for thinking about why features of well-known learning rules (like momentum) might be useful. The most important consequence of our framework is that, given an assumed objective (e.g., with a certain amount of temporal discounting, and a particular model of landscape shape belief updating), any two learning rules can be compared, and one or more learning rules can be shown to be optimal. Whether and how to link this framework to specific empirical circumstances is a different question which we expect to be more difficult.