ο‚›GitHub

How to run stochastic gradient descent

Ronny Bergmann

This tutorial illustrates how to use the stochastic_gradient_descent solver and different DirectionUpdateRules to introduce the average or momentum variant, see Stochastic Gradient Descent.

Computationally, we look at a very simple but large scale problem, the Riemannian Center of Mass or FrΓ©chet mean: for given points $p_i ∈\mathcal M$, $i=1,…,N$ this optimization problem reads

\[\operatorname*{arg\,min}_{p∈\mathcal M} \frac{1}{2N}\sum_{i=1}^{N} \operatorname{d}^2_{\mathcal M}(p,p_i),\]

which of course can be (and is) solved by a gradient descent, see the introductory tutorial or Statistics in Manifolds.jl. If $N$ is very large, evaluating the complete gradient might be quite expensive. A remedy is to evaluate only one of the terms at a time and choose a random order for these.

We first initialize the packages

using Manifolds, Manopt, Random, BenchmarkTools, ManifoldDiffusing ManifoldDiff: grad_distanceRandom.seed!(42);

We next generate a (little) large(r) data set

n = 5000Οƒ = Ο€ / 12M = Sphere(2)p = 1 / sqrt(2) * [1.0, 0.0, 1.0]data = [exp(M, p, Οƒ * rand(M; vector_at = p)) for i in 1:n];

Note that due to the construction of the points as zero mean tangent vectors, the mean should be very close to our initial point p.

In order to use the stochastic gradient, we now need a function that returns the vector of gradients. There are two ways to define it in Manopt.jl: either as a single function that returns a vector, or as a vector of functions.

The first variant is of course easier to define, but the second is more efficient when only evaluating one of the gradients.

For the mean, the gradient is

\[\operatorname{grad}f(p) = \frac{1}{N}\sum_{i=1}^N \operatorname{grad}f_i(p) \quad\text{ where }\quad \operatorname{grad}f_i(p) = -\log_p p_i\]

which we define in Manopt.jl in two different ways: either as one function returning all gradients as a vector (see gradF), or, maybe more fitting for a large scale problem, as a vector of small gradient functions (see gradf)

F(M, p) = 1 / (2 * n) * sum(map(q -> distance(M, p, q)^2, data))gradF(M, p) = [grad_distance(M, q, p) for q in data]gradf = [(M, p) -> grad_distance(M, q, p) for q in data];DL = DecreasingLength(M; length = 1.0) # stochastic gradient descent needs a decreasing step sizep0 = 1 / sqrt(3) * [1.0, 1.0, 1.0]
3-element Vector{Float64}:
 0.5773502691896258
 0.5773502691896258
 0.5773502691896258

The calls are only slightly different, but notice that accessing the second gradient element requires evaluating all logs in the first function, while we only call one of the functions in the second array of functions. So while you can use both gradF and gradf in the following call, the second one is (much) faster:

p_opt1 = stochastic_gradient_descent(M, gradF, p0; stepsize = DL)
3-element Vector{Float64}:
  0.7084813270817236
 -0.006829271064941041
  0.7056965142561222
@benchmark stochastic_gradient_descent($M, $gradF, $p0; stepsize = $DL)
BenchmarkTools.Trial: 1 sample with 1 evaluation per sample.
 Single result which took 5.806 s (17.15% GC) to evaluate,
 with a memory estimate of 7.83 GiB, over 200260802 allocations.
p_opt2 = stochastic_gradient_descent(M, gradf, p0; stepsize = DL)
3-element Vector{Float64}:
  0.6999851090058944
 -0.009965127289917282
  0.7140879101399926
@benchmark stochastic_gradient_descent($M, $gradf, $p0; stepsize = $DL)
BenchmarkTools.Trial: 903 samples with 1 evaluation per sample.
 Range (min … max):  4.687 ms … 16.016 ms  β”Š GC (min … max):  0.00% … 66.68%
 Time  (median):     4.826 ms              β”Š GC (median):     0.00%
 Time  (mean Β± Οƒ):   5.540 ms Β±  2.517 ms  β”Š GC (mean Β± Οƒ):  11.92% Β± 16.58%

  β–ˆβ–…β–β–‚β–                                                       
  β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–β–„β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–†β–‡β–‡β–†β–„β–‡β–‡β–‡ β–‡
  4.69 ms      Histogram: log(frequency) by time     15.5 ms <

 Memory estimate: 5.21 MiB, allocs estimate: 131779.

This result is reasonably close. We can also modify the search direction by using a DirectionUpdateRule, namely:

On the one hand MomentumGradient, which keeps track of the iterate and vector transports the last direction to the current iterate. The necessary vector_transport_method= keyword is set to a suitable default on every manifold, see default_vector_transport_method. We get

p_opt3 = stochastic_gradient_descent(    M, gradf, p0; stepsize = DL, direction = MomentumGradient(; direction = StochasticGradient()))
3-element Vector{Float64}:
  0.7056414705569414
 -0.003917305187978256
  0.7085582331398751
MG = MomentumGradient(; direction = StochasticGradient());@benchmark stochastic_gradient_descent($M, $gradf, $p0; stepsize = $DL, direction = $MG)
BenchmarkTools.Trial: 652 samples with 1 evaluation per sample.
 Range (min … max):  6.202 ms … 17.894 ms  β”Š GC (min … max):  0.00% … 61.51%
 Time  (median):     6.514 ms              β”Š GC (median):     0.00%
 Time  (mean Β± Οƒ):   7.678 ms Β±  3.237 ms  β”Š GC (mean Β± Οƒ):  14.81% Β± 19.14%

  β–ˆβ–‡                                                          
  β–ˆβ–ˆβ–‡β–…β–„β–ƒβ–‚β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–‚β–ƒβ–ƒβ–‚β–‚β–‚β–‚β–ƒβ–ƒβ–‚ β–‚
  6.2 ms         Histogram: frequency by time        17.7 ms <

 Memory estimate: 9.03 MiB, allocs estimate: 231930.

And on the other hand the AverageGradient computes an average of the last n gradients. This is done by

p_opt4 = stochastic_gradient_descent(    M, gradf, p0; stepsize = DL,    direction = AverageGradient(; n = 10, direction = StochasticGradient()),)
3-element Vector{Float64}:
  0.7094044023644667
 -0.006431061655364132
  0.7047723287359534
AG = AverageGradient(; n = 10, direction = StochasticGradient());@benchmark stochastic_gradient_descent($M, $gradf, $p0; stepsize = $DL, direction = $AG)
BenchmarkTools.Trial: 340 samples with 1 evaluation per sample.
 Range (min … max):  11.412 ms … 23.851 ms  β”Š GC (min … max):  0.00% … 46.09%
 Time  (median):     12.234 ms              β”Š GC (median):     0.00%
 Time  (mean Β± Οƒ):   14.705 ms Β±  4.435 ms  β”Š GC (mean Β± Οƒ):  17.76% Β± 20.17%

    β–‡β–ˆβ–„                                                        
  β–ƒβ–†β–ˆβ–ˆβ–ˆβ–†β–ƒβ–‚β–‚β–β–β–‚β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–‚β–ƒβ–„β–ƒβ–ƒβ–ƒβ–ƒβ–ƒβ–ƒβ–„β–ƒβ–„β–‚ β–ƒ
  11.4 ms         Histogram: frequency by time        23.4 ms <

 Memory estimate: 20.48 MiB, allocs estimate: 531944.

Note that the default StoppingCriterion here is StopAfterIteration(10000) combined with StopWhenGradientNormLess(1e-9).

For both update rules we have to internally specify that we are still in the stochastic setting, since both rules can also be used with the IdentityUpdateRule within gradient_descent.

For this not-that-large-scale example we can of course also use a gradient descent with ArmijoLinesearch,

fullGradF(M, p) = 1 / n * sum(grad_distance(M, q, p) for q in data)p_opt5 = gradient_descent(M, F, fullGradF, p0; stepsize = ArmijoLinesearch())
3-element Vector{Float64}:
  0.7050420976839262
 -0.006374163322665689
  0.7091368066426853

but in general it is expected to be a bit slow.

AL = ArmijoLinesearch();@benchmark gradient_descent($M, $F, $fullGradF, $p0; stepsize = $AL)
BenchmarkTools.Trial: 78 samples with 1 evaluation per sample.
 Range (min … max):  54.477 ms … 66.702 ms  β”Š GC (min … max):  0.00% … 14.88%
 Time  (median):     64.396 ms              β”Š GC (median):    14.75%
 Time  (mean Β± Οƒ):   64.358 ms Β±  1.286 ms  β”Š GC (mean Β± Οƒ):  14.60% Β±  1.71%

                                                   β–ƒβ–‚β–ˆβ–ƒ β–„β–‚     
  β–ƒβ–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–β–…β–…β–†β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–†β–„β–„ ▁
  54.5 ms         Histogram: frequency by time        65.7 ms <

 Memory estimate: 74.26 MiB, allocs estimate: 1716685.
Technical Details

This tutorial is cached. It was last run on the following package versions.

Status `~/work/Manopt.jl/Manopt.jl/tutorials/Project.toml`
  [47edcb42] ADTypes v1.24.0
  [6e4b80f9] BenchmarkTools v1.8.0
  [5ae59095] Colors v0.13.1
  [a0c0ee7d] DifferentiationInterface v0.7.21
  [31c24e10] Distributions v0.25.131
  [26cc04aa] FiniteDifferences v0.12.34
  [f6369f11] ForwardDiff v1.4.6
  [8ac3fa9e] LRUCache v1.6.2
  [af67fdf4] ManifoldDiff v0.4.5
  [1cead3c2] Manifolds v0.11.29
  [3362f125] ManifoldsBase v2.5.1
  [0fc0a36d] Manopt v0.6.8 `.`
  [91a5bcdd] Plots v1.41.7
  [731186ca] RecursiveArrayTools v4.5.1
  [37e2e46d] LinearAlgebra v1.12.0
  [9a3f8284] Random v1.11.0

This tutorial was last rendered September 20, 2026, 7:53:25.