Support for specialized matrix types (e.g., BlockBandedMatrices) with optimized solvers

[rsenne]

Hello, and thank you for the excellent package.

I was wondering if there are plans to support additional matrix types that have specialized solving routines or optimized numerical algorithms. For instance, in my current workflow I am using a Hessian matrix that is block-tridiagonal. I pass it as a SparseMatrixCSC for performance, but I realized that I could instead represent it using BlockBandedMatrices.jl, which provides efficient linear solvers for such structures.

This would be particularly helpful in Newton-based optimization routines involving expressions like:
$\theta^{new} = \theta^{old} - H \setminus \nabla$

Having native support for block-banded matrices—or even a hook to dispatch to a custom solver—would be a great enhancement.

Thanks again for your work!

[rsenne]

i poked around a bit and I think this is as easy as modifying update_state! like follows:

function update_state!(d, state::NewtonState, method::Newton)
    # Search direction is always the negative gradient divided by
    # a matrix encoding the absolute values of the curvatures
    # represented by H. It deviates from the usual "add a scaled
    # identity matrix" version of the modified Newton method. More
    # information can be found in the discussion at issue #153.
    T = eltype(state.x)
    H = NLSolversBase.hessian(d)
    g = gradient(d)

    if H isa AbstractSparseMatrix || H isa AbstractBlockBandedMatrix # Added BlockBandedMatrix support here
        # General fallback to backslash (dense RHS is fine)
        state.s .= -(H \ convert(Vector{T}, g))
    else
        state.F = cholesky!(Positive, H)
        if g isa Array
            ldiv!(state.s, state.F, -g)
        else
            gv = Vector{T}(undef, length(g))
            copyto!(gv, -g)
            copyto!(state.s, state.F \ gv)
        end
    end

I then ran the following MWRE:

using Optim, BlockBandedMatrices, LinearAlgebra

# test for BlockBandedMatrix Newton method
function test_blockbanded_newton()
    println("Testing Newton method with BlockBandedMatrices...")
    
    # Create a simple 4x4 block-banded matrix with 2x2 blocks
    n = 4
    λ = [2, 2]  # row block sizes
    μ = [2, 2]  # column block sizes
    l, u = 1, 1  # lower and upper bandwidth
    
    # Create and populate the BlockBandedMatrix
    H = BlockBandedMatrix(zeros(n, n), λ, μ, (l, u))
    
    # Fill the blocks manually to ensure it's positive definite
    H[Block(1,1)] = [3.0 1.0; 1.0 3.0]  # First diagonal block
    H[Block(2,2)] = [4.0 0.5; 0.5 4.0]  # Second diagonal block
    H[Block(1,2)] = [0.5 0.2; 0.2 0.5]  # Super-diagonal block
    H[Block(2,1)] = [0.5 0.2; 0.2 0.5]  # Sub-diagonal block
    
    println("BlockBandedMatrix H:")
    display(Matrix(H))  # Convert to dense for display
    
    # Simple quadratic: f(x) = 0.5 * x'Hx - b'x + c
    b = [1.0, 2.0, 3.0, 4.0]
    c = 5.0
    
    println("RHS vector b: ", b)
    
    # Objective function and derivatives
    f(x) = 0.5 * dot(x, H * x) - dot(b, x) + c
    
    function g!(storage, x)
        mul!(storage, H, x)
        storage .-= b
    end
    
    function h!(storage, x)
        copyto!(storage, H)
    end
    
    # Starting point
    x0 = [0.5, 0.5, 0.5, 0.5]
    
    # Check initial conditions
    g0 = similar(x0)
    g!(g0, x0)
    println("Starting point x0: ", x0)
    println("Initial gradient: ", g0)
    println("Initial gradient norm: ", norm(g0))
    println("Initial f-value: ", f(x0))
    
    # Verify Hessian type
    println("Hessian type: ", typeof(H))
    
    # Solve analytically for comparison
    x_analytical = H \ b
    println("Analytical solution: ", x_analytical)
    println("Analytical f-value: ", f(x_analytical))
    
    # Optimize using Newton method
    println("\nRunning Newton optimization...")
    result = optimize(f, g!, h!, x0, Newton(), 
                     Optim.Options(show_trace=true, g_tol=1e-8))
    
    println("\nOptimization result:")
    println("  Converged: ", Optim.converged(result))
    println("  Iterations: ", Optim.iterations(result))
    println("  Final f-value: ", Optim.minimum(result))
    
    x_optim = Optim.minimizer(result)
    println("  Final solution: ", x_optim)
    println("  Solution error: ", norm(x_optim - x_analytical))
    
    return result
end

# Run the test
test_blockbanded_newton()

And this worked:

Optimization result:
Converged: true
Iterations: 1
Final f-value: 1.8658389128064923
Final solution: [0.01402547727533815, 0.4771224382015322, 0.6165672541057787, 0.8625875145978192]
Solution error: 2.520129940731374e-16

• Status: success

• Candidate solution
  Final objective value: 1.865839e+00

• Found with
  Algorithm: Newton's Method

• Convergence measures
  |x - x'| = 4.86e-01 ≰ 0.0e+00
  |x - x'|/|x'| = 5.63e-01 ≰ 0.0e+00
  |f(x) - f(x')| = 6.09e-01 ≰ 0.0e+00
  |f(x) - f(x')|/|f(x')| = 3.26e-01 ≰ 0.0e+00
  |g(x)| = 0.00e+00 ≤ 1.0e-08

• Work counters
  Seconds run: 1 (vs limit Inf)
  Iterations: 1
  f(x) calls: 3
  ∇f(x) calls: 3
  ∇²f(x) calls: 1

[pkofod]

I support making this more modular / improving the existing code. I'll have to look at the exact changes you made.

[rsenne]

Great! I'm happy to take charge on this for a PR. Right now my solution from above is pretty naively implemented--I'm sure there is a better way to handle this to improve overall genericity when special matrix types can be exploited. I can also work on a non-trivial use-case to confirm there is a speed up when doing this?

[pkofod]

That would be great. If you prefer, you can sketch a solution out. If you accept that I may ask for modifications, you can also try out something you think first and make a PR and we can discuss.

[rsenne]

Sounds good! I'll go through the code this week to see how this can be most elegantly handled. I'll give you an update if I come up with what I view is a good idea and then you can give me feedback. In parallel, I'll also confirm this does give a performance speed up; it definitely should in theory, but def want to make sure

[rsenne]

Hey @pkofod hoping you can give me a little insight on a few things. In your code, specifically the update_state! function I've been playing around with:

function update_state!(d, state::NewtonState, method::Newton)
    # Search direction is always the negative gradient divided by
    # a matrix encoding the absolute values of the curvatures
    # represented by H. It deviates from the usual "add a scaled
    # identity matrix" version of the modified Newton method. More
    # information can be found in the discussion at issue #153.
    T = eltype(state.x)
    if typeof(NLSolversBase.hessian(d)) <: AbstractSparseMatrix
        state.s .= .-(NLSolversBase.hessian(d) \ convert(Vector{T}, gradient(d)))
    else
        state.F = cholesky!(Positive, NLSolversBase.hessian(d))
        if typeof(gradient(d)) <: Array
            # is this actually StridedArray?
            ldiv!(state.s, state.F, -gradient(d))
        else
            # not Array, we can't do inplace ldiv
            gv = Vector{T}(undef, length(gradient(d)))
            copyto!(gv, -gradient(d))
            copyto!(state.s, state.F \ gv)
        end
    end
    # Determine the distance of movement along the search line
    lssuccess = perform_linesearch!(state, method, d)
    # Update current position # x = x + alpha * s
    @. state.x = state.x + state.alpha * state.s
    return !lssuccess # break on linesearch error
end

It basically does one of two things: 1.) If it's sparse, do a linear solve or 2.) if it's not sparse, do a factorization that enforces it is positive definite via PositiveFactorizations.jl to ensure Newton's method always goes in the correct direction.

So here are my questions: why do we ignore the possibility of incorrect ascent in the sparse case? Is it because the linear solve for sparse is so much faster that we can take an incorrect direction once in awhile? Or is there something else I'm missing?

Because otherwise, it seems like we could just allow any AbstractMatrix and just simply do for any matrix type:

state.s .= .-(NLSolversBase.hessian(d) \ convert(Vector{T}, gradient(d)))

From here you could just check for a Matrix in which case you do this cholesky factorization that has been implemented.

Does this make sense?

[rsenne]

Non sequitur: a long term advancement could be to use LinearSolvers.jl as well. I could see this being useful for anywhere in the package a linear solve is performed and regardless of what matrix types are used. Obviously out of the scope of the current proposed enhancement, but could be a nice follow up down the line.

[pkofod]

Yeah, we do the PositiveFactorization to ensure a descent step in the Matrix case. I think someone wanted sparse support at some point and as you rightly point out that one is not safeguarded against ascent directions. I'd argue that if your Hessian is not PSD then you should maybe use a trust region approach. Line search methods are generally most useful if you can assure PSD Hessian (approximations).

[rsenne]

Okay got it. To me this begs the question if it makes sense to keep this functionality? I would agree with you that a trust region method would be more suitable in the case of not being able to guarantee PSD. This would then also mean that my proposed solution of just using the standard linear solve would be optimal as If we're always solving via A \ b, Julia's dispatch system can handle the specifics—whether sparse, dense, or otherwise. What do you think?

[rsenne]

Hey @pkofod just wanted to check in on this!

Here are my two basic ideas for moving forward:

Option 1: Linear Solve Dispatch

We simplify the code path by allowing any AbstractMatrix and rely on Julia's multiple dispatch to handle the best solver for the matrix type:

state.s .= .-(H \ convert(Vector{T}, g))

This avoids the current bifurcation between AbstractSparseMatrix and dense Matrix and delegates the solving strategy to Julia’s dispatch system (e.g., BlockBandedMatrix, SparseMatrixCSC, etc). It assumes the user ensures their Hessian is PSD, and thus appropriate for line search–based Newton.

Option 2: Allow optional user-defined solvers

We could allow users to optionally define their own preferred method of solving. Potentially adding a kwarg to Optim.Options like so:

optimize(f, g!, h!, x0, Newton(); solve = (H, g) -> H \ g)

Then inside update_state!, use this hook:

state.s .= -method.solve(H, convert(Vector{T}, g))

In the event nothing is provided, we could default to the current method of using PositiveFactorizations.jl etc.

Do you like either of these options? Is there something else you would suggest?

[pkofod]

optimize(f, g!, h!, x0, Newton(); solve = (H, g) -> H \ g)

I think it would more common to put it inside the Newton method, but yes we could allow for that. The question is: does the type uniquely define how you want to solve? I don't think so, so probably we want to provide a factorization / solve interface instead of relying solely on the type. Does that make sense?

[rsenne]

I think it would more common to put it inside the Newton method

Yes, I agree. That makes better sense.

Does that make sense?

I think so. What about something like this:

Newton(;
    factorize = H -> cholesky!(Positive, H),
    solve = nothing  # or (H, g) -> H \ g
)

Then integrated to update_state! as follows:

function update_state!(d, state::NewtonState, method::Newton)
    T = eltype(state.x)
    H = NLSolversBase.hessian(d)
    g = gradient(d)

    if method.solve !== nothing
        # User provided full custom solve routine
        state.s .= -method.solve(H, convert(Vector{T}, g))
    else
        # Use factorization (default: Cholesky for dense, user-customizable)
        F = method.factorize(H)
        if g isa StridedVector
            ldiv!(state.s, F, -g)
        else
            gv = Vector{T}(undef, length(g))
            copyto!(gv, -g)
            copyto!(state.s, F \ gv)
        end
    end

    # Line search and position update
    lssuccess = perform_linesearch!(state, method, d)
    @. state.x = state.x + state.alpha * state.s
    return !lssuccess
end

This should preserve the current default behavior of using cholesky!(Positive, H), but would allow users to override the factorization strategy if they know their matrix is always PD and want to skip the check—or plug in something like lu, ldlt, or a structure-specific routine. Alternatively, users can supply their own solve method directly, which could internally manage caching, preconditioning, or heuristics. I tried to mirror the original structure as closely as possible. Let me know if you'd like anything changed—happy to adjust.

[pkofod]

I think it would be good to see it in a PR at this point :) Please remember to add some tests to make sure that everything is exercised properly.