torch-levenberg-marquardt

A PyTorch implementation of the Levenberg-Marquardt algorithm for mini batch training, with GPU support, pluggable losses via the square root trick, and damping strategies, aimed at models small enough to form a Jacobian.

First order methods like SGD and Adam rule large scale training because they scale, and for models with millions or billions of parameters they are the only realistic option. The README's argument is that smaller models are a different story: second order methods can converge faster there, and sometimes succeed where first order methods give up. Levenberg-Marquardt sits between the two, building on Gauss-Newton for second order information and adding adaptive damping for stability.

The implementation leans on a square root trick so LM works with any PyTorch loss, not just least squares. Supported out of the box are MSELoss, L1Loss, HuberLoss, CrossEntropyLoss, BCELoss, and BCEWithLogitsLoss. Three damping strategies ship: the standard form of the Jacobian product plus lambda times identity, the Fletcher variant using the diagonal of the Jacobian product, and custom strategies. Jacobian computation can be split into sub batches to cut memory, and you can choose a subset of parameters to update. Example training loops cover plain PyTorch and PyTorch Lightning.

Two examples show where it helps. Fitting y equals sinc(10x) with a 61 parameter network, Adam fails to converge while LM converges quickly with very low loss. On MNIST with a 1026 parameter CNN, both methods reach roughly the same accuracy, but LM stops at epoch 8. The README notes that batch size and parameter count need choosing carefully to get the best out of the algorithm.

torch-levenberg-marquardt brings second order optimization to PyTorch for small models, with loss flexibility through the square root trick and split Jacobian computation. The examples show it converging where Adam stalls.

• Source 1 · Project repository
• Source 2 · README