Nonlinearsolve.jl: Fast and Robust Solvers for Nonlinear Equations in Julia

Consider claim:

> The feasible region {x | f(x) = 0} is nonconvex no matter whether f is convex.

For some positive integer n and for the set of real numbers R, consider closed (in the usual topology for R^n), convex set A a subset of R^n. Define function f: R^n --> R so that f(x) = 0 for all x in A, f(x) > 0 for all x not in A, and for all x in R^n f infinitely differentiable at x. It is a theorem that such an f exists.

Then { x | f(x) = 0 } = A and is both closed and convex in contradiction to the claim.

No, there are two types of GPU usage: the ensemble form and the array form. In the array form, you simply make your state variable a GPU and can run all internal operations on the GPU. This is used for GPU acceleration of `f` functions where you have a large state and a relatively "regular" function in `f`, such as neural networks or PDE discretizations. Thus you can see this in action for example on things like the DeepEquilibriumNetworks.jl https://docs.sciml.ai/DeepEquilibriumNetworks/stable/tutoria... use cases, where it's big neural networks with nonlinear solvers in them.

The other use case is if you're solving tons of f(u,p)=0 equations for many different p, but all the same `f`. This shows up for example when doing parameter sweeps. You might think that you can just use the former for this: if you originally had `u` a vector, then you can make `u` a matrix and solve `f(u[:,i],p[:,i])` all simultaneously. This is in fact what machine learning libraries tend to do with their `vmap` approach. However, that hides a choice: how do you change `f` to build the composite `f`? You can either do the matrix operations on each operation of `f`, or you can build a specific GPU kernel for `f`. Of course you can do the former, that is exactly the same as the use case above, reusing the same machinery. However, we also offer the latter. And the reason for that is because calling GPU kernels has a lot of overhead and inhibits interprocedural optimizations, and therefore the former is a lot slower than the latter when you can exploit this known repeated structure. We were able to demonstrate this in a recent publication which showed using the latter approach for ODEs was about 20x-100x faster than PyTorch, Jax, and our own implementation of the former, while our approach to the latter is just about as fast as MPGOS which is a CUDA kernel, indicating that it's some limitation of the design (see https://www.sciencedirect.com/science/article/abs/pii/S00457..., or the arxiv version https://arxiv.org/abs/2304.06835). The kernel building approach is thus simply more efficient for things like parameter sweeps, and it's a pretty unique way of using the GPU as most machine learning libraries are not built around kernel generation but instead only give you the linear algebra kernel calling approaches.

tl;dr, DiffEqGPU.jl's documentation is very clear on the two types of GPU usage https://docs.sciml.ai/DiffEqGPU/stable/getting_started/, but is specifically for ODEs. We need to write some similar tutorials for GPU usage on the nonlinear solvers to this effect.

those are indeed very popular tools but for a slightly different kind of problem than this package is attempting to solve

You can also put it as a minimization problem to minimize ||g(u)||, but I assume that does not magically make anything more efficient.

That makes your Newton method require calculating the Hessian of the objective function, which is a second derivative. A normal Newton method for nonlinear systems uses the Jacobian, which is the first derivative (the connection is that the Hessian is the Jacobian of the gradient). But indeed, formulating it so that you have discontinuities (absolute values) and require second derivatives instead of the first for the same algorithm is not the good kind of magic :-/. You could then try using only gradient-based methods, like Adam, but you can see that is actually using less problem structure. BFGS for example is related to Broyden's method and is thus using simple Hessian approximations, so for nonlinear system solving this is effectively less stable than just using Newton.

That's not to say that there's no reason to use this trick though. I have seen that some optimization tools can improve robustness in some cases, and this is a way to make use of heuristic methods (genetic algorithms, differential evolution, etc.) for globalizing the process, but it's probably not the first trick you'd want to do in most cases.

While raising derivative order is one way to understand why this should be less efficient, but another is to understand that it loses some specific structure. In the nonlinear system problem f(u)=0, any u that solves this is a solution. All solutions are the same, if you find one you're good. In the optimization problem, you want to find "the" u that minimizes f(u). Now if you are minimizing f(u)=||g(u)|| and you know there is a u that causes g(u) to be zero, then all u that causes g(u) to be zero is a solution, all match the global optima. But solvers don't necessarily know that. Solvers need to use second order information to know that they have not hit a saddle point or other behavior. They don't know that "near zero = done", and they can't specialize on this information. This is another way in which you're just giving sophisticated more solvers more work to do where you already know the answer by problem design.