# Solver and AD advice for nonlinear LS problem
**URL:** https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014
**Category:** Optimization (Mathematical)
**Tags:** question, multithreading, ad, preallocation
**Created:** [January 21, 2025, 12:28pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014 "2025-01-21T12:28:34Z")
**Posts on this page:** 6
**Page:** 1
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [January 21, 2025, 12:28pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/1 "2025-01-21T12:28:34Z")
(This is a vague open-ended question for which I cannot provide an MWE.)
I have a nonlinear least squares problem
\min\_x \| f(x) \|\_2^2
where x \in \mathbb{R}^N, N \approx 20...40, and f(x) has about 100–200 elements. f should be treated as a black box by the solver. There are no constraints on x.
I have f implemented in Julia, but it still needs optimization, specifically preallocation of buffers (but it has multithreading, oh my aching head). Currently it takes about 40s on a desktop.
I am using an implementation of a trust region method I wrote for a smaller problem, which calculates the Jacobian, using ForwardDiff.jl as a backend. I wonder an algorithm that calculates the Jacobian-vector product would make sense, but I am not sure which packages have a robust implementation (I could write one given a reference).
f allocates like crazy. I could pre-allocate the buffers but, not knowing the types `ForwardDiff` calls it with, I find this challenging. What are the best practices for preallocation combined with ForwardDiff?
Any other AD backend I should consider?
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [January 21, 2025, 12:59pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/2 "2025-01-21T12:59:51Z")
> [@Tamas\_Papp](#):
>
> I have fff implemented in Julia, but it still needs optimization, specifically preallocation of buffers (but it has multithreading, oh my aching head).
OhMyThreads.jl has nice functionality and documentation for this, if you want to give it a try.
> [@Tamas\_Papp](#):
>
> I am using an implementation of a trust region method I wrote for a smaller problem, which calculates the Jacobian, using [ForwardDiff.jl](https://juliaregistries.github.io/General/packages/redirect_to_repo/ForwardDiff) as a backend.
Is there a reason not to use existing packages like NonlinearSolve.jl?
Also, is the Jacobian sparse? If so, you may achieve significant savings by exploiting it.
> [@Tamas\_Papp](#):
>
> I wonder an algorithm that calculates the Jacobian-vector product would make sense, but I am not sure which packages have a robust implementation (I could write one given a reference).
Do you mean a robust implementation of the JVP? You can always try them all with DifferentiationInterface.jl, the relevant operator is `pushforward`.
> [@Tamas\_Papp](#):
>
> ff allocates like crazy. I could pre-allocate the buffers but, not knowing the types `ForwardDiff` calls it with, I find this challenging. What are the best practices for preallocation combined with ForwardDiff?
Probably PreallocationTools.jl?
> [@Tamas\_Papp](#):
>
> Any other AD backend I should consider?
Have you tried Enzyme.jl?
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [January 21, 2025, 1:20pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/3 "2025-01-21T13:20:26Z")
Yup all of this. NonlinearSolve’s methods are what would be most robust and efficient here, and it ties into DifferentiationInterface.jl so you should just be able to define the ADType and let it fly. The example recreating some of the PETSc examples should be a good starting point:
> **[PETSc SNES Example 2 · NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/tutorials/snes_ex2/)**
>
> Documentation for NonlinearSolve.jl.
For nonlinear least squares, you simplify need to tell it what the output sizing is. A tutorial on that is found here:
> **[Getting Started with Nonlinear Rootfinding in Julia · NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/tutorials/getting_started/#Problem-Type-4:-Solving-Nonlinear-Least-Squares-Problems)**
>
> Documentation for NonlinearSolve.jl.
You can make your system completely non-allocating, see the full discussion in:
> **[Efficiently Solving Large Sparse Ill-Conditioned Nonlinear Systems in Julia ·...](https://docs.sciml.ai/NonlinearSolve/stable/tutorials/large_systems/)**
>
> Documentation for NonlinearSolve.jl.
That doesn’t discuss PreallocationTools.jl, but the docs there should be relatively straightforward.
If you really need performance, you can mix in some of the ModelingToolkit structural simplification tricks, see:
> **[Symbolic Nonlinear System Definition and Acceleration via ModelingToolkit ·...](https://docs.sciml.ai/NonlinearSolve/stable/tutorials/modelingtoolkit/)**
>
> Documentation for NonlinearSolve.jl.
Whether that is applicable depends on the type of equations you have, it will increase compliation but in most cases should give a straight multiple order of magnitude performance gain if you can split to SCCs and all of the fancy stuff.
Finally the list of solvers is in:
> **[Nonlinear Least Squares Solvers · NonlinearSolve.jl](https://docs.sciml.ai/NonlinearSolve/stable/solvers/nonlinear_least_squares_solvers/)**
>
> Documentation for NonlinearSolve.jl.
`LevenbergMarquardt()` is not done by the book and there’s some changes that make it more stable than the standard form for ill-conditioned cases, I’d probably recommend that. But of course there is a `TrustRegion` method there.
These schemes are benchmarked in quite a lot of detail to make sure we have full performance and robustness, for the proof of that see:
> **[NonlinearSolve.jl: High-Performance and Robust Solvers for Systems of...](https://arxiv.org/abs/2403.16341)**
>
> Efficiently solving nonlinear equations underpins numerous scientific and engineering disciplines, yet scaling these solutions for complex system models remains a challenge. This paper presents NonlinearSolve.jl - a suite of high-performance...
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [January 21, 2025, 1:23pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/4 "2025-01-21T13:23:13Z")
> [@Tamas\_Papp](#):
>
> I wonder an algorithm that calculates the Jacobian-vector product would make sense, but I am not sure which packages have a robust implementation (I could write one given a reference).
Probably not for this size. Making sure you fill the matrix with sparse AD is probably a big deal, but using a sparse matrix in the LU or an iterative solver is likely not a good idea at this size.
---
### Author:  [@Tamas\_Papp](https://discourse.julialang.org/u/Tamas_Papp)
#### Post date: [January 21, 2025, 4:20pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/5 "2025-01-21T16:20:36Z")
> [@gdalle](#):
>
> [OhMyThreads.jl](https://juliaregistries.github.io/General/packages/redirect_to_repo/OhMyThreads) has nice functionality and documentation for this, if you want to give it a try.
I am already using that package for the threaded computation, but I have not implemented the preallocation part because I am still not sure how to do it.
Conceptually, this is how it works: the computation is about 100–150 “recipes” (recipes are always the same), each needs a number of temporary matrices (varies for each recipe, but constant given the recipe) for a two-pass algorithm.
Probably the best option for me would be [`Channel`](https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#The-safe-way:-Channel), but I do not know in advance how many matrices I will need in advance.
To clarify with an MWE, consider the following:
```julia
using DifferentiationInterface, ForwardDiff
struct MyComputation
N::Int
M::Int
recipes::Vector{Int}
end
function compute_recipe(α::T, N, M, recipe) where T
matrices = Vector{Matrix{T}}(undef, recipe)
# “backward pass” (for MWE)
for i in recipe:(-1):1
if i == recipe
# this line should grab an empty matrix from a Channel
matrices[i] = fill(exp(α), N, M)
else
# ditto
matrices[i] = matrices[i + 1] .+ α
end
end
# “forward pass” (for MWE)
mapfoldl(sum, +, matrices)
end
function (mc::MyComputation)(α)
(; N, M, recipes) = mc
mapreduce(r -> compute_recipe(α, N, M, r), +, recipes)
end
# runtime code
mc = MyComputation(4, 5, sample(1:9, 20, replace = true))
@code_warntype mc(0.5)
DifferentiationInterface.value_and_derivative(mc, AutoForwardDiff(; chunksize = 1), 0.5)
```
Now recipes could come in any order depending on the scheduler, and I may need a variable number depending on the threads. Or can I add to the `Channel` “on demand”?
---
### Author:  [@cgeoga](https://discourse.julialang.org/u/cgeoga)
#### Post date: [January 21, 2025, 5:26pm UTC](https://discourse.julialang.org/t/solver-and-ad-advice-for-nonlinear-ls-problem/125014/6 "2025-01-21T17:26:51Z")
If it is possible to get a trial license, [Artelys KNITRO](https://www.artelys.com/solvers/knitro/) has a least squares mode that uses Gauss-Newton and has always worked very well for my problems (solving the score equations of a GP parameter estimation problem, which is a very ugly but low-dimensional problem).
I don’t think the interface is exposed in JuMP, but it is implemented in `KNITRO.jl`. I have my own little wrapper demonstrating the boilerplate [here](https://git.sr.ht/~cgeoga/StandaloneKNITRO.jl/tree/master/item/src/knitro_optimize.jl#L59).
Maybe worth a shot!
Quick edit: Maybe `Bumper.jl` could be useful for reducing allocations without too much pre-allocation suffering?