🚀 The feature, motivation and pitch

Motivation: sparse LM optimizer relies on a sparse Ax = b solver

Hi! We are working on a sparse Levenberg–Marquardt optimizer, and we have already sparsified the Jacobian matrix and A matrix (derived from J_T @ J) into COO tensors. By default, the pypose LM optimizer uses Cholesky solver to solve for the the param update for LM, which does not support sparse PyTorch tensors.

We tried to implement a sparse CG solver ourselves, but it did not work very well

Therefore, we implemented a Conjugate Gradient solver where the A matrix can be sparsed. We reference this open source project. This solver works for optimizing BAL problems - we can see the loss decreasing - but the terminal loss is still very high compared to GTSAM's LM and Scipy's non-linear least square optimizer (which used trf optimizer):

Figure: Our terminal least square loss with CG solver

Figure: GTSAM's least square loss on the same dataset

Experiment shows that scipy's CG solver can achieve better end-to-end result than our CG solver

Going back to the CG solver. We suspected that the loss difference is due to the inaccuracy of our CG solver. Luckily, scipy also has official CG solver implementation, so we compared our CG with scipy's CG. It turns out, with scipy's CG solver we achieved much better least square loss, so our original CG solver may be doing something wrong here...

Figure: Our LM using scipy's CG solver

Towards a more accurate CG solver

Therefore, we came to the conclusion that we need a more accurate sparse linear system solver. For example, JAX hand-crafted its CG solver that is bit-accurate with scipy.

Figure: JAX's CG solver

However, we could really use some help for the sparse CG solver from the pypose community :)

I am Huan Xu in the pypose Slack, shoot me a message if you want to collaborate with us on this... Thanks!