Scalable Levenberg-Marquardt Training Framework for Deep Learning
Xianghao Wang, Johan Barthélemy, John Taylor
Rigorous, well-analyzed systems engineering that makes exact LMA feasible at 1B parameters, but with no convergence or accuracy evidence on realistic tasks and weak baselines, so practical value remains unproven.
Abstract
The Levenberg-Marquardt (LMA) algorithm is a widely used nonlinear least-squares solver known for its rapid convergence. For small-scale neural networks, LMA provides better convergence and architectural robustness than first-order optimizers like Adam. However, LMA imposes prohibitive memory and computational costs. These costs are primarily from the Jacobian matrix, whose dimensions scale with both the number of samples and parameters. This bottleneck prevents its use in large-scale neural networks. Recent efforts scale LMA using preconditioned conjugate gradient (PCG) solvers to avoid materializing the Jacobian. However, PCG only approximates the solution, requiring many network propagations and a task-specific preconditioner. Other studies optimize LMA to reduce costs, but they restrict their methods to specific tasks or network architectures. As a result, LMA remains impractical for training general large-scale neural networks. In this paper, we propose LeMA, a distributed LMA algorithm for large-scale neural network training on multi-GPU platforms. Our contributions are twofold: First, an exact dual formulation of LMA removes the memory dependence on the parameter dimension. Second, a Jacobian sharding strategy computes Jacobian shards on the fly and in parallel to avoid materializing the full Jacobian. For large-scale neural networks, LeMA reduces the memory footprint of LMA from quadratic in the number of parameters to quadratic in the batch size. This makes exact LMA steps feasible for billion-parameter models without approximate PCG solves. Furthermore, for a convolutional neural network with 1,000,022,632 parameters on an 8-GPU node, LeMA achieves a 7.97x speedup over a single GPU, corresponding to a 99.7% strong scaling efficiency.
AI Impact Assessments
(1 model)Core Contribution
LeMA is a systems-oriented framework that makes exact Levenberg-Marquardt (LMA) steps feasible for billion-parameter networks on multi-GPU nodes. It has three ingredients:
1. An exact dual formulation of full Marquardt LMA, using damping with diag(JᵀJ) rather than the identity. This moves the linear solve from the M×M parameter space to the N×N batch space. Prior dual forms (Di Marco) covered only the Levenberg variant. The proof is a short push-through identity, so the mathematical novelty is modest but correct and useful.
2. Jacobian sharding. Shards are computed on the fly, accumulated into the Gram matrix, and freed. This removes the O(NM) Jacobian footprint, leaving O(min(M,N)² + SM + LA) memory.
3. Engineering optimizations: vmap-vectorized per-sample VJPs, a split-K batched GEMM for skinny S×M block products, separate NCCL/Gloo channels, and a reusable Gram matrix across rejected steps.
The dual form trades memory for recomputation, since each Jacobian shard is re-evaluated N/S times. The paper analyzes this trade-off explicitly.
Methodological Rigor
The analysis is careful. Table I gives memory, FLOP and communication costs, and the measured values match the predictions closely (for example, 3SM + 3M elements for CNN-1B, within 0.05 GB). The numerical exactness study compares against an FP64 reference. It shows the dual form is far more stable than the standard form in FP32 at small λ, with the standard form's error reaching 21–280% against 1.3×10⁻² for the dual form. This is a genuinely interesting side finding. The ablations on S, L, and split-K are clean and well explained. The split-K microbenchmark and the full-step 2.39× gain are convincing.
The evaluation has serious limits:
Potential Impact
The practical significance is bounded by an unresolved question: do exact LMA/Gauss-Newton steps at batch sizes of a few hundred to a few thousand beat first-order or approximate second-order methods on real tasks? Per-step cost is O(N²M) plus re-evaluation overhead. For CNN-1B, a step takes 39.8 s on 8 GPUs with N=256, which is extremely expensive compared with an Adam step. Without evidence of dramatically fewer steps or better final quality, adoption in mainstream deep learning is unlikely.
More plausible impact lies in scientific machine learning, PINNs, regression and curve fitting, and small-data and high-precision regimes. There LMA is already favored, and a scalable exact solver could enable larger models. The dual-form observation, memory scaling with batch size rather than parameter count, is reusable by other Gauss-Newton and natural-gradient work. The numerical-stability finding may also be of interest.
Timeliness
Interest in second-order and curvature-aware optimizers (Shampoo, Muon, SOAP) is high, and a scalable exact Gauss-Newton primitive is topical. However, the field's focus is on LLM-scale stochastic training, where this approach (small batch, per-sample Jacobians, no BatchNorm, MSE-style residuals from CE via √ℓ) is not obviously competitive. The residual formulation guideline for CE and MAE is a practical but minor contribution.
Strengths
Limitations
Overall, this is a solid, well-executed systems contribution with limited demonstrated downstream value. Its influence will depend on follow-up work showing training benefits.
Generated Oct 19, 2026
Comparison History (0)
Rigorous, well-analyzed systems engineering that makes exact LMA feasible at 1B parameters, but with no convergence or accuracy evidence on realistic tasks and weak baselines, so practical value remains unproven.