Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning Rules
Authors: John Vastola
Samuel J. Gershman
Kanaka Rajan
Venue: NeurIPS 2025 poster ()
Keywords: gradient descent
optimization
momentum
natural gradient
Adam
learning rules
theory
optimal control
normative framework
variational principles
continual learning
Primary area: theory
TL;DR: We derive well-known learning rules from an objective that casts learning rules as policies for navigating uncertain loss landscapes

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.

Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning Rules
John J. Vastola 12.3
Samuel Gershman
Kanaka Rajan 1,3+
Department Neurobiology, Harvard Medical School
Department of Psychology and Center for Brain Science, Harvard University Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University (john_vastole janaka_rajes Ahhms harvard.edee gershman@fas harvard edu
Abstract
Learning rules- press-rresrritioss for updating model parameters improve performance- are typically assumed rather than derived. Why do some learn- ing rules work better than others, and under what assumptions can a given rule be considered optimal? We propose theoretical framework that casts learning rules as policies for navigating (partially observable) loss landscapes, and identi- fies optimal rules solutions to associated optimal control problem. A range well-known rules emerge naturally within this framework under different as- sumptions: gradient descent from short-horizon optimization, momentum from longer- honger-horizon planning, natural gradients from accounting for parameter space geometry, non -oon-rgndeent 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 be understood as optimal responses to task uncertainty. By unifying these phenomena under a single objective, framework clarifies the computational structure of learning and offers principled foundation for designing adaptive algorithms.
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 digh--imensional parameter spaces characteristic of models based on artificial neural networks, optimization typically involves seeking iterative improvements performance rather than directly searching for a global optimum. Procedures for iterative parameter improvement, or leurning rules, are most commonly some variant of gradient descent 1 2), with the backpropagation algorithm being a notable example. In biological neural networks, the plausibility of gradient descent is hotly debated
3 4
[), and alternative rules that do not follow gradients have been proposed [78]
Variants of gradient descent can largely be classified in terms of the presence or absence of three elements: momentum 9 [0], an adaptive learning rate [TH-3]] and loss approximation. From 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 flat; and computing the loss its gradients approximately, for example over mini-batches, both improves efficiency and may add noise useful for generalization [64
Conresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

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 wide variety of contexts [1] It would be helpful to have a normative framework for answering them in principled fashion, which, given 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 (2T]. 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 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 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 learning rules- includinn gradient descent with momentum (Sec. 3.. natural gradient descent (Sec. 4. and rules with adaptive learning rates (Sec. [ -an all be derived from our framework. Finally, Sec. uses our framework to justify recently-identifer 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, is useful to observe that gradient descent and Newton's 's method (a second- order analogue (23) minimize a certain objective. Given a loss [(0) that depends on a parameter vector O, and local first- or second-order approximation of it, we would like an update AO 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 || is not too large. We can do this by minimizing a combination of the loss and step-size-related regularization term:
2 27
|A0 2
+ [(0) + (oC(())
((20)=
C(O + A0) =
2n
|AP 2n
+ L()) + (Vo(()) + (())))A)))+(()))HHI))) ((AA))
where > 0 is the learning rate (which determines the size of the 'trust region'), ) and H()) is the Hessian of L at O. Given the first-order loss approximation, the solution is A0 (O),, i.e., gradient descent with a learning rate n) Geometrically, this step is down the 'steepest' part of the loss landscape near O Given the second-order loss approximation, the solution is AO -1 I + H (0)) VeC(O), i.c., gradient descent with a curvature- sewssiiieeessiiivveeeeeeeeeeeee effective learning rate Men := 1 + H() One still takes step down the steepest part of the loss landscape, but takes a larger step if the landscape 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 n) limit, where 7et = H (0) -1; more generally, we obtain a regularized version, which (given a specific Hessian approximation) is sometimes called the Levenberg-Marquare 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. D 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?

---

a. single-step b.
C.
|012
kinetic
transient increase inkintti
one-step
multi-step
20
grefer minimum
time
multi-step
potential E
tranclent increase
in potential Eolos
kC(
global
B2) - (08 1)2 81 2
8
minimum
local min
global min
Figure 1: Basic idea of our framework and simple example. Single-step approaches (top) optimize over short-term changes to the loss, while multi-step approach (bottom) optimizes over longer-term changes. b. Gradient descent vs multi-step optimization for double-well loss, with the optimal multi-step trajectory computed by directly minimizing the objective (see Appendix for details). Note that the multi-step rule converges to the global rather than local minimum. 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-ting time optimal control problem 25 over learning trajectories:
J((0,)) = EEooss landscappe helier
{ parameter change cost + [ loss cost }
= (c.) {%000
( 1 1. 2n
f(0)) + e - di
where 17 > 0 is the learning rate, k > 0 weights the influence of the loss, and 2 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 en This objective also effectively turns learning into a physics problem (28), with these two terms analogous kinetic and potential energy. and G (the parameter space metric) and 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. [b).
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 L. of L 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 the sheer size of parameter space, it is difficult to plan using more than small part of the loss landscape 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 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 entire trajectory in one shot, in the optimal direction, re-estimate your belief given new observations, and repeat.
any given time.
instead work iteratively: take step
Deriving learning rules. Our objective (Eq. D takes as input four pieces of data: loss function L,, a parameter space metric G. drift term f. and random variable Le, which models the leamerr' 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 (Ej), that minimizes J. Given an initial state On and a desired time step At > O, an optimal learning rule is the difference 0 Pac-00 Oo along the optimal trajectory, which is proportional to the initial velocity O when At is small. This defines navigation policy and continuous-tin 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 al. (29 use this type of approach to numerically solve classical mechanics problems, and we followed their approach to generate Fig. Dp. с (see Appendix B for details). But there a powerful alternative approach: framing learning rule optimization in terms of continuous-ting
analogue of the optimal 'next step'. The naive
learning trajectories allows us to leverage a result from the calculus of variations 288 30 to identify optimal dynamics with solutions of the Euler-Lagrange (EL) equutions (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

---

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 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
(( /
1 27
1
k.((:))
dt.
(2)
This objective is the simplest version of Eq. (G I, f 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 D
= P:
Pr = Pp + nk Ve C(0.)
where we defined momentum as p, == 0:. These equations have two interesting consequences. First, obtain momentum essentially for free, simply from going from one-step multi-step optimization. Second, the temporal discounting rate y allows one to interpolate between a momentum-bast rule and standard gradient descent, since in the 'overdamped" (large 7) limit these equations become
0. np Vo,
This objective has a straightforward analogy with a mechanical system 28 a particle with a mass 1/n moves in a potential k.L and experiences amount of friction proportional y. When 'friction' is sufficiently high, dynamics become non -oneriatia,,, and dominated by the shape of the potential/loss.
Deriving first- and second-order learning rules with and without momentum. We derive first- and second-order learning rules by tweaking Eq. to locally approximate L near 0g as
(00)) L(O,) + 9 "0 (0: 0o) 00. 0o)² (O. -0o)
where g and H are the gradient and Hessian of L at Og. For any At 0, we find (Appendix B that
0 - 0, = -I 3 - 1 H H-2
is the optimal update. See Fig. 2 for example Og traces and Fig. 2 for example loss traces. Importantly, we can derive three learning rules from this result. When the step size At is sufficiently large, we recover the Hessian-comee conditiooed rule (i.e., Newton's method) that 'jumps' straight to the minimum at - -H -H g. When A/ is somewhat smaller than the other characteristic time scales, we either get gradient descent (in the overdamped large y limit) or ballistic learning rule:
(I -----Ar- V TInRRHA H H Newton's method; At large Gradient descent; At small and 7 1 (7)
0 Oar On
1k ghtt
V 1/2 gAi
Ballistie; At small and 0
The last learning rule is not usually considered, but has square-root normalization similar to Adam [3]. We call such rule ballistic because 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 's method does. Because depends on the square root of the Hessian, differences in how quickly different parameter directions converge
reduced somewhat relative to gradient descent (Fig. 2 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.
Thhi convention matches machine learning practice. To better parallel physics, would choose P; 0./n-

---

a.
b.0.5
C.
d. 4
0t
C
7
0
2
3
7
1
71
10
0.5
0.25
7
10
10
72 2
1
10
0
5
10
time
0
10 time
01
1
H11/ H2
e.
train (MNIST, MLP)
test acc. (MNIST, MLP)
f.
train loss (CIFAR-10, CNN)
acc. (CIFAR-10, CNN)
100
2:
100
SGD
SGD
0.3
Ballistic Adam
Ballistic
train loss
test
train loss
Adam
test
SGD Ballistic Adam
75
SGD Ballistic
0.1
Adam
epoch
5 epoch
0 10 15 epach
500
10 epoch
Figure 2: Effeet of modulating temporal discounting rate. a. Example optimal 0, traces for a ID quadratic loss, assuming different values of the temporal discounting rate (7 = 0, 1, 10). b. Loss over time given the 0g from (a), same values of y. Note that lower values of ('longer' planning horizon) produce loss curves that converge more quickly. c. Shape of trajectories for different " given a 2D anisotropic loss. In the mendient-descent like regime (7 > 1), 02 converges much more quickly than 0, due to the anisotropy. In the ballistic (y = 0) regime, the difference in convergence rates is not as extreme. d. Ratio of convergence rates r; == />2/4 H 7/2 assuming diagonal regime, directions with four times as much curvature converge
Hessian. In the gradient-descent times faster; in the ballistic regime, they only converge / times faster. e. An Adam- like implementation of the ballistic (y x 0) rule was used to train 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 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).
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 parameter space follow non- -un--ddide metric G(0) (e.g., the Fisher information matrix), the objective we used to derive gradient descent must modified to involve parameter change penalty 2n
1 (A))" (AO), which causes the optimal first-order (one-step) rule to become A0 - , We can implement Amari's insight in the continuous- time,, multi-step setting by adding a metric to Eq. 2
( 211
((O,))) ((10)))
G(O.)0, + ) e- dt. dt.
(8)
The corresponding EL equations that describe optimal learning dynamics read (see Appendix
p: = Pp + 1jk I Vo.
= Pt
and by argument analogous the one in the previous section, recover three types of rules:
(I - ad v TIA H ' Newton's method; At large
0s1 00
2k. G gàt
Nat. gradient; At small and 7 > (10) Nat. ballistic; At small and 0
T 'At
These follow from making the replacements - G and H in a model that assumes G is locally 0 O-independent (i.c.. here, G refers to G(O))). In addition to recovering natural gradient descent (middle), which modifies gradient descent to account for oon-Euclidean parameter space geometry (21). we also recover other learning rules. The first is Newton's 's method with metric- and deessian-depende learning rate, and the last is a ballistic method which looks like Newton's
method, but involves - HHH-1 'G''HH- instead of H H1 This method compromises between using G-1
5

---

a. no G, no f
b.
anisotropic G
C. rotational f
d.
decay f
1
0
-1
1
0
- 0 0 -1 0
-1 0 1
Figure 3: Parameter space geometry affects optimal learning trajectories. a. Optimal trajectory through 01-02 space for isotropic quadratic loss, assuming nontrivial G' and 0. The heatmap and contours show the value of the loss at each (01, 82) value. Black line: optimal trajectory, red dot: global minimum of loss. Note that, because the loss is isotropic, the optimal trajectory is b. Same as (a), but given strongly anisotropic constant metric G. Note that the optimal trajectory is longer the same along each direction, but converges much more quickly along the 0 direction. Same as (a), but given f that corresponds to purely rotational dynamics. Note two differences: spirals about the origin, and no longer converges to the global minimum of the loss, but a different point closer to the origin (orange dot). d. Same as (a), but given f that corresponds to weight decay, There is no anisotropy, but the trajectory does not converge the minimum of the loss.
and H as preconditioners, or equivalently between moving based on parameter space and loss landscape geometry. While prior work has attempted combine natural gradients and momentum heuristically [32), our framework shows how these elements arise jointly from principled objective. As is well-known, using an anisotropic G can add anisotropy learning dynamics (Fig, 31. b).
Natural gradient descent is not second-order optimization in disguise. It is often argued that natural gradient descent behaves like second-order method, with the metric G acting a surrogate for the Hessian H. Martens (33 explores this view in detail, and shows that natural gradient descent sometimes equivalent
Generalized Gauss-Newton method. But Martens also notes a variety of problems with this view, like the fact that existing theory (c.g., convergence rates) portrays the Hessian 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 G and H play fundamentally different roles: G governs the parameter velocity penalty, while 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. also clear when one examines the different roles of G and 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" state-dependent fashion, like a 'wind' that determines which directions
pushes on the parameter more or less difficult travel in. This feature can arise from partial controllability; if our control of parameter is partial, e.g., 0, = (0t), controlling the parameter via involves not just moving 0, in direction that improves the loss, but also fighting against the 'default' dynamics f(0). Here, we show that this feature produce non con-gradient rules. If we add a drift term f to Eq. 8
J({0,)) ( ( (2)10 f(0.))
+(0.)++c(()) ) c- dt.
- (100))-((([00
(11)
The presence of and of can greatly complexify the EL equations, especially if they are state- dependent. Consider the effect of the drift by itself (i.e., assume G I; see Appendix G
0. - h + ((,))
((0))
(J,0)) (O,) vVVoo(()))
where ((,) is the Jacobian of f at 0, One can interpret J as contributing to dynamics in two ways: contributes an effective "potential-like term (J(0)) I (0.), or equivalently an effective loss; and it produces Note that this reduces to the usual discounting factor [ = -((0) for all which true if and only if ff is the gradient of some function.
effective state-dependent discounting factor Tsn (0,) == 7I J(O.) ((0))*

---

In Appendix G we consider two examples in more detail: the case where f corresponds to rotational dynamics (Fig. ) and the case where f corresponds to weight decay (Fig. Bld. The former case is 'nonconservative (e.g., J( & ( but the second is not. Interestingly, rotational dynamics can make 'spiraling' trajectories optimal, and both choices of 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 includ- ing a drift term f is that can cause the optimal learning rule to involve non con-gradient dynamics, which cannot be described as gradient descent (with or without momentum) down any objective. In our G = I example, this happens if and only if f itself is not the gradient of any function (or equivalently. if J (0) is not symmetric for all 0). In terms of Helmholtz decomposition [35] f(0) = VoV( 1)+++++ + R(O), where V is some non- unique 'potential' function and R is divergence-fre (i.e., Va R = 0). R is the interesting component of f. See Appendix 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 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 and RMSprop 12 use 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 work well in practice, with Adam variants outperforming most other known optimizers [1720 in typical settings, but it is unclear why, or under what assumptions Adam optimal. this section, we show that an Adam-like strategy is optimal under 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 nav- igating partially observable loss landscape is to maintain rumning estimates of quantities that characterize its local shape, like the gradient and curvature. Motivated by this idea, assume that the learner maintains local model of loss landscape shape near their current location ou, i.e.,
mm((O: B(C(O.)) = L(0o) + m.. (0. - -nn) + 2 (0. - 0o) -o) V V.,, - 0o)
(12)
where the gradient estimate m. and Hessian estimate VV ; assumed to be time- Here, k > is a fixed (known) scaling factor. Since m, and V. 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 9.- Assume that the learner has a Bayesian model of how local landscape geometry evolves with two components: observation model, which connects g, to m, and V,; ; and a prior belief about how landscape geometry evolves.
Motivated by simple model of gradient drift and diffusion in a quadratic loss (see Appendix H). assume that 9t ~ N (m, V./ Ve/An At),), which implies that E(g,] = m, and Cov(g,) V.. To make inference simpler, we will use a crude approximation 3 associated with methodd-of-momett based strategies, and assume 9, and (g, - m.)gg m provide independent observations of m, and V.. i.e. 9t ~ N(m,, (o I) and (g, m.) mm," ~ "m)) and 03 > are 'observation noise' parameters.
NV (o At) Here, of > 0
One reasonable prior, which assumes that both parameters tend towards zero and become more uncertain in the absence of observations, an Ornstein- Uhheebbbk prior. Assume
m m, + E1711
V. 022 V. + E272
(13)
where a 2 and 0g 2 parameterize decay, & and & > 0 parameterize the uncertainty growth rate, and Th and 7/2 are Gaussian white noise terms. This means m.+1 M(m, (f and V,, 141 NVV V, (()..
'See Appendix H for a discussion of why this is reasonable, and for description of the somewhat more complicated rule one gets if' one does not assume this.
7

---

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):
%
| log logp(m++1. V my, Vi) kB(C(O.)) dt
J((O,)) = lim At-
2n
Al
where p(m++1, V ++1 9r, m. V.) is the learner's posterior belief about local landscape shape dy- namics. This objective has well-defined At + 0 limit (see Appendix H with EL equations
0 -0- = nk m + x xV
0o)l
m. - ( + ammm = a, + +
k(0; - 0u) + oi
(m 9.)
V. (V. = ag + 2
K 2 (0) 00)² + 02
V - (g. - (g: - mm,)
These equations somewhat complicated, but considerably simplify if one makes assumptions about the relative sizes of parameters. If we assume that parameter changes are ballistic (17 >>>>> 7) but landscape beliefs change somewhat more slowly (5t <2 < 7),
0- syls mm + (o. 0o)
1 k(O, 0o) + Mm 91)
m mh
7
01
V. = 0z V.
k K (0. (0) 00)
1 [ (g. mm))²)
To good approximation, this means that the optimal learning rule has
O++A 0, = -VnkV, 1/2 m,
(14)
Note that the effective learning rate goes like V/-/² rather than V Our framework identifies the square root as a consequence of assuming highly inertial parameter changes, but gradient-descent- like landscape shape estimate dynamics. Moreover, contrary to other ideas about Adam (36), the square root is feature rather than a bug, and need not be 'fixed'. Other features of Adam, like estimating , using the uncentered averages of 9.9" and approximating V. as diagonal, appear to be approximations that improve efficiency and scalability, as is usually believed.
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 (c.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.
learning dynamics. Assume that the learner uses distri- butional estimates, rather than point estimates, of the model parameters O. For simplicity, assume each 8g is associated with a normal distribution A (ps, v.). where μ4 denotes the mean estimate of 6,, and e denotes the posterior variance. Instead of penalizing abrupt weight changes, in this setting makes more sense to penalize abrupt changes in weight distribution:
Modeling fright-uncertainty.oo
DKL(p(O) M 11+1))pp(0 (000)..000 v.)) H((p(0)), v()) + KC(H, vi) e e 7t dt
J((0.)) lim Ar
(At)2
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. I with nontrivial metric G' (see
8

---

Appendix D. equations read
Given a local (quadratic) loss approximation with gradient g and Hessian H, the EL
14
+ 7 ie; = nyko
+
Σ
Hig (P, 1300)
(15)
5
1: 2
0/
84 -ú; = 2n
1/4
+
(16)
2
The equation for μ4 (Eq. 15p has interesting feature: it involves variance-dependent effective discounting rate "ent := 7 Ex and effective learning rate Mar := 1pt; One can interpret the former as saying that planning should become more short-term when weight variances are changing quickly (i.e., 0 high), and that learning should speed up when uncertainty (i.e., v.) is high.
v,
Variance reflects both parsimony and task-driven instability. The equation for oz (Eq. 16 specifies how variance ought to increase decrease along the optimal learning trajectory. The two terms on the right-hand side tell us how this happens: variance decreases the nearby loss landscape very curved (intuitively, the nearest minimum casier 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 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 transition between tasks. If the variance is changing quickly but the mean not changing, decreases variance somewhat. See Appendix foor more discussion.
The entropic term enforces behavior analogous the weight resetting of Dohare et al. [22]: when the loss is not particularly sensitive to a given weight's 's value (i.e., Ha 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 parsimony, the other volatility- explain the behavior of continual learning algorithms that inject noise or reset unused weights, like Dohare et al.'s method.
7 Discussion
We proposed unifying framework that treats learning rules solutions 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 shift in how learning rules should be evaluated: rather than focusing only on convergence rates optimization guarantees, we should consider what assumptions rule implicitly encodes and what problem it is actually solving (c.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
Connection to physics. Our optimal control formulation of learning dynamics can be precisely mapped to a classical mechanics problem, which may allow physics tools be adapted to study learning algorithms. For example, Noether's theorem [3] implies that (quasi-) symmetries of objectives like Eq. [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 466 have already begun to be explored.
Biological relevance. Neural circuits in the brain are thought to implement learning as Hebbian plasticity, homeostatic mechanisms, or ppike-timing-deeendeee plasticity-that often lack clear gradient-descent interpretation 7 47]. These rules may include time-asymmetric

---

or rotational dynamics, operate under metabolic or architectural constraints, or adjust weights in response 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 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 unify learning rules under broader frameworks. Khan and Rue [50] portray various optimizers as natural gradient descent with respect to 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. 12 which relates accelerated optimization methods to a continuous-time variational objective. However, there 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") 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 classical mechanics, 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 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 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 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 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 assumed objective (c.g., with certain amount of temporal discounting, and a particular model of landscape shape belief updating), any two learning rules can be compared, and one more learning rules be shown be optimal. Whether and how link this framework to specific empirical circumstances is a different question which we expect to be more difficult.
Acknowledgments and Disclosure of Funding
SJG was funded by the Kempner Institute for the Study of Natural and Artificial Intelligence. and a Polymath Award from Schmidt Sciences. KR was funded by the NIH (RFIDA056403, U01NS136507), James S. McDonnell Foundation (220020466), Simons Foundation (Pilot Extension- 00003332-02) McKnight Endowment Fund, CIFAR Azrieli Global Scholar Program, NSF (2046583), a Harvard Medical School Neurobiology Lefler Small Grant Award, and a Harvard Medical School Dean's 's Innovation Award. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence at Harvard University.
10

---

References
[1] Haskell B. Curry. The method of steepest descent for non-linear minimization problems. Quarterly of Applied Mathematics, 2(3):258 2615, 1944.
[2] Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv e-prints, page arXiv:1609.04447 September 2016.
[3] D. E. Rumelhart, E. Hinton, and R. J. Williams. Learning internal representations by error propagation, page 318-362. MIT Press, Cambridge, MA, USA, 1986.
[4] David E. Rumelhart, Geoffrey E. Hinton, and Ronald Williams. Learning representations by back-propagating errors. Nature, (206388):53333333333 536-636, October 1986.
[5] Timothy Lillicrap. Daniel Cownden, Douglas B. Tweed, and Colin Akerman. Random synaptic feedback weights support error backpropagation for deep learning, Nature Communi- cations, 7(1):13276, November 2016.
[6] Timothy P. Lillicrap. Adam Santoro, Luke Marris, Colin J. Akerman, and Geoffrey Hinton. Backpropagation and the brain. Nature Reviews Neuroscience, 21(6):335-346, June 2020.
[7] Donald O. Hebb. The organization of behavior: A neuropsychological theory. New York: Wiley and Sons, 1949.
[8] Wulfram Gerstner and Werner M. Kistler. Mathematical formulations of Hebbian learning. Biological Cybernetics, 8755)4004441, December 2002.
[9] B.T. Polyak. Some methods speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964.
[10] Πya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. the importance of initialization and momentum in deep learning. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 1139-11444 Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR.
[11] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12661):2111-2159, 201
[12] T. Tieleman and G. Hinton. Rmsprop: Divide the gradient by running average of its recent magnitude. Coursera: Neural Networks for Machine Learning. 2012.
[13] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Third International Conference on Learning Representations, 2015.
[14] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[15] Samuel L. Smith and Quoe V. Le. A bayesian perspective generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
[16] Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th Inter- national Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9058-9067, PMLR, 13-18 Jul 2020.
[17] Dami Choi, Christopher J. Shallue, Zachary Nado, Jachoon Lec, Chris J. Maddison, and George E. Dahl. On Empirical Comparisons of Optimizers for Deep Learning. arXiv e-prints, page 1rXiv:1910.05446 October 2019.
[18] Derya Soydaner. A comparison of optimization algorithms for deep learning. International Journal of Pattern Recognition and Artificial Intelligence, 34(13):2052013, 2020.
11

---

[19] Frederik Kunstner, Alan Milligan, Robin Yadav, Mark Schmidt, and Alberto Bietti. Heavy- tailed class imbalance and why adam outperforms gradient descent on language models. In A. Globerson, I.. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang, editors, Advances Neural Information Processing Systems, volume 37, pages 30106-30148, Curran Associates, Inc., 2024.
[20] Nathaniel Morrison and Eric Y. Ma. Efficiency of machine learning optimiz- and meta-optimization for nanophotonic inverse design tasks. APL Machine Learning, 3(1):016101, January 2025, _eprint: https:/pubbs.aip: wwwpuusss.i..rgg....... aipp.ooggaai//mm//atiiiccl
[21] Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-
276, 1998.
[22] Shibhansh Dohare, J. Fernando Hermandez-Garcia, Qingfeng Lan, Parash Rahman, A. Rupam Mahmood, and Richard S. Sutton. Loss of plasticity in deep continual learning. Nature, 632(8026):768-774 August 2024.
[23] Peter Deuflhard. Newton Methods for Nonlinear Problems: Affine Invariance and Adaptive Algorithms. Springer Berlin Heidelberg, Berlin, Heidelberg, 2011.
[24] Stephen J. Wright Jorge Nocedal. Numerical Optimization. Springer New York, New York, NY, 2006.
[25] Dimitri Bertsekas. Dynamic programming and optimal control: Volume Athena scientific, 4th edition, 2012.
[26] Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction. MIT Press, Cambridge, MA, 2nd edition, 2018.
[27] Kenji Doya. Reinforcement Learning in Continuous Time and Space. Neural Com- pulation, 12(1):219-245. January 2000. _eprint:
5961..ddf.
[28] Herbert Goldstein. Classical mechanics. Addison-Wesley, 1980.
[29] Tim Strang, Isabella Caruso, and Sam Greydanus. Nature's Cost Function: Simulating Physics by Minimizing the Action. arXiv e-prints, page 1xXiv:2303.02115, March 2023.
[30] Lev D. Elsgolc. Calculus of variations. Dover, 2007.
[31] Kevin P. Murphy. Probabilistic Machine Learning: An introduction. MIT Press, 2022.
[32] Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Bayesian deep learning by perturbaett in Adam. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2611-2620. PMLR, 10-15 Jul 2018.
[33] James Martens. New insights and perspectives the natural gradient method. Journal of Machine Learning Research. 21(146):1-76, 2020.
[34] S.M. Carroll. Spacetime and Geometry. Cambridge University Press, 2019.
[35] D.J. Griffiths. Introduction to Electrodynamics. Cambridge University Press, 2023.
[36] Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? A second-order perspective. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 29949-29973. PMLR. 21-27 Jul 2024.
12

---

[37] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Ttrassa-BBawwinsse Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell, Overcoming catas- trophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 1 114(13)3521-3322 2017.
[38] Ronald Kemker, Mare McClure, Angelina Abitino, Tyler Hayes, and Christopher Kanan. Measuring catastrophic forgetting in neural networks. Proceedings of the AAAI Conference on Artificial Intelligence, 32(1), Apr. 2018.
[39] Mohamed Elsayed and A. Rupam Mahmood. Addressing loss of plasticity and catastrophic forgetting in continual learning. In The Tweifth International Conference on Learning Repre- sentations, 2024.
[40] Mason Kamb and Surya Ganguli. An analytic theory of creativity in convolutional diffusion models, arXiv e-prints, page #fXiv2241..0029) December 2024.
[41] Binxu Wang and John J. Vastola. The unreasonable effectiveness of Gaussian score approxima- tion for diffusion models and its applications. Transactions on Machine Learning Research, 2024.
[42] John J. Vastola. Generalization through variance: how noise shapes inductive biases in diffusion models. In The Thirteenth International Conference on Learning Representations, 2025.
[43] E. Noether. Invariante variationsprobleme. Nachrichten von der Gesellschafi der Wissenschafter zu Göutüngen, Kaabermutisch--hyssikklll
Klasse, 1918:235-2577 1918.
[44] John J. Vastola. Dynamical symmetries in the Auctuation-driver regime: application of Noether's theorem to noisy dynamical systems. In NeurIPS 2024 Workshop on Symmetry and Geometry in Neural Representations. 2025.
[45] Hidenori Tanaka and Daniel Kunin. Noether's learning dynamics: Role of symmetry breaking in neural networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 25646-25660. Curran Associates, Inc., 2021.
[46] John J. Vastola, Samuel J. Gershman, and Kanaka Rajan. A variational manifold embedding framework for nonlinear dimensionality reduction. In NeurlPS 2025 Workshop on Symmetry and Geometry in Neural Representations, 2025.
[47] Friedemann Zenke and Wulfram Gerstner. The temporal paradox of hebbian learning and homeostatic plasticity. Current Opinion in Neurobiology, 43:166-176, 2017.
[48] Benjamin Scellier and Yoshua Bengio. Equilibrium propagation: Bridging the gap between energy-based models and backpropagation. Frontiers in Computational Neuroscience, Volume 11 2017, 2017.
[49] Riley Simmons- -enmen,-EEdler, Ryan P. Badman, Felix Baastad Berg, Raymond Chua, John Vastola, Joshua Lunger, William Qian, and Kanaka Rajan. Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Ended Environments. arXiv e-prints, page (tw:20060069888 June 2025.
[50] Mohammad Emtiyaz Khan and Havard Rue. The Bayesian learning rule. Journal of Machine Learning Research, 24(281):1-46 2023.
[51] Lucas Shoji, Kenta Suzuki, and Leo Kozachkov. Is All Learning (Natural) Gradient Descent? arXiv e-prints, page arXiv:2409.1 16422, September 2024.
[52] Andre Wibisono, Ashia C. Wilson, and Michael I. Jordan. variational perspective on acceler- ated methods in optimization. Proceedings of the National Academy of Sciences, E7355) E7358, 2016.
[53] Antonio Orvieto and Robert Gower. In Search of Adam's 's Secret Sauce. arXiv e-prints, page htl:22505.21829 May 2025.
13

---

[54] Shaden Naif Alshammari, Mark Hamilton, Axel Feldmann, John R. Hershey, and William T. Freeman. A unifying framework for representation learning. In The Thirteenth International Conference Learning Representations, 2025.
[55] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, (6(11))2278-2333 1998.
[56] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. Technical
Report.
14

---

A Motivating our continuous-time objective function
In this appendix, we describe in detail how the most general form of our objective function (Eq. can be viewed as a direct generalization the objective we used to justify gradient descent (see Sec.
1
2
A.1
From single-step objective to a multi-step objective
Recall that the objective we used to justify gradient descent (and Newton's method) had the form
|002 L(O + A0)
J(20)
(17)
27y
where have neglected to Taylor expand C to keep objective more general. Note that J involves two terms: one which penalizes large parameter changes, and another which penalizes high values of the loss. The former term is especially necessary if consider local approximation (e.g., an expansion in powers of A of L, since the approximation may no longer be valid if we consider sufficiently large steps.
would like to go from this single-step objective to an analogous multi-step objective. The most obvious way to do this is to define the objective over sum of K terms, each of which involves parameter change penalty and loss term:
K-1 Σ t=
Jmulte A0.. A 40K-1)==
1 L().)
(18)
21
Note the philosophy of including the loss at each step: it implies that we would like a path through parameter space that involves decreases to the loss each step, rather than just at the end. It is also possible to penalize the loss only at the end, but the resulting learning rules would look somewhat different than, e.g., variants of gradient descent.
A.2 From discrete time to continuous time
Although we could directly study the optimization of Eq. 18 we can instead exploit the fact that objectives like this tend to be easier to analyze in continuous time, since determining the optimal sequence A A becomes a well-studied calculus of variations [28 30 problem. If each step takes an amount of 'time' At, the cost of each step is scaled to be proportional to At, and we adjust ? + of A1)2 for dimensional reasons, we obtain a continuous-time objective that directly generalizes Eq. 18
J( ((O0)) = t=
0 24
+ c(0,)
At Ar-30
/
T
1 2n
=10.1 2 ((O.) dt (19)
where T' is the total optimization 'time". Note that our optimization problem now looks like an optimal control problem (with control cost x 1/1 and state cost L) or classical mechanics problem (where particle with mass m == 1/n moves in potential determined by L).
To make it slightly clearer how our solution depends on L, and to more explicitly control the influence of the loss, we will add a prefactor & > 0 (which has units of inverse time) in front of it:
/ ] 10,11 kC(O.) dt
J((0))
(20)
21j
This change does not meaningfully affect our optimization problem; it only amounts change in the units of J.
A.3
Incorporating temporal discounting
Our objective (Eq. 21 now has the form of an optimal control problem. Motivated by this observation, we incorporate temporal discounting factor ent which causes the learner to overemphasize rewards and costs that are nearer in time:
(((0.))
1.11 2 + kC(O.)
e-78 dt.
(21)
2n
15

---

We could have included a more general discounting factor, as some prior work does [2]. We restrict ourselves to exponential discounting because it is canonical choice [26 27 and yields relatively simple El. equations. As an aside, allowing a more general discounting factor would allow to derive learning rules that involve Nesterov momentum rather than just standard (Polyak) momentum.
Since we are not generally interested in any specific learning time T and since taking the T + co substantially simplifies some EL-equation-related math, we will consider the version the objective in everything that follows:
10,11 k=((,)) e-x dt. 2n
-
J((O.) =
(22)
This objective is infinitely forward-looking in the sense that trajectories which optimize depend on considering 0, at arbitrary distant future times.
A.4 Incorporating parameter space geometry
Our original single-step objective (Eq. [7) implicitly assumes that parameter space is Euclidean, or at least that Euclidean distance metric most appropriate for penalizing large parameter changes. As researchers like Amari [11] have observed, this not necessarily true. More generally, we might want to penalize distances according to less trivial metric like the Fisher information metric.
We can account for this fact by modifying Eq. 17 to involve a metric G:
1 (())) (( + A0)
J(A0) = 2n)
(23)
Note that, for all 0, we will assume G(Ø) is symmetric and positive definite (and hence invertible). To account for partial controllability (see Sec. 4. we could make the similarly minor change of penalizing not the size of A0, but the size of A f0) (i.e., the size of deviations from the 'default' dynamics determined by the drift function f):
((A0) = 2m
A0 - f(0)) f(0) + C(7 + A0)
(24)
We can generalize this objective something which operates in continuous- time over multiple steps by the same argument as behavior. We must only make the small change - f so that the continuous-timg limit is well-defined. We obtain
J((0.))
0 f(0.)) f(0.)) + +C((O)) e- dt.
(25)
2m
A.5 Accounting for loss approximation
Lastly, must account for assumption that the true loss L is generally only partially observable. Since we have already posed learning as a reinforcement learning and optimal control problem, this change is easy to make. In those settings, accounts for random variables by averaging the objective over them; we will do the same here. We finally have
J((,,) = B(c,) {o-00 (1, f(0)) -
f(0.) kε((0)) ) e- dt } (26)
16

---

B Experiment details
In this appendix, we provide details relevant to understanding the numerical experiments mentioned in the main text. See http://wwwwo.o/oo/o//////////aaa reproduces all figures.
for code that
Direct optimization of the objective. In Fig. objective
we numerically estimate the minimizer of the
J((O.)) = S/ ( 27 11 kC(O,))
6- dt {00
01
62
kL 02) (2)
e e dt
2m
2n
(27)
given a double-well loss
L(01,0g)=0 - (0) - -b)² + c 02 + d 1
(28)
with a =,, 1, and d d= -1/2, assuming = 1 and 7 = 0.1. We do this optimization directly (rather than via the EL equations) by discretizing (t) and (2(8), i.e.
[0r(L441) (1e)) [02(1=-1) 02(12)1 + kC (4x), (1x). 2n))
e - (At)
(n)22
(29) where we choose N + 1 equally spaced time points tqp tq...., ty for simplicity. This means that 4 = k&t for all I: = 0, .... N, where At == T/N. Here, the cutoff time T > 0 is chosen to be large enough that the optimal trajectory is near the global minimum of the loss at the final time point i.- (This means that, even though we do not integrate all the way until t co, not much interesting happens beyond time T.)
By fixing the initial point ((4)) (01 (ln), 02(l)))² final point ((Ix) (0) (Ix), (2(Lx))² and the cutoff time T, we can vary the remaining 2(N 1) degrees freedom to determine the optimal trajectory. Following Strang et al. (29]. minimize J with respect to these variables using PyTorch-based gradient descent approach.
Given the solution, we can estimate the 'kinetic energy' throughout trajectory by computing
KE(ta): = (10) (10)(tu.1) - 0 (tx))
(101(6 2n) At)2
+ (02(tx+1) ) - 02 (tx))
(30)
2n)
and the 'potential energy' by computing
PE(ta) = (tx).
(31)
Application of ballistic rule to MNIST and CIFAR-10 image classification. For Fig. 2 we
implement the 'ballistic' rule that emerges from one of our exact solutions in the 7 limit (see Eq. 7. and use it to train classifiers on the MNIST (55) and CIFAR-10 [56) image datasets. At each step, the ballistic rule prescribes parameter update proportional to
AO x H 1/2
(32)
where g is the current loss gradient and H is the current Hessian of the loss. Since Hessian computation is expensive and difficult to scale, we implement the ballistic rule by using an Adam-like [13] approach: at each step, we update running average of (uncentered) gradient variances along each direction. This involves three crude approximations: we use this running average instead of directly computing the Hessian; we only compute the diagonal entries of the Hessian proxy; and we do not center the variance estimates. Despite these approximations, we still believe that this heuristic approach captures the spirit of the ballistic rule. Furthermore, using an Adam-like approach is theoretically reasonable given the link we identify between Adam and our ballistic rule (Sec. 5)
Let v; denote the running gradient variance associated with the parameter 0,, For each i, the precise updates per step are
V; +- B28 + (1 - B2)9
(33)
A0; = 19: o
17

---

where 9i denotes the current gradient along the 6 direction, 77 denotes the learning rate, and 32 E [0, 1] controls the time scale on which gradient observations are averaged. Note that this usage of 3g is intended to match Adam's s; like in Adam, values like 32 0.9 and 32 0.999 (i.e., values close to one) appear work well.
We consider only two simple architectures to illustrate the ballistic rule's behavior:
multilayer perceptron (MLP) with hidden layers, which we trained on MNIST; and pooling
small convolutional neural network (CNN) with four convolutional layers, layers, and final fully-connected layer, which we trained on CIFAR-10.
We generally found that, especially since this implementation of the ballistic rule is similar the standard implementation of Adam, hyperparameter settings that work well for Adam also work well for it. For example, learning rates around 1e-3 and 1e-4 worked well. For MNIST, 82 0.999 performed reasonably well, but for CIFAR- 10 lower value (32 = 0.9) appeared to be necessary for good performance.
18

---

C Deriving learning rules via the Euler-Lagrange equations
Suppose that the parameter vector 0 is D-dimensional. After averaging over the loss landscape belief' model in Eq. | our objective has the form
J((0)))
10. f(0.)) -((G)))))-_((()))) + + e e 7e dt
(34)
2m
for some L. G, and so on Given an objective like this, how do we go about (analytically) deriving learning rules?
We are looking for a trajectory (ejo.c)) which minimizes J. This trajectory is not necessarily unique, for example due to symmetry, but often unique in practice. By the calculus of variations, subject to the relevant boundary conditions (and smoothness-related technical conditions), the optimal trajectory can be shown [28 30 to satisfy the Euler- Lagrange (EL) equations.
The idea to use the objective to define Lagrangian
L(O, 0,
1 0 f(0)) 27
f(0)) +c(0)) e-
(35)
The Euler-Lagrange equations are the D equations
d 0LL) = de
OL
d 0L
aL 00,
i.e.
for i = I,.... D.......
(36)
dt 00
dt 00;
More explicitly, we have second-order ODEs
d dt (GO) -()) = nkk V
(1 00 ( 1 G(0)0 f(0)
0
(37)
Since they are second-order, specifying solutions to these ODEs (up to symmetry-related degen- eracies) requires two pieces of boundary data. In classical mechanics (28), one often looks for solutions with a specified initial position and initial momentum (or equivalently, initial velocity). In our context, this does not quite make sense. We do want a solution with specified initial parameter vector 0. but requiring some particular initial velocity or momentum is less obviously meaningful.
The second piece of boundary data follows from the fact that we would like J to be minimized. This equivalent to requiring that the asymptotic (t - oo) endpoint of the trajectory corresponds to the global minimum of the loss, at least if f O. If f is nonzero, the drift contributes term to the "effective' loss, and it is the global minimum this function which must be asymptotically approached instead.
In short, we can find optimal trajectories (O.heeo,ss) by solving the EL equations (Eq. 37 subject to the constraints that:
Og is fixed, and corresponds the 'current' parameter vector:
The remaining degree of freedom (e.g., the trajectory endpoint or initial momentum) is chosen so that J is minimized.
Below, we provide simple but instructive one-dimensional examples that provide intuition about how this process plays out practice.
C.1 Example: quadratic loss
Suppose D = 1, G I. 0, and
(()) = (0 00)2
(38)
"If the loss model à involves parameters (e.g., estimated gradients and curvature) which are themselves dynamic, these parameters should be considered components of é, and can hence contribute terms to an 'effective' metric G. an effective drift f. etc. This is easiest to see in the context of specific examples, like the ones that appear in Sec.
19

---

i.., the loss is quadratic with global minimum at 0.. Eq.
34 becomes
kh
J((0,))= Ju
62
(0. 0.) )
e r de.
(39)
27
2
The corresponding Lagrangian is
1
kh (A 0.)² e-
L(O,0, 0,
82
(40)
2n
and the corresponding EL equation is
0 - -0= = nkh(0, 0.)
= - nkh0, = -nkhb.
(41)
This is linear, second-order ODE with constant coefficients, and so can be solved in the usual way This ODE's characteristic equation is
p2 ---7 . - -nkh 0
(42)
and has roots
7= 7
72 + nkh..
(43)
4
Importantly, one of these roots is positive and one negative, fact which we will return shortly. The full solution can hence be written
0. c4e -c
(44)
where 0, is the (obvious) particular solution. Enforcing the initial condition,
Oo
0, + c+ +
=0 0. C+
(45)
Hence,
0. = 0. + c+e + (0p 0. C. )--
(46)
0 c+r4er+ + (0l - 0, c+)r c
We can determine c+ by substituting these into J and minimizing with respect to c+. Note,
2 + 2c, (0o 0, C+ ctfeff+ (r-+ (0. 0.)2 2. 2r+ + (0o - - 2
(47)
8: 2 + (80- 0, -c_ - 2r + rc, (0o - 0, - -c.)el++t
Since
r -r+ + nkh
r = -r- + nkh
(48)
74r- -nkh
the 6? expression can be simplified to
03 njth [2] 2r (0o 0. 2 10000000 -0-- -e
(49)
+ 7 [c2 + (fn - 0. - c. -2+)2 2
After some algebra, the integrand of is
kh +
21 r+
7
c, e. (2r, -290 + kh + 2n
-r- (0o 0. - (2
(50)
Integrating from / = 0 to 1 T.J. equals
-1) r+] 2 (e(2r+ (eler++)11-1
kh 7 7 (Oo (00-0-- -24) (cl2r--))) - 2m
J == lim T
kh + 2n
2r+ 7
2r -r (51)
'Physically, it is analogous to a damped harmonic oscillator, but with the "wrong" sign on the 0 coefficient. This difference makes sense, since a learning trajectory ought eventually settle down into a global minimum rather than oscillate.
20

---

Recall that r+ is positive. The quantity
222 2r+ - V 4
+ nkh
(52)
also positive, which means that the term with e (2)+ approaches infinity in the 7 - op limit we're interested in. Inspecting Eq. 51] the only way to remove the offending term is set c4 O.
But this outcome was foreseeable if we note that the positive-root term e 'blows up" in the long-time limit, whereas the other term doesn't. For this reason, in future derivations we will bypass this argument and simply set c+ (or its higher-dimensional analogue) zero.
Incidentally, if our objective involved finite time horizon E |0, T rather than t E [0, 20), in general we would have C+ * O, which would somewhat complicate our expressions for optimal learning trajectories.
Setting C+ = 0, finally find that the optimal learning trajectory has
0: 0- + (0. - 0,)e"
(53)
0, (0n 0.)r__c-t
i... it approaches the global minimum exponentially quickly, at a rate r_.
C.2
Example: double-well loss
While the previous example is instructive, the loss function we considered was convex, and involved only one (local/global) minimum. It is somewhat more interesting to see what happens with double-well loss
(()) = - 62)² -
9
Eman
(54)
where h >> 0 and q > O, and where the additive constant funin is chosen so that L equals zero its global minimum. This loss has two local minima, as is clear from its derivative:
C'( h@( 82) q0²
(55)
In particular, since
C'(0) = he 82
q
62
(56)
h
the minima
q2
0+ = =- +
02
+
(57)
4h2
2h
and are near +0,, if y is small. The lower minimum is at 0+; the h term is identical at both 0, and 0_, but the g term is negative at 0_..
For this loss, the objective J is
J({O.)) 0
1
07 +
h (0
0,)2 02)2-990 983 e -t dt
(58)
2m
and the EL equation is
- - = (),) = nkhh,, -
(59)
This second-order ODE is highly nonlinear, and it is not obvious if it is analytically solvable. If we assume 7 = (i.e., no temporal discounting). the EI. equation becomes
8, = (),)
(60)
This can be simplified somewhat by using a trick. If we multiply the left-hand side and right-hand side by 8, we can integrate both with respect to time; we obtain
0 = C()) + 2
(61)
where E is a constant. Since the left-hand side is nonnegative, E> -k.((mmm)) where Ownn is the argument for which L achieves its global minimum. (Here, due the additive offset we included, L(0min) = 0.)
21

---

We label this constant E since it corresponds to this setting's notion of energy, as we could figure out from an analysis of this problem's Hamiltoniar
Eq. 61 implies that
{"
1 0 +kC(()))) 27
["
lim ["
J
lim T-+00
dt lim T T-000 0
2 2k:((0)) dt ET (62)
T-000
along the optimal trajectory. But this a problem, since ET + 00 as T + 0o for most choices of E. We do best by setting energy equal to its minimum possible value km0mmm))) O, since the problematic term vanishes and the loss term asymptotically approaches zero (since the value of the loss the global minimum zero). That is,
J = lim
2k(((e;) dt < 00.
(63)
T T-000
Going back to Eq. 61 this means that
0. = - 2
(64)
This defines an optimal learning trajectory, and also an optimal learning rule.
What does this equation mean? To understand this expression, suppose 0p 0_. i.e., that the learner begins in the shallower minimum. Clearly, the optimal learning dynamics must move right (towards 04); ); the optimal trajectory in this case has
à = / 2nk
(65)
i.e., we take the plus sign. If the learner begins a parameter 04, we would instead take the minus sign.
Although this example is somewhat complicated, there two important takeaways:
Long-horizon optimal learning trajectories approach the global minimum rather than any local minima.
Even if the functional form of L is more complicated than quadratic, it is in some cases possible to derive optimal learning rules and trajectories.
In the rest of this paper, partly because analyses like the above difficult, and partly because goal to derive well-known learning rules, we essentially only consider quadratic (local) approximations of the loss.
"See Vastola 44 for the details of how to follow this line of thought in an analogous theoretical setting. At least if 7 0, this notion of energy is conserved along the optimal trajectory.
22

---

D More on the boundary conditions of the Euler-Lagrange equations
The boundary conditions for the objective minimization problem we describe Sec. are reasonably clear: of all the possible bell-behaved (e.g., smooth, bounded) trajectories ( (o) through the D-dimensional parameter space which begin at a prescribed initial parameter vector Qg. the 'optimal' trajectory is the one which makes J smallest. 'Smallest' here is well-defined notion, since J is bounded from below as long as the loss L is bounded from below. If multiple trajectories make J as small as possible, then each of them is optimal.
However, the boundary conditions for the Buler-Lagrange (EL) equations are less obvious. The EL equations answer the following question: given all possible smooth trajectories that begin at do and reach (T)) at time T >> O, which of them makes the Lagrangian (i.c., the integrand of the objective) stationary? This is three steps removed from our original minimization problem, since (i) stationary points may not correspond to local minima, (ii) local minima may not be global minima, and (iii) given (T) may not correspond to an optimal trajectory.
We think about the boundary conditions of the EL equations encounter in this paper in the following way. First, we assume that the trajectory of interest has prescribed initial point Og. since this is boundary condition of the original minimization problem. Second, since (according to ( standard calculus of variations results, under mild assumptions") the Lagrangian J. is stationary at the global minimizer of J, satisfying the EL equations a necessary (but not sufficient) condition for trajectory to be optimal. Third, any two trajectories can be compared (i.c., which corresponds to lower value of J?).
Together, the first two insights tell us that the EL equations must be satisfied for a specific initial point Ou and some other point OTT). The third insight tells that, if we do not know which ((T) to use, we can compare any two possibilities by comparing the corresponding values of J. This yields a two-level optimization strategy: for many endpoints (T)), solve the EL equations; then, choose the solution whose corresponding J is smallest.
This strategy is not circular, and is useful because it reduces an optimization problem defined over an infiite-dimensiona. function space (i.c., the infiiitediimensional space of all possible parameter trajectories) to optimization problem over D-dimensional space (since each possible solution corresponds to a particular choice of (T) E P).
Finally, it is worth noting that the rather abstract conditions under which calculus of variations problem is interesting and well-defined are not that useful for understanding many of the extremely simple objectives we consider in this paper. Especially given quadratic loss, we can often verify directly (through algebra rather than numerics) that solutions of the EL equations correspond to global minimizers of J, and even directly compute J as function of any 0g and ((T).
FFrr example, we may want to assume that the kinetic term is convex and coercive. These two properties hold for the simple quadratic kinetic terms consider throughout.
23

---

E Deriving learning rules with momentum: more details
In this appendix, we derive the results presented in Sec.
The relevant objective is
J((0,)) /
1 kc(0)) ) 0
e e dt.
(66)
2m
The corresponding Lagrangian is
L(0,0,1)
0011 + kc(())
em
(67)
21
and the corresponding EL equations are
0. 00 vVVo,, ((0.))
(68)
We can rewrite these equations a suggestive form by defining 'momentum' as p; := Oz, so that we have
= Pt
P: - PP + Vo. ((.)).
(69)
Note that convention for momentum matches machine learning practice, but does not necessarily match the physics convention for momentum, which has P. 8e,
dh.
E.
Solution to the Euler-Lagrange equations for a quadratic loss
If we assume quadratic loss
Ž(0.) == C(0,) + g" (0.-0)) (0, -0)-0))" - -00)
(70)
where is the local gradient and H is the (symmetric, positive definite) local Hessian, the EL equations (Eq. 688 become
0, 0 ni i H(, 0n) = -0 m)H nn Ig HO,)
(71)
The above represents a system of coupled second-order ODEs. We can decouple it by exploiting an cigendecomposition of H. Since H = QAQ" for some orthogonal matrix Q and diagonal matrix A (whose diagonal entries are nonnegative), we can premultiply both sides of this equation by Q" to obtain
3. ~ nkAØ, nk [g A@æl
(72)
where define Φ, = 0, and 9 Q" We have many linear, second-order ODEs identical form to the one from the first example in Appendix C
For a given di, we have
- n = sp [g: >;
(73)
where X; = As 2 0. The corresponding characteristic equation reads
2 7r nbλ; =0
(74)
and has roots
11 +
4 + nk:
(75)
By the same argument we used Appendix C we throw away the positive root (since does not minimize J) and keep the negative one. Denote the (negative) root that keep r,. We have
(x(t) = Th 6s0
A 0:
(76)
where Φıu A T. is the particular solution of Eq. enforcing the initial condition. Doing so, we find
73 and is a constant. We can determine c, by
(x d. x; 0z etif).
(77)
Transforming back to 0 space via the relationship 0, = Q we have
0t( Σ Quoio - 1, (1+
(78)
24

---

or equivalently
= -200-A--0 QI Rg) QΦu QA-
Q Q'Q R Q"" Qg
(79)
where we define the diagonal matrix R whose diagonal entries are the ri. Noting that
Q - - exp {
V + nkH - 2
2 11))
(80)
we can simplify our result to
0. = 0n -0-0(eexx
1I-exp-
V
72
n
2
11)
H H
(81)
2
This implies that, for any At O.
A Pu
cxp
2 I n
I
At sr)
H g.
H
(82)
2
E.2 Special cases of the quadratic loss solution
At is large, the exponential asymptotically vanishes, so
lim - - Ar-+00
(83)
and we recover Newton's 's method (i.e., one 'jumps' to the minimum). For large but not infinite At. Eq. 82 says that one ought follow a Newton-like method, but with possibly asymmetric learning rate in different directions, whose precise form depends on the argument of the exponential.
For small At, obtain (to first order in At)
1 nk
21 H H
0 Oa: Oo
(84)
We can simplify this further if we make an assumption about the relative sizes of H and yy If y VnkkA for all eigenvalues X of H, then
4nk
7 2mph
n
~
H
(85)
2
2
so in this limit (the 'overdamped" limit) we obtain the learning rule
Ast O, TyK HH 9 At 1/k
? g
(86)
which just corresponds to gradient descent.
Meanwhile, if At is small but < ViJkAz for all eigenvalues X; of H (or if 0), then
Oat Oo v /I g At / V 1 AL.
(87)
This 'ballistic' learning rule operates in regime where momentum dominates learning dynamics.
Lastly, it's worth noting that if is larger than some eigenvalues of H but smaller than others, one can obtain learning rule that looks like gradient descent along some directions, and looks ballistic along other directions.
25

---

Deriving learning rules in non-Euclidean parameter spaces
In this appendix, we derive the results presented in the first part of Sec. The relevant objective
J((0.)) - [ [ 0 G(O,)0++ + +k(0,) n' dt
(88)
2n
where G(O,) is a symmetric and positive definite matrix for all 0,. The Lagrangian
L(0,0,t)
0' D'G( +c(()) e-
(89)
21
and the EL equations
d (G(O.)0.)- - - = n/ Vo, + Vo. (bOTc10)j..)
(90)
di
If the metric is approximately independent of O., we have the special form
G(O.)O, -((0,), ((.,,
0, - Vo,
(91)
that appears in the main text. We can rewrite this in terms the 'momentum' p, == 0, to exactly reproduce the expression from Sec.
F.1 Solution to the Euler-Lagrange equations for quadratic loss
Assume G is independent of 0, and that the loss is quadratic, i.c.,
c(O.) :- _(0o) + g (0. 0o) = (0, -
0o)
(92)
where g is the local gradient and H is the (symmetric, positive definite) local Hessian. The EL equations become
-.- = ny Ig + H(O; - 0o)]
(93)
This is identical to what we considered in Appendix E. up to the changes - G 9 and H - G 1 H. The solution, then, is also identical up to these changes.
26

---

G Non-gradient learning rules from partial controllability
In this appendix, we derive results related to the second part Sec. which concerns the influence of 'drift' term (e.g., due to partial controllability) on optimal learning trajectories. The relevant objective is
( [ ( 1/-
(0.--(())'' f(0.)) [(0.))++((0)) k.((.)) ) 2n
]((0)))
e de.
(94)
The Lagrangian is
00,,0,1 - (100 e
(95)
21j
and the EL equations are
d dt (G(O))0, f(0:))) f(o)) - ( -(000 (f(.)))) s(0.)))
Trivial metric.
In the special case that G I, we have
0. dt f(0,) - - f(0.) = oV c(0.) + Vo. L1 110.--(01)))
(96)
Recall that the Jacobian J of f is defined a
0f((0)
Ju
(97)
08,
Using it, we can more explicitly write our expression as
i - ((,)), - -(0, f(0,)) = v,. - J(0,)=10: J(0,) f(0.)
(98)
Rearranging this, obtain
0, [ + ((,) - -, (J(0,)" - + VA C(O,)
(99)
the equation that appears in the main text.
Trivial metric and linear drift. In the special case that f(0) -O,
0, - 17/ + J 0, = [J" -1]J0, V L(O.). In the that L is approximated locally quadratic (as in, e.g., Appendix this becomes
(100)
0,
JJ [ -I.0 + gg H 0u)]
(101)
In the next two subsections, we will consider two explicit examples of dynamics of this form.
G.1 Example: drift term with rotational dynamics
Setup. For simplicity, assume that the Hessian isotropic, i.e., hI. Assume that the default dynamics are rotational the sense that (@) J0, where J is skew-symmetric, i.e., J --" If skew-symmetric, solutions of the default dynamics Ô = f(0) are purely rotational, since they have the form
0. = ="0
(102)
where exp(Jt) is rotation matrix. The EL equations have the form
0, 171 + 2J 0, -1 - +JJ10 + 1ı [g + h(O, - 0o)l
(103)
'Note that using boldfaced J denote the Jacobian of and to denote the objective.
27

---

All skew-symmetrie matrices are unitarily diagonalizable over the complex numbers, so we can write
J UAU
(104)
where U is unitary (i.e., UU-= = - = I) and A is diagonal with complex entries.
Our second-order ODEs become decoupled the space of eigenvectors of Q. The quantity Φ, U Og changes according to
の 17] + 2A] o. - + >Aj0; 1/ + h(0, - +.)l
(105)
where define U g. Happily, each component Φ, of Φ, evolves independently according to linear second-order ODE, and each of these can be solved in the usual way.
Decoupled ODEs.
Each d, evolves according to
- - 2 & Inkh- X2 >Ailo - 7/ 10: hovoll
(106)
where X; = Aug The characteristic equation of this ODE is
y2 - r 22.Jr [nkhh X - ->1=-= = -0
(107)
and its solution is
7 2A:
(++ + 21,)²
r+
+
+ 1yKh - X2- - >>>;
2
4
(108)
72
7
Ax +
+ +nkh.
As before (see, e.g., Appendix C Appendix E}. we ignore the positive root since its associated solution asymptotically blows up. Define
Ts A; +
+
nkh
(109)
2
4
and R as the diagonal matrix with Rs = ri. Combining the general and particular solutions of Eq.
106 yields
nkh
(x(() w/hh (A; +
[0:0 11 --1ag + ce
(110)
some constant c. Enforcing the initial condition.
whh (t) 1ykh 3(+++))
10:0 h-1g, (1 - e1t) + 0s
(111)
Define the diagonal matrix B via
n/h
Bu nkh A= + 7)
(112)
so that the o solution becomes
dc(t) (1 err1) (0r0 h --1gg)) ct.t0aa
(113)
or in vector form,
Φ, (I - e fu B ( - h 1- it $
(114)
Solution. The relationship 0, = UΦ, tells us that
0: = n 0, + (I - 12 B - h-11) g)
(115)
where we define the matrix R as
72
R U = J+
+ mkh I
(116)
2
4
28

---

and B as
B:=UBUT = sykh [nkh - J[ +
(117)
Interestingly, in the long-time limit this solution doesn't converge to the global minimum, but to slightly different location controlled by the 'bias' matrix I::
lim 0. B(0, -
(118)
even more interesting feature of this solution is that, since the ri are complex-va valued (through their dependence on the A; which are pure imaginary for akew-symmetrie matrices), they approach their asymptotic value a spiraling-in fashion. To see this explicitly, consider the particular skew- symmetric matrix
(119)
which is the infinitesimal generator of a counterclockwise rotation in two-dimensional plane. Then
(cost
sin cos t
e cos(() + sin(t)
(120)
sin t
ie., obtain matrix that performs counterclockwis rotation by an angle t. The optimal learning trajectory approaches its final value according to
Rl c (V
(121)
which combines decay with rotation.
G.2 Example: drift term representing weight decay
Setup. this example, assume a general Hessian and that f(0) = =-j for some weight decay rate j > 0. The EL equations are
0, - = [j 7)- 0. + sple [g + H(,, [[))]
(122)
or equivalently
0, 0, - [(j - -) 1 1jk 0: = 1/k (g HO,)
(123)
We can solve this system of second-order ODEs using the strategy from Appendix by using an eigendecomposition H = QAQ The only differences are that the relevant rates are
7
2 + +j - 7)5 + nkAs
r: =
(124)
2
4
and that we once again get biases
nk b, = nバλ + [j 7)
(125)
Define the diagonal matrix R using the Tú, and the diagonal matrix using the b,. The solution has
0, e
R 0u + (I (I- (B) - -H'9)
(126)
where we define the matrix R as
- -0-=-11- 1 + -)) +11 I nhH
R = QRQ
(127)
2
and as
nkH
B QBQ
(128)
kk + (j - 7)j
As in the previous example, the form of J (here, just the scalar j) contributes to both drift-related bias B and the rate R at which the optimal trajectory approaches its asymptotic value.
29

---

G.3
Optimal learning trajectories generally do not follow gradients
The G EL equations above can be written
p: = = J(0,) J(0,))) P: + - + ((0.)
(129)
where we define momentum (as before) as P; := 0r. Assume Helmholtz decomposition f(0) (0) + R(0), where V is some non-unique 'potential' function and is divergence-free (i.e.,
Vo R = 0). This decomposition implies that the Jacobian of f has entries
(R)
Juy = 02 ( +
(130)
00,
Note that the only possible source of asymmetry comes from R. Moreover, note that
0Rj 00,-
(0,V)
R
Σ 5J - Σ 02 V
,,V + R2] = 0
Σ
+ (0,V) +
2
j
(131)
L0, R2}²
0
Σ
If J n denotes the Jacobian of R, the EL equations can be written
p: + (R(0,) P. ~V, -V, - R((0)
f(0)) + -
= + - RR())) P: -((,) + Vo, (o,Ve,)(0.)
where we define the effective loss/potential
(e,) == mkL((,)) + Vo,) +R| -
(132)
Hence, it is clear that a nontrivial R contributes non- con-gradient dynamics to learning in two ways: first, through determining an asymmetric effective temporal discounting rate; and second, through the y R term.
30

---

H
Deriving adaptive learning rules from dynamic loss landscape beliefs
In this appendix, we motivate and derive the Adam-like adaptive learning rule discussed in Sec.
H.1 Motivation: why gradient variance relates the Hessian
First, we briefly motivate that the (average) loss landscape can be locally approximated as
B[((0,)) = ((0o) + m.(0. (0: Ou) 00 0o)
(0. Vo 0a) -0)
(133)
where my is the average observed gradient and V, is the covariance of these gradient observations. Why is it reasonable to suppose that V, proportional to the local Hessian H?
Assume that the true loss landscape has a gradient ! and Hessian H. The idea is to view noise in gradients as due to unobservable, noisy parameter vector 0, that explores the local loss landscape according to stochastic process. Since 0, is noisy, its time derivative (i.e., gradients, which are observable) will also be noisy. We can view this change in perspective. Rather than assuming that O, remains fixed but the landscape changes dynamically (and partly randomly) around us, can assume that the landscape is fixed, but that where we are is randomly changing.
The simplest process we can assume is one that is unbiased (i.e., noisy gradients equal true gradients on average), and has structureless (i.e., white and state- (ntteeddddeeeent) noise term. Consider
0, = - on - H(, - Ov) + om, =- - 0: + tan,
(134)
dt
where r > 0 a decay time scale, a > controls the amount of noise, is Gaussian white noise term, and
μ = 0g - H
(135)
At steady state (or at least, on time scales somewhat longer than 7), we have
o² Ø ~~ NN,, o "
(136)
Near steady state, this implies that
E (V () HE (() 0.) -;) H o² 27 H
1
(137)
or equivalently that the variance of observed gradients is proportional to H. This result can be interpreted in the following way. For sharp/curved minima, H is by definition large, so small parameter changes can produce relatively large changes in gradients; on the other hand, in flat loss landscape regions, H is small, so even large parameter changes do not change gradients much.
H.2 Deriving the Euler-Lagrange equations
More complex observation model.
Assume that g. ~ ((m,, v/) The relevant objective is
|0,10
logp(m++1 V Ve+lgg, m,, V.) + KE(L(0,)) e- ddt
J((O,, m,, v,}) = lim At-+
27
At
where p( pmM+11, V V++119 19, mi V.) is the learner's posterior belief about local landscape shape dy- namics. The posterior term contains two types of terms: terms from the observation model, and terms from the prior. In particular, contains the terms
(g.
m2)
I|m + aamm
|6 + a.v, + const.
Σ
log(2 +
+
(138)
20;
2
2
where we neglect unimportant additive constants, and assume that V, diagonal. Here, is denotes Vi, and v denotes the vector containing the 04. The corresponding Lagrangian is
Σ
02
(g:
m2
1
log(2mus) + (mi, + amm))2
(6: + a2v))²
ABI2(O.)) e
21
20
2
2
31

---

and the corresponding EL equations are
0. - 17h (9s + (6; Bux)
k(0; Oni) + (m, 9b)
m; - a. a,m,) aAm; + E
Vi
(139)
1
(gi m;)
0, 7(0 + a2u;) a²2; + 3 k (0. 0ox)
2 0
20
The equation for v; is somewhat more complicated than the presented in the main text, but note that it can be written as
8 7(0 + a20;) = azu; + 22 K
2 (0: 0o.) )
+ 2 u
(u. - (g. (a--m)33) m;)
(140)
which matches the main text form up the prefactor 1/(2 If the variance e; fairly stable (for example, because a reasonable amount of evidence about gradients has already been accumulated), then this prefactor is approximately constant, and the forms are identical.
Simplified observation model. we instead treat gradients and squares of gradients as consisting of separate observations, i.c., 9, N(mm, (o)/At)I) and (g. m.) (9.m;;)) (o3 At) then the posterior contains the terms
(g) m2)² 20
(gi m2)
| + aamml 1
| + 020+ll 1|
const.
(141)
20
257
2
The corresponding Lagrangian is
02
(g: - -2
[0s (g: m2)
(rin; + aim,))
(ú: + a20;)
L L: Σ
+ NE[[((,)]
e-
2m
207
203
257
2
and the corresponding EL equations are
0 - 70: np [gs + kv (6, -Bus)]
k(0; - -Oos) + (m; 9i)
m4 - (m + a)m:) =amm + 5
of
(142)
6; - >00 + a22)) az42 + E
k R. (0, 0or)² + o2
(v; (gi mz)
2
32

---

I Deriving learning rules sensitive to weight uncertainty
learning rule discussed Sec. 6 Recall that assume each model parameter 0, (for .....; D) is associated with normal distribution Wi,, Us), where 14, is the average value of 8, and U; is its variance. The advantage of this setup is that it allows the learner not just estimate what their parameters are, but also how certain they about them. In principle, we could consider a model with a more general covariance matrix, but we restrict ourselves diagonal covariance for simplicity.
In this appendix, we derive the weight-unceertaint-ssennn
L1 Simplifying the objective
The relevant objective is
DRI(p(O) (10000 U 1lp1)lp(00 ve)) H(p(0)pe + kC vi) e e m(tt)²
J(H,, v1)) = lim A1-0
where the first term does not penalize abrupt parameter changes, but abrupt changes in parameter distribution. Note that the Kullback-Leibler (KI.) divergence term be written
(xz(p(0) 14, + St, v2 + 8, b HL ve))
Eo { log p( O |4 + ju At, + i. At) log p(0 p4; We)
0,
A1)
1
[0, - pp:
1
Σ Eo,
log [T (U: + [;,ti]
+
log [2rrss]
2(1; A.)
2
20;
2
Σ 1
27 + út
p
1
log[2n [2 (vr + ا;
+
A1)2
log [og[2
20%
20s
2
At
11 (At)²
log (1 1+ 02
Σ
At
20;
20;
2
Σ
(At)²
2 14
4 2
and that the entropy term is
Hlp(0)m (1)))
Σ 2
log 2Tee:;
(143)
which means that our objective is effectively
a(H and the corresponding Lagrangian is
1 2 1 0?
1
log(2rre;) + kC vi) e-
2n
1/3
41j
1
2
02
U,,, 2n U/
log(2)re;;) + kL vi) e-
4n
2
Here, the effective metric is the Fisher information metric for normal distribution.
I.2 Simplifying the Euler-Lagrange equations Taking derivatives, the EL equations read
d j4
tL(u, v)
dt
v1
Op
(144)
d
1 0;
ac(u, v)
17
di
2
2
Iv:
2 1/4
2v
20
33

---

Simplifying, these become
U/
ac(H. v)
? 144 = -nku;
Op.. (L((,,)) v)
(145)
ü; - - 2n
ko?
+
dv:
2
Consider local (quadratic) approximation of the loss, i.e.
1 ((0.) [(0o) + g" (0. 0o) / (0,- -00)² -0o).
(146)
where g is the local gradient and H is the local Hessian. Averaging this quantity over 0, and Og.
1 C(u, v) == Eg,,, (2(0.)) :- g" (m, 1o) 1 j(u, -00" H(,, 14) Σ Han + const.
where 'const.' denotes terms we can ignore. Using this particular L(us, v). the EL equations become
i4
1/ pi
0;
+ ?
nkvn
+
Σ
Hij(1s
1430)
j
(147)
1:
0?
8 - --y00 = 21
1/3
+
2
1/3
which are the equations that appear in the main text.
1.3 The Euler-Lagrange equations in the overdamped limit
The EL equations simplify somewhat in the overdamped (large 7) limit?
nku; i4
9i +
Hig(1s 1500
7
(148)
ú; 20 Vb
kHav?
Note that, in this limit, E2 influences how μ4 evolves (by modulating the effective learning rate), but v; evolves independently of μu, at least time scales where the quadratic loss landscape approximation remains valid. The entropic term U; tends to increase the variance, and the loss gradient term tends to decrease the variance; they balance when
1
U/
(149)
H
which implies that narrower loss landscape basins (high H&)) produce small uncertainties, and broader basins (low H;s) produce high uncertainties. This is intuitively reasonable, and reflects the selection of the 'simplest" model compatible with the data. Meanwhile, when the entropic term is absent, in this limit 0; 0.
Suppose that the learner has converged to a 'good' (i.c., deep) local or global minimum. When there task transition, learner is probably no longer in a good local minimum), which means that H;; suddenly decreases. Eq. 148ssas that optimal learning dynamics involves sudden increase in variance, which persists until nú equilibrates to its new value. By Eq. 149 this value is proportional to the new 1/H
generally expects the local landscape to longer be curved (since the
"T justify this reduction more rigorously, we could have performed singular perturbation analysis. This analysis not particularly enlightening, merely report the result.
34

---

1.4 Qualitative behavior outside of the overdamped limit
Outside of the overdamped limit, the EL equations for μi and D influence each other: changes in μa contribute an effective 'force' that affects 2/ and affects both the effective discounting rate and learning rate the μ4, equation. Because the ju term has the same sign the entropic term, it plays the same qualitative role, and can produce increases in variance.
Again suppose that there is task transition after the learner has converged to 'good' minimum, that H suddenly decreases. Optimal learning dynamics again involves a sudden increase variance, this time driven by both the entropic term and the juţ term.
Unlike the overdamped case, if His does not change after a task transition but gi does (i.e., the location, but not size, of the basin has changed), the jáž produces transient increase in variance. After μr equilibrates, the j term becomes zero, and v; approaches the same equilibrium value as in the overdamped case (Eq. 149).
35

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We claim we develop normative framework, and do so.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Discussed at the end.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
36

---

Justification:
See appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We do not do many experiments, but include all relevant code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
37

---

Answer:
[Yes]
Justification:
We do not do many experiments, but include all relevant code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7/nips.com wwtpcccpppllccccccc guwdes/ddeeeeee Cooeeeoodemsssseeeeee /ttp:///b///uudddd///eeeeeeeesss ssiniioiscyp for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Yes (and we include code), but it is worth noting that parameter details are a minor part of the paper mainly relevant Fig. 2.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[NA]
Justification: Statistics not really relevant here.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
38

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, implicitly. (The relevant experiments can be done on a laptop.)
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer: [Yes]
Justification:
No one was harmed in the making of this theory paper.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
We did not feel it was relevant here.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
39

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Not relevant for fairly abstract pure theory paper like this.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The MNIST and CIFAR-10 papers are cited.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
40

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
41

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [No]
Justification: No LLMs used for nontrivial tasks.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy fattps://neurii: http://s//////ooeeeeeee///////// for what should or should not be described.
42


Reviews:
[{"review_id": "6nkHyXVQZH", "reviewer": "Reviewer_BRpX", "summary": "In this work, the authors develop a framework for describing existing learning algorithms. Across a variety of algorithms, they work towards identifying rules for which each particular algorithm is the optimal algorithm.", "questions": "how do the results relate to practice? are there experiments that would made this paper more relevant to practice?", "limitations": "yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n* the authors study the important problem of characterizing learning algorithms\n* normative rules for learning algorithms is an interesting way to characterize learning algorithms \n\nWeaknesses:\n\n* no experiments for identifying whether or not the identified rules are actually realistic in practice (as a basic example: do continuous time approximations actually match what happens in practice? one would think a priori that infinitely many extraordinarily small learning rates would actually significantly change navigation of the loss landscape for DL optimization tasks). \n* the paper is rather confusing, treating prior work as well known and obvious (e.g., the amari insight on line 141) and considered theory as obvious (e.g., line 118 or line 80). ideally, the authors would work to guide the reader better through their work.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback, and think that it has improved the paper. We would first like to clarify a minor but conceptually important point.\n\n**Our theory does not require infinitesimally small learning rates.**\n\n> (as a basic example: do continuous time approximations actually match what happens in practice? one would think a priori that infinitely many extraordinarily small learning rates would actually significantly change navigation of the loss landscape for DL optimization tasks)\n\nThe learning rules we derive in this paper, including well-known rules like gradient descent with momentum, and Adam, do not require infinitesimally small parameter updates. See, for example, our derivation of gradient descent with momentum (Sec. 2; see especially Eq. 7). In the final learning rules, the step size $\\\\Delta t$ can be arbitrarily large. \n\n(As a caveat, the precise *form* of the learning rule definitely depends on the step size. For example, in Eq. 7, small step sizes look more like gradient descent, and larger step sizes look more like Newton's method.)\n\nWe agree that, if our theory proposed only learning rules with extraordinarily small learning rates, it would not be practical. However, it does not propose this. Rather, the continuous-time formulation should be mostly viewed as a trick for getting concrete math results (since it enables the use of the Euler-Lagrange equations, our central tool). In the final learning rules, updates do not have to be infinitesimally small. For concrete examples with finite $\\Delta t$, see (for example) Eq. 7, Eq. 10, and Eq. 14.\n\n**On guiding the reader through our results.**\n\n> the paper is rather confusing, treating prior work as well known and obvious (e.g., the amari insight on line 141) and considered theory as obvious (e.g., line 118 or line 80). ideally, the authors would work to guide the reader better through their work.\n\nThis is a fair point, and we could do a better job clarifying our approach throughout the paper. While we do think certain concepts may be familiar to some readers---for example, the continuous-time objectives we discuss on lines 118 and 80 may be familiar to readers familiar with control theory---we acknowledge that this will not be true for all readers. We have already tried to add context where possible (for example, when we discuss Amari's insight around line 141, we also go on to describe that insight using the terminology of this paper), but we will add more context to further improve clarity.\n\nTo address this concern, we plan to improve the final version of the paper in two ways. First, we can introduce some of the math and ideas a bit more gently, and add additional context where possible. It is certainly true that some sections are quite dense in the current version. Second, we will add more figures and examples to provide readers another way to understand what is going on. We think these two changes will greatly improve the clarity of the paper.\n\n**On practical consequences of our framework.**\n\n> how do the results relate to practice? are there experiments that would made this paper more relevant to practice?\n\nThis is a good question. First, our main aim with this paper is to provide a framework for unifying and better understanding existing learning rules. This aim is admittedly more modest than trying to derive new state-of-the-art learning rules, but we think we have achieved it reasonably well. Other recent machine learning work also tries to do this (see, for example, \"I-Con: A Unifying Framework for Representation Learning\" by Alshammari et al., which we cite). \n\nHowever, it is definitely true that one can be more confident in a framework like ours if it derives novel learning rules that perform well on nontrivial machine learning tasks. We have derived a few novel learning rules (for example, in Eq. 7 we mention a novel 'ballistic' learning rule which involves the square root of the Hessian), but have not seriously evaluated their performance in the current version of the paper.\n\nWe are happy to perform new experiments to test the novel learning rules suggested by our framework, and think these experiments would improve the quality of the paper and more clearly demonstrate the power of the framework. We are open to suggestions for what experiments may be most interesting, but we think that since the main goal of the paper is to establish the framework and rederive known learning rules, we will try to keep things simple. We will show some simple toy examples (e.g., we will compare different learning rules for a convex, quadratic loss) and slightly more complex examples related to image classification (e.g., on simple data sets like MNIST and CIFAR-10). These are the kinds of examples used in a variety of benchmarking papers, for example to study the performance of Adam. \n\nWe have not yet run these experiments, partly since the rebuttal period does not allow explicitly showing new results, but do not expect they will be difficult to perform."}, {"review_id": "23dYouTQxg", "reviewer": "Reviewer_WGtm", "summary": "The paper presents a framework to derive generally learning rules retrieving known ones in the process. Namely, the authors propose to define learning rules by finding continuous-time parameter curves minimizing ${}^{(1)}$ some average future discounted costs comprising the loss but also some parameters-specific costs enforcing e.g. smoothness of parameter changes or beliefs in the process changes. \nThe authors illustrate their framework to retrieve (i) momentum SGD, (ii) natural gradient descent, (iii) adaptive algorithms à la Adam. They also apply their framework to get insights on continual learning approaches. A long appendix provides detailed computations and explanations for each setting\n\n\n${}^{(1)}$ the authors solve the Euler-Lagrange equations associated to their objective, but the Euler-Lagrange equations only give stationary points for the objective in terms of continuous time curves.", "questions": "- How can the framework explain why Adam could work better than SGD with momentum? More generally the choice of objective to solve for with Euler-Lagrange equations appear as arbitrary as choosing an algorithm, and does not provide much insights on the choice of the algorithm.\n- Could the authors use their framework to\n  - derive Nesterov's accelerated method?\n  - derive learning-rate tuners such as line-searches?\n  - derive algorithms using exponential moving average?\n  - derive sign-sgd (a variant of sgd that was linked to adam)\n- A current difficulty of usual frameworks is to be able to take into account phenomena like edge of stability [1]. Does the proposed framework help to integrate such effects?\n- Could the authors test some new algorithms they propose?\n\n\n[1] Cohen, Jeremy M., et al. \"Gradient descent on neural networks typically occurs at the edge of stability.\" arXiv preprint arXiv:2103.00065 (2021).", "limitations": "See weaknesses and questions.", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The idea of considering \"long term\" optimization horizons could be interesting. The authors should integrate algorithms like Lookahead [1] in their framework since it bears similar motivations. \n- The derivations in the appendix are clear. \n- The paper is well written and pleasant to read (see small remark about tone below though).\n- Some new algorithms are derived, though not tested. \n- While the momentum and natural gradient viewpoints are very well known from a usual variational viewpoint, the viewpoint for Adam appears new.\n\n**Weaknesses**\n\n*Main*\n- The field of optimization may be older than the authors themselves. There already exists frameworks in optimization to design algorithms in a principled way. An essential example is the one developed by e..g A. Nemirovisky and D. Rudin [2] that influenced optimization for years. Their framework defines what could the best possible algorithm for a given class of problems, read for example [3] for a gentler introduction. Such a framework had the benefits to understand how different algorithms compare given some preliminary information on it. It also bore fruits: the accelerated algorithm of Nesterov was motivated by finding the best algorithm for smooth and convex problems. In comparison, the framework proposed by the authors does not help us understand how these algorithms compare for a given class of problems. The framework provides ad-hoc explanations to well-known algorithms.\n- \"it it makes it easier to generate new learning rules from a principled starting point, and to justify them without empirical guesswork\" Actually an algorithm is made to be applied. Experiments are the gold standard, the truth to understand the relevance of the theory. The ultimate test for the proposed framework is to get a new algorithm that perform better in practice (like sharpness aware minimization was designed after studying sharpening effects). The authors actually have some new algorithms like the \"ballistic one\", why aren't they tested empirically? At least each intuition that the framework provides could be tested empirically to see how relevant the framework is. Why aren't each intuition tested empirically to analyze their relevance? Is it left for other people to do this work?\n- Numerous approximations are sometimes needed to find the desired algorithm in the given framework (see Adam for example).\n- All these algorithms had previous justifications: momentum is well known to be motivated from a physics viewpoint (it's called Polyak's heavy ball), Adam is justified as a form of Adagrad, Adagrad itself is justified a s a form of a normalized gradient, there are many studies from those viewpoints that gave new intuitions like understanding how normalized gradient can adapt to the edge of stability phenomenon. The authors need to justify what brings their approach. \"One rule to rule them all\" is not a good motivation (especially given that their framework won't encompass all possible algorithms).\n\nThe paper presents many algorithms through the given framework. It could be much better to focus on one algorithm each time and provide numerical experiments illustrating either new algorithms that the framework gave rise to (like the ballistic one) or to illustrate empirically some intuitions that the framework gave. Even the \"Newton's method\" that the framework gives in section 3 is not the actual Newton method. Maybe it gives a much better algorithm, maybe not at all. Probably this framework completely misses the need to adjust for the learning rate once a learning rate rule is found. A Newton method needs a priori to be damped for example, is it the case for what the authors propose? There could be interesting fruits from the framework but currently trying to frame each algorithm through one lens does not add much relevance to the approach. Experiments will prove the relevance.\n\n*Minor*\n- \"Natural gradient descent is not second-order optimization in disguise\" The authors are presenting generic preconditioning methods. Indeed they are not presenting natural gradient descent which is one particular instance. Updates using pre-conditioners are well motivated from a variational viewpoint. Actually one can go beyond using simple Mahalanobis distances and consider Bregman divergences, or any gradient mapping coming from a partial linearization of the objective. \n\n*Additional notes*:\n- A humble tone may be best to invite the reader to understand the relevance of the idea rather than a forceful tone affirming for example that \"[this] framework is significant\". A paper is written to present an idea not to argue or convince about its significance, let the reader judge by themselves. \n- There is a missing identity matrix in the \"effective learning rate\" of the Newton step and it's not an \"effective learning rate\", learning rates are scalars, a Newton method does more than multiplying with a scalar.\n\n[1] Zhang, Michael, et al. \"Lookahead optimizer: k steps forward, 1 step back.\" Advances in neural information processing systems 32 (2019).\n[2] Nemirovskij, Arkadij Semenovič, and David Borisovich Yudin. \"Problem complexity and method efficiency in optimization.\" (1983).\n[3] Nesterov, Yurii. Introductory lectures on convex optimization: A basic course. Vol. 87. Springer Science & Business Media, 2013.", "was_revised": true, "final_justification": "I fully acknowledge that I am not the right audience for this paper, while it appears that other reviewers like the paper. A fresh perspective on optimization is needed. This one may be wrong, and experiments (even simple ones like the snippet I proposed) would have tremendously helped. I cannot recommend this paper for acceptance, but I shall not block it. So I will increase my score to 3.", "rebuttal": "We deeply appreciate the reviewer's time and detailed feedback, and think that your comments have improved the paper. We will respond to some of the major concerns first.\n\n**What is the value of a normative framework for learning rules?**\n\n> \"it makes it easier to generate new learning rules from a principled starting point, and to justify them without empirical guesswork\" Actually an algorithm is made to be applied. Experiments are the gold standard, the truth to understand the relevance of the theory. \n\n> Experiments will prove the relevance.\n\nThis touches on an important philosophical point, which we will clarify in a revised version of the paper. In the end, most machine learning practitioners would like an algorithm that helps their models learn as quickly as possible. What's the point of theory here?\n\nWe think neither theory nor experiment is the endpoint. Rather, the process is best viewed as a 'virtuous cycle': experiments identify which learning rules and hyperparameters work well in practice; theory can help clarify why those algorithms work as well as they do, and perhaps suggest improvements; those improvements can be tested in further experiments; and so on. Said differently, the ideal workflow is not theory -> experiment, nor is it experiment -> theory; it is theory -> experiment -> theory -> experiment ...\n\nOur statement \"it makes it easier ...\" does *not* mean that there is nothing more to be said, or that no experiments should be done. It means that we provide a principled framework for justifying learning rules from a theoretical perspective, which can *then* be tested in experiments. It is not the only framework possible, but it is fairly broad (for example, it encompasses a number of known learning rules, as we show in the paper) and has a number of advantages.\n\nOne of these advantages is that the 'goodness' of a given learning rule, given a set of assumptions, can be quantified by a single scalar, rather than more complicated quantities like convergence rates. This means that there is an unambiguous way to compare learning rules; one rule is better than another if the corresponding number is lower (i.e., lower loss and more efficient navigation). Another advantage is that it generally identifies a unique rule and set of hyperparameters as optimal. \n\nWe emphasize that the idea of this variational approach is not completely new, and has been successfully used in previous literature to derive interesting insights about optimization methods. The Wibisono et al. 2016 paper we cite is a good example.\n\n**Our framework is not ad-hoc.**\n\n> the framework proposed by the authors does not help us understand how these algorithms compare for a given class of problems. The framework provides ad-hoc explanations to well-known algorithms.\n\n> ... the choice of objective to solve for with Euler-Lagrange equations appear as arbitrary as choosing an algorithm, and does not provide much insights on the choice of the algorithm.\n\nWe think it is inaccurate to describe our framework as ad-hoc. It is 'normative' in the sense that, given a set of assumptions, it identifies one or a family of learning rules as 'optimal'. Moreover, these assumptions are fairly reasonable, involving (among other things) minimizing the loss, efficiently navigating through the loss landscape, and temporal discounting. The framework says, *given some assumptions*, that one or a family of rules is best in a precise sense. \n\nOne potential objection to this approach is, why these assumptions and not others? Equivalently, where does the objective function come from? There is admittedly infinite freedom here, which is why we use fairly simple assumptions throughout. For example, we only consider quadratic kinetic terms like $\\dot{x}^2$, and an exponential temporal discounting function. One becomes more confident that these assumptions are reasonable by rederiving known rules. \n\nOn the other hand, it is worth emphasizing that any approach of this type has to assume *something*, and even in classical work (e.g., by Nesterov) this is true. The assumptions are taken as primitive, and we can use theory tools to derive the consequences. Whether the resulting rules apply in more realistic situations is indeed a question best clarified by experiments, and by proceeding with the 'virtuous cycle' described above.\n\n**Not all algorithms are encompassed by the framework.**\n\n> (especially given that their framework won't encompass all possible algorithms).\n\nIt is absolutely true that not all optimization algorithms can be derived via our framework, and we think this is a feature rather than a bug. In particular, this fact can be used to suggest changes to well-known algorithms. For example, although you point out that:\n\n> Numerous approximations are sometimes needed to find the desired algorithm in the given framework (see Adam for example)\n\n...a different interpretation is that the framework identifies appropriate ways Adam could be modified. While it remains to be seen whether these modifications perform better in practice, these hypotheses are at least interesting consequences of the framework. The framework does not provide 'just-so' stories that can be used to 'derive' absolutely any algorithm.\n\nAs a minor point, algorithms like line search (which do not necessarily involve smooth changes in a parameter, since it 'jumps' around to look for a minimum) are probably not encompassed by our framework, which only concerns optimizers that involve smooth parameter changes.\n\n**How is this framework's justifications different than other justifications?**\n\n> All these algorithms had previous justifications\n\nThis is absolutely true. But one feature of our framework that we think is nice is that it provides quantitative guidance about how good an algorithm is through the aforementioned 'goodness' score (essentially, the objective $J$). This is somewhat different than how momentum is usually justified: while the physics analogy suggests including momentum is intuitively *reasonable*, our framework shows that it is *quantitatively optimal* under certain assumptions. A similar statement holds for Adam and the other algorithms we discuss. In short, the framework offers quantitative insight (that algorithm X is optimal, and *moreover* Y amount better than algorithm Z) that complements existing insight. Not all existing insight is qualitative---for example, convergence rates are not---but a single performance scalar, and a claim regarding optimality, is arguably somewhat simpler, or at the very least helpful.\n\n**Where are the experiments?**\n\n> The authors actually have some new algorithms like the \"ballistic one\", why aren't they tested empirically? \n\n> Could the authors test some new algorithms they propose?\n\nAlthough our aim is primarily to develop a unifying theoretical framework for various existing learning rules (in the spirit of other recent work in machine learning, like Alshammari et al. 2025), we agree that additional experiments will help demonstrate the utility of the framework and learning rules derived from it. \n\nIn a revised version of the paper, we plan to explicitly test some of these algorithms (e.g., the ballistic one) in a few simple cases, including a toy convex quadratic loss and in simple image classification problems (e.g., MNIST, CIFAR-10). These types of problems are already used in many papers benchmarking algorithms like Adam, so although they are simple, they seem appropriate here. Our goal is not necessarily to derive a state-of-the-art algorithm that performs better than existing algorithms on a frontier problem. \n\nWe cannot show the results here, since it is not possible to include links or PDFs during the rebuttal, but do not anticipate that these experiments will be difficult to perform. \n\n**Extensions of the framework.**\n\n> Could the authors use their framework to derive ...\n\nThere isn't space to go into details, but in short, the framework definitely includes Nesterov's accelerated method and exponential moving average as examples. The accelerated method can be derived assuming a similar objective to the one from Sec. 3, and algorithms involving an EMA can be derived using an objective similar to Adam. (The insight is that EMA is optimal when past gradients are informative about current and future gradients.) We will add both derivations to the paper.\n\nAs mentioned above, the framework probably doesn't encompass line searches or algorithms where parameter updates are 'jumpy', like sign-SGD (although maybe sparsity or compression constraints can help). \n\n**Miscellaneous and minor concerns.**\n\n- In the revised version, we will add many more numerical experiments and figures. This both helps show the framework is useful, and makes the paper easier to follow.\n- We are happy to modify the tone somewhat, and agree that the lack of humbleness in certain places could put off some readers.\n- We will add the identity matrix to our expression for $\\eta_{eff}$. This wasn't a typo, just slightly sloppy notation. It's true that the term \"learning rate\" generally refers to a scalar, but we think it is a helpful and not-too-egregious abuse of terminology to speak about a 'non-scalar learning rate' in the context of preconditioning methods.\n- We agree that the Lookahead optimizer provides an interesting test case for our framework. We will cite it and show how to derive it using our framework.\n- We agree that the edge of stability phenomenon is extremely interesting, but think it would require a completely separate analysis to link our framework to it. The problem is that most learning rules are derived using simple assumptions (e.g., a convex and often quadratic loss); these simple assumptions are probably insufficient for modeling a phenomenon associated with the highly non-convex loss landscapes of neural networks. We will cite this paper and note this as a limitation of our work, but not necessarily as a limitation of the framework overall."}, {"review_id": "4ENDa4W6QW", "reviewer": "Reviewer_AiaD", "summary": "The paper provides a general framework to model learning dynamics. The focus is not a parameter that minimize a loss, but rather a trajectory (on parameter space) that minimize a novel objective: an integral over the trajectory of a combination of loss and parameter change. The optimal solution is given by the Euler Lagrange equations providing two boundary conditions: an initial parameter, and the integral to be minimized.\n\nDifferent hyperparameters characterize the weight of these two components, as well as their specific form. Importantly, for some specific choices of hyperparameters, the optimal solution can be expressed analitically, and the discrete time difference along such optimal trajectory corresponds to learning rules prescribed by established optimizers. \nThis motivates the relevance of this novel framework: it can recover and generalize multiple optimizers.", "questions": "Can the authors elaborate on the second boundary condition of the Euler Lagrange solutions?\n\nIn Line 1099 you state that Equation (133) \"matches the main text form up to the prefactor $1/(2v_i^2$\". Referring to the equation before line 231 in the main text. Can the authors expand on this mismatch? Why is there such difference? And how is the equation in the main text derived?\n\nCan the authors provide an example where rotational forces exists?", "limitations": "yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "The paper is well written an easy to read. The various elements are properly discussed and the Appendix provides several useful examples and a clear derivation of the results.\n\nThe proposed framework is, to my knowledge, novel. The derivation are correct and multiple learning rules are derived from the same framework. For this reason I consider this work very relevant and useful for the community.\n\nWhat I personally find extremely interesting is the analysis in Section 4, where the loss is assumed to be quadratic and the parameter geometry is given by the Fisher. This lead to a proper distinction of Hessian and Generalized Gauss Newton which, as the authors also notice, are too often simply viewed as surrogates of each other. The authors framework provides a rigorous derivation of their role.\n\nI found the assumption in Section 5 on gradient noise to be a bit arbitrary. It would be very helpful if the authors can show some references motivating (at least partially) this assumption. Aren't there other works that motivates Adam through this kind of assumption? \n\nI like the presence of the $f$ term in the general form, but I personally don't see how the rotational term can be of any use in practice. Is there a practical use case that can be shown?\n\nOverall, I really like the proposed framework and the approach of finding solutions through Euler Lagrange. My only concern (which is actually my only main concern about the paper) is on the boundary conditions. First of all, this part is a bit swiped under the carpet in the Appendix, and I don't really like that. The authors should be more clear about this. \nThen, more specifically, no problem about the initial parameter condition, that's very reasonable. My concern is about \"The remaining degree of freedom\" that \"is chosen so that J is minimized\". This feels circular. Isn't the whole point of Euler Lagrange solution to minimize J? My understanding is that Euler Lagrange provides LOCALLY optimal solutions, i.e. optimal among small perturbations. How is this compatible with your choice of boundary condition? Please elaborate on this.\n\nSome minor details:\n\n-In Line 42 I'd recommend changing \"sensitive to structure in parameter space\" to \"sensitive to non-local structure in parameter space\" to better clarify what is actually being generalized of the natural gradient descent methods\n\n-In Line 855 there is a typo: \"J to minimized\" should be changed in \"J to be minimized\"\n\n-In Equations 123 and 124, isn't that $J$ supposed to be a $\\mathbf{J}$? Is it a typo or am I missing something", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their support, close reading, and detailed comments. We think these comments have greatly helped strengthen the paper. Let's begin by addressing your most major concern.\n\n**On boundary conditions.**\n\n> Can the authors elaborate on the second boundary condition of the Euler Lagrange solutions?\n\n> Overall, I really like the proposed framework and the approach of finding solutions through Euler Lagrange. My only concern (which is actually my only main concern about the paper) is on the boundary conditions. First of all, this part is a bit swiped under the carpet in the Appendix, and I don't really like that. The authors should be more clear about this. Then, more specifically, no problem about the initial parameter condition, that's very reasonable. My concern is about \"The remaining degree of freedom\" that \"is chosen so that J is minimized\". This feels circular. Isn't the whole point of Euler Lagrange solution to minimize J? My understanding is that Euler Lagrange provides LOCALLY optimal solutions, i.e. optimal among small perturbations. How is this compatible with your choice of boundary condition? Please elaborate on this.\n\nThis is indeed a subtle point, and we agree that it would be helpful to clarify it in the paper and SI. **We will add an extra appendix section to discuss this issue explicitly.** \n\nOur most general objective (Eq. 2) involves expectations, but once these are evaluated, one ends up with an objective of the form\n\n$$ J[ \\mathbf{x}(t) ] = \\int_0^{\\infty} F[ \\mathbf{x}(t), \\dot{\\mathbf{x}}(t), t ] \\ dt = \\int_0^{\\infty} \\frac{1}{2} \\left[ \\dot{\\mathbf{x}}(t) - \\mathbf{f}(\\mathbf{x}(t), t) \\right]^T \\mathbf{G}(\\mathbf{x}(t), t) \\left[ \\dot{\\mathbf{x}}(t) - \\mathbf{f}(\\mathbf{x}(t), t) \\right] e^{- \\gamma t} + k \\mathcal{L}(\\mathbf{x}(t), t) e^{- \\gamma t} \\ dt$$\nwhere the scalar-valued function $F$ is smooth, and moreover equals a 'kinetic' term plus a 'loss' term. We would like to minimize $J$ over all possible smooth trajectories $\\mathbf{x}(t)$ with a fixed initial point $\\mathbf{x}(0) = \\mathbf{x}_0$.\n\nThe idea is to optimize $J$ via a *two-step* optimization: we find the function $\\mathbf{x}(t)$ that minimizes $J$ for a specified initial point $\\mathbf{x}(0)$ and final point $\\mathbf{x}(\\infty)$, and *then* (since our initial problem only prescribes $\\mathbf{x}(0)$) we choose the solution whose final point corresponds to the smallest value of $J$. Note that the latter minimization—over a finite set of $D$ numbers $\\mathbf{x}(\\infty) = ( x_1(\\infty), ..., x_D(\\infty) )^T$—is different than the more challenging optimization over a function space that we started with. The argument is not circular, since it reduces a minimization over functions to a (much easier) minimization over a finite number of scalars.\n\nHow can we minimize $J$ for prescribed start- and endpoints? A standard result in the calculus of variations says that the stationary points of $J$ given these boundary conditions correspond to the solutions of the Euler-Lagrange equations. But a priori, these stationary points might be local maxima or saddle points. It turns out that if the matrix $\\mathbf{G}$ is positive definite for all $\\mathbf{x}$ and $t$, and the other functions appearing in $J$ are smooth in their arguments, *additional* results from the calculus of variations say that stationary points correspond to local minima, and that the global minimum is among them. This is true even if the loss is non-convex, as long as it is smooth.  \n\n(If the loss is non-smooth, because of functions like ReLUs, the result still holds as long as the non-smoothness is not too bad.)\n\nIt is possible that there is not a unique solution to the EL equations, as mentioned, e.g., in Appendix B (for example, one can navigate around an object either to the left or to the right, and these paths usually involve an equal amount of effort). But we usually have uniqueness, and uniqueness implies that the EL equation solution is the global minimum. If we do *not* have uniqueness, we simply choose the solution(s) with the lowest $J$ value, which must correspond to the global minimum.\n\n**In summary, the roadmap is this:** we start by hoping to minimize $J$ over a space of functions that correspond to trajectories through parameter space. This space of functions is only constrained by the condition that $\\mathbf{x}(0)$ is fixed. Using some calculus of variations results, we can argue that (because of the special form of the objective) stationary points of $J$ correspond to minima, that the global minimum is one of these. Finally, getting this solution required us to specify $\\mathbf{x}(\\infty)$, so we actually obtained a family of possible solutions, each corresponding to a different value of $J$. The solution we ultimately care about is the one which makes $J$ smallest, and we can find it by directly minimizing $J$ with respect to the $D$ scalars contained in $\\mathbf{x}(\\infty)$. \n\nIn the aforementioned new appendix, we will spell out the required technical conditions (e.g., the convexity and coercivity of the kinetic term) more precisely. But we absolutely agree that getting all of these details straight is crucial for being confident in our framework, and thank you for pointing out that we could be clearer about this.\n\n**On rotational forces.**\n\n> Can the authors provide an example where rotational forces exists?\n\n> I like the presence of the $\\mathbf{f}$ term in the general form, but I personally don't see how the rotational term can be of any use in practice. Is there a practical use case that can be shown?\n\nThis is a good question, and we will address it in a revised version of the paper. We think the $\\mathbf{f}$ term is probably not that useful for modeling machine learning optimization, but may be useful for modeling optimization processes in biological neural networks. Processes like weight decay are known to occur, but more generally one can imagine dynamics that tend to drive a system towards some non-loss-related fixed point. In the simplest case, we can imagine that each parameter decays to zero at the same rate, but a more complicated case could involve each parameter going to zero (or some other value) with different rates, and possibly in a coupled fashion. \n\nAn explicit but tractable model of this has \n$$\\mathbf{f}(\\mathbf{x}) = - \\mathbf{A} \\mathbf{x}$$ \nwhere $\\mathbf{A}$ has eigenvalues whose real parts are positive, so that in the long run $\\mathbf{f}$ drives the system towards zero. In general, $\\mathbf{A}$ is not necessarily symmetric (i.e., the effect of $x_j$ on $x_i$ may not be the same as the effect of $x_i$ on $x_j$), which means that it is possible for $\\mathbf{A}$ to have complex eigenvalues, and hence it is possible that its dynamics include a rotational component.\n\nThe pure rotation case that we use as an example is probably unrealistic (although we think it is cool), but a case involving a mix of decay and rotation appears to be more biologically plausible.\n\n**Gradient noise and Adam.**\n\n> I found the assumption in Section 5 on gradient noise to be a bit arbitrary. It would be very helpful if the authors can show some references motivating (at least partially) this assumption. Aren't there other works that motivates Adam through this kind of assumption?\n\nIt is funny that you ask this. After the NeurIPS paper submission deadline, another paper appeared that proposes a way to motivate Adam's effectiveness: \"In Search of Adam’s Secret Sauce\" by Orvieto and Gower. It *also* motivates Adam as effectively implementing an online inference algorithm by assuming noisy observations of gradients. On the other hand, they are not as explicit about a connection between the local curvature of the loss landscape, and the size of the gradient observation noise.\n\nWe will cite this paper and find some other references to support our approach.\n\nIn any case, we definitely agree that this section can get confusing, and will think hard about how we might be able to simplify/clarify the derivation.\n\n**Miscellaneous and minor concerns.**\n\n> In Line 42 I'd recommend changing \"sensitive to structure in parameter space\" to \"sensitive to non-local structure in parameter space\"\n\nWe can make this change, thanks for the suggestion.\n\n> In Line 855 there is a typo: \"J to minimized\" should be changed in \"J to be minimized\"\n\nGood catch, will fix.\n\n> In Equations 123 and 124, isn't that $J$\n supposed to be a $\\mathbf{J}$? Is it a typo or am I missing something\n\nGood catch. In those equations, we use $J_{ij}$ to denote the components of the Jacobian matrix $\\mathbf{J}$. Although it is standard to use non-boldface to indicate components, we admit this may be confusing, and will try to clarify this better in the text. Sorry for the confusion.\n\n> In Line 1099 you state that Equation (133) \"matches the main text form up to the prefactor $1/(2 v_i^2)$. Referring to the equation before line 231 in the main text. Can the authors expand on this mismatch? Why is there such difference? And how is the equation in the main text derived?\n\nEq. 133 in the SI corresponds to the more conceptually simple but more technically annoying gradient noise model, which has\n$$ \\mathbf{g}_t \\sim \\mathcal{N}( \\mathbf{m}_t, \\mathbf{V}_t / \\Delta t) .$$\nIt is a little annoying because the corresponding log-likelihood depends non-quadratically on $\\mathbf{V}_t$, which complicates the EL equations a bit. In the main text, we consider the ad-hoc assumption that one effectively has *separate/uncoupled* noisy observations of the mean $\\mathbf{m}_t$ and the variance $\\mathbf{V}_t$. The point of the argument around Eq. 133 is that the ad-hoc assumption yields almost the same EL equations as the non-ad-hoc assumption, since the difference is \n\n$$\\frac{[ v_i - (g_i - m_i)^2]}{2 v_i^2} \\text{ versus } \\frac{[ v_i - (g_i - m_i)^2]}{2 \\sigma_2^2} . $$\n\nThey approximately agree if $v_i$ equilibrates quickly and $v_i \\approx \\sigma_2$. We can clarify this point in the text."}, {"review_id": "OjSirww0Jl", "reviewer": "Reviewer_KkuV", "summary": "Different learning rules, such as variants of gradient descent, are often compared empirically. Theoretical analysis generally looks at convergence rates or guarantees. This paper takes a slightly different perspective and tries to understand what problem each rule is implicitly solving. By framing parameter updates as a multi-step, continuous-time optimal control problem, the authors find they can derive many popular learning rules from a common framework. The objective for deriving the parameter updates is composed of two terms: a parameter change cost and the loss function. Different choices of the parameter change cost and approximations of the loss function give us different learning rules. Common tricks, like momentum, naturally fall out of the framework due to the multi-step optimization. Using a second-order approximation of the loss can give Netwon's method, gradient descent, or an in-between \"ballistic\" learning rule depending on the choice of time step and temporal discount factor. Incorporating different metrics in the parameter change cost can respect parameter geometry. For instance, using the Fisher information matrix gives us natural gradient descent. An important take-away from this analysis is that natural gradient descent is not actually second-order optimization. Rather, the metric governs the parameter velocity penalty and the Hessian measures loss landscape curvature. Making assumptions about partial controllability of parameters, such as existing external \"forces\" on parameters that may act in a state-dependent way, can produce non-gradient rules, such as rotational dynamics, and account for weight decay. Next, assuming the true loss is only partially observable and must be estimated, they show that adaptive optimizers, like Adam and RMSprop, can be derived. Finally, they show that some common continual learning strategies can fall out of modeling weight uncertainty and making the learning dynamics sensitive to this uncertainty.", "questions": "- Are there more subtle differences between Wibisono et al. [49] and the formulation considered in this paper, when not considering partial observability or controllability?\n- What are some situations in which one would want to use the ballistic version of the learning rule? What are some trade-offs of these different choices besides Hessian computation?", "limitations": "- The authors are clear that their analysis does not consider comptuational requirements and efficiency in its definition of optimal.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- Building a better theoretical understanding of what implicit assumptions different learning rules make and when to use them is an important problem. This paper provides a new perspective on this problem, deriving many different learning rules under a common framework.\n- New learning rules can possibly be derived from this framework, making different assumptions on parameter geometry, partial controllability, and partial observability. In fact, the ballistic learning rules appear to be novel, although it is not evaluated empirically.\n- The insight on the relationship between natural gradient descent and second order methods appears to be novel and makes a valuable distinction between the two.\n- The view of adaptive learning rules as modeling partial observability also appears novel and provides a nice explanation for what these algorithms are doing under the hood.\n- The discussion of continual learning algorithms, which may inject noise or reset unused weights, as being tied to a weight-uncertainty-aware learning rule is a nice justification of choices that are often empirically motivated.\n\nWeaknesses:\n- While the paper is primarily focused on theory, and is deriving known algorithms from a new perspective, there are some new variants of common learning rules put forth. It would be interesting and valuable to compare these different choices on some simple examples to better understand what trade-offs they make.\n- The framing of momentum as falling out of a continuous-time formulation appears in Wibisono et al. [49], which is briefly mentioned in the related work section. While this work does not consider partial observability, partial controllability, weight uncertainty, or second-order methods, it does account for momentum falling out of a multi-step perspective on optimization. Moreover, it uses the Bregman divergence as a distance metric. Natural gradient descent can be viewed as gradient descent with a specific choice of Bregman divergence [1, 2]. Therefore, while I do think the contributions of this paper are substantial and novel, the paper could do a better job relating its innovations to this prior work and making the distinctions clear.\n\nReferences:\n[1] G. Raskutti and S. Mukherjee, \"The information geometry of mirror descent,\" IEEE Transactions on Information Theory, 2015.\n[2] N. Wagener, C. Cheng, J. Sacks, and B. Boots, \"An online learning approach to model predictive control,\" RSS, 2019.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their support and useful feedback, which we think has greatly strengthened the paper. See below for detailed responses.\n\n**What is the relationship between our framework and the results from Wibisono et al.?**\n\n> The framing of momentum as falling out of a continuous-time formulation appears in Wibisono et al. [49], which is briefly mentioned in the related work section. While this work does not consider partial observability, partial controllability, weight uncertainty, or second-order methods, it does account for momentum falling out of a multi-step perspective on optimization. Moreover, it uses the Bregman divergence as a distance metric. Natural gradient descent can be viewed as gradient descent with a specific choice of Bregman divergence [1, 2]. Therefore, while I do think the contributions of this paper are substantial and novel, the paper could do a better job relating its innovations to this prior work and making the distinctions clear.\n\n> Are there more subtle differences between Wibisono et al. [49] and the formulation considered in this paper, when not considering partial observability or controllability?\n\nWe agree that the relationship between our work and the work of Wibisono et al. could be made clearer, and **we will explicitly address this issue in the revised version by including a new appendix section**. In brief:\n- There is a *key conceptual difference* between our work and that of Wibisono et al. Our variational problem can be viewed as optimal loss landscape navigation, while theirs can't.\n- Their use of the Bregman divergence makes their Lagrangian more general in some ways, although, as you point out, our framework is more general in different ways (partial observability, partial controllability, weight uncertainty, second-order methods...).\n\nThe first difference we mention is fairly interesting. To understand it, recall that our objective (ignoring partial controllability and observability, and explicit time-dependence in the loss) involves the Lagrangian\n\n$$ L(\\mathbf{x}, \\dot{\\mathbf{x}}, t) = e^{- \\gamma t} \\left\\\\{  \\frac{1}{2} \\dot{\\mathbf{x}}^T \\mathbf{G}(\\mathbf{x}, t) \\dot{\\mathbf{x}} + k \\mathcal{L}(\\mathbf{x})  \\right\\\\} .$$\n\nMeanwhile, if we replace some of their variable labels with ours so that comparison is easier, their objective involves the Lagrangian\n\n$$L(\\mathbf{x}, \\dot{\\mathbf{x}}, t) = e^{\\alpha_t + \\gamma_t} \\left\\\\{  D_h(  \\mathbf{x} + e^{-\\alpha_t} \\dot{\\mathbf{x}}, \\mathbf{x} )- e^{\\beta_t} \\mathcal{L}(\\mathbf{x}) \\right\\\\} $$\n\nwhere $D_h(\\mathbf{y}, \\mathbf{x})$ denotes a Bregman divergence, and where $\\gamma_t$ has a slightly different meaning than in our Lagrangian. The use of the Bregman divergence is one obvious difference between our framework and theirs, although it is not crucial, since we could easily have *also* used it. (In the new aforementioned appendix, we will talk about what changes when we do this.) \n\nThe more interesting difference is that their objective involves a *minus sign* (i.e., it has the form kinetic minus potential) whereas ours involves a *plus sign* (i.e., its form is kinetic plus potential). \n\nWhy does this matter? With the plus sign, each increment $L( \\mathbf{x}_t, \\dot{\\mathbf{x}}_t, t) \\Delta t$ that contributes to the overall objective $J$ is can be assumed to be nonnegative (if we shift the loss so that it equals zero at its global minimum), and more importantly these increments have a simple interpretation. The first term represents the 'cost' of movement, while the second represents the 'cost' of the current loss value. The total cost of any single parameter update is the sum of these two.\n\nThis fact implies that our framework's objective has the appealing interpretation that one is trying to navigate a loss landscape as efficiently as possible. The Wibisono et al. Lagrangian does not have the same interpretation, since the cost of movement and the cost incurred by the loss have different signs. \n\nOne might expect that this subtle sign difference does not matter, and that one can recover our Lagrangian from theirs by some reparameterization, but this is not true. It turns out that the sign difference fundamentally changes the nature of the solutions one obtains. This is easiest to see in the simplest possible case, where \n- we take $\\alpha_t = \\beta_t = \\gamma_t = 0$ in their Lagrangian, and we use $D_h(\\mathbf{y},\\mathbf{x}) = \\Vert \\mathbf{y} - \\mathbf{x} \\Vert_2^2$; and\n- we take $\\gamma_t = 0$ and $\\mathbf{G}$ to be the identity in our Lagrangian.\n\nAssuming a quadratic loss, one obtains two simple Lagrangians (quadratic kinetic term, quadratic potential) that differ only via the aforementioned sign difference. In this case, their Lagrangian yields sinusoidal solutions, while ours yields hyperbolic sines and cosines (or equivalently $e^{r t}$ and $e^{- r t}$ for some real $r$). These are very different!\n\n\n**When would one want to use the ballistic learning rule? More generally, when are different versions of the proposed learning algorithms useful?**\n\n> While the paper is primarily focused on theory, and is deriving known algorithms from a new perspective, there are some new variants of common learning rules put forth. It would be interesting and valuable to compare these different choices on some simple examples to better understand what trade-offs they make.\n\n> What are some situations in which one would want to use the ballistic version of the learning rule? What are some trade-offs of these different choices besides Hessian computation?\n\nThis is a great question, and we agree that we have not clearly addressed this in the current version of the paper. **In the revised version, we will include simple experiments that empirically compare different algorithms, including the ballistic one.** In the simplest case, we will compare the algorithms for a toy quadratic loss. We will go on to compare them for simple image classification problems (like MNIST and CIFAR-10). We think this is appropriate, since these kinds of simple tasks are already used to benchmark algorithms like Adam.\n\nWhat do we expect to see? Theoretically, the ballistic algorithm is optimal when one prioritizes good performance in the distant future (low temporal discounting, $\\gamma \\approx 0$). Hence, we might expect the algorithm to take steps that are costlier in the short-term, but better (in terms of minimizing the loss) in the long-term.\n\nConsider the quadratic loss case. If $\\lambda$ denotes an eigenvalue of the Hessian, Newton's method scales like $1/\\lambda$ along a given eigendirection. The ballistic method scales like $1/\\sqrt{\\lambda}$ along the same direction, which can be both bigger (if $\\lambda > 1$) and smaller (if $\\lambda < 1$). Hence, it can make the learning rate both faster and slower along a given direction than if one followed Newton's method. The larger consequence of this is that although some directions will converge more slowly, overall, the global minimum will be reached more quickly.\n\nOn the other hand, it is not obvious how this rule (which was derived as optimal assuming a simple quadratic loss) will perform given the highly non-convex loss landscapes associated with more realistic problems. One of the papers we cite (\"Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspective\" by Lin et al.) claims that removing the square root of Adam (which our framework identifies as related to the ballistic rule) actually *boosts* performance, which is contrary to the intuition we present here. There are probably many factors that make the real story more complicated. \n\nIt is worth emphasizing that, although we find our framework theoretically interesting and helpful for thinking about learning rules, we agree with the claim of another reviewer that the ultimate test for any learning rule is a real experiment on the kind of task one is actually interested in (e.g., training some neural network to do image classification)."}]

Author rebuttal:

Meta-review: I have the following concerns:

(1) The main contribution is a unifying framework, but I am not convinced of its significance. For a fixed set of phenomena (or, a fixed set of learning rules), one can find infinitely many ways to unify them. The theory space is much larger than the space of reality. To actually propose a unifying framework, I think the authors are obliged to predict/design/do something new with it. But the only thing the authors did is to explain existing ones. This also relates to some reviewers' criticism that it lacks practical relevance and experiments

(2) The paper does not achieve what it suggests to achieve. A key question raised by the authors is "Why do some learning rules work better than others?" But this question is unanswered, and I do not feel a hint of answer after reading the paper. I think those who support the acceptance of this paper should explain to me using authors' logic the answer to this question. This point also links to two other problems: (1) lack of significance and (2) the paper is poorly written (in terms of overclaiming)

(3) Novelty. I find the paper of limited novelty. As other reviewers have also pointed out, the Lagrangian formalism has been proposed by Wibisono et al. The derivation of the second-order rules and Natural GD is very straightforward and involves no novel technique (given the standard derivations of them and other related works pointed out by reviewers). Not to mention that these rules are not really used. I think the primary novelty is about the partial observability, but this part is way too short and shallow. I would happier if this section is 5 page long with 2 experiments on transformers revealing novel phenomena about them (exaggerated but you get my point).

Also, I partially disagree with the authors' reply to KkuV. The authors claim that a disadvantage of Wibisono et al is that "that their objective involves a minus sign (i.e., it has the form kinetic minus potential)." This is not a valid criticism. First of all, the Lagrangian in physics has kinetic minus potential, thus this criticism applies to physics. Thus, the same logic implies that the authors' formalism is also a better theory for physics. This cannot be true. Second, if we transform the Lagrangian to a Hamiltonian, the Hamiltonian indeed has kinetic plus potential (usually). Thus, the problem is trivially avoided by well known results. In fact, the equivalent Bregman Hamiltonian has indeed been pointed out by Wibisono et al.
Decision: Accept (poster)