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.
Figures & tables
Component
Standard Form
Dual Form
Memory (elements)
Gram matrix
3M2
3N2
Jacobian shards
2SM
3SM
Slice activations
LA
LA
Vectors
O(M)
O(M+N)
Computation (FLOPs)
TABLE I: Per-Device Cost of Algorithms 3 and 4
Fig. 1: Peak memory of one trial step on one GPU with N=256 and S=L=4 . (a) Crosses mark the first model that runs out of memory. (b) Breakdown of each variant at its largest model.
Model
N
λ
Standard form
Dual form
LeNet-5
1,024
10−3
2.8
8.8×10−6
1
3.7×10−3
5.9×10−6
103
4.3×10−6
3.6×10−7
8,192
10−3
2.0
2.4×10−4
1
2.3×10−3
1.6×10−5
103
4.4×10−6
3.2×10−7
TABLE II: Relative Error of the Update Against an FP64 Reference
Fig. 2: Standard and dual forms for LeNet-5 with N=8,192<M=61,706 . Dashed lines show ideal scaling.
Fig. 3: Strong scaling of LeMA with a fixed global batch. (a) Step time, labeled with the speedup and efficiency on 8 GPUs. (b) Phase breakdown of CNN-1B.
Fig. 4: Effect of the shard size S and the slice size L for CNN-100M with N=256 on one GPU.
Fig. 5: Gram block products with torch.mm and with split-K. (a) One product of two S×M shards on one GPU. The shaded range is where LeMA applies split-K. (b) Step time of CNN-1B on 8 GPUs.
The Levenberg-Marquardt (LM) algorithm is a well-known second-order method for rapid convergence and strong robustness when training small- to medium-sized neural networks (NNs). However, its computational and memory costs increase significantly as the number of parameters in an NN grows. To address this limitation, subspace methods have been proposed, such as the Krylov subspace LM (KSLM) and the hybrid subspace LM (HSLM), making second-order algorithms more efficient. In this work, we evaluate the subspace Levenberg-Marquardt algorithms for regression and classification tasks in neural networks. We compare the performance of subspace LM variants with the classical LM method, as well as other popular first-order algorithms, such as stochastic gradient descent (SGD) and Adam.
M. Duc Hoang
Department of Mathematics, University of California, Davis, Davis, CA, USA
Efficient and stable training of large language models (LLMs) remains a core challenge in modern machine learning systems. To address this challenge, Reparameterized Orthogonal Equivalence Training (POET), a spectrum-preserving framework that optimizes each weight matrix through orthogonal equivalence transformation, has been proposed. Although POET provides strong training stability, its original implementation incurs high memory consumption and computational overhead due to intensive matrix multiplications. To overcome these limitations, we introduce POET-X, a scalable and memory-efficient variant that performs orthogonal equivalence transformations with significantly reduced computational cost. POET-X maintains the generalization and stability benefits of POET while achieving substantial improvements in throughput and memory efficiency. In our experiments, POET-X enables the pretraining of billion-parameter LLMs on a single Nvidia H100 GPU, and in contrast, standard optimizers such as AdamW run out of memory under the same settings.
Zeju Qiu, Lixin Liu, Adrian Weller +2
The Chinese University of Hong Kong · MPI for Intelligent Systems · University of Cambridge +1
Matrix-based optimizers have demonstrated immense potential in training Large Language Models (LLMs), however, designing an ideal optimizer remains a formidable challenge. A superior optimizer must satisfy three core desiderata: efficiency, achieving Muon-like preconditioning to accelerate optimization; stability, strictly adhering to the scale-invariance inherent in neural networks; and speed, minimizing computational overhead. While existing methods address these aspects to varying degrees, they often fail to unify them, either incurring prohibitive computational costs like Muon, or allowing radial jitters that compromise stability like RMNP. To bridge this gap, we propose Nora, an optimizer that rigorously satisfies all three requirements. Nora achieves training stability by explicitly stabilizing weight norms and angular velocities through row-wise momentum projection onto the orthogonal complement of the weights. Simultaneously, by leveraging the block-diagonal dominance of the Transformer Hessian, Nora effectively approximates structured preconditioning while maintaining an optimal computational complexity of O(mn). Furthermore, we prove that Nora is a scalable optimizer and establish its corresponding scaling theorems. With a streamlined implementation requiring only two lines of code, our preliminary experiments validate Nora as an efficient and highly promising optimizer for large-scale training.
Jinghui Yuan, Jiaxuan Zou, Shuo Wang +2
1. School of Artificial Intelligence, Optics and Electronics (iOPEN), Northwestern Polytechnical University. · 2. School of Mathematics and Statistics, Xi’an Jiaotong University. · 4. Gaoling School of Artificial Intelligence, Renmin University of China. +1