Gradient descent

Gradient descent is a simple algorithm for optimization, and empirically works extremely well for modern deep learning.

The idea is very simple:

  • At any given point in parameter space, compute the gradient of the loss with respect to parameters – the direction of steepest ascent. Since we want to decrease loss, we step in the negative gradient direction – the direction of steepest descent.
  • This brings us to a new point in parameter space, and we repeat.

A key hyperparameter for gradient descent is the step size (i.e., the learning rate), commonly denoted as η\eta. This hyperparameter dictates how large of a step we take along the direction of steepest descent.

Formally, we can write gradient descent as follows:

θt+1=θt−η∇θL(θt),\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta \nabla_{\boldsymbol{\theta}} \mathcal{L}(\boldsymbol{\theta}_t),

where θt\boldsymbol{\theta}_t denotes the parameters at iteration tt, η\eta denotes the step size, and L\mathcal{L} represents a loss function mapping parameters to a scalar.

Gradient flow

Gradient flow is an idealized, continuous version of gradient descent, taking η→0\eta \rightarrow 0.

We can write gradient flow as follows:

dθdt=−∇L(θ),\frac{d\boldsymbol{\theta}}{dt} = -\nabla \mathcal{L}(\boldsymbol{\theta}),

where we think of θ\boldsymbol{\theta} evolving over continuous time tt.

Approximating gradient descent with gradient flow

While gradient flow is always continuously moving in the direction of steepest descent, gradient descent takes discrete steps and does not account for the change in loss landscape between each step. This can lead to a mismatch between the paths taken by gradient flow and gradient descent.

One natural question is: can we understand finite-step gradient descent as gradient flow on a slightly different objective?

More precisely, suppose the original loss is L(θ)\mathcal{L}(\boldsymbol{\theta}). We can write a single step of gradient descent as

θGD=θ0−η∇L(θ0).\boldsymbol{\theta}_{\mathrm{GD}} = \boldsymbol{\theta}_0 - \eta \nabla \mathcal{L}(\boldsymbol{\theta}_0).

We want to find a nearby loss Lη\mathcal{L}_\eta whose gradient flow, started at the same θ0\boldsymbol{\theta}_0 and run for time η\eta, ends up near the endpoint of the discrete gradient descent step. Formally, we want to find a loss Lη\mathcal{L}_\eta such that the gradient flow trajectory θ(t)\boldsymbol{\theta}(t), defined by

dθ(t)dt=−∇Lη(θ(t)),θ(0)=θ0,\frac{d\boldsymbol{\theta}(t)}{dt} = -\nabla \mathcal{L}_\eta(\boldsymbol{\theta}(t)), \qquad \boldsymbol{\theta}(0)=\boldsymbol{\theta}_0,

has endpoint θflow:=θ(η)\boldsymbol{\theta}_{\mathrm{flow}} := \boldsymbol{\theta}(\eta) satisfying

θflow≈θGD.\boldsymbol{\theta}_{\mathrm{flow}} \approx \boldsymbol{\theta}_{\mathrm{GD}}.

This adjusted loss Lη\mathcal{L}_{\eta} has different names in different fields; in this note we’ll refer to it as the “modified loss”.

A quadratic warmup

We’ll start with a simple quadratic loss:

L(θ)=λ2θ2,\mathcal{L}(\theta) = \frac{\lambda}{2}\theta^2,

where θ∈R\theta \in \mathbb{R} is our learned parameter, and λ∈R+\lambda \in \mathbb{R}^{+} is a constant that controls the curvature.

The gradient is ∇L(θ)=λθ\nabla \mathcal{L}(\theta) = \lambda\theta, and so one step of gradient descent gives

θGD=θ0−η∇L(θ0)=θ0−η(λθ0)=(1−ηλ)θ0.\begin{aligned} \theta_{\mathrm{GD}} &= \theta_0 - \eta \nabla \mathcal{L}(\theta_0) \\ &= \theta_0 - \eta (\lambda\theta_0) \\ &= (1 - \eta \lambda)\theta_0. \end{aligned}

Now we’ll approximate the endpoint of gradient flow on the same loss. Recall that gradient flow is defined by the differential equation

dθdt=−∇L(θ)=−λθ.\begin{aligned} \frac{d\theta}{dt} &= -\nabla \mathcal{L}(\theta) \\ &= -\lambda\theta. \end{aligned}

The solution to this differential equation is

θ(t)=e−λtθ0.\theta(t) = e^{-\lambda t}\theta_0.

After flowing for time η\eta,

θflow=e−ληθ0.\theta_{\mathrm{flow}} = e^{-\lambda\eta}\theta_0.

Recalling the Taylor expansion of exe^x, we can write

e−λη=1−λη+λ2η22+O(η3).\begin{aligned} e^{-\lambda\eta} &= 1 - \lambda\eta + \frac{\lambda^2\eta^2}{2} + O(\eta^3). \end{aligned}

Using this approximation, we can write the endpoint of gradient flow as

θflow=(1−λη+λ2η22+O(η3))θ0.\theta_{\mathrm{flow}} = \left(1 - \lambda\eta + \frac{\lambda^2\eta^2}{2} + O(\eta^3)\right)\theta_0.

Comparing this to the result of discrete gradient descent, we see that the two agree to first order in η\eta, but disagree at order η2\eta^2 and higher.

Left: the quadratic bowl L(θ)=λ2θ2\mathcal{L}(\theta) = \frac{\lambda}{2}\theta^2 and the starting point θ0=1.2\theta_0=1.2. Right: parameter trajectories over time: GF (blue, continuous) and discrete GD (black dots, plotted at times t=kηt = k\eta with η=0.5\eta=0.5). Note that the two paths do not match. In the case of our quadratic loss, GD approaches the minimum more quickly than GF – this is because GD locks in the starting gradient for the entire step, while GF’s gradient continuously shrinks as ∣θ∣\lvert\theta\rvert decreases.

The modified quadratic

From the above analysis, we see that gradient descent lands at θGD=(1−ηλ)θ0\theta_{\mathrm{GD}} = (1 - \eta \lambda)\theta_0, while gradient flow lands at θflow=(1−λη+λ2η22+O(η3))θ0\theta_{\mathrm{flow}} = \left(1 - \lambda\eta + \frac{\lambda^2\eta^2}{2} + O(\eta^3)\right)\theta_0. The mismatch primarily stems from the extra second-order term in gradient flow: λ2η22θ0\frac{\lambda^2\eta^2}{2}\theta_0. Note that this extra (positive) term causes gradient flow to move slightly less far down the quadratic bowl than gradient descent.

So, in order to construct a modified gradient flow that behaves similarly to gradient descent here, we can try making the bowl effectively steeper:

Lη(θ)=12(λ+η2λ2)θ2.\mathcal{L}_\eta(\theta) = \frac{1}{2}\left(\lambda+\frac{\eta}{2}\lambda^2\right)\theta^2.

We can now analyze how gradient flow behaves on this modified loss.

At the start of the flow, the initial velocity is

dθdt∣t=0=−dLη(θ)dθ∣θ=θ0=−(λ+η2λ2)θ0=−λθ0−η2λ2θ0.\begin{aligned} \left.\frac{d\theta}{dt}\right|_{t=0} &= \left.-\frac{d\mathcal{L}_\eta (\theta)}{d\theta} \right|_{\theta=\theta_0}\\ &= -\left(\lambda+\frac{\eta}{2}\lambda^2\right)\theta_0 \\ &= -\lambda\theta_0-\frac{\eta}{2}\lambda^2\theta_0. \end{aligned}

To get the acceleration, we differentiate the flow equation once more with respect to time:

d2θdt2=−(λ+η2λ2)dθdt.\begin{aligned} \frac{d^2\theta}{dt^2} &= -\left(\lambda+\frac{\eta}{2}\lambda^2\right)\frac{d\theta}{dt}. \end{aligned}

At t=0t=0, this gives

d2θdt2∣t=0=[−(λ+η2λ2)dθdt]t=0=−(λ+η2λ2)[−(λ+η2λ2)θ0]=(λ+η2λ2)2θ0=λ2θ0+O(η).\begin{aligned} \left.\frac{d^2\theta}{dt^2}\right|_{t=0} &= \left[-\left(\lambda+\frac{\eta}{2}\lambda^2\right)\frac{d\theta}{dt}\right]_{t=0} \\ &= -\left(\lambda+\frac{\eta}{2}\lambda^2\right) \left[-\left(\lambda+\frac{\eta}{2}\lambda^2\right)\theta_0\right] \\ &= \left(\lambda+\frac{\eta}{2}\lambda^2\right)^2\theta_0 \\ &= \lambda^2\theta_0+O(\eta). \end{aligned}

Note that the O(η)O(\eta) part of the acceleration will end up contributing only to an order η3\eta^3 term, because acceleration will be multiplied by η2/2\eta^2/2.

Now we can approximate the endpoint of gradient flow on the modified loss:

θflow=θ(0)+η⋅dθdt∣t=0+η22⋅d2θdt2∣t=0+O(η3)=θ0+η(−λθ0−η2λ2θ0)+η22(λ2θ0+O(η))+O(η3)=θ0−ηλθ0−η22λ2θ0+η22λ2θ0+O(η3)=θ0−ηλθ0+O(η3).\begin{aligned} \theta_{\mathrm{flow}} &= \theta(0) + \eta \cdot \left.\frac{d\theta}{dt}\right|_{t=0} + \frac{\eta^2}{2} \cdot \left.\frac{d^2\theta}{dt^2}\right|_{t=0} + O(\eta^3) \\ &= \theta_0 + \eta\left(-\lambda\theta_0-\frac{\eta}{2}\lambda^2\theta_0\right) + \frac{\eta^2}{2}\left(\lambda^2\theta_0+O(\eta)\right) + O(\eta^3) \\ &= \theta_0 -\eta\lambda\theta_0 -\frac{\eta^2}{2}\lambda^2\theta_0 +\frac{\eta^2}{2}\lambda^2\theta_0 +O(\eta^3) \\ &= \theta_0-\eta\lambda\theta_0+O(\eta^3). \end{aligned}

The order-η2\eta^2 terms cancel, so gradient flow on the modified loss matches one gradient descent step on the original loss (up to order η2\eta^2). Nice!

Rewriting the modified loss in terms of gradient norm

Let’s look back at the modified loss that we used. The modified loss is just the original loss plus one extra correction term:

Lη(θ)=12λθ2⏟original loss L(θ)+η4λ2θ2⏟finite-step correction.\mathcal{L}_\eta(\theta) = \underbrace{\frac{1}{2}\lambda\theta^2}_{\text{original loss } \mathcal{L}(\theta)} + \underbrace{\frac{\eta}{4}\lambda^2\theta^2}_{\text{finite-step correction}}.

Since ∇L(θ)=λθ\nabla \mathcal{L}(\theta)=\lambda\theta, we can write the extra term as

η4λ2θ2=η4∥∇L(θ)∥2.\frac{\eta}{4}\lambda^2\theta^2 = \frac{\eta}{4}\|\nabla\mathcal{L}(\theta)\|^2.

We will see that, in general, the modified loss can be written as

Lη(θ)=L(θ)+η4∥∇L(θ)∥2+O(η2).\mathcal{L}_\eta(\boldsymbol{\theta}) = \mathcal{L}(\boldsymbol{\theta}) + \frac{\eta}{4}\|\nabla \mathcal{L}(\boldsymbol{\theta})\|^2 + O(\eta^2).

Left: the original loss L(θ)\mathcal{L}(\theta) (blue) and the modified loss Lη(θ)\mathcal{L}_\eta(\theta) (orange), with λ=1\lambda=1 and η=0.5\eta=0.5. Note that the modified loss is a slightly steeper bowl. Right: parameter trajectories over time: GF on the original loss (blue), GF on the modified loss (orange), and discrete GD (black dots). Note that GF on the modified loss now tracks the discrete GD steps more closely, compared to GF on the original loss.

The general case

Now we’ll try to generalize the modified loss construction.

We’ll denote the initial gradient as g:=∇L(θ0)\mathbf{g} := \nabla\mathcal{L}(\boldsymbol{\theta}_0) and the initial Hessian as H:=∇2L(θ0)\mathbf{H} := \nabla^2\mathcal{L}(\boldsymbol{\theta}_0).

We can write the endpoint of one gradient descent step as

θGD=θ0−ηg.\boldsymbol{\theta}_{\mathrm{GD}} = \boldsymbol{\theta}_0-\eta \mathbf{g}.

We can similarly write the endpoint of gradient flow on the original loss as

θflow=θ0−ηg+η22Hg+O(η3).\boldsymbol{\theta}_{\mathrm{flow}} = \boldsymbol{\theta}_0-\eta \mathbf{g}+\frac{\eta^2}{2}\mathbf{H}\mathbf{g}+O(\eta^3).

Note that, just like in our toy quadratic case, gradient descent and gradient flow agree up to first order in η\eta, but disagree at order η2\eta^2 and higher. Gradient flow has an extra second-order term: η22Hg\frac{\eta^2}{2}\mathbf{H}\mathbf{g}.

Now we try a nearby loss whose first correction is proportional to the step size:

Lη(θ)=L(θ)+ηR(θ)+O(η2),\mathcal{L}_\eta(\boldsymbol{\theta}) = \mathcal{L}(\boldsymbol{\theta})+\eta R(\boldsymbol{\theta})+O(\eta^2),

where RR is the correction term.

For gradient flow on this modified loss, initial velocity is

dθdt∣t=0=−∇Lη(θ0)=−∇L(θ0)−η∇R(θ0)+O(η2)=−g−η∇R(θ0)+O(η2).\begin{aligned} \left.\frac{d\boldsymbol{\theta}}{dt}\right|_{t=0} &= -\nabla \mathcal{L}_\eta(\boldsymbol{\theta}_0) \\ &= -\nabla \mathcal{L}(\boldsymbol{\theta}_0)-\eta\nabla R(\boldsymbol{\theta}_0)+O(\eta^2) \\ &= -\mathbf{g}-\eta\nabla R(\boldsymbol{\theta}_0)+O(\eta^2). \end{aligned}

Initial acceleration is

d2θdt2∣t=0=−∇2Lη(θ0)dθdt∣t=0=−(∇2L(θ0)+O(η))(−∇L(θ0)+O(η))=∇2L(θ0)∇L(θ0)+O(η)=Hg+O(η).\begin{aligned} \left.\frac{d^2\boldsymbol{\theta}}{dt^2}\right|_{t=0} &= -\nabla^2\mathcal{L}_\eta(\boldsymbol{\theta}_0) \left.\frac{d\boldsymbol{\theta}}{dt}\right|_{t=0} \\ &= -\left(\nabla^2\mathcal{L}(\boldsymbol{\theta}_0)+O(\eta)\right) \left(-\nabla\mathcal{L}(\boldsymbol{\theta}_0)+O(\eta)\right) \\ &= \nabla^2\mathcal{L}(\boldsymbol{\theta}_0)\nabla\mathcal{L}(\boldsymbol{\theta}_0)+O(\eta) \\ &= \mathbf{H}\mathbf{g}+O(\eta). \end{aligned}

Note that the RR-dependent pieces in the acceleration are only order η\eta, so after multiplying by η2/2\eta^2/2 they only matter at order η3\eta^3.

We can now write the approximation for the endpoint of gradient flow on the modified loss:

θflow=θ0+η⋅dθdt∣t=0+η22⋅d2θdt2∣t=0+O(η3)=θ0+η(−g−η∇R(θ0)+O(η2))+η22(Hg+O(η))+O(η3)=θ0−ηg+η2(12Hg−∇R(θ0))+O(η3).\begin{aligned} \boldsymbol{\theta}_{\mathrm{flow}} &= \boldsymbol{\theta}_0 +\eta \cdot \left.\frac{d\boldsymbol{\theta}}{dt}\right|_{t=0} +\frac{\eta^2}{2} \cdot \left.\frac{d^2\boldsymbol{\theta}}{dt^2}\right|_{t=0} +O(\eta^3) \\ &= \boldsymbol{\theta}_0 +\eta\left(-\mathbf{g}-\eta\nabla R(\boldsymbol{\theta}_0)+O(\eta^2)\right) +\frac{\eta^2}{2}\left(\mathbf{H}\mathbf{g}+O(\eta)\right) +O(\eta^3) \\ &= \boldsymbol{\theta}_0-\eta \mathbf{g} +\eta^2\left(\frac{1}{2}\mathbf{H}\mathbf{g}-\nabla R(\boldsymbol{\theta}_0)\right) +O(\eta^3). \end{aligned}

To match θGD=θ0−ηg\boldsymbol{\theta}_{\mathrm{GD}} = \boldsymbol{\theta}_0-\eta \mathbf{g} through order η2\eta^2, we need the η2\eta^2 term to vanish; i.e., we need

∇R(θ0)=12Hg.\nabla R(\boldsymbol{\theta}_0)=\frac{1}{2}\mathbf{H}\mathbf{g}.

A function with this gradient is

R(θ)=14∥∇L(θ)∥2.R(\boldsymbol{\theta})=\frac{1}{4}\|\nabla\mathcal{L}(\boldsymbol{\theta})\|^2.
Proof

Write RR coordinate-wise as

R(θ)=14∑i(∂L(θ)∂θi)2.R(\boldsymbol{\theta}) = \frac{1}{4}\sum_i \left(\frac{\partial \mathcal{L}(\boldsymbol{\theta})}{\partial \theta_i}\right)^2.

Differentiate with respect to coordinate θj\theta_j:

∂R(θ)∂θj=12∑i∂L(θ)∂θi∂2L(θ)∂θj ∂θi.\begin{aligned} \frac{\partial R(\boldsymbol{\theta})}{\partial \theta_j} &= \frac{1}{2}\sum_i \frac{\partial \mathcal{L}(\boldsymbol{\theta})}{\partial \theta_i} \frac{\partial^2 \mathcal{L}(\boldsymbol{\theta})}{\partial \theta_j\,\partial \theta_i}. \end{aligned}

At θ0\boldsymbol{\theta}_0, this becomes

∂R(θ0)∂θj=12∑iHjigi.\frac{\partial R(\boldsymbol{\theta}_0)}{\partial \theta_j} = \frac{1}{2}\sum_i H_{ji}g_i.

That is the jjth coordinate of 12Hg\frac{1}{2}\mathbf{H}\mathbf{g}. Therefore ∇R(θ0)=12Hg\nabla R(\boldsymbol{\theta}_0)=\frac{1}{2}\mathbf{H}\mathbf{g}.

Thus

Lη(θ)=L(θ)+η4∥∇L(θ)∥2+O(η2).\mathcal{L}_\eta(\boldsymbol{\theta}) = \mathcal{L}(\boldsymbol{\theta}) + \frac{\eta}{4}\|\nabla \mathcal{L}(\boldsymbol{\theta})\|^2 + O(\eta^2).

References

References are listed in alphabetical order.

  1. Implicit gradient regularization [link]
    David Barrett and Benoit Dherin. International Conference on Learning Representations. 2021.
  2. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics [link]
    Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. International Conference on Learning Representations. 2021.
  3. On the origin of implicit regularization in stochastic gradient descent [link]
    Samuel L Smith, Benoit Dherin, David Barrett, and Soham De. International Conference on Learning Representations. 2021.
  4. Neural thermodynamics: Entropic forces in deep and universal representation learning [link]
    Liu Ziyin, Yizhou Xu, and Isaac L. Chuang. Advances in Neural Information Processing Systems. 2025.