Add Distance over Gradients Stepsize

docs/src/references.bib

@@ -5,6 +5,13 @@
 % --- A
 %
 %
+@article{dodd2024learning,
+  title={Learning-rate-free stochastic optimization over Riemannian manifolds},
+  author={Dodd, Daniel and Sharrock, Louis and Nemeth, Christopher},
+  journal={arXiv preprint arXiv:2406.02296},
+  year={2024}
+}
+
 @book{AbsilMahonySepulchre:2008,
     AUTHOR    = {Absil, P.-A. and Mahony, R. and Sepulchre, R.},
     DOI       = {10.1515/9781400830244},

src/Manopt.jl

@@ -522,7 +522,12 @@ export SmoothingTechnique, LinearQuadraticHuber, LogarithmicSumOfExponentials
 #
 # Stepsize
 export Stepsize
-export AdaptiveWNGradient, ConstantLength, DecreasingLength, Polyak
+export AdaptiveWNGradient,
+    ConstantLength,
+    DecreasingLength,
+    Polyak,
+    DistanceOverGradients,
+    DistanceOverGradientsStepsize
 export ProximalGradientMethodBacktracking
 export ArmijoLinesearch, Linesearch, NonmonotoneLinesearch
 export get_stepsize, get_initial_stepsize, get_last_stepsize

src/plans/stepsize.jl

@@ -1687,3 +1687,193 @@ end
 function get_last_stepsize(step::WolfePowellBinaryLinesearchStepsize, ::Any...)
     return step.last_stepsize
 end
+
+@doc raw"""
+    DistanceOverGradientsStepsize{R<:Real} <: Stepsize
+
+Implements the Riemannian Distance over Gradients (RDoG) stepsize schedule, a learning-rate-free 
+optimization method for Riemannian manifolds introduced by [@dodd2024learning].
+
+This stepsize adapts automatically without hyperparameter tuning by tracking the maximum distance 
+from the initial point and accumulating gradient norms.
+
+# Fields
+
+* `initial_distance::R`: initial distance estimate ``ϵ > 0``
+* `max_distance::R`: maximum distance from initial point ``\bar{r}_t = \max_{s≤t} d(p_0, p_s)``
+* `gradient_sum::R`: accumulated squared gradient norms ``G_t = \sum_{s=0}^t \|∇f(p_s)\|^2``
+* `initial_point`: the initial point ``p_0``
+* `use_curvature::Bool`: whether to use sectional curvature in the stepsize
+* `sectional_curvature_bound::R`: lower bound on sectional curvature ``κ`` (only used if `use_curvature=true`)
+
+The stepsize at iteration ``t`` is computed as:
+
+```math
+η_t = \frac{\bar{r}_t}{\sqrt{ζ_κ(\bar{r}_t)} \sqrt{G_t}}
+```
+
+where ``ζ_κ`` is the geometric curvature function:
+- ``ζ_κ(d) = 1`` if ``κ ≥ 0`` or `use_curvature=false`
+- ``ζ_κ(d) = \frac{\sqrt{|κ|} d}{\tanh(\sqrt{|κ|} d)}`` if ``κ < 0``
+
+# Constructor
+
+    DistanceOverGradientsStepsize(M::AbstractManifold; kwargs...)
+
+Initialize the RDoG stepsize on manifold `M`.
+
+## Keyword arguments
+
+* `initial_distance=1e-3`: initial distance estimate ``ϵ``
+* `use_curvature=false`: whether to incorporate sectional curvature
+* `sectional_curvature_bound=0.0`: lower bound on sectional curvature (if known)
+$(_var(:Keyword, :p; add="initial point, used to track distances"))
+
+# See also
+
+[`AdaptiveWNGradientStepsize`](@ref), [`gradient_descent`](@ref)
+"""
+mutable struct DistanceOverGradientsStepsize{R <: Real, P} <: Stepsize
+    initial_distance::R
+    max_distance::R
+    gradient_sum::R
+    initial_point::P
+    use_curvature::Bool
+    sectional_curvature_bound::R
+    last_stepsize::R
+end
+
+function DistanceOverGradientsStepsize(
+        M::AbstractManifold;
+        p = rand(M),
+        initial_distance::R = 1.0e-3,
+        use_curvature::Bool = false,
+        sectional_curvature_bound::R = 0.0,
+    ) where {R <: Real}
+    return DistanceOverGradientsStepsize{R, typeof(p)}(
+        initial_distance,
+        initial_distance,  # max_distance starts at initial_distance
+        zero(R),          # gradient_sum starts at 0
+        copy(M, p),       # store initial point
+        use_curvature,
+        sectional_curvature_bound,
+        NaN,              # last_stepsize
+    )
+end
+
+"""
+    geometric_curvature_function(κ::Real, d::Real)
+
+Compute the geometric curvature function ``ζ_κ(d)`` for sectional curvature bound ``κ``.
+
+```math
+ζ_κ(d) = \\begin{cases}
+    \\frac{\\sqrt{|κ|} d}{\\tanh(\\sqrt{|κ|} d)} & \\text{if } κ < 0 \\\\
+    1 & \\text{if } κ ≥ 0
+\\end{cases}
+```
+"""
+function geometric_curvature_function(κ::Real, d::Real)
+    if κ < 0 && d > 0
+        sqrt_abs_κ = sqrt(abs(κ))
+        arg = sqrt_abs_κ * d
+        # Avoid numerical issues for small arguments
+        if arg < 1.0e-8
+            return 1.0 + arg^2 / 3.0  # Taylor expansion
+        else
+            return arg / tanh(arg)
+        end
+    else
+        return 1.0
+    end
+end
+
+function (rdog::DistanceOverGradientsStepsize)(
+        mp::AbstractManoptProblem,
+        s::AbstractManoptSolverState,
+        i,
+        args...;
+        gradient = nothing,
+        kwargs...,
+    )
+    M = get_manifold(mp)
+    p = get_iterate(s)
+    grad = isnothing(gradient) ? get_gradient(mp, p) : gradient
+
+    # Compute gradient norm
+    grad_norm_sq = norm(M, p, grad)^2
+
+    if i == 0
+        # Initialize on first call
+        rdog.gradient_sum = grad_norm_sq
+        rdog.initial_point = copy(M, p)
+        rdog.max_distance = rdog.initial_distance
+
+        # Initial stepsize
+        if rdog.use_curvature
+            ζ = geometric_curvature_function(
+                rdog.sectional_curvature_bound, rdog.max_distance
+            )
+            stepsize = rdog.initial_distance / (sqrt(ζ) * sqrt(max(grad_norm_sq, eps())))
+        else
+            stepsize = rdog.initial_distance / sqrt(max(grad_norm_sq, eps()))
+        end
+    else
+        # Update gradient sum
+        rdog.gradient_sum += grad_norm_sq
+
+        # Update max distance
+        current_distance = distance(M, rdog.initial_point, p)
+        rdog.max_distance = max(rdog.max_distance, current_distance)
+
+        # Compute stepsize
+        if rdog.use_curvature
+            ζ = geometric_curvature_function(
+                rdog.sectional_curvature_bound, rdog.max_distance
+            )
+            stepsize = rdog.max_distance / (sqrt(ζ) * sqrt(rdog.gradient_sum))
+        else
+            stepsize = rdog.max_distance / sqrt(rdog.gradient_sum)
+        end
+    end
+
+    rdog.last_stepsize = stepsize
+    return stepsize
+end
+
+get_initial_stepsize(rdog::DistanceOverGradientsStepsize) = rdog.last_stepsize
+get_last_stepsize(rdog::DistanceOverGradientsStepsize) = rdog.last_stepsize
+
+function show(io::IO, rdog::DistanceOverGradientsStepsize)
+    s = """
+    DistanceOverGradients(;
+      initial_distance=$(rdog.initial_distance),
+      use_curvature=$(rdog.use_curvature),
+      sectional_curvature_bound=$(rdog.sectional_curvature_bound)
+    )
+
+    Current state:
+      max_distance = $(rdog.max_distance)
+      gradient_sum = $(rdog.gradient_sum)
+      last_stepsize = $(rdog.last_stepsize)
+    """
+    return print(io, s)
+end
+
+"""
+    DistanceOverGradients(; kwargs...)
+    DistanceOverGradients(M::AbstractManifold; kwargs...)
+
+Creates a factory for the [`DistanceOverGradientsStepsize`](@ref).
+
+## Keyword arguments
+
+* `initial_distance=1e-3`: initial distance estimate
+* `use_curvature=false`: whether to use sectional curvature
+* `sectional_curvature_bound=0.0`: lower bound on sectional curvature
+
+$(_note(:ManifoldDefaultFactory, "DistanceOverGradientsStepsize"))
+"""
+function DistanceOverGradients(args...; kwargs...)
+    return ManifoldDefaultsFactory(Manopt.DistanceOverGradientsStepsize, args...; kwargs...)
+end

test/plans/test_stepsize.jl

@@ -128,6 +128,121 @@ using ManoptTestSuite
             "Polyak()\nA stepsize with keyword parameters\n   * initial_cost_estimate = 0.0\n"
         @test ps(dmp, sgs, 1) == (f(M, p) - 0 + 1) / (norm(M, p, X)^2)
     end
+    @testset "Distance over Gradients Stepsize" begin
+        @testset "does not use sectional cuvature (Eucludian)" begin
+            M = Euclidean(2)
+            f(M, p) = sum(p .^ 2)
+            grad_f(M, p) = sum(2 .* p)
+            dmp = DefaultManoptProblem(M, ManifoldGradientObjective(f, grad_f))
+            p = [2.0, 2.0]
+            gds = GradientDescentState(M; p = p)
+            ds = DistanceOverGradientsStepsize(
+                M; p = p, initial_distance = 1.0, use_curvature = false
+            )
+            @test ds.gradient_sum == 0
+            @test ds.max_distance == 1.0
+            @test ds.initial_point == p
+            @test ds.last_stepsize === NaN
+            lr = ds(dmp, gds, 0)
+            @test lr == 0.125
+        end
+        @testset "use sectional cuvature (Euclidian)" begin
+            M = Euclidean(2)
+            f(M, p) = sum(p .^ 2)
+            grad_f(M, p) = sum(2 .* p)
+            dmp = DefaultManoptProblem(M, ManifoldGradientObjective(f, grad_f))
+            p = [2.0, 2.0]
+            gds = GradientDescentState(M; p = p)
+            ds = DistanceOverGradientsStepsize(
+                M;
+                p = p,
+                initial_distance = 1.0,
+                use_curvature = true,
+                sectional_curvature_bound = 0.0,
+            )
+            @test ds.gradient_sum == 0
+            @test ds.max_distance == 1.0
+            @test ds.initial_point == p
+            @test ds.last_stepsize === NaN
+            lr = ds(dmp, gds, 0)
+            @test lr == 0.125
+        end
+        @testset "do not use sectional cuvature (Sphere)" begin
+            M = Sphere(1)
+            f(M, p) = sum(p .^ 2)
+            grad_f(M, p) = sum(2 .* p)
+            dmp = DefaultManoptProblem(M, ManifoldGradientObjective(f, grad_f))
+            p = [1, 0]
+            gds = GradientDescentState(M; p = p)
+            ds = DistanceOverGradientsStepsize(
+                M; p = p, initial_distance = 1.0, use_curvature = false
+            )
+            @test ds.gradient_sum == 0
+            @test ds.max_distance == 1.0
+            @test ds.initial_point == p
+            @test ds.last_stepsize === NaN
+            lr = ds(dmp, gds, 0)
+            @test lr == 0.5
+        end
+        @testset "use sectional cuvature (Sphere)" begin
+            M = Sphere(1)
+            f(M, p) = sum(p .^ 2)
+            grad_f(M, p) = sum(2 .* p)
+            dmp = DefaultManoptProblem(M, ManifoldGradientObjective(f, grad_f))
+            p = [1, 0]
+            gds = GradientDescentState(M; p = p)
+            ds = DistanceOverGradientsStepsize(
+                M;
+                p = p,
+                initial_distance = 1.0,
+                use_curvature = true,
+                sectional_curvature_bound = 1.0,
+            )
+            @test ds.gradient_sum == 0
+            @test ds.max_distance == 1.0
+            @test ds.initial_point == p
+            @test ds.last_stepsize === NaN
+            lr = ds(dmp, gds, 0)
+            @test lr == 0.5
+        end
+        @testset "use sectional curvature (Hyperbolic)" begin
+            M = Hyperbolic(2)              # Lorentz model in R^3
+
+            t = 0.5
+            p = [cosh(t), sinh(t), 0.0]
+
+            v_ambient = [0.0, 1.0, 0.0]
+            g = project(M, p, v_ambient)
+            gnorm = norm(M, p, g)
+            f(M, q) = 0.0
+            grad_f(M, q) = g
+
+            dmp = DefaultManoptProblem(M, ManifoldGradientObjective(f, grad_f))
+            gds = GradientDescentState(M; p = p)
+            ds = DistanceOverGradientsStepsize(
+                M;
+                p = p,
+                initial_distance = 1.0,
+                use_curvature = true,
+                sectional_curvature_bound = -1.0,
+            )
+
+            @test ds.gradient_sum == 0
+            @test ds.max_distance == 1.0
+            @test ds.initial_point == p
+            @test ds.last_stepsize === NaN
+
+            # Expected initial step:
+            # ζκ(1) = (√|κ| * 1) / tanh(√|κ| * 1) = 1 / tanh(1)
+            # lr = initial_distance / (sqrt(ζκ(1)) * sqrt(gnorm^2)) = 1 / (sqrt(ζκ(1)) * gnorm)
+            zeta = 1 / tanh(1.0)
+            expected_lr = 1.0 / (sqrt(zeta) * gnorm)
+
+            # Call with the known gradient to avoid any objective/gradient mismatches
+            lr = ds(dmp, gds, 0; gradient = g)
+            @test isapprox(lr, expected_lr; rtol = 1.0e-12, atol = 0)
+        end
+    end
     @testset "max_stepsize fallbacks" begin
         M = ManoptTestSuite.DummyManifold()
         @test isinf(Manopt.max_stepsize(M))