Stochastic Gradient Descent
This tutorial illustrates how to use the stochastic_gradient_descent
solver and different DirectionUpdateRule
s in order 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}_{x∈\mathcal M} \frac{1}{2}\sum_{i=1}^{N} \operatorname{d}^2_{\mathcal M}(x,p_i),$$
which of course can be (and is) solved by a gradient descent, see the introductionary tutorial or the Statistics in Manifolds.jl. If $N$ is very large it might be quite expensive to evaluate the complete gradient. 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 BenchmarkTools, Colors, Manopt, Manifolds, PlutoUI, Random
and we define some colors from Paul Tol
begin
black = RGBA{Float64}(colorant"#000000")
TolVibrantOrange = RGBA{Float64}(colorant"#EE7733") # Start
TolVibrantBlue = RGBA{Float64}(colorant"#0077BB") # a path
TolVibrantTeal = RGBA{Float64}(colorant"#009988") # points
end;
We next generate a (little) large(r) data set
begin
n = 5000
σ = π / 12
M = Sphere(2)
x = 1 / sqrt(2) * [1.0, 0.0, 1.0]
Random.seed!(42)
data = [exp(M, x, random_tangent(M, x, Val(:Gaussian), σ)) for i in 1:n]
localpath = join(splitpath(@__FILE__)[1:(end - 1)], "/") # files folder
image_prefix = localpath * "/stochastic_gradient_descent"
@info image_prefix
render_asy = false # on CI or when you do not have asymptote, this should be false
end
false
render_asy && asymptote_export_S2_signals(
image_prefix * "/center_and_large_data.asy";
points=[[x], data],
colors=Dict(:points => [TolVibrantBlue, TolVibrantTeal]),
dot_sizes=[2.5, 1.0],
camera_position=(1.0, 0.5, 0.5),
)
false
render_asy && render_asymptote(image_prefix * "/center_and_large_data.asy"; render=2);
PlutoUI.LocalResource(image_prefix * "/center_and_large_data.png")
Note that due to the construction of the points as zero mean tangent vectors, the mean should be very close to our initial point x
.
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
: as one function, that returns a vector or a vector of funtions.
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 we have as a gradient
$$ gradF(x) = \sum_{i=1}^N \operatorname{grad}f_i(x) \quad \text{where} \operatorname{grad}f_i(x) = -\log_x 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, x) = 1 / (2 * n) * sum(map(p -> distance(M, x, p)^2, data))
F (generic function with 1 method)
gradF(M, x) = [grad_distance(M, p, x) for p in data]
gradF (generic function with 1 method)
gradf = [(M, x) -> grad_distance(M, p, x) for p in data];
The calls are only slightly different, but notice that accessing the 2nd gradient element requires evaluating all logs in the first function. So while you can use both gradF
and gradf
in the following call, the second one is (much) faster:
x_opt1 = stochastic_gradient_descent(M, gradF, x)
3-element Vector{Float64}: 0.7071067811865475 0.0 0.7071067811865475
@benchmark stochastic_gradient_descent($M, $gradF, $x)
BenchmarkTools.Trial: 3170 samples with 1 evaluation. Range (min … max): 967.301 μs … 27.029 ms ┊ GC (min … max): 0.00% … 93.80% Time (median): 1.286 ms ┊ GC (median): 0.00% Time (mean ± σ): 1.564 ms ± 1.919 ms ┊ GC (mean ± σ): 10.62% ± 8.27% ▇█▂ ▄▇███▆▄▃▅▅▄▃▂▂▂▂▂▂▂▂▂▂▂▁▁▂▂▂▁▂▂▂▁▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▁▂ ▃ 967 μs Histogram: frequency by time 6.21 ms < Memory estimate: 861.52 KiB, allocs estimate: 10031.
x_opt2 = stochastic_gradient_descent(M, gradf, x)
3-element Vector{Float64}: 0.7071067811865475 0.0 0.7071067811865475
@benchmark stochastic_gradient_descent($M, $gradf, $x)
BenchmarkTools.Trial: 10000 samples with 6 evaluations. Range (min … max): 6.000 μs … 1.787 ms ┊ GC (min … max): 0.00% … 99.13% Time (median): 9.883 μs ┊ GC (median): 0.00% Time (mean ± σ): 14.850 μs ± 70.544 μs ┊ GC (mean ± σ): 22.73% ± 4.82% ▃▆█▆ ▁▁▁▃▇████▇▄▃▄▄▃▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂ 6 μs Histogram: frequency by time 34.7 μs < Memory estimate: 41.08 KiB, allocs estimate: 26.
x_opt2
3-element Vector{Float64}: 0.7071067811865475 0.0 0.7071067811865475
This result is reasonably close. But we can improve it by using a DirectionUpdateRule
, namely:
On the one hand MomentumGradient
, which requires both the manifold and the initial value, in order to keep track of the iterate and parallel transport the last direction to the current iterate. You can also set a vector_transport_method
, if ParallelTransport()
is not available on your manifold. Here we simply do
x_opt3 = stochastic_gradient_descent(
M, gradf, x; direction=MomentumGradient(M, x, StochasticGradient(zero_vector(M, x)))
)
3-element Vector{Float64}: 0.7071067811865475 0.0 0.7071067811865475
MG = MomentumGradient(M, x, StochasticGradient(zero_vector(M, x)));
@benchmark stochastic_gradient_descent($M, $gradf, $x; direction=$MG)
BenchmarkTools.Trial: 10000 samples with 4 evaluations. Range (min … max): 6.850 μs … 2.662 ms ┊ GC (min … max): 0.00% … 98.68% Time (median): 9.350 μs ┊ GC (median): 0.00% Time (mean ± σ): 15.104 μs ± 89.791 μs ┊ GC (mean ± σ): 23.10% ± 3.94% ▂▅▇███▆▃▁▂▃▄▄▄▃▁ ▂▂▁▁ ▂ ████████████████████████▇▇▆▆▇▆▆▆▇▆█████████▇▆▅▆▅▆▆▆▄▅▅▄▆▅▆▅ █ 6.85 μs Histogram: log(frequency) by time 37.8 μs < Memory estimate: 40.98 KiB, allocs estimate: 24.
And on the other hand the AverageGradient
computes an average of the last n
gradients, i.e.
x_opt4 = stochastic_gradient_descent(
M, gradf, x; direction=AverageGradient(M, x, 10, StochasticGradient(zero_vector(M, x)))
);
AG = AverageGradient(M, x, 10, StochasticGradient(zero_vector(M, x)));
@benchmark stochastic_gradient_descent($M, $gradf, $x; direction=$AG)
BenchmarkTools.Trial: 10000 samples with 4 evaluations. Range (min … max): 6.900 μs … 2.455 ms ┊ GC (min … max): 0.00% … 99.01% Time (median): 9.475 μs ┊ GC (median): 0.00% Time (mean ± σ): 15.152 μs ± 88.396 μs ┊ GC (mean ± σ): 23.03% ± 3.94% ▂▅▇██▇▅▂▂▂▃▄▅▄▂▁ ▁▁▁▁▁ ▁▁▁ ▂ ███████████████████████▇▇▆▇▇▄▆▆▆▇███████▇▇█▇▅▇▇▅▆▅▆▆▅▅▄▄▃▃▅ █ 6.9 μs Histogram: log(frequency) by time 38.6 μs < Memory estimate: 40.98 KiB, allocs estimate: 24.
Note that the default StoppingCriterion
is a fixed number of iterations which helps the comparison here.
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
, but it will be a little slower usually
fullGradF(M, x) = sum(grad_distance(M, p, x) for p in data)
fullGradF (generic function with 1 method)
x_opt5 = gradient_descent(M, F, fullGradF, x; stepsize=ArmijoLinesearch())
3-element Vector{Float64}: 0.7071067811865475 -4.886637652164662e-17 0.7071067811865475
AL = ArmijoLinesearch();
@benchmark gradient_descent($M, $F, $fullGradF, $x; stepsize=$AL)
BenchmarkTools.Trial: 4 samples with 1 evaluation. Range (min … max): 1.342 s … 1.417 s ┊ GC (min … max): 10.95% … 10.59% Time (median): 1.387 s ┊ GC (median): 11.15% Time (mean ± σ): 1.384 s ± 31.775 ms ┊ GC (mean ± σ): 11.16% ± 0.47% █ █ █ █ █▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁ 1.34 s Histogram: frequency by time 1.42 s < Memory estimate: 711.53 MiB, allocs estimate: 9034671.