# torch-levenberg-marquardt: Second-order training for small PyTorch models > A PyTorch package that implements the Levenberg-Marquardt algorithm with mini-batch support, custom damping, and loss-agnostic residuals. It targets small models where first-order optimizers stall, but it is not a replacement for Adam at scale. **fabiodimarco/torch-levenberg-marquardt** — PyTorch implementation of Levenberg-Marquardt training algorithm - Repository: https://github.com/fabiodimarco/torch-levenberg-marquardt - Stars: 101 · Forks: 5 - Language: Python - License: MIT - Published: 2026-08-17 · Updated: 2026-08-18 · Language: en - Canonical page: https://hysenlabs.com/projects/fabiodimarco-torch-levenberg-marquardt ## What this optimizer is for The Levenberg-Marquardt algorithm is a second-order optimization method that combines the Gauss-Newton approach with adaptive damping. This package brings that method to PyTorch, with support for mini-batch training and both regression and classification. The intended user is someone training small neural networks, where first-order methods like SGD or Adam are not the only option. The README states that second-order methods can converge faster and sometimes succeed where first-order methods fail, but it also acknowledges that first-order methods remain the only viable choice for models with millions or billions of parameters. So the target is a narrow but real niche: small models where convergence speed matters more than scaling to massive parameter counts. ## The square root trick for arbitrary losses Levenberg-Marquardt is designed for least-squares problems, which normally means mean squared error. This package extends it to any loss that is always non-negative by redefining residuals as the square root of the loss. The README shows the derivation: if the loss is l(y, f(x, W)) >= 0, then the residual r_i becomes sqrt(l(y, f(x, W))). This is a simple but effective trick. For MSELoss, the residual simplifies to y_pred - y_true, which is the standard form. For CrossEntropyLoss, the residual is the square root of the per-sample cross-entropy. The package ships with six built-in losses: MSELoss, L1Loss, HuberLoss, CrossEntropyLoss, BCELoss, and BCEWithLogitsLoss. If you need a different loss, you must implement a custom residual definition, which means writing a class with a forward method and a residuals method. That is a concrete extension point, but it also means the package is not plug-and-play for arbitrary PyTorch losses. ## How the algorithm is implemented The core is the LevenbergMarquardtModule, which wraps a PyTorch model and a loss function. The algorithm computes the Jacobian of the residuals with respect to the model parameters, then forms the Gauss-Newton Hessian approximation as J^T J. The damping term is added to this approximation, and the resulting linear system is solved to get the parameter update. The README mentions a 'split Jacobian matrix computation' to reduce memory usage, which is important because the full Jacobian for a model with many parameters can be enormous. The package also supports custom parameter selection strategies, so you can choose to update only a subset of the model's parameters during a training step. This is a practical feature for models where some layers should stay frozen or where you want to limit the computational cost. The solve method is configurable; the examples use 'qr', which suggests that the linear solve is done via QR decomposition, though the README does not detail other options. ## Damping strategies and their trade-offs The damping factor is what makes Levenberg-Marquardt stable. The package implements three strategies: Standard, which adds lambda times the identity matrix to J^T J; Fletcher, which adds lambda times the diagonal of J^T J; and Custom, which lets you define your own. The Standard approach is the classic one, but it can be less effective when the Hessian approximation has very different scales across parameters. Fletcher's diagonal scaling is often better in practice because it normalizes the curvature per parameter. The README does not provide guidance on which to choose, so you will have to experiment. The damping factor is updated adaptively, and the number of attempts per step is controlled by the attempts_per_step parameter in the training module. A higher value allows more tries to find an acceptable update, but it also increases computation time. This is a genuine trade-off that the documentation leaves for the user to tune. ## Getting it running Installation is straightforward: pip install torch-levenberg-marquardt. The README shows a minimal training loop using the utils.fit function, which takes a LevenbergMarquardtModule and a DataLoader. You set the learning_rate, attempts_per_step, and solve_method. For example, tlm.training.LevenbergMarquardtModule(model=model, loss_fn=tlm.loss.MSELoss(), learning_rate=1.0, attempts_per_step=10, solve_method='qr'). There is also a PyTorch Lightning wrapper, utils.CustomLightningModule, so you can integrate with the Trainer. For development, the repository provides conda environment files for CUDA, CPU-only, and macOS. The PyPI badge indicates the package is published, but the README does not show the exact version or release history. The examples include a Colab notebook, which is a good way to try it without setting up a local environment. ## Memory and scalability limits The biggest limitation is memory. The Gauss-Newton method requires computing the Jacobian of the residuals with respect to all parameters. For a model with P parameters and a batch of size B, the Jacobian is B times P, and the J^T J matrix is P times P. That is quadratic in the parameter count. The package tries to mitigate this with split Jacobian computation and parameter selection, but it does not change the fundamental scaling. The README explicitly says that first-order methods are the only viable option for models with millions or billions of parameters. So if your model has more than a few hundred thousand parameters, this package will likely run out of memory or be impractically slow. There is also a numerical caveat: the square root trick requires the loss to be non-negative, and for losses like cross-entropy, the gradient of the square root can become unstable near zero. The README does not discuss this failure mode, but it is a real risk. ## Alternatives and what makes this different The direct alternative is the TensorFlow implementation, tf-levenberg-marquardt, by the same author. It offers the same algorithm but for TensorFlow, so the choice is largely about which framework you use. Beyond that, you could use a first-order optimizer like Adam, which is what most PyTorch users do. The difference is fundamental: Adam uses only gradients and adaptive per-parameter learning rates, while Levenberg-Marquardt uses second-order information from the Jacobian. That means LM can take larger, more informed steps, but at a much higher per-step cost. There are also other second-order methods in PyTorch, such as LBFGS, which is built into the torch.optim module. LBFGS approximates the Hessian using a limited memory history, which is more memory-efficient than forming J^T J explicitly. The trade-off is that LBFGS is designed for full-batch optimization and can be finicky with mini-batches. This package's mini-batch support is a real advantage over LBFGS in that regard. ## Maintenance and license The repository is licensed under MIT, which allows commercial use, modification, and redistribution with attribution. The README does not mention a specific version or a changelog, and the last push date is unknown from the material provided. There are no recent releases listed, which could mean the project is stable or that it is not actively maintained. The README links to a Colab notebook and provides development environment files, which suggests the author cares about reproducibility. However, the lack of release history and the fact that the repository is not archived means you should check the commit log and issue tracker before depending on it. The package also depends on PyTorch and optionally PyTorch Lightning, so you need to keep those dependencies in sync. The maintenance cost is moderate: the API is small, but you may need to implement custom losses or damping strategies, which requires understanding the internals. ## Conclusion Adopt torch-levenberg-marquardt if you train small models, typically under a few hundred thousand parameters, on regression or classification tasks where first-order methods converge slowly or fail. The square root trick and split Jacobian computation make it practical for mini-batches, and the MIT license permits commercial use. Do not use it for large-scale networks; the memory cost of the Gauss-Newton approximation grows with the square of the parameter count, and the documentation does not report benchmarks for models beyond toy examples. Before adopting, verify that your loss function is non-negative and that you can define residuals, or use one of the six built-in losses. Also confirm that your damping strategy and solve method ('qr' or otherwise) work on your hardware; the README shows 'qr' but does not compare it to alternatives. Test on a small subset of your data first, because convergence behavior with adaptive damping can be sensitive to the learning rate and attempts_per_step settings. ## FAQ ### Can you explain the Levenberg-Marquardt algorithm in a simple way? It optimises a sum of squared residuals using second-order information, building on the Gauss-Newton method, and adds adaptive damping to fix the stability problems of plain Gauss-Newton. The README's own guidance is that it suits models small enough to afford second-order work, since first-order methods are the only viable choice at millions or billions of parameters. ### Is PyTorch a library? In this project it is used as one, declared as a dependency at version 2.0.0 or newer alongside NumPy, PyTorch Lightning, torchmetrics and tqdm. The library adds an optimiser built on second-order information and says it leverages GPU acceleration, and it offers a Lightning module so the same training step can be used inside that framework's trainer. ### Which loss functions does torch-levenberg-marquardt support? Six out of the box: MSELoss, L1Loss, HuberLoss, CrossEntropyLoss, BCELoss and BCEWithLogitsLoss. Anything else is added by implementing custom residual definitions, and each loss class provides both a forward and a residuals method so the optimiser can work with a per-sample quantity rather than a mean. ### What damping strategies are available? Three: a standard one that adds a scaled identity to the product of the Jacobian transpose with itself, a Fletcher one that scales the diagonal of that product instead, and a custom option for anything else. The Fletcher variant is the interesting one when parameters are badly scaled relative to one another. ## Sources - [Official README](https://github.com/fabiodimarco/torch-levenberg-marquardt#readme) - [Project repository](https://github.com/fabiodimarco/torch-levenberg-marquardt) - [README](https://github.com/fabiodimarco/torch-levenberg-marquardt/blob/main/README.md) - [Releases](https://github.com/fabiodimarco/torch-levenberg-marquardt/releases) --- Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/fabiodimarco-torch-levenberg-marquardt