Guard optimizers and linesearch against non-descent directions and zero steps

src/OptimKit.jl

@@ -15,6 +15,14 @@ const GRADTOL = ScopedValue(1.0e-8)
 const MAXITER = ScopedValue(1_000_000)
 const VERBOSITY = ScopedValue(1)
 
+# `!isnegative(x)` is not equivalent to `x >= 0`: it is also `true` for NaN
+@static if !isdefined(Base, :isnegative) # added to Base in Julia 1.13
+    isnegative(x::Real) = x < 0
+end
+@static if !isdefined(Base, :ispositive) # added to Base in Julia 1.13
+    ispositive(x::Real) = x > 0
+end
+
 # Default values for the manifold structure
 _retract(x, d, α) = (add(x, d, α), d)
 _invretract(x, y) = add(y, x, -1)

src/cg.jl

@@ -125,8 +125,33 @@ function optimize(
                     )
                 end
             )
-            η = add!(η, ηprev, β)
+            # e.g. 0/0 after a zero step, where g == gprev
+            if isfinite(β)
+                η = add!(η, ηprev, β)
+            else
+                β = zero(α)
+            end
         end
+        dϕ = inner(x, g, η)
+        # not equivalent to `dϕ >= 0`: a NaN slope must also count as non-descent
+        if !isnegative(dϕ) && !iszero(β)
+            verbosity >= 2 &&
+                @info "CG: not a descent direction, restarting with the preconditioned gradient"
+            β = zero(α)
+            η = scale!(deepcopy(Pg), -1)
+            dϕ = inner(x, g, η)
+        end
+        # not equivalent to `dϕ >= 0`: a NaN slope must also count as non-descent
+        if !isnegative(dϕ)
+            verbosity >= 1 &&
+                @warn @sprintf(
+                "CG: preconditioned gradient is not a descent direction (dϕ = %.2e), stopping",
+                dϕ
+            )
+            break
+        end
+        # after a zero step, doubling would leave the initial guess at zero
+        iszero(α) && (α = 1 / sqrt(-dϕ))
 
         # store current quantities as previous quantities
         xprev = x
@@ -174,6 +199,12 @@ function optimize(
         end
         ηprev = transport!(deepcopy(ηprev), xprev, ηprev, α, x)
 
+        if iszero(α) && iszero(β)
+            verbosity >= 1 &&
+                @warn "CG: linesearch made no progress along the preconditioned gradient, stopping"
+            break
+        end
+
         # increase α for next step
         α = 2 * α
     end

src/gd.jl

@@ -86,6 +86,16 @@ function optimize(
         # compute new search direction
         Pg = precondition(x, deepcopy(g))
         η = scale!(Pg, -1) # we don't need g or Pg anymore, so we can overwrite it
+        dϕ = inner(x, g, η)
+        # not equivalent to `dϕ >= 0`: a NaN slope must also count as non-descent
+        if !isnegative(dϕ)
+            verbosity >= 1 &&
+                @warn @sprintf(
+                "GD: preconditioned gradient is not a descent direction (dϕ = %.2e), stopping",
+                dϕ
+            )
+            break
+        end
 
         # perform line search
         _xlast[] = x # store result in global variables to debug linesearch failures
@@ -118,6 +128,11 @@ function optimize(
             numiter, format_time(Δt), f, normgrad, α, nfg
         )
 
+        if iszero(α)
+            verbosity >= 1 && @warn "GD: linesearch made no progress, stopping"
+            break
+        end
+
         # increase α for next step
         α = 2 * α
     end

src/lbfgs.jl

@@ -95,11 +95,28 @@ function optimize(
                 H(g, ξ -> precondition(x, ξ), (ξ1, ξ2) -> inner(x, ξ1, ξ2), add!, scale!)
             end
             η = scale!(Hg, -1)
-        else
+            # not equivalent to `inner(x, g, η) >= 0`: a NaN slope must also count as non-descent
+            if !isnegative(inner(x, g, η))
+                verbosity >= 2 &&
+                    @info "LBFGS: not a descent direction, resetting the inverse Hessian approximation"
+                empty!(H)
+            end
+        end
+        if length(H) == 0
             Pg = precondition(x, deepcopy(g))
             normPg = sqrt(inner(x, Pg, Pg))
             η = scale!(Pg, -0.01 / normPg) # initial guess: scale invariant
         end
+        dϕ = inner(x, g, η)
+        # not equivalent to `dϕ >= 0`: a NaN slope must also count as non-descent
+        if !isnegative(dϕ)
+            verbosity >= 1 &&
+                @warn @sprintf(
+                "LBFGS: preconditioned gradient is not a descent direction (dϕ = %.2e), stopping",
+                dϕ
+            )
+            break
+        end
 
         # store current quantities as previous quantities
         xprev = x
@@ -139,6 +156,17 @@ function optimize(
             numiter, format_time(Δt), f, normgrad, α, length(H), nfg
         )
 
+        if iszero(α)
+            if length(H) == 0
+                verbosity >= 1 &&
+                    @warn "LBFGS: linesearch made no progress along the preconditioned gradient, stopping"
+                break
+            end
+            # the same direction would be proposed again
+            empty!(H)
+            continue
+        end
+
         # transport gprev, ηprev and vectors in Hessian approximation to x
         gprev = transport!(gprev, xprev, ηprev, α, x)
         for k in 1:length(H)

src/linesearches.jl

@@ -124,8 +124,12 @@ function (ls::HagerZhangLineSearch)(
     )
     (f₀, g₀) = fg₀
     ϕ₀ = f₀
+    if !(isfinite(initialguess) && ispositive(initialguess))
+        throw(ArgumentError("initial guess for the step length should be positive and finite, got $initialguess"))
+    end
     dϕ₀ = inner(x₀, g₀, η₀)
-    if dϕ₀ >= zero(dϕ₀)
+    # not equivalent to `dϕ₀ >= 0`: a NaN slope must also count as non-descent
+    if !isnegative(dϕ₀)
         @warn "Linesearch was not given a descent direction: returning zero step length"
         return x₀, f₀, g₀, η₀, zero(one(f₀)), 0
     end
@@ -425,13 +429,23 @@ function bracket(iter::HagerZhangLineSearchIterator{T}, c::LineSearchPoint) wher
                 c.α, c.dϕ, c.ϕ - p₀.ϕ
             )
         end
+        if !(isfinite(c.ϕ) && isfinite(c.dϕ))
+            verbosity >= 1 &&
+                @warn "  Linesearch bracket: no finite function value or slope within the allowed function evaluations"
+            return a, a, numfg
+        end
         c.dϕ >= 0 && return a, c, numfg # B1
         # from here: c.dϕ < 0
         if c.ϕ > fmax # B2
             a, b, nfg = bisect(iter, iter.p₀, c)
             return a, b, numfg + nfg
         else # B3
             a = c
+            if numfg >= iter.parameters.maxfg
+                verbosity >= 1 &&
+                    @warn "  Linesearch bracket: slope still negative after the allowed function evaluations"
+                return a, a, numfg
+            end
             α *= iter.parameters.ρ
             c = takestep(iter, α)
             numfg += 1

test/runtests.jl

@@ -98,6 +98,46 @@ algorithms = (GradientDescent, ConjugateGradient, LBFGS)
     @test f < 1.0e-12
 end
 
+# an indefinite preconditioner eventually produces a non-descent direction
+@testset "Non-descent direction $algtype" for algtype in algorithms
+    fg = quadraticproblem(Matrix(1.0I, 2, 2), zeros(2))
+    precondition(x, g) = [g[1], -g[2] / 2]
+    alg = algtype(; verbosity = 1, gradtol = 1.0e-12, maxiter = 100)
+    x, f, g, numfg, history = @test_logs (:warn, r"not a descent direction") (:warn, r"not converged") optimize(
+        fg, [1.0, 0.1], alg; precondition
+    )
+    @test all(isfinite, x)
+    @test size(history, 1) - 1 < 100
+end
+
+struct ConstantFlavor <: OptimKit.CGFlavor
+    β::Float64
+end
+(flavor::ConstantFlavor)(args...) = flavor.β
+
+# a bad β should restart from the preconditioned gradient instead of breaking the optimization
+@testset "ConjugateGradient restarts for β = $β" for β in (NaN, -1.0e3)
+    n = 10
+    y = randn(n)
+    A = randn(n, n)
+    A = A' * A + I
+    fg = quadraticproblem(A, y)
+    alg = ConjugateGradient(; flavor = ConstantFlavor(β), verbosity = 0, gradtol = 1.0e-8, maxiter = 10_000)
+    x, f, g, numfg, history = optimize(fg, randn(n), alg)
+    @test all(isfinite, x)
+    @test x ≈ y rtol = cond(A) * 1.0e-8
+end
+
+@testset "Linesearch rejects invalid initial guess $α" for α in (0.0, -1.0, NaN, Inf)
+    fg = x -> (x^2, 2 * x)
+    @test_throws ArgumentError HagerZhangLineSearch()(fg, 1.0, -2.0; initialguess = α)
+end
+
+@testset "Linesearch with non-finite slope" begin
+    fg = x -> (x^2, 2 * x)
+    @test_logs (:warn, r"not given a descent direction") match_mode = :any HagerZhangLineSearch()(fg, 1.0, NaN)
+end
+
 include("sphere.jl")
 
 @testset "Manifold correctness" begin