面向深度学习的可扩展 Levenberg-Marquardt 训练框架

# Scalable Levenberg-Marquardt Training Framework for Deep Learning
Source: [https://arxiv.org/html/2610.22925](https://arxiv.org/html/2610.22925)
Xianghao WangAffiliation:School of Computing Australian National University Canberra, Australia xianghao\.wang@anu\.edu\.auJohn TaylorAffiliation:School of Computing Australian National University Canberra, Australia john\.taylor@anu\.edu\.au

###### Abstract

The Levenberg\-Marquardt \(LMA\) algorithm is a widely used nonlinear least\-squares solver in scientific computing, known for its rapid convergence\. Recent studies show that for small\-scale neural networks, LMA provides better convergence and architectural robustness than first\-order optimizers like Adam\. However, as a second\-order method, 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 limits LMA in its traditional applications and 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, and an accurate approximation requires 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\. The contributions of LeMA 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 advancement 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 a7\.97×7\.97\\timesspeedup over a single GPU, corresponding to a 99\.7% strong scaling efficiency\.

###### Index Terms:

Levenberg\-Marquardt algorithm, distributed training, deep learning, multi\-GPU

## IIntroduction

The Levenberg\-Marquardt \(LMA\) algorithm\[[1](https://arxiv.org/html/2610.22925#bib.bib1),[2](https://arxiv.org/html/2610.22925#bib.bib2)\]is a standard solver for nonlinear least\-squares problems in scientific computing\[[3](https://arxiv.org/html/2610.22925#bib.bib3),[4](https://arxiv.org/html/2610.22925#bib.bib4)\]\. Each LMA step solves a damped Gauss\-Newton system\. A damping factorλ\\lambdamoves the step between the Gauss\-Newton method and gradient descent, allowing LMA to converge quickly near a minimum while remaining robust far from it\. Since Hagan and Menhaj\[[5](https://arxiv.org/html/2610.22925#bib.bib5)\]applied LMA to neural networks, many studies have used it to train small models containing up to several thousand parameters\[[6](https://arxiv.org/html/2610.22925#bib.bib6),[7](https://arxiv.org/html/2610.22925#bib.bib7),[8](https://arxiv.org/html/2610.22925#bib.bib8),[9](https://arxiv.org/html/2610.22925#bib.bib9),[10](https://arxiv.org/html/2610.22925#bib.bib10),[11](https://arxiv.org/html/2610.22925#bib.bib11),[12](https://arxiv.org/html/2610.22925#bib.bib12),[13](https://arxiv.org/html/2610.22925#bib.bib13)\]\.

Despite its fast convergence, LMA is difficult to scale to large neural networks\. The bottleneck lies in the linear system solved at each step, which involves anM×MM\\times Mcoefficient matrix𝐉T​𝐉\+λ​diag⁡\(𝐉T​𝐉\)\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\operatorname\{diag\}\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)and anN×MN\\times MJacobian matrix𝐉\\mathbf\{J\}, whereNNis the batch size andMMis the number of parameters\. This formulation creates an𝒪⁡\(M​N\+M2\)\\mathcal\{O\}\(MN\+M^\{2\}\)memory footprint\. For a model with only one million parameters, the coefficient matrix alone contains a trillion elements\. The computational cost is equally high\. Evaluating𝐉\\mathbf\{J\}requiresNNbackward propagations, computing the coefficient matrix takes𝒪⁡\(M2​N\)\\mathcal\{O\}\(M^\{2\}N\)time, and solving the linear system takes𝒪⁡\(M3\)\\mathcal\{O\}\(M^\{3\}\)time\. These extreme memory and computation demands prevent the extension of LMA to large\-scale training\.

To address these memory bottlenecks, several studies and libraries\[[14](https://arxiv.org/html/2610.22925#bib.bib14),[15](https://arxiv.org/html/2610.22925#bib.bib15),[16](https://arxiv.org/html/2610.22925#bib.bib16),[17](https://arxiv.org/html/2610.22925#bib.bib17),[18](https://arxiv.org/html/2610.22925#bib.bib18)\]use preconditioned conjugate gradient \(PCG\) solvers\[[19](https://arxiv.org/html/2610.22925#bib.bib19)\]to avoid materializing the large matrices\. PCG computes an approximate solution to the linear system𝐀𝐱=𝐛\\mathbf\{A\}\\mathbf\{x\}=\\mathbf\{b\}using only matrix\-vector products𝐀𝐯\\mathbf\{A\}\\mathbf\{v\}, without explicitly constructing𝐀\\mathbf\{A\}\. For instance, both the optimization toolkitPETSc\[[17](https://arxiv.org/html/2610.22925#bib.bib17)\]and the Python libraryJAXopt\[[14](https://arxiv.org/html/2610.22925#bib.bib14)\]provide LMA implementations with a PCG option\. However, applying PCG to LMA requires numerous network propagations to ensure a robust approximation\. Additionally, PCG relies on a suitable preconditioner matrix, and selecting an effective one is highly task\-specific\[[20](https://arxiv.org/html/2610.22925#bib.bib20)\]\.

As an exact alternative to PCG, Bilski et al\.\[[21](https://arxiv.org/html/2610.22925#bib.bib21)\]decompose LMA into independent tasks where each neuron computes its own Jacobian and weight updates\. However, this approach assumes each neuron produces a single output\. It is therefore incompatible with convolutional layers, which share parameters to generate multiple spatial outputs\.

Our goal is to design an exact LMA framework for general deep learning tasks, without relying on PCG approximations or specific network architectures\. We focus strictly on the systems engineering challenges to make each exact LMA step feasible and fast at scale, as the numerical convergence properties of LMA fall beyond the scope of this work\. Di Marco\[[13](https://arxiv.org/html/2610.22925#bib.bib13),[12](https://arxiv.org/html/2610.22925#bib.bib12)\]utilizes a dual form of the simpler Levenberg method, which appliesλ​𝐈\\lambda\\mathbf\{I\}as the damping term\. This approach successfully reduces the memory footprint from𝒪⁡\(M​N\+M2\)\\mathcal\{O\}\(MN\+M^\{2\}\)to𝒪⁡\(M​N\+N2\)\\mathcal\{O\}\(MN\+N^\{2\}\), making it feasible to train large neural networks whereM≫NM\\gg N\. Inspired by this, we derive and implement an exact dual formulation specifically for LMA\. To further reduce memory consumption, we design a partitioning strategy that avoids storing the full Jacobian𝐉\\mathbf\{J\}\. This design allows us to parallelize LMA and accelerate training across multiple GPUs\.

We presentLeMA, a multi\-GPU Levenberg\-Marquardt training framework for neural networks\. We summarize our main contributions as follows:

- •We derive an exact dual formulation for LMA, which shifts the linear system from the parameter dimension to the batch dimension \(Section[III\-A](https://arxiv.org/html/2610.22925#S3.SS1)\)\.
- •We provide guidelines for designing robust residual functions tailored for LMA \(Section[III\-B](https://arxiv.org/html/2610.22925#S3.SS2)\)\.
- •We design a partitioning and parallelization strategy to further reduce the memory footprint and accelerate computation \(Sections[IV\-B](https://arxiv.org/html/2610.22925#S4.SS2)and[IV\-C](https://arxiv.org/html/2610.22925#S4.SS3)\)\.

In addition to these core contributions, we introduce several supplementary optimizations detailed in Sections[III](https://arxiv.org/html/2610.22925#S3)and[IV](https://arxiv.org/html/2610.22925#S4)\.

We evaluate LeMA on image classification tasks across various neural networks\. LeMA scales LMA to models with up to1,000,022,6321\{,\}000\{,\}022\{,\}632parameters and achieves a99\.7%99\.7\\%strong scaling efficiency on 8 GPUs\. For LeNet\-5, which is small enough for the standard LMA form to fit in GPU memory, LeMA achieves an8\.3×8\.3\\timesspeedup and a33×33\\timesreduction in memory footprint on a single GPU compared to a standard LMA implementation\. When running on 8 GPUs, it delivers a37\.8×37\.8\\timesspeedup and a41×41\\timesmemory reduction over a standard multi\-GPU LMA baseline\.

The rest of this paper is organized as follows\. Section[II](https://arxiv.org/html/2610.22925#S2)reviews LMA and Jacobian computation within deep learning frameworks\. Section[III](https://arxiv.org/html/2610.22925#S3)derives the exact dual formulation and provides residual representations for common loss functions\. Section[IV](https://arxiv.org/html/2610.22925#S4)details the distributed implementation of LeMA and analyzes its associated costs\. Finally, Section[V](https://arxiv.org/html/2610.22925#S5)evaluates the performance of LeMA\.

## IIBackground

This section reviews the Levenberg\-Marquardt algorithm \(Section[II\-A](https://arxiv.org/html/2610.22925#S2.SS1)\) and the computation of Jacobian matrices within deep learning frameworks \(Section[II\-B](https://arxiv.org/html/2610.22925#S2.SS2)\)\.

### II\-AThe Levenberg\-Marquardt Algorithm

The Levenberg\-Marquardt \(LMA\) algorithm solves nonlinear least\-squares problems\[[1](https://arxiv.org/html/2610.22925#bib.bib1),[2](https://arxiv.org/html/2610.22925#bib.bib2),[3](https://arxiv.org/html/2610.22925#bib.bib3),[4](https://arxiv.org/html/2610.22925#bib.bib4)\]\. Consider a modelf⁡\(⋅,𝜽\)f\(\\cdot;\\boldsymbol\{\\theta\}\)with parameters𝜽∈ℝM\\boldsymbol\{\\theta\}\\in\\mathbb\{R\}^\{M\}and a batch ofNNsamples with inputs𝐗∈ℝN×din\\mathbf\{X\}\\in\\mathbb\{R\}^\{N\\times d\_\{\\text\{in\}\}\}and targets𝐘∈ℝN×dout\\mathbf\{Y\}\\in\\mathbb\{R\}^\{N\\times d\_\{\\text\{out\}\}\}, where theii\-th rows correspond to𝐱i\\mathbf\{x\}\_\{i\}and𝐲i\\mathbf\{y\}\_\{i\}\. LMA minimizes the sum of squared residuals‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}, where𝐫∈ℝN\\mathbf\{r\}\\in\\mathbb\{R\}^\{N\}contains one scalar residual per sample\. In classical curve fitting, the model output is a scalar, making the residual simplyri=f⁡\(𝐱i,𝜽\)−yir\_\{i\}=f\(\\mathbf\{x\}\_\{i\};\\boldsymbol\{\\theta\}\)\-y\_\{i\}\. However, neural networks typically produce vector outputs, such as a probability distribution over classes\. Following Hagan and Menhaj\[[5](https://arxiv.org/html/2610.22925#bib.bib5)\], recent works\[[12](https://arxiv.org/html/2610.22925#bib.bib12),[9](https://arxiv.org/html/2610.22925#bib.bib9),[6](https://arxiv.org/html/2610.22925#bib.bib6),[7](https://arxiv.org/html/2610.22925#bib.bib7),[10](https://arxiv.org/html/2610.22925#bib.bib10),[11](https://arxiv.org/html/2610.22925#bib.bib11)\]therefore define each residual asri=r⁡\(f⁡\(𝐱i,𝜽\),𝐲i\)r\_\{i\}=r\(f\(\\mathbf\{x\}\_\{i\};\\boldsymbol\{\\theta\}\);\\mathbf\{y\}\_\{i\}\), where the residual functionrrmaps thedoutd\_\{\\text\{out\}\}\-dimensional output of a sample and its target to a scalar\. Section[III\-B](https://arxiv.org/html/2610.22925#S3.SS2)derivesrrfor common loss functions\.

Each LMA step computes the parameter updateΔ​𝜽\\Delta\\boldsymbol\{\\theta\}by solving

\(𝐉T​𝐉\+λ​𝐃\)​Δ​𝜽=−𝐉T​𝐫,𝐃=diag⁡\(𝐉T​𝐉\)\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)\\Delta\\boldsymbol\{\\theta\}=\-\\mathbf\{J\}^\{T\}\\mathbf\{r\},\\quad\\mathbf\{D\}=\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)\(1\)where𝐉=∂𝐫/∂𝜽∈ℝN×M\\mathbf\{J\}=\\partial\\mathbf\{r\}/\\partial\\boldsymbol\{\\theta\}\\in\\mathbb\{R\}^\{N\\times M\}is the Jacobian with itsii\-th row defined as\(∇𝜽ri\)T\(\\nabla\_\{\\boldsymbol\{\\theta\}\}r\_\{i\}\)^\{T\}, andλ\>0\\lambda\>0is thedamping factor\. Decreasingλ\\lambdamoves LMA towards the Gauss\-Newton method for faster convergence, whereas increasingλ\\lambdamoves it towards gradient descent for robustness\.

Implementations of LMA differ in their update criteria, damping strategies, and iteration logic\. We adopt the logic oftorch\-levenberg\-marquardt\[[12](https://arxiv.org/html/2610.22925#bib.bib12)\], a Python library for training neural networks using LMA, as detailed in Algorithm[1](https://arxiv.org/html/2610.22925#alg1)\. A step is accepted only if it reduces the sum of squared residuals, specifically when‖𝐫′‖22<‖𝐫‖22\\left\\\|\\mathbf\{r\}^\{\\prime\}\\right\\\|\_\{2\}^\{2\}<\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}\. In this case,λ\\lambdais decreased; otherwise,λ\\lambdais increased and the step is recomputed\.

Algorithm 1The Levenberg\-Marquardt Algorithm0:Model

f⁡\(⋅,𝜽\)f\(\\cdot;\\boldsymbol\{\\theta\}\), initial parameters

𝜽0\\boldsymbol\{\\theta\}\_\{0\}, data

\(𝐗,𝐘\)\(\\mathbf\{X\},\\mathbf\{Y\}\), initial damping factor

λ0\\lambda\_\{0\}, adjustment factors

λ↑\>1\\lambda\_\{\\uparrow\}\>1,

0<λ↓<10<\\lambda\_\{\\downarrow\}<1, maximum iterations

TT
0:Trained parameters

𝜽\\boldsymbol\{\\theta\}
1:

𝜽←𝜽0\\boldsymbol\{\\theta\}\\leftarrow\\boldsymbol\{\\theta\}\_\{0\},

λ←λ0\\lambda\\leftarrow\\lambda\_\{0\}
2:for

t=1t=1to

TTdo

3:

𝐫←𝐫⁡\(f⁡\(𝐗,𝜽\),𝐘\)\\mathbf\{r\}\\leftarrow\\mathbf\{r\}\(f\(\\mathbf\{X\};\\boldsymbol\{\\theta\}\);\\mathbf\{Y\}\)\{residuals\}

4:

𝐉←∂𝐫/∂𝜽\\mathbf\{J\}\\leftarrow\\partial\\mathbf\{r\}/\\partial\\boldsymbol\{\\theta\}\{Jacobian\}

5:

𝐃←diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}\\leftarrow\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)
6:repeat

7:Solve

\(𝐉T​𝐉\+λ​𝐃\)​Δ​𝜽=−𝐉T​𝐫\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)\\Delta\\boldsymbol\{\\theta\}=\-\\mathbf\{J\}^\{T\}\\mathbf\{r\}for

Δ​𝜽\\Delta\\boldsymbol\{\\theta\}
8:

𝜽′←𝜽\+Δ​𝜽\\boldsymbol\{\\theta\}^\{\\prime\}\\leftarrow\\boldsymbol\{\\theta\}\+\\Delta\\boldsymbol\{\\theta\}
9:

𝐫′←𝐫⁡\(f⁡\(𝐗,𝜽′\),𝐘\)\\mathbf\{r\}^\{\\prime\}\\leftarrow\\mathbf\{r\}\(f\(\\mathbf\{X\};\\boldsymbol\{\\theta\}^\{\\prime\}\);\\mathbf\{Y\}\)
10:if

‖𝐫′‖22<‖𝐫‖22\\\|\\mathbf\{r\}^\{\\prime\}\\\|^\{2\}\_\{2\}<\\\|\\mathbf\{r\}\\\|^\{2\}\_\{2\}then

11:

𝜽←𝜽′\\boldsymbol\{\\theta\}\\leftarrow\\boldsymbol\{\\theta\}^\{\\prime\},

λ←λ⋅λ↓\\lambda\\leftarrow\\lambda\\cdot\\lambda\_\{\\downarrow\}\{accept step\}

12:else

13:

λ←λ⋅λ↑\\lambda\\leftarrow\\lambda\\cdot\\lambda\_\{\\uparrow\}\{reject step\}

14:endif

15:untilstep accepted

16:endfor

17:return

𝜽\\boldsymbol\{\\theta\}

Many studies\[[10](https://arxiv.org/html/2610.22925#bib.bib10),[21](https://arxiv.org/html/2610.22925#bib.bib21),[5](https://arxiv.org/html/2610.22925#bib.bib5)\]and applications\[[12](https://arxiv.org/html/2610.22925#bib.bib12),[6](https://arxiv.org/html/2610.22925#bib.bib6),[9](https://arxiv.org/html/2610.22925#bib.bib9),[7](https://arxiv.org/html/2610.22925#bib.bib7)\]confuse LMA with the earlier Levenberg algorithm\[[1](https://arxiv.org/html/2610.22925#bib.bib1)\], which damps all parameters equally with the identity matrix𝐈M\\mathbf\{I\}\_\{M\}:

\(𝐉T​𝐉\+λ​𝐈M\)​Δ​𝜽=−𝐉T​𝐫\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{I\}\_\{M\}\)\\Delta\\boldsymbol\{\\theta\}=\-\\mathbf\{J\}^\{T\}\\mathbf\{r\}\(2\)Marquardt\[[2](https://arxiv.org/html/2610.22925#bib.bib2)\]replaced𝐈M\\mathbf\{I\}\_\{M\}with𝐃=diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}=\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\), which scales the damping of each parameter according to its curvature\.

### II\-BJacobian Matrix Computation

Deep learning frameworks\[[22](https://arxiv.org/html/2610.22925#bib.bib22),[23](https://arxiv.org/html/2610.22925#bib.bib23),[24](https://arxiv.org/html/2610.22925#bib.bib24)\]provide two primitives to compute products involving the Jacobian without explicitly materializing the full matrix: thevector\-Jacobian product \(VJP\)and theJacobian\-vector product \(JVP\)\. A VJP computes𝐉T​𝐯\\mathbf\{J\}^\{T\}\\mathbf\{v\}for a vector𝐯∈ℝN\\mathbf\{v\}\\in\\mathbb\{R\}^\{N\}; setting𝐯\\mathbf\{v\}to the standard basis vector𝐞i\\mathbf\{e\}\_\{i\}extracts theii\-th row of𝐉\\mathbf\{J\}\. A JVP computes𝐉𝐯\\mathbf\{J\}\\mathbf\{v\}for a vector𝐯∈ℝM\\mathbf\{v\}\\in\\mathbb\{R\}^\{M\}; setting𝐯=𝐞j\\mathbf\{v\}=\\mathbf\{e\}\_\{j\}extracts thejj\-th column\. Equation \([3](https://arxiv.org/html/2610.22925#S2.E3)\) illustrates both operations\. Internally, a VJP relies on reverse\-mode automatic differentiation, requiring a forward pass that stores intermediate activations followed by a backward pass\. A JVP instead relies on forward\-mode automatic differentiation, requiring only a single forward pass without storing activations\.

𝐉=\[∂r1∂θ1⋯∂r1∂θj⋯∂r1∂θM⋱∂ri∂θ1⋯∂ri∂θj⋯∂ri∂θM←VJP⋱∂rN∂θ1⋯∂rN∂θj⋯∂rN∂θMJVP↑\\CodeAfter​\]∈ℝN×M\\mathbf\{J\}=\\begin\{bNiceMatrix\}\\frac\{\\partial r\_\{1\}\}\{\\partial\\theta\_\{1\}\}&\\cdots&\\frac\{\\partial r\_\{1\}\}\{\\partial\\theta\_\{j\}\}&\\cdots&\\frac\{\\partial r\_\{1\}\}\{\\partial\\theta\_\{M\}\}&\\\\ \\vdots&\\ddots&\\vdots&&\\vdots&\\\\ \\frac\{\\partial r\_\{i\}\}\{\\partial\\theta\_\{1\}\}&\\cdots&\\frac\{\\partial r\_\{i\}\}\{\\partial\\theta\_\{j\}\}&\\cdots&\\frac\{\\partial r\_\{i\}\}\{\\partial\\theta\_\{M\}\}&\{\\color\[rgb\]\{0,0,1\}\\leftarrow\\text\{VJP\}\}\\\\ \\vdots&&\\vdots&\\ddots&\\vdots&\\\\ \\frac\{\\partial r\_\{N\}\}\{\\partial\\theta\_\{1\}\}&\\cdots&\\frac\{\\partial r\_\{N\}\}\{\\partial\\theta\_\{j\}\}&\\cdots&\\frac\{\\partial r\_\{N\}\}\{\\partial\\theta\_\{M\}\}&\\\\ &&\{\\color\[rgb\]\{1,0,0\}\\overset\{\\uparrow\}\{\\text\{JVP\}\}\}&&&\\CodeAfter\\hbox to0\.4pt\{\\vbox to0\.4pt\{\\pgfpicture\\makeatletter\\hbox\{\\hskip 0\.2pt\\lower\-0\.2pt\\hbox to0\.0pt\{\\lxSVG@begingroup@\{\_scopebegin=1\} \\lxSVG@begingroup@\{stroke=\#000000\} \\lxSVG@begingroup@\{fill=\#000000\} \\lxSVG@setlinewidth\{\\the\\pgflinewidth\}\\lxSVG@begingroup@\{stroke\-width=0\.4pt\} \\lx@inpgf@ignorespaces\\nullfont\\hbox to0\.0pt\{\\lxSVG@begingroup@\{\_scopebegin=1\} \{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\{\}\\lx@inpgf@ignorespaces\{\}\}\{\}\{\{\}\}\{\} \{\\lx@inpgf@ignorespaces\{\}\\lx@inpgf@ignorespaces\{\}\}\{\}\{\{\}\}\{\}\{\}\{\}\{\}\{\{\}\}\{\}\\lxSVG@begingroup@\{\_scopebegin=1\} \\color\[rgb\]\{0,0,1\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\}\\lxSVG@stroke\\lxSVG@drawpath@unclipped\{M 0 0 M 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 Z M 0 0\}\{fill:none\} \\lx@inpgf@ignorespaces \\lxSVG@closescope \\lxSVG@closescope \{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\}\\hss\}\\lxSVG@discardpath\\lxSVG@closescope \\hss\}\}\\lxSVG@closescope\\endpgfpicture\}\}\\hbox to0\.4pt\{\\vbox to0\.4pt\{\\pgfpicture\\makeatletter\\hbox\{\\hskip 0\.2pt\\lower\-0\.2pt\\hbox to0\.0pt\{\\lxSVG@begingroup@\{\_scopebegin=1\} \\lxSVG@begingroup@\{stroke=\#000000\} \\lxSVG@begingroup@\{fill=\#000000\} \\lxSVG@setlinewidth\{\\the\\pgflinewidth\}\\lxSVG@begingroup@\{stroke\-width=0\.4pt\} \\lx@inpgf@ignorespaces\\nullfont\\hbox to0\.0pt\{\\lxSVG@begingroup@\{\_scopebegin=1\} \{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\{\}\\lx@inpgf@ignorespaces\{\}\}\{\}\{\{\}\}\{\} \{\\lx@inpgf@ignorespaces\{\}\\lx@inpgf@ignorespaces\{\}\}\{\}\{\{\}\}\{\}\{\}\{\}\{\}\{\{\}\}\{\}\\lxSVG@begingroup@\{\_scopebegin=1\} \\color\[rgb\]\{1,0,0\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\{\}\{\}\{\{\}\}\}\{\{\}\{\}\{\{\}\}\}\{\}\{\}\{\}\\lxSVG@stroke\\lxSVG@drawpath@unclipped\{M 0 0 M 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 L 5\.53 0 C 2\.48 0 2\.48 0 5\.53 0 Z M 0 0\}\{fill:none\} \\lx@inpgf@ignorespaces \\lxSVG@closescope \\lxSVG@closescope \{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\}\{\\lx@inpgf@ignorespaces\}\\hss\}\\lxSVG@discardpath\\lxSVG@closescope \\hss\}\}\\lxSVG@closescope\\endpgfpicture\}\}\\end\{bNiceMatrix\}\\in\\mathbb\{R\}^\{N\\times M\}\(3\)
Our framework assumes that each sample in a batch is processed independently, which excludes intra\-batch operations such as Batch Normalization\. Under this assumption,rir\_\{i\}depends only on theii\-th sample\. Therefore, one single\-sample VJP yields one row of𝐉\\mathbf\{J\}, and assembling the full Jacobian requiresNNsingle\-sample VJPs\. In contrast, a column involves allNNsamples, meaning its extraction requiresNNsingle\-sample JVPs, or equivalently, oneNN\-sample JVP\. Assembling the full Jacobian column\-wise thus requiresM​NMNsingle\-sample JVPs\. Consequently, we always construct the Jacobian row\-wise using VJPs\.

## IIIAlgorithm Design

This section introduces the core algorithm of LeMA\. Specifically, Section[III\-A](https://arxiv.org/html/2610.22925#S3.SS1)derives an exact dual formulation of LMA to bypass theM×MM\\times Mmatrix𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}\. Section[III\-B](https://arxiv.org/html/2610.22925#S3.SS2)formulates residuals for common loss functions, and Section[III\-C](https://arxiv.org/html/2610.22925#S3.SS3)outlines the complete LeMA algorithm\.

### III\-ADual Transformation

When training large\-scale models whereM\>NM\>N, the denseM×MM\\times MGram matrix𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}in \([1](https://arxiv.org/html/2610.22925#S2.E1)\) dominates the memory footprint\. We therefore transform the LMA equation into its dual form, which only involves theN×NN\\times Nmatrix𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}and thus moves the linear system from the parameter dimension to the batch dimension:

\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)​𝐯=−𝐫,Δ​𝜽=𝐃−1​𝐉T​𝐯\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)\\mathbf\{v\}=\-\\mathbf\{r\},\\quad\\Delta\\boldsymbol\{\\theta\}=\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\\mathbf\{v\}\(4\)where𝐃=diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}=\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)as in \([1](https://arxiv.org/html/2610.22925#S2.E1)\)\. We assume that𝐃\\mathbf\{D\}is invertible\. Section[IV\-B1](https://arxiv.org/html/2610.22925#S4.SS2.SSS1)describes how we ensure this in practice\.

We now prove that \([4](https://arxiv.org/html/2610.22925#S3.E4)\) yields the same update as \([1](https://arxiv.org/html/2610.22925#S2.E1)\)\. Sinceλ\>0\\lambda\>0and𝐃\\mathbf\{D\}has a positive diagonal, both𝐉T​𝐉\+λ​𝐃\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}and𝐉𝐃−1​𝐉T\+λ​𝐈N\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}are positive definite and thus invertible\. Expanding the products shows that

\(𝐉T​𝐉\+λ​𝐃\)​𝐃−1​𝐉T\\displaystyle\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}=𝐉T​𝐉𝐃−1​𝐉T\+λ​𝐉T\\displaystyle=\\mathbf\{J\}^\{T\}\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{J\}^\{T\}=𝐉T​\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)\\displaystyle=\\mathbf\{J\}^\{T\}\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)Multiplying both sides by\(𝐉T​𝐉\+λ​𝐃\)−1\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)^\{\-1\}on the left and by−\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)−1​𝐫\-\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)^\{\-1\}\\mathbf\{r\}on the right turns the update of \([1](https://arxiv.org/html/2610.22925#S2.E1)\) into that of \([4](https://arxiv.org/html/2610.22925#S3.E4)\):

Δ​𝜽\\displaystyle\\Delta\\boldsymbol\{\\theta\}=−\(𝐉T​𝐉\+λ​𝐃\)−1​𝐉T​𝐫\\displaystyle=\-\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)^\{\-1\}\\mathbf\{J\}^\{T\}\\mathbf\{r\}=−𝐃−1​𝐉T​\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)−1​𝐫=𝐃−1​𝐉T​𝐯\\displaystyle=\-\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)^\{\-1\}\\mathbf\{r\}=\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\\mathbf\{v\}
Notably, the dual formulation in recent implementations\[[12](https://arxiv.org/html/2610.22925#bib.bib12),[13](https://arxiv.org/html/2610.22925#bib.bib13)\]corresponds to the Levenberg algorithm \([2](https://arxiv.org/html/2610.22925#S2.E2)\), i\.e\.,𝐃=𝐈M\\mathbf\{D\}=\\mathbf\{I\}\_\{M\}in \([4](https://arxiv.org/html/2610.22925#S3.E4)\), and thus lacks Marquardt’s defining modification𝐃=diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}=\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)\[[2](https://arxiv.org/html/2610.22925#bib.bib2)\]\. In contrast, we derive and prove an exact dual formulation of the complete LMA algorithm\.

### III\-BResidual Formulation

Deep learning models minimize a standard lossℒ=1N​∑i=1Nℓi\\mathcal\{L\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\ell\_\{i\}, whereas LMA minimizes the sum of squared residuals‖𝐫‖22=∑i=1Nri2\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}=\\sum\_\{i=1\}^\{N\}r\_\{i\}^\{2\}\[[5](https://arxiv.org/html/2610.22925#bib.bib5)\]\. As a practical heuristic to guide the migration of existing code to the LMA framework, Equation \([5](https://arxiv.org/html/2610.22925#S3.E5)\) formulates the exact residual representations for three common loss functions\. Here,y^i\\hat\{y\}\_\{i\}denotes the network output for theii\-th sample, whiley^i,c\\hat\{y\}\_\{i,c\}andyi,cy\_\{i,c\}correspond to thecc\-th ofCCclasses:

LossℓiriMSE\(y^i−yi\)2y^i−yiCE−∑c=1Cyi,clogy^i,cℓi\+ϵMAE\|y^i−yi\|sign⁡\(y^i−yi\)​ℓi\+ϵ\\begin\{array\}\[\]\{l@\{\\quad\}l@\{\\quad\}l\}\\text\{Loss\}&\\ell\_\{i\}&r\_\{i\}\\\\ \\text\{MSE\}&\(\\hat\{y\}\_\{i\}\-y\_\{i\}\)^\{2\}&\\hat\{y\}\_\{i\}\-y\_\{i\}\\\\ \\text\{CE\}&\-\\sum\_\{c=1\}^\{C\}y\_\{i,c\}\\log\\hat\{y\}\_\{i,c\}&\\sqrt\{\\ell\_\{i\}\+\\epsilon\}\\\\ \\text\{MAE\}&\\left\|\\hat\{y\}\_\{i\}\-y\_\{i\}\\right\|&\\sign\(\\hat\{y\}\_\{i\}\-y\_\{i\}\)\\sqrt\{\\ell\_\{i\}\+\\epsilon\}\\end\{array\}\(5\)
The Mean Squared Error \(MSE\) directly yields a residual\. For the Cross\-Entropy \(CE\) loss, the natural residual would beℓi\\sqrt\{\\ell\_\{i\}\}\. However, the derivative ofx\\sqrt\{x\}is12​x\\frac\{1\}\{2\\sqrt\{x\}\}, so a perfectly predicted sample \(ℓi=0\\ell\_\{i\}=0\) causes a division by zero in the Jacobian computation\. A small constantϵ\>0\\epsilon\>0inside the square root prevents this singularity\.

The Mean Absolute Error \(MAE\) requiresϵ\>0\\epsilon\>0for the same reason, and additionally the sign of the errorei=y^i−yie\_\{i\}=\\hat\{y\}\_\{i\}\-y\_\{i\}\. Without the sign, the derivative of\|ei\|\+ϵ\\sqrt\{\\left\|e\_\{i\}\\right\|\+\\epsilon\}with respect toeie\_\{i\}issign⁡\(ei\)2​\|ei\|\+ϵ\\frac\{\\sign\(e\_\{i\}\)\}\{2\\sqrt\{\\left\|e\_\{i\}\\right\|\+\\epsilon\}\}, which jumps between−12​ϵ\-\\frac\{1\}\{2\\sqrt\{\\epsilon\}\}and12​ϵ\\frac\{1\}\{2\\sqrt\{\\epsilon\}\}aseie\_\{i\}crosses zero\. These values are large for a smallϵ\\epsilon, so the Jacobian becomes discontinuous, and the optimizer oscillates around the minimum\. Multiplying the residual bysign⁡\(ei\)\\sign\(e\_\{i\}\)cancels the sign in the derivative, which becomes12​\|ei\|\+ϵ\\frac\{1\}\{2\\sqrt\{\\left\|e\_\{i\}\\right\|\+\\epsilon\}\}, i\.e\., strictly positive and continuous atei=0e\_\{i\}=0\.

In summary, matching a loss to the LMA objective often introduces a square root into the residual and hence into the denominator of its derivative\. Residual functions must therefore avoid near\-zero denominators and sign oscillations as the error approaches zero\.

### III\-CThe LeMA Algorithm

Algorithm[2](https://arxiv.org/html/2610.22925#alg2)presents the complete LeMA procedure\. We collectively refer to𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}in the standard form \([1](https://arxiv.org/html/2610.22925#S2.E1)\) and𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}in the dual form \([4](https://arxiv.org/html/2610.22925#S3.E4)\) as the Gram matrix𝐆\\mathbf\{G\}\. LeMA selects the dual form whenM\>NM\>Nand the standard form otherwise, so𝐆\\mathbf\{G\}hasmin⁡\(M,N\)2\\min\(M,N\)^\{2\}elements\. The memory footprint thus drops from𝒪⁡\(M​N\+M2\)\\mathcal\{O\}\(MN\+M^\{2\}\)in Algorithm[1](https://arxiv.org/html/2610.22925#alg1)to𝒪⁡\(M​N\+min⁡\(M,N\)2\)\\mathcal\{O\}\(MN\+\\min\(M,N\)^\{2\}\), and Section[IV\-B](https://arxiv.org/html/2610.22925#S4.SS2)further eliminates the𝒪⁡\(M​N\)\\mathcal\{O\}\(MN\)term of the Jacobian\. In addition, LeMA forms𝐆\\mathbf\{G\}once per iteration and reuses it across rejected steps, which only changeλ\\lambda, and carries the residuals of an accepted step over to the next iteration\.

Algorithm 2The Levenberg\-Marquardt Algorithm of LeMA0:Model

f⁡\(⋅,𝜽\)f\(\\cdot;\\boldsymbol\{\\theta\}\)with

MMparameters, per\-sample residual function

rr, initial parameters

𝜽0\\boldsymbol\{\\theta\}\_\{0\}, data

\(𝐗,𝐘\)\(\\mathbf\{X\},\\mathbf\{Y\}\)with

NNsamples, initial damping factor

λ0\\lambda\_\{0\}, adjustment factors

λ↑\>1\\lambda\_\{\\uparrow\}\>1,

0<λ↓<10<\\lambda\_\{\\downarrow\}<1, maximum iterations

TT
0:Trained parameters

𝜽\\boldsymbol\{\\theta\}, loss value

ℒ\\mathcal\{L\}
1:

𝜽←𝜽0\\boldsymbol\{\\theta\}\\leftarrow\\boldsymbol\{\\theta\}\_\{0\},

λ←λ0\\lambda\\leftarrow\\lambda\_\{0\}
2:

𝐫←𝐫⁡\(f⁡\(𝐗,𝜽\),𝐘\)\\mathbf\{r\}\\leftarrow\\mathbf\{r\}\(f\(\\mathbf\{X\};\\boldsymbol\{\\theta\}\);\\mathbf\{Y\}\)\{residuals, Sec\.[III\-B](https://arxiv.org/html/2610.22925#S3.SS2)\}

3:for

t=1t=1to

TTdo

4:

𝐉←∂𝐫/∂𝜽\\mathbf\{J\}\\leftarrow\\partial\\mathbf\{r\}/\\partial\\boldsymbol\{\\theta\}\{Jacobian\}

5:

𝐃←diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}\\leftarrow\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)
6:if

M\>NM\>Nthen

7:

𝐆←𝐉𝐃−1​𝐉T\\mathbf\{G\}\\leftarrow\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\{dual form,

N×NN\\times N\}

8:else

9:

𝐆←𝐉T​𝐉\\mathbf\{G\}\\leftarrow\\mathbf\{J\}^\{T\}\\mathbf\{J\},

𝐠←𝐉T​𝐫\\mathbf\{g\}\\leftarrow\\mathbf\{J\}^\{T\}\\mathbf\{r\}\{standard form,

M×MM\\times M\}

10:endif

11:repeat

12:if

M\>NM\>Nthen

13:Solve

\(𝐆\+λ​𝐈N\)​𝐯=−𝐫\(\\mathbf\{G\}\+\\lambda\\mathbf\{I\}\_\{N\}\)\\mathbf\{v\}=\-\\mathbf\{r\}for

𝐯\\mathbf\{v\}
14:

Δ​𝜽←𝐃−1​𝐉T​𝐯\\Delta\\boldsymbol\{\\theta\}\\leftarrow\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\\mathbf\{v\}
15:else

16:Solve

\(𝐆\+λ​𝐃\)​Δ​𝜽=−𝐠\(\\mathbf\{G\}\+\\lambda\\mathbf\{D\}\)\\Delta\\boldsymbol\{\\theta\}=\-\\mathbf\{g\}for

Δ​𝜽\\Delta\\boldsymbol\{\\theta\}
17:endif

18:

𝜽′←𝜽\+Δ​𝜽\\boldsymbol\{\\theta\}^\{\\prime\}\\leftarrow\\boldsymbol\{\\theta\}\+\\Delta\\boldsymbol\{\\theta\}
19:

𝐫′←𝐫⁡\(f⁡\(𝐗,𝜽′\),𝐘\)\\mathbf\{r\}^\{\\prime\}\\leftarrow\\mathbf\{r\}\(f\(\\mathbf\{X\};\\boldsymbol\{\\theta\}^\{\\prime\}\);\\mathbf\{Y\}\)
20:if

‖𝐫′‖2<‖𝐫‖2\\\|\\mathbf\{r\}^\{\\prime\}\\\|^\{2\}<\\\|\\mathbf\{r\}\\\|^\{2\}then

21:

𝜽←𝜽′\\boldsymbol\{\\theta\}\\leftarrow\\boldsymbol\{\\theta\}^\{\\prime\},

𝐫←𝐫′\\mathbf\{r\}\\leftarrow\\mathbf\{r\}^\{\\prime\},

λ←λ⋅λ↓\\lambda\\leftarrow\\lambda\\cdot\\lambda\_\{\\downarrow\}\{accept step\}

22:else

23:

λ←λ⋅λ↑\\lambda\\leftarrow\\lambda\\cdot\\lambda\_\{\\uparrow\}\{reject step\}

24:endif

25:untilstep accepted

26:endfor

27:

ℒ←1N​‖𝐫‖22\\mathcal\{L\}\\leftarrow\\frac\{1\}\{N\}\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}
28:return

𝜽\\boldsymbol\{\\theta\},

ℒ\\mathcal\{L\}

## IVLeMA Implementation

This section describes how LeMA executes Algorithm[2](https://arxiv.org/html/2610.22925#alg2)onPPGPUs\. Section[IV\-A](https://arxiv.org/html/2610.22925#S4.SS1)covers the distributed setup, Section[IV\-B](https://arxiv.org/html/2610.22925#S4.SS2)the partitioned computation of both forms, Section[IV\-C](https://arxiv.org/html/2610.22925#S4.SS3)the evaluation of Jacobian shards, and Section[IV\-D](https://arxiv.org/html/2610.22925#S4.SS4)the resulting cost\.

### IV\-ADistributed Setup

LeMA runsPPprocesses in the Single Program, Multiple Data \(SPMD\) model\[[25](https://arxiv.org/html/2610.22925#bib.bib25)\], one per NVIDIA GPU, and communicates throughtorch\.distributed\[[22](https://arxiv.org/html/2610.22925#bib.bib22)\]\. Each process belongs to two communication groups\. Large tensors move between GPUs over NCCL\[[26](https://arxiv.org/html/2610.22925#bib.bib26)\], NVIDIA’s collective communication library, and small scalars such as local loss values move between hosts over Meta’s Gloo library\[[27](https://arxiv.org/html/2610.22925#bib.bib27)\]\. Keeping the scalars on the host avoids many small host\-to\-device and device\-to\-host copies that would otherwise interrupt the GPU streams\.

All parameters are stored in one flat tensor, and the weights of each layer are mapped to their positions in it\. A VJP therefore returns a singleMM\-dimensional gradient, and no per\-layer gradients need to be concatenated\.

Each process loads its own mini\-batch ofN/PN/Psamples instead of receiving part of a scattered global batch, and thePPmini\-batches together form the global batch\. PyTorch’sDistributedSampler\[[22](https://arxiv.org/html/2610.22925#bib.bib22)\]keeps the mini\-batches disjoint\. The dual form also needs the global batch on every process \(Section[IV\-B2](https://arxiv.org/html/2610.22925#S4.SS2.SSS2)\)\.

### IV\-BPartition and Parallelization

TheN×MN\\times MJacobian dominates the memory of Algorithm[2](https://arxiv.org/html/2610.22925#alg2), so we never store it in full\. We split𝐉\\mathbf\{J\}by rows intoN/SN/Sshards𝐉=\[𝐉1T,…,𝐉N/ST\]T\\mathbf\{J\}=\[\\mathbf\{J\}\_\{1\}^\{T\},\\dots,\\mathbf\{J\}\_\{N/S\}^\{T\}\]^\{T\}, where𝐉k∈ℝS×M\\mathbf\{J\}\_\{k\}\\in\\mathbb\{R\}^\{S\\times M\}and𝐫k∈ℝS\\mathbf\{r\}\_\{k\}\\in\\mathbb\{R\}^\{S\}correspond to thekk\-th group ofSSsamples\. Each shard is computed on the fly, added to the Gram matrix𝐆\\mathbf\{G\}, and freed, so the Jacobian occupies𝒪⁡\(S​M\)\\mathcal\{O\}\(SM\)memory instead of𝒪⁡\(N​M\)\\mathcal\{O\}\(NM\)\. We split by rows because a row shard needs onlySSsingle\-sample VJPs\. The subroutineJacevaluates a shard and returns𝐉k\\mathbf\{J\}\_\{k\}and𝐫k\\mathbf\{r\}\_\{k\}, and we drop𝐫k\\mathbf\{r\}\_\{k\}from the notation when it is not used\. Internally,Jacworks onslicesofLLsamples \(Section[IV\-C](https://arxiv.org/html/2610.22925#S4.SS3)\)\. On processpp,\(𝐗p,i,𝐘p,i\)\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)and𝐉p,i\\mathbf\{J\}\_\{p,i\}denote theii\-th of theN/\(P​S\)N/\(PS\)shards of its mini\-batch\(𝐗p,𝐘p\)\(\\mathbf\{X\}\_\{p\},\\mathbf\{Y\}\_\{p\}\)\. We assume thatP​SPSdividesNNandLLdividesSS\.

#### IV\-B1Standard Form

Under the row partition, both terms of \([1](https://arxiv.org/html/2610.22925#S2.E1)\) become sums over shards:

𝐉T​𝐉=∑k=1N/S𝐉kT​𝐉k,𝐉T​𝐫=∑k=1N/S𝐉kT​𝐫k\\mathbf\{J\}^\{T\}\\mathbf\{J\}=\\sum\_\{k=1\}^\{N/S\}\\mathbf\{J\}\_\{k\}^\{T\}\\mathbf\{J\}\_\{k\},\\qquad\\mathbf\{J\}^\{T\}\\mathbf\{r\}=\\sum\_\{k=1\}^\{N/S\}\\mathbf\{J\}\_\{k\}^\{T\}\\mathbf\{r\}\_\{k\}\(6\)The summands are independent, so each process sums the shards of its own mini\-batch locally\. An NCCL all\-reduce then gives every process𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}and𝐉T​𝐫\\mathbf\{J\}^\{T\}\\mathbf\{r\}, and a Gloo all\-reduce gives‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}\. Parameters without curvature, such as those of inactive ReLU units, have a zero diagonal entry and would make𝐃\\mathbf\{D\}singular\. We therefore clamp each\(𝐉T​𝐉\)j​j\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)\_\{jj\}in𝐃\\mathbf\{D\}to at least1%1\\%of the mean diagonal entry, and the dual form uses the same clamp\. Every process then solves theM×MM\\times Msystem on its own\. Algorithm[3](https://arxiv.org/html/2610.22925#alg3)lists the steps on processpp\. If LeMA rejects a step, only the final solve is redone with the newλ\\lambda\.

Algorithm 3Distributed Standard Form Solve on Processpp0:Local mini\-batch

\(𝐗p,𝐘p\)\(\\mathbf\{X\}\_\{p\},\\mathbf\{Y\}\_\{p\}\)with shards

\(𝐗p,i,𝐘p,i\)\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)of

SSsamples, parameters

𝜽\\boldsymbol\{\\theta\}, damping factor

λ\\lambda
0:Parameter update

Δ​𝜽\\Delta\\boldsymbol\{\\theta\}, loss

‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}
1:

𝐆p←𝟎M×M\\mathbf\{G\}\_\{p\}\\leftarrow\\mathbf\{0\}^\{M\\times M\},

𝐠p←𝟎M\\mathbf\{g\}\_\{p\}\\leftarrow\\mathbf\{0\}^\{M\},

ρp←0\\rho\_\{p\}\\leftarrow 0
2:for

i=1i=1to

N/\(P​S\)N/\(PS\)do

3:

𝐉p,i,𝐫p,i←Jac​\(𝐗p,i,𝐘p,i\)\\mathbf\{J\}\_\{p,i\},\\mathbf\{r\}\_\{p,i\}\\leftarrow\\textsc\{Jac\}\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)\{

𝐉p,i∈ℝS×M\\mathbf\{J\}\_\{p,i\}\\in\\mathbb\{R\}^\{S\\times M\}\}

4:

𝐆p←𝐆p\+𝐉p,iT​𝐉p,i\\mathbf\{G\}\_\{p\}\\leftarrow\\mathbf\{G\}\_\{p\}\+\\mathbf\{J\}\_\{p,i\}^\{T\}\\mathbf\{J\}\_\{p,i\}
5:

𝐠p←𝐠p\+𝐉p,iT​𝐫p,i\\mathbf\{g\}\_\{p\}\\leftarrow\\mathbf\{g\}\_\{p\}\+\\mathbf\{J\}\_\{p,i\}^\{T\}\\mathbf\{r\}\_\{p,i\}
6:

ρp←ρp\+𝐫p,iT​𝐫p,i\\rho\_\{p\}\\leftarrow\\rho\_\{p\}\+\\mathbf\{r\}\_\{p,i\}^\{T\}\\mathbf\{r\}\_\{p,i\}
7:endfor

8:

𝐉T​𝐉←AllReduce​\(𝐆p\)\\mathbf\{J\}^\{T\}\\mathbf\{J\}\\leftarrow\\textsc\{AllReduce\}\(\\mathbf\{G\}\_\{p\}\),

𝐉T​𝐫←AllReduce​\(𝐠p\)\\mathbf\{J\}^\{T\}\\mathbf\{r\}\\leftarrow\\textsc\{AllReduce\}\(\\mathbf\{g\}\_\{p\}\)\{NCCL\}

9:

‖𝐫‖22←AllReduce​\(ρp\)\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}\\leftarrow\\textsc\{AllReduce\}\(\\rho\_\{p\}\)\{Gloo\}

10:

𝐃←diag⁡\(𝐉T​𝐉\)\\mathbf\{D\}\\leftarrow\\diag\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)
11:Solve

\(𝐉T​𝐉\+λ​𝐃\)​Δ​𝜽=−𝐉T​𝐫\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\+\\lambda\\mathbf\{D\}\)\\Delta\\boldsymbol\{\\theta\}=\-\\mathbf\{J\}^\{T\}\\mathbf\{r\}for

Δ​𝜽\\Delta\\boldsymbol\{\\theta\}
12:return

Δ​𝜽\\Delta\\boldsymbol\{\\theta\},

‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}

#### IV\-B2Dual Form

The dual form \([4](https://arxiv.org/html/2610.22925#S3.E4)\) first solves\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)​𝐯=−𝐫\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)\\mathbf\{v\}=\-\\mathbf\{r\}and then setsΔ​𝜽=𝐃−1​𝐉T​𝐯\\Delta\\boldsymbol\{\\theta\}=\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\\mathbf\{v\}\. To avoid forming anyM×MM\\times Mmatrix, we keep only the diagonal of𝐃−1\\mathbf\{D\}^\{\-1\}as a vector𝐝∈ℝM\\mathbf\{d\}\\in\\mathbb\{R\}^\{M\}withdj=1/\(𝐉T​𝐉\)j​jd\_\{j\}=1/\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)\_\{jj\}\. Multiplying by𝐃−1\\mathbf\{D\}^\{\-1\}on the right then scales columnjjbydjd\_\{j\}, so𝐉~k=𝐉k​𝐃−1=𝐉k⊙𝐝T\\tilde\{\\mathbf\{J\}\}\_\{k\}=\\mathbf\{J\}\_\{k\}\\mathbf\{D\}^\{\-1\}=\\mathbf\{J\}\_\{k\}\\odot\\mathbf\{d\}^\{T\}, where⊙\\odotis element\-wise multiplication and𝐝T\\mathbf\{d\}^\{T\}is broadcast across rows\. Under the row partition, the Gram matrix splits intoS×SS\\times Sblocks:

𝐉𝐃−1​𝐉T=\[𝐉~1​𝐉1T⋯𝐉~1​𝐉N/ST⋱𝐉~N/S​𝐉1T⋯𝐉~N/S​𝐉N/ST\]\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}=\\begin\{bmatrix\}\\tilde\{\\mathbf\{J\}\}\_\{1\}\\mathbf\{J\}\_\{1\}^\{T\}&\\cdots&\\tilde\{\\mathbf\{J\}\}\_\{1\}\\mathbf\{J\}\_\{N/S\}^\{T\}\\\\ \\vdots&\\ddots&\\vdots\\\\ \\tilde\{\\mathbf\{J\}\}\_\{N/S\}\\mathbf\{J\}\_\{1\}^\{T\}&\\cdots&\\tilde\{\\mathbf\{J\}\}\_\{N/S\}\\mathbf\{J\}\_\{N/S\}^\{T\}\\end\{bmatrix\}\(7\)
Algorithm[4](https://arxiv.org/html/2610.22925#alg4)computes the dual form on processppin three phases\. The first phase computes𝐝\\mathbf\{d\}\. Because\(𝐉T​𝐉\)j​j\(\\mathbf\{J\}^\{T\}\\mathbf\{J\}\)\_\{jj\}is the squared norm of columnjjof𝐉\\mathbf\{J\}, the vector𝐝\\mathbf\{d\}can be accumulated shard by shard and all\-reduced in the same way as𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}in Algorithm[3](https://arxiv.org/html/2610.22925#alg3)\. The diagonal entries are clamped as in the standard form before taking reciprocals\. In the second phase, processppcomputes theN/PN/Prows of \([7](https://arxiv.org/html/2610.22925#S4.E7)\) that belong to its mini\-batch\. For each local shard𝐉~p,i\\tilde\{\\mathbf\{J\}\}\_\{p,i\}, it evaluates every global shard𝐉j\\mathbf\{J\}\_\{j\}and the block𝐉~p,i​𝐉jT\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\mathbf\{J\}\_\{j\}^\{T\}, and the global residuals𝐫j\\mathbf\{r\}\_\{j\}come out of the same evaluation\. An all\-gather then assembles the full Gram matrix on every process\. In the third phase, every process solves for𝐯\\mathbf\{v\}, andVjpcomputes𝐉T​𝐯\\mathbf\{J\}^\{T\}\\mathbf\{v\}slice by slice over the local mini\-batch without forming the Jacobian again\. Here𝐯p,l\\mathbf\{v\}\_\{p,l\}is the part of𝐯\\mathbf\{v\}for thell\-th local slice\. An all\-reduce of the partial sums, followed by scaling with𝐝\\mathbf\{d\}, givesΔ​𝜽\\Delta\\boldsymbol\{\\theta\}\. If LeMA rejects a step, only the solve and the third phase are repeated with the newλ\\lambda\. The standard form evaluates each Jacobian row once, whereas the dual form re\-evaluates every global shard for each local shard and thus spends extra computation to save memory \(Section[IV\-D](https://arxiv.org/html/2610.22925#S4.SS4)\)\.

Algorithm 4Distributed Dual Form Solve on Processpp0:Global batch

\(𝐗,𝐘\)\(\\mathbf\{X\},\\mathbf\{Y\}\)with shards

\(𝐗j,𝐘j\)\(\\mathbf\{X\}\_\{j\},\\mathbf\{Y\}\_\{j\}\)of

SSsamples, local mini\-batch

\(𝐗p,𝐘p\)\(\\mathbf\{X\}\_\{p\},\\mathbf\{Y\}\_\{p\}\)with shards

\(𝐗p,i,𝐘p,i\)\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)of

SSsamples and slices

\(𝐗p,l,𝐘p,l\)\(\\mathbf\{X\}\_\{p,l\},\\mathbf\{Y\}\_\{p,l\}\)of

LLsamples, parameters

𝜽\\boldsymbol\{\\theta\}, damping factor

λ\\lambda
0:Parameter update

Δ​𝜽\\Delta\\boldsymbol\{\\theta\}, loss

‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}
1:

𝐝p←𝟎M\\mathbf\{d\}\_\{p\}\\leftarrow\\mathbf\{0\}^\{M\}
2:for

i=1i=1to

N/\(P​S\)N/\(PS\)do

3:

𝐉p,i←Jac​\(𝐗p,i,𝐘p,i\)\\mathbf\{J\}\_\{p,i\}\\leftarrow\\textsc\{Jac\}\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)
4:

𝐝p←𝐝p\+𝟏ST​\(𝐉p,i⊙𝐉p,i\)\\mathbf\{d\}\_\{p\}\\leftarrow\\mathbf\{d\}\_\{p\}\+\\mathbf\{1\}\_\{S\}^\{T\}\(\\mathbf\{J\}\_\{p,i\}\\odot\\mathbf\{J\}\_\{p,i\}\)
5:endfor

6:

𝐝←𝟏⊘AllReduce​\(𝐝p\)\\mathbf\{d\}\\leftarrow\\mathbf\{1\}\\oslash\\textsc\{AllReduce\}\(\\mathbf\{d\}\_\{p\}\)\{NCCL\}

7:

𝐆p←𝟎\(N/P\)×N\\mathbf\{G\}\_\{p\}\\leftarrow\\mathbf\{0\}^\{\(N/P\)\\times N\}\{local rows of

𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\}

8:for

i=1i=1to

N/\(P​S\)N/\(PS\)do

9:

𝐉~p,i←Jac​\(𝐗p,i,𝐘p,i\)⊙𝐝T\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\leftarrow\\textsc\{Jac\}\(\\mathbf\{X\}\_\{p,i\},\\mathbf\{Y\}\_\{p,i\}\)\\odot\\mathbf\{d\}^\{T\}
10:for

j=1j=1to

N/SN/Sdo

11:

𝐉j,𝐫j←Jac​\(𝐗j,𝐘j\)\\mathbf\{J\}\_\{j\},\\mathbf\{r\}\_\{j\}\\leftarrow\\textsc\{Jac\}\(\\mathbf\{X\}\_\{j\},\\mathbf\{Y\}\_\{j\}\)
12:

𝐆p​\[i,j\]←𝐉~p,i​𝐉jT\\mathbf\{G\}\_\{p\}\[i,j\]\\leftarrow\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\mathbf\{J\}\_\{j\}^\{T\}
13:endfor

14:endfor

15:

𝐉𝐃−1​𝐉T←AllGather​\(𝐆p\)\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\\leftarrow\\textsc\{AllGather\}\(\\mathbf\{G\}\_\{p\}\)\{NCCL\}

16:

𝐫←\[𝐫1;…;𝐫N/S\]\\mathbf\{r\}\\leftarrow\[\\mathbf\{r\}\_\{1\};\\dots;\\mathbf\{r\}\_\{N/S\}\]
17:Solve

\(𝐉𝐃−1​𝐉T\+λ​𝐈N\)​𝐯=−𝐫\(\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\+\\lambda\\mathbf\{I\}\_\{N\}\)\\mathbf\{v\}=\-\\mathbf\{r\}for

𝐯\\mathbf\{v\}
18:

𝐮p←𝟎M\\mathbf\{u\}\_\{p\}\\leftarrow\\mathbf\{0\}^\{M\}
19:for

l=1l=1to

N/\(P​L\)N/\(PL\)do

20:

𝐮p←𝐮p\+Vjp​\(𝐗p,l,𝐘p,l,𝐯p,l\)\\mathbf\{u\}\_\{p\}\\leftarrow\\mathbf\{u\}\_\{p\}\+\\textsc\{Vjp\}\(\\mathbf\{X\}\_\{p,l\},\\mathbf\{Y\}\_\{p,l\},\\mathbf\{v\}\_\{p,l\}\)
21:endfor

22:

Δ​𝜽←𝐝⊙AllReduce​\(𝐮p\)\\Delta\\boldsymbol\{\\theta\}\\leftarrow\\mathbf\{d\}\\odot\\textsc\{AllReduce\}\(\\mathbf\{u\}\_\{p\}\)\{NCCL\}

23:return

Δ​𝜽\\Delta\\boldsymbol\{\\theta\},

‖𝐫‖22\\left\\\|\\mathbf\{r\}\\right\\\|\_\{2\}^\{2\}

Each block product𝐉~p,i​𝐉jT\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\mathbf\{J\}\_\{j\}^\{T\}multiplies twoS×MS\\times Mshards and produces anS×SS\\times Sblock\. A general matrix multiplication \(GEMM\) usually maps one output tile to one thread block\. WhenS≪MS\\ll M, the output has few tiles, so only a few thread blocks run, each looping over the full inner dimensionMM, and most streaming multiprocessors \(SMs\) of the GPU sit idle\. We address this with split\-K\[[28](https://arxiv.org/html/2610.22925#bib.bib28)\]\. LeMA cuts the inner dimension into chunks, multiplies matching chunks of the two shards into partialS×SS\\times Sblocks with one batched GEMM, and sums the partial blocks\. The number of chunks is a multiple of the SM count, so the batched GEMM keeps all SMs busy\. The chunks are strided views of the shards, so no data is copied, and the partial blocks are small\. Shards with more rows already give the GEMM enough parallelism, so LeMA uses split\-K only for shards with few rows\.

### IV\-CJacobian Shard Computation

This section describes the functionsJacandVjpused in Algorithms[3](https://arxiv.org/html/2610.22925#alg3)and[4](https://arxiv.org/html/2610.22925#alg4)\. Samples are independent, so theSSrows of a shard come fromSSsingle\-sample VJPs\. Running them one after another sends a single sample through the network at a time and leaves most of the GPU idle\.Jacinstead vectorizes the VJPs of a slice ofLLsamples withvmap\[[22](https://arxiv.org/html/2610.22925#bib.bib22)\]and obtains anL×ML\\times MJacobian slice per call \(Algorithm[5](https://arxiv.org/html/2610.22925#alg5)\)\. The per\-sample function also returns the residual as an auxiliary output, so the residuals need no extra forward pass\. VectorizingLLVJPs, however, also storesLLtimes as many activations, so the slice sizeLLtrades execution speed against memory\.

We slice the computation of𝐉T​𝐯\\mathbf\{J\}^\{T\}\\mathbf\{v\}in the dual form in the same way\. One VJP over allN/PN/Plocal samples would compute it, but its forward pass would keep the activations of all of them at once\.Vjptherefore handles one slice ofLLsamples per call, and Algorithm[4](https://arxiv.org/html/2610.22925#alg4)sums the results\. UnlikeJac,Vjpdoes not build a Jacobian slice\. A single backward pass on the scalar𝐯lT​𝐫l\\mathbf\{v\}\_\{l\}^\{T\}\\mathbf\{r\}\_\{l\}gives the vector𝐉lT​𝐯l∈ℝM\\mathbf\{J\}\_\{l\}^\{T\}\\mathbf\{v\}\_\{l\}\\in\\mathbb\{R\}^\{M\}directly \(Algorithm[6](https://arxiv.org/html/2610.22925#alg6)\)\.

Algorithm 5Jacobian Shard EvaluationJac​\(𝐗k,𝐘k\)\\textsc\{Jac\}\(\\mathbf\{X\}\_\{k\},\\mathbf\{Y\}\_\{k\}\)0:Shard

\(𝐗k,𝐘k\)\(\\mathbf\{X\}\_\{k\},\\mathbf\{Y\}\_\{k\}\)of

SSsamples with slices

\(𝐗k,l,𝐘k,l\)\(\\mathbf\{X\}\_\{k,l\},\\mathbf\{Y\}\_\{k,l\}\)of

LLsamples, model

f⁡\(⋅,𝜽\)f\(\\cdot;\\boldsymbol\{\\theta\}\), residual function

rr
0:Jacobian shard

𝐉k∈ℝS×M\\mathbf\{J\}\_\{k\}\\in\\mathbb\{R\}^\{S\\times M\}, residuals

𝐫k∈ℝS\\mathbf\{r\}\_\{k\}\\in\\mathbb\{R\}^\{S\}
1:

g⁡\(𝐱,𝐲\)←\(∇𝜽r​\(f⁡\(𝐱,𝜽\),𝐲\),r⁡\(f⁡\(𝐱,𝜽\),𝐲\)\)g\(\\mathbf\{x\},\\mathbf\{y\}\)\\leftarrow\\big\(\\nabla\_\{\\boldsymbol\{\\theta\}\}r\(f\(\\mathbf\{x\};\\boldsymbol\{\\theta\}\);\\mathbf\{y\}\),\\ r\(f\(\\mathbf\{x\};\\boldsymbol\{\\theta\}\);\\mathbf\{y\}\)\\big\)
2:for

l=1l=1to

S/LS/Ldo

3:

𝐉k,l,𝐫k,l←vmap​\(g\)​\(𝐗k,l,𝐘k,l\)\\mathbf\{J\}\_\{k,l\},\\mathbf\{r\}\_\{k,l\}\\leftarrow\\texttt\{vmap\}\(g\)\(\\mathbf\{X\}\_\{k,l\},\\mathbf\{Y\}\_\{k,l\}\)\{

𝐉k,l∈ℝL×M\\mathbf\{J\}\_\{k,l\}\\in\\mathbb\{R\}^\{L\\times M\}\}

4:endfor

5:return

\[𝐉k,1;…;𝐉k,S/L\]\[\\mathbf\{J\}\_\{k,1\};\\dots;\\mathbf\{J\}\_\{k,S/L\}\],

\[𝐫k,1;…;𝐫k,S/L\]\[\\mathbf\{r\}\_\{k,1\};\\dots;\\mathbf\{r\}\_\{k,S/L\}\]

Algorithm 6Vector\-Jacobian ProductVjp​\(𝐗l,𝐘l,𝐯l\)\\textsc\{Vjp\}\(\\mathbf\{X\}\_\{l\},\\mathbf\{Y\}\_\{l\},\\mathbf\{v\}\_\{l\}\)0:Slice

\(𝐗l,𝐘l\)\(\\mathbf\{X\}\_\{l\},\\mathbf\{Y\}\_\{l\}\)of

LLsamples, vector

𝐯l∈ℝL\\mathbf\{v\}\_\{l\}\\in\\mathbb\{R\}^\{L\}, model

f⁡\(⋅,𝜽\)f\(\\cdot;\\boldsymbol\{\\theta\}\), residual function

rr
0:

𝐉lT​𝐯l∈ℝM\\mathbf\{J\}\_\{l\}^\{T\}\\mathbf\{v\}\_\{l\}\\in\\mathbb\{R\}^\{M\}
1:

𝐫l←r⁡\(f⁡\(𝐗l,𝜽\),𝐘l\)\\mathbf\{r\}\_\{l\}\\leftarrow r\(f\(\\mathbf\{X\}\_\{l\};\\boldsymbol\{\\theta\}\);\\mathbf\{Y\}\_\{l\}\)\{forward pass\}

2:return

∇𝜽\(𝐯lT​𝐫l\)\\nabla\_\{\\boldsymbol\{\\theta\}\}\(\\mathbf\{v\}\_\{l\}^\{T\}\\mathbf\{r\}\_\{l\}\)\{backward pass\}

### IV\-DMemory Footprint and Performance Analysis

We now analyze the per\-device cost of the standard form \(Algorithm[3](https://arxiv.org/html/2610.22925#alg3)\) and the dual form \(Algorithm[4](https://arxiv.org/html/2610.22925#alg4)\)\. LetAAbe the number of activation elements stored for one sample in a forward pass, andFFthe FLOPs of one single\-sample VJP, which consists of one forward and one backward pass\. Because samples are independent, a vectorized pass overLLsamples inJacorVjpstoresL​ALAactivation elements and costsL​FLFFLOPs\. We leave out the input mini\-batch, which every training method needs\. Table[I](https://arxiv.org/html/2610.22925#S4.T1)lists the remaining components\. In the asymptotic bounds below, we drop a term only if a retained term strictly dominates it\.

TABLE I:Per\-Device Cost of Algorithms[3](https://arxiv.org/html/2610.22925#alg3)and[4](https://arxiv.org/html/2610.22925#alg4)#### IV\-D1Memory Footprint

The memory of both forms has three main parts\. The first is the Gram matrix, which is𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}withM2M^\{2\}elements in the standard form and the gathered𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}withN2N^\{2\}elements in the dual form\. LeMA keeps the Gram matrix for reuse after a rejected step, so the solver applies damping to a separate copy and factorizes it into LU, which triples this cost\. The second part is the Jacobian shards\.Jaccomputes the per\-sample gradients slice by slice and then assembles them into the rows of the shard, so the gradients and their assembled copy briefly coexist and take2​S​M2SMelements\. The standard form evaluates one shard at a time\. The nested loop of the dual form also holds𝐉~p,i\\tilde\{\\mathbf\{J\}\}\_\{p,i\}while evaluating𝐉j\\mathbf\{J\}\_\{j\}, which raises this part to3​S​M3SMelements\. The third part is the activations of one slice,L​ALAelements, sinceJacandVjpnever run at the same time\. The remaining vectors, such as𝜽\\boldsymbol\{\\theta\},𝐝\\mathbf\{d\}, and𝐫\\mathbf\{r\}, are bounded byS​MSMandN2N^\{2\}\. Algorithm[2](https://arxiv.org/html/2610.22925#alg2)selects the dual form only whenM\>NM\>N, so the memory per device is

𝒪⁡\(min⁡\(M,N\)2\+S​M\+L​A\)\\mathcal\{O\}\\big\(\\min\(M,N\)^\{2\}\+SM\+LA\\big\)\(8\)The first term is quadratic in the smaller ofMMandNN\. The other two grow linearly with the model size and are set by the shard sizeSSand the slice sizeLL, regardless of the batch sizeNN\. Algorithm[1](https://arxiv.org/html/2610.22925#alg1), by contrast, stores the fullN×MN\\times MJacobian\.

#### IV\-D2Computation

In the standard form, each Jacobian row is evaluated exactly once, so each device runsN/PN/Psingle\-sample VJPs\. Accumulating𝐉p,iT​𝐉p,i\\mathbf\{J\}\_\{p,i\}^\{T\}\\mathbf\{J\}\_\{p,i\}costs2​S​M22SM^\{2\}FLOPs per shard, or2​N​M2/P2NM^\{2\}/PFLOPs over theN/\(P​S\)N/\(PS\)local shards\. In the dual form, the inner loop re\-evaluates allN/SN/Sglobal shards for each of theN/\(P​S\)N/\(PS\)local shards\. This amounts toN2/\(P​S\)N^\{2\}/\(PS\)single\-sample VJPs and outweighs the3​N/P3N/PVJPs spent on𝐝\\mathbf\{d\},𝐉~p,i\\tilde\{\\mathbf\{J\}\}\_\{p,i\}, and𝐉T​𝐯\\mathbf\{J\}^\{T\}\\mathbf\{v\}\. TheN2/\(P​S2\)N^\{2\}/\(PS^\{2\}\)block products𝐉~p,i​𝐉jT\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\mathbf\{J\}\_\{j\}^\{T\}cost2​S2​M2S^\{2\}MFLOPs each and2​N2​M/P2N^\{2\}M/PFLOPs in total\. Solving theM×MM\\times MorN×NN\\times Nsystem by LU factorization costs𝒪⁡\(M3\)\\mathcal\{O\}\(M^\{3\}\)or𝒪⁡\(N3\)\\mathcal\{O\}\(N^\{3\}\)FLOPs\. This term does not shrink withPP, since every device solves the same system\. The remaining matrix\-vector and element\-wise operations, such as𝐉p,iT​𝐫p,i\\mathbf\{J\}\_\{p,i\}^\{T\}\\mathbf\{r\}\_\{p,i\}and𝐉p,i⊙𝐝T\\mathbf\{J\}\_\{p,i\}\\odot\\mathbf\{d\}^\{T\}, cost𝒪⁡\(N​M/P\)\\mathcal\{O\}\(NM/P\)FLOPs and are bounded by the shard products\. The per\-device computational complexities are therefore

Standard:\\displaystyle\\text\{Standard:\}𝒪⁡\(N​FP\+N​M2P\+M3\)\\displaystyle\\mathcal\{O\}\\Big\(\\frac\{NF\}\{P\}\+\\frac\{NM^\{2\}\}\{P\}\+M^\{3\}\\Big\)\(9\)Dual:\\displaystyle\\text\{Dual:\}𝒪⁡\(N2​FP​S\+N2​MP\+N3\)\\displaystyle\\mathcal\{O\}\\Big\(\\frac\{N^\{2\}F\}\{PS\}\+\\frac\{N^\{2\}M\}\{P\}\+N^\{3\}\\Big\)WhenM\>NM\>N, the block products of the dual form cost a factor ofM/NM/Nless than the accumulation of𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}in the standard form\. The extra cost of the dual form comes from evaluating the JacobianN/SN/Stimes, which is what it pays to remove the𝒪⁡\(M2\)\\mathcal\{O\}\(M^\{2\}\)memory\.

#### IV\-D3Communication

The standard form all\-reduces𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}and𝐉T​𝐫\\mathbf\{J\}^\{T\}\\mathbf\{r\}\. The dual form all\-gathers𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}and all\-reduces the twoMM\-dimensional vectors𝐝\\mathbf\{d\}and𝐉T​𝐯\\mathbf\{J\}^\{T\}\\mathbf\{v\}\. As with memory, the Gram matrix dominates the communication volume, which is𝒪⁡\(min⁡\(M,N\)2\+M\)\\mathcal\{O\}\(\\min\(M,N\)^\{2\}\+M\)elements\. The scalar loss sent over Gloo is negligible\.

Besides the model sizeMMand the batch sizeNN, the shard sizeSSand the slice sizeLLdecide the balance between memory and speed\. A largerSSreduces the Jacobian re\-evaluation of the dual form but costs𝒪⁡\(S​M\)\\mathcal\{O\}\(SM\)memory\. A largerLLimproves the utilization of vectorized VJPs at a cost of𝒪⁡\(L​A\)\\mathcal\{O\}\(LA\)memory, with the same FLOP count\.SSandLLshould therefore be tuned together to fill the device memory that the Gram matrix leaves free\.

## VEvaluation

LeMA changes how an LMA step is computed, not which step is taken\. As long as it computes the same update as LMA, it inherits the convergence behavior of LMA, which prior work has studied for neural networks\[[5](https://arxiv.org/html/2610.22925#bib.bib5),[6](https://arxiv.org/html/2610.22925#bib.bib6),[12](https://arxiv.org/html/2610.22925#bib.bib12)\]\. We therefore evaluate correctness and performance rather than training accuracy\. We first verify numerically that the dual form matches an FP64 reference update \(Section[V\-B](https://arxiv.org/html/2610.22925#S5.SS2)\)\. We then measure how far LeMA extends the model size on one GPU \(Section[V\-C](https://arxiv.org/html/2610.22925#S5.SS3)\), compare the two forms \(Section[V\-D](https://arxiv.org/html/2610.22925#S5.SS4)\), and evaluate strong scaling \(Section[V\-E](https://arxiv.org/html/2610.22925#S5.SS5)\)\. Finally, we study the shard and slice sizes \(Section[V\-F](https://arxiv.org/html/2610.22925#S5.SS6)\) and split\-K \(Section[V\-G](https://arxiv.org/html/2610.22925#S5.SS7)\)\.

Although LeMA implements both the standard and dual forms of LMA, we treat the standard form as the baseline for most experiments because the number of parameters typically exceeds the batch size in practice\.

Fig\. 1:Peak memory of one trial step on one GPU withN=256N=256andS=L=4S=L=4\. \(a\) Crosses mark the first model that runs out of memory\. \(b\) Breakdown of each variant at its largest model\.### V\-AExperiment Setup

#### V\-A1Hardware and Software

All experiments are conducted on a server equipped with two Intel Xeon Platinum 8275CL CPUs and eight NVIDIA A100\-SXM4 GPUs\. Each GPU features 80 GB of memory, a peak memory bandwidth of 2,039 GB/s, and a peak FP32 performance of 19\.5 TFLOPS\. The GPUs are fully connected via NVLink 3\.0, providing 600 GB/s of bidirectional bandwidth per GPU\.

We implement LeMA in Python 3\.14 using PyTorch 2\.14, CUDA 12\.6, and NCCL 2\.29\. All computations are performed in single precision \(FP32\)\.

#### V\-A2Datasets and Models

We use MNIST\[[29](https://arxiv.org/html/2610.22925#bib.bib29)\], which has 60,000 training and 10,000 test images of28×2828\\times 28pixels\. The whole dataset is pinned in host memory\[[30](https://arxiv.org/html/2610.22925#bib.bib30)\]\.

To reach the billion\-parameter scale on MNIST, we take a convolutional neural network \(CNN\) with two3×33\\times 3convolutional layers of 32 and 64 channels, a2×22\\times 2max pooling, and two fully connected layers, with ReLU after each hidden layer, and widen its hidden fully connected layer tohhneurons\. The model hasM=18,826\+9,227​hM=18\{,\}826\+9\{,\}227hparameters\. We call the models withh=10,836h=10\{,\}836and108,378108\{,\}378CNN\-100M and CNN\-1B, with 100,002,598 and 1,000,022,632 parameters\. Experiments that need the standard LMA form use smaller models: LeNet\-5\[[31](https://arxiv.org/html/2610.22925#bib.bib31)\]with ReLU and max pooling \(M=61,706M=61\{,\}706\) and a 784\-16\-10 MLP \(M=12,730M=12\{,\}730\)\. All models use the CE residual in \([5](https://arxiv.org/html/2610.22925#S3.E5)\) withϵ=10−9\\epsilon=10^\{\-9\}\.

#### V\-A3Measurements

We split the step time into four phases: Jacobian shard evaluation \(Jac\), the Gram block products𝐉~p,i​𝐉jT\\tilde\{\\mathbf\{J\}\}\_\{p,i\}\\mathbf\{J\}\_\{j\}^\{T\}, the NCCL and Gloo collectives including the wait for the slowest process, and the rest, such asVjp, the linear solve, and the loss evaluation of the trial step\. To analyze the peak memory footprint, we track all memory allocations during execution\. At the exact moment of maximum memory usage, we categorize every live tensor as the Gram matrix, the Jacobian or its shards,MM\-dimensional vectors, activations, or other buffers\.

### V\-BExactness of the Dual Form

Section[III\-A](https://arxiv.org/html/2610.22925#S3.SS1)proves that the dual form yields the same update as the standard form in exact arithmetic\. To verify this equivalence in FP32, we execute a single LeMA step with each form across different models and damping factors\.

As a reference independent of LeMA, we explicitly form the full Jacobian𝐉\\mathbf\{J\}and𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}, clamp𝐃\\mathbf\{D\}with the same rule as both forms \(Section[IV\-B1](https://arxiv.org/html/2610.22925#S4.SS2.SSS1)\), and solve the standard form \([1](https://arxiv.org/html/2610.22925#S2.E1)\) directly in FP64 to obtainΔ​𝜽ref\\Delta\\boldsymbol\{\\theta\}\_\{\\text\{ref\}\}\. We evaluate LeNet\-5 with batch sizesN=1,024N=1\{,\}024and8,1928\{,\}192\(both smaller thanMM\), and an MLP withN=4,096<MN=4\{,\}096<MandN=16,384\>MN=16\{,\}384\>M\. For each setup, we test damping factorsλ=10−3\\lambda=10^\{\-3\},11, and10310^\{3\}\.

Table[II](https://arxiv.org/html/2610.22925#S5.T2)reports the relative error‖Δ​𝜽−Δ​𝜽ref‖2/‖Δ​𝜽ref‖2\\left\\\|\\Delta\\boldsymbol\{\\theta\}\-\\Delta\\boldsymbol\{\\theta\}\_\{\\text\{ref\}\}\\right\\\|\_\{2\}/\\left\\\|\\Delta\\boldsymbol\{\\theta\}\_\{\\text\{ref\}\}\\right\\\|\_\{2\}\. Forλ=103\\lambda=10^\{3\}, the errors of both forms are at most4\.4×10−64\.4\\times 10^\{\-6\}\. Asλ\\lambdadecreases, the linear systems become ill\-conditioned, and the standard form degrades rapidly\. Atλ=10−3\\lambda=10^\{\-3\}, its error reaches 21% to 280%, rendering the update unusable\. In contrast, the error of the dual form stays within1\.3×10−21\.3\\times 10^\{\-2\}\. To check the reference itself, we also run the dual form of LeMA in FP64\. It differs from the reference by at most5\.3×10−95\.3\\times 10^\{\-9\}, and mostly by10−1510^\{\-15\}to10−1210^\{\-12\}, so the FP64 reference stays accurate even atλ=10−3\\lambda=10^\{\-3\}, and the dual transformation introduces no error beyond rounding\. We conclude that the errors in Table[II](https://arxiv.org/html/2610.22925#S5.T2)arise from the limited FP32 precision, to which the standard form is far more sensitive than the dual form\.

TABLE II:Relative Error of the Update Against an FP64 Reference
### V\-CMemory Footprint

We next evaluate the maximum number of parameters each variant supports on a single GPU\. We compare four configurations: the standard and dual forms of LMA, each implemented with either a full or a sharded Jacobian\. By increasinghhfrom 1 to 108,378, we scale the parameter countMMfrom2\.8×1042\.8\\times 10^\{4\}to1\.0×1091\.0\\times 10^\{9\}until each configuration exhausts the GPU memory\. All configurations useN=256N=256andL=4L=4, while the sharded implementations additionally setS=4S=4\.

Figure[1](https://arxiv.org/html/2610.22925#S5.F1)shows the peak memory and its breakdown\. The standard forms encounter theM2M^\{2\}bottleneck first\. For the implementation with the full Jacobian,𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}and the intermediate data for LU factorization consume all but 0\.02 GB of the 83\.1 GB footprint atM≈105M\\approx 10^\{5\}\. As a result, the solver exhausts memory atM≈3×105M\\approx 3\\times 10^\{5\}\. The sharded standard form fails even earlier atM≈105M\\approx 10^\{5\}because it retains𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}for reuse after a rejected step, requiring a thirdM×MM\\times Mmatrix\. The dual form with the full Jacobian avoids theM2M^\{2\}memory limit, but itsN×MN\\times MJacobian footprint grows linearly withMM, reaching 61\.6 GB atM=3\.0×107M=3\.0\\times 10^\{7\}\.

LeMA resolves both of these memory bottlenecks\. AtM=3\.0×107M=3\.0\\times 10^\{7\}, it requires only 1\.82 GB, achieving a34×34\\timesreduction compared to the dual form with the full Jacobian\. For the CNN\-1B model, the peak memory of 60\.0 GB consists of 48\.0 GB for shards and 12\.0 GB for vectors\. The shards exactly match the3​S​M=12​M3SM=12Melements predicted in Table[I](https://arxiv.org/html/2610.22925#S4.T1), and the vectors amount to threeMM\-dimensional vectors, giving15​M15Melements in total\. Consequently, LeMA successfully trains a model on a single GPU that is33×33\\timeslarger than the maximum capacity of the full\-Jacobian dual form, and approximately104×10^\{4\}\\timeslarger than that of the standard LMA\.

### V\-DStandard Form vs\. Dual Form

Fig\. 2:Standard and dual forms for LeNet\-5 withN=8,192<M=61,706N=8\{,\}192<M=61\{,\}706\. Dashed lines show ideal scaling\.We compare the memory footprint and execution performance of the standard and dual forms\. We use LeNet\-5 withN=8,192<MN=8\{,\}192<M\. Both forms useS=1,024S=1\{,\}024andL=256L=256on identical samples\.

Figure[2](https://arxiv.org/html/2610.22925#S5.F2)shows that the dual form achieves a smaller memory footprint and faster execution\. It outperforms the standard form across allPPconfigurations, delivering an8\.3×8\.3\\timesspeedup on a single GPU and a37\.8×37\.8\\timesspeedup on 8 GPUs\. The standard form exhibits poor scaling efficiency\. Specifically, on a single GPU it spends 9\.23 s of its 12\.65 s step time solving theM×MM\\times Mlinear system\. Because every GPU solves this system redundantly without parallelization, this 73% share bounds the speedup to1\.37×1\.37\\times, and the standard form reaches only1\.28×1\.28\\timeson 8 GPUs\. The dual form scales significantly better\. Even on 8 GPUs, it requires only 0\.05 s to solve theN×NN\\times Nsystem out of a 0\.26 s step time, so this unparallelized overhead limits scaling far less, and the dual form reaches5\.84×5\.84\\timeson 8 GPUs\.

The dual form achieves a33×33\\timesreduction in memory footprint compared to the standard form\. The standard form consumes 45\.8 GB per GPU regardless ofPP, dominated by threeM×MM\\times Mmatrices:𝐉T​𝐉\\mathbf\{J\}^\{T\}\\mathbf\{J\}, its damped copy, and the LU factors\. In contrast, the dual form requires only 1\.37 GB on a single GPU and 1\.11 GB on 8 GPUs, which is33×33\\timesto41×41\\timesless\. This footprint decreases asPPincreases because each GPU computes onlyN/PN/Prows of the Gram matrix before the all\-gather operation\.

### V\-EStrong 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\.We scale CNN\-100M and CNN\-1B from 1 to 8 GPUs with a fixed global batch\. Both use the dual form, withN=2,048N=2\{,\}048andS=L=32S=L=32for CNN\-100M andN=256N=256andS=L=4S=L=4for CNN\-1B\. These settings are the same for everyPP, so each GPU processesN/PN/Psamples\.

LeMA scales almost ideally \(Figure[3](https://arxiv.org/html/2610.22925#S5.F3)\(a\)\)\. On 8 GPUs, the step time of CNN\-1B drops from 317\.2 s to 39\.8 s, a speedup of7\.97×7\.97\\timesand a parallel efficiency of 99\.7%\. CNN\-100M reaches7\.96×7\.96\\times\(99\.5%\), and the efficiency stays at or above 99\.4% for everyPP\. The peak memory per GPU stays at 60\.0 GB for CNN\-1B and 39\.7 GB for CNN\-100M, since \([8](https://arxiv.org/html/2610.22925#S4.E8)\) does not depend onPP\.

Figure[3](https://arxiv.org/html/2610.22925#S5.F3)\(b\) shows the reason\. The Jacobian evaluation and the Gram block products take 68% and 29% of the CNN\-1B step, and both are split evenly across GPUs\. Communication takes only 0\.06 s on 8 GPUs, 0\.16% of the step, although it all\-reduces two 4 GB vectors\. The dual form re\-evaluates each Jacobian rowN/SN/Stimes but communicates itsN2\+2​MN^\{2\}\+2Melements only once per step \(Section[IV\-D](https://arxiv.org/html/2610.22925#S4.SS4)\)\. The redundant linear solve takes less than 0\.03% of the step\.

### V\-FShard and Slice Sizes

Section[IV\-D](https://arxiv.org/html/2610.22925#S4.SS4)predicts thatSStrades Jacobian re\-evaluation against memory, whileLLonly affects hardware utilization\. We train CNN\-100M withN=256N=256on one GPU forSSfrom 2 to 64 andLLfrom 1 toSS\(Figure[4](https://arxiv.org/html/2610.22925#S5.F4)\)\.

Fig\. 4:Effect of the shard sizeSSand the slice sizeLLfor CNN\-100M withN=256N=256on one GPU\.The step time falls almost in proportion to1/S1/S\. WithL=SL=S, it drops from 77\.1 s atS=2S=2to 4\.75 s atS=32S=32,16\.2×16\.2\\timesfor a16×16\\timeslarger shard, andS=64S=64withL=8L=8reaches 3\.70 s\. Larger shards help both dominant phases: the Jacobian is re\-evaluated fewer times, and larger blocks raise the arithmetic intensity of the block products at the same FLOP count\. The slice size matters much less\. AtS=32S=32, raisingLLfrom 1 to 32 speeds up the step by only1\.5×1\.5\\times\.

The peak memory matches3​S​M\+3​M3SM\+3Melements, i\.e\., the3​S​M3SMJacobian shard elements of Table[I](https://arxiv.org/html/2610.22925#S4.T1)plus threeMM\-dimensional vectors, within 0\.05 GB and does not depend onLL, because the activations of a slice are negligible for this model\. The only exception isS=L=64S=L=64, which runs out of memory although it needs the same amount asS=64S=64with smallerLL, because about 24 GiB of the reserved memory are too fragmented for its large slice allocations\. Since the time scales with1/S1/Sand the memory withSS,SSshould be as large as the GPU memory allows\.

### V\-GSplit\-K of the Gram Block Products

Fig\. 5:Gram block products withtorch\.mmand with split\-K\. \(a\) One product of twoS×MS\\times Mshards on one GPU\. The shaded range is where LeMA applies split\-K\. \(b\) Step time of CNN\-1B on 8 GPUs\.We evaluate split\-K \(Section[IV\-B2](https://arxiv.org/html/2610.22925#S4.SS2.SSS2)\) on a single block product and on a full step\. The baseline computes each block product with onetorch\.mmcall, which uses cuBLAS\. In the microbenchmark, we multiply two randomS×MS\\times Mshards withM≈108M\\approx 10^\{8\}or10910^\{9\}, matching CNN\-100M and CNN\-1B, andSSfrom 1 to 64, skipping shapes that do not fit in memory\. Both variants write into a strided block of a larger matrix, as in Algorithm[4](https://arxiv.org/html/2610.22925#alg4), and split\-K uses its default of 16 chunks per SM \(1,728 chunks on an A100\)\. After two warm\-up runs, we repeat each product for at least 0\.5 s and report the mean time from CUDA events\. A block product reads8​S​M8SMbytes in FP32, so a bandwidth\-bound kernel should take time proportional toSS\.

A plain GEMM does not \(Figure[5](https://arxiv.org/html/2610.22925#S5.F5)\(a\)\)\. FromS=2S=2to88, the time oftorch\.mmstays at 12\.9–13\.6 ms forM≈108M\\approx 10^\{8\}and 130–135 ms forM≈109M\\approx 10^\{9\}while the data grow4×4\\times, so the GEMM is limited by its parallelism\. Split\-K follows the data volume and is1\.271\.27–11\.0×11\.0\\timesfaster forS≤8S\\leq 8\. For the CNN\-1B shards withS=4S=4, it cuts the time from 131\.3 ms to 22\.9 ms \(5\.7×5\.7\\times\), which corresponds to 1\.40 TB/s, or 69% of the peak memory bandwidth\. ForS≥16S\\geq 16, the GEMM already has enough parallelism and split\-K is up to9\.8×9\.8\\timesslower, so LeMA applies split\-K only forS≤8S\\leq 8\. The result is not sensitive to the chunk count, which changes the time by at most 12% beyond 8 chunks per SM\. Split\-K changes the products by a relative difference of at most4\.5×10−54\.5\\times 10^\{\-5\}because of the different FP32 summation order\.

For the full step, we train CNN\-1B on 8 GPUs with the configuration of Section[V\-E](https://arxiv.org/html/2610.22925#S5.SS5)\. Each GPU computesN2/\(P​S2\)=512N^\{2\}/\(PS^\{2\}\)=512block products per step\. Without split\-K, they take 66\.9 s, 70% of the 95\.0 s step, which matches 512 products of 131 ms each in the microbenchmark\. Split\-K reduces them by5\.7×5\.7\\timesto 11\.8 s and speeds up the step by2\.39×2\.39\\timesto 39\.8 s \(Figure[5](https://arxiv.org/html/2610.22925#S5.F5)\(b\)\)\. The peak memory stays at 60\.0 GB per GPU\.

## VIConclusion and Future Work

We presented LeMA, a multi\-GPU neural network training framework based on the Levenberg\-Marquardt \(LMA\) algorithm\. By utilizing an exact dual formulation that shifts the linear system from the parameter dimension to the batch dimension, alongside efficient partitioning strategies and Split\-K optimizations, we successfully scale LMA to train models with billions of parameters\. Consequently, LeMA achieves significant speedups over standard LMA implementations\.

We identify several limitations in the current design of LeMA, which direct our future work\. First, our assumption that each input sample is processed independently prohibits the use of intra\-batch operations like Batch Normalization\. We plan to extend LeMA to support these layers without introducing significant overhead\. Second, the dual form repeatedly evaluates all global Jacobian shards for each local block of the Gram matrix𝐉𝐃−1​𝐉T\\mathbf\{J\}\\mathbf\{D\}^\{\-1\}\\mathbf\{J\}^\{T\}\. This re\-evaluation makes Jacobian computation a primary bottleneck\. Since inter\-process communication time is currently negligible, future implementations could exchange evaluated shards across GPUs via fast NVLink to avoid this redundant computation\. Third, LeMA gathers the full coefficient matrix on every GPU and solves the linear system redundantly\. Integrating a distributed solver, such as NVIDIA cuSOLVERMp\[[32](https://arxiv.org/html/2610.22925#bib.bib32)\], would allow us to partition the coefficient matrix and solve the system in parallel, yielding further memory and speed improvements\. Finally, this work focuses strictly on systems engineering to optimize memory and execution time\. We use a basic damping strategy and acceptance criterion, and leave refining them and evaluating convergence on practical deep learning tasks to future work\.

## Acknowledgment

This work was supported by the NVIDIA Academic Grant Program\. The authors used generative AI tools, including Claude and Gemini, to check grammar, improve clarity, and place figures near the text that discusses them\. The authors reviewed all changes and take full responsibility for the content\.

## References

- \[1\]K\. Levenberg, “A method for the solution of certain non\-linear problems in least squares,”*Quarterly of Applied Mathematics*, vol\. 2, no\. 2, pp\. 164–168, July 1944\. \[Online\]\. Available:http://dx\.doi\.org/10\.1090/qam/10666
- \[2\]D\. W\. Marquardt, “An algorithm for least\-squares estimation of nonlinear parameters,”*Journal of the Society for Industrial and Applied Mathematics*, vol\. 11, no\. 2, pp\. 431–441, June 1963\. \[Online\]\. Available:http://dx\.doi\.org/10\.1137/0111030
- \[3\]J\. Nocedal and S\. J\. Wright, “Least\-squares problems,” in*Numerical Optimization*, 2nd ed\. New York, NY: Springer, 2006, pp\. 245–269\. \[Online\]\. Available:https://doi\.org/10\.1007/978\-0\-387\-40065\-5\_10
- \[4\]H\. P\. Gavin, “The levenberg\-marquardt algorithm for nonlinear least squares curve\-fitting problems,” Department of Civil and Environmental Engineering Duke University, Tech\. Rep\., 2019\. \[Online\]\. Available:https://people\.duke\.edu/~hpgavin/ce281/lm\.pdf
- \[5\]M\. Hagan and M\. Menhaj, “Training feedforward networks with the marquardt algorithm,”*IEEE Transactions on Neural Networks*, vol\. 5, no\. 6, pp\. 989–993, 1994\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/72\.329697
- \[6\]S\. Ali, S\. Riaz, Safoora, X\. Liu, and G\. Wang, “A levenberg\-marquardt based neural network for short\-term load forecasting,”*Computers, Materials & Continua*, vol\. 75, no\. 1, pp\. 1783–1800, 2023\. \[Online\]\. Available:http://dx\.doi\.org/10\.32604/cmc\.2023\.035736
- \[7\]J\. Zhang, J\. Tian, A\. M\. Alcaide, J\. I\. Leon, S\. Vazquez, L\. G\. Franquelo, H\. Luo, and S\. Yin, “Lifetime extension approach based on the levenberg\-marquardt neural network and power routing of dc\-dc converters,”*IEEE Transactions on Power Electronics*, vol\. 38, no\. 8, pp\. 10 280–10 291, Aug\. 2023\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/tpel\.2023\.3275791
- \[8\]J\. A\. Taylor, W\. Wang, B\. Bala, and T\. Bednarz, “Optimizing the optimizer for data driven deep neural networks and physics informed neural networks,”*CoRR*, vol\. abs/2205\.07430, 2022\. \[Online\]\. Available:https://doi\.org/10\.48550/arXiv\.2205\.07430
- \[9\]M\. Revanesh, S\. S\. Gundal, J\. R\. Arunkumar, P\. J\. Josephson, S\. Suhasini, and T\. K\. Devi, “Artificial neural networks\-based improved levenberg\-marquardt neural network for energy efficiency and anomaly detection in wsn,”*Wireless Networks*, vol\. 30, no\. 6, pp\. 5613–5628, Mar\. 2023\. \[Online\]\. Available:http://dx\.doi\.org/10\.1007/s11276\-023\-03297\-6
- \[10\]M\. Troiano, E\. Nobile, F\. Mangini, M\. Mastrogiuseppe, C\. Conati Barbaro, and F\. Frezza, “A comparative analysis of the bayesian regularization and levenberg\-marquardt training algorithms in neural networks for small datasets: A metrics prediction of neolithic laminar artefacts,”*Information*, vol\. 15, no\. 5, p\. 270, May 2024\. \[Online\]\. Available:http://dx\.doi\.org/10\.3390/info15050270
- \[11\]S\. S\. Koshy and L\. J\. Anbarasi, “Lmhistnet: Levenberg\-marquardt based deep neural network for classification of breast cancer histopathological images,”*IEEE Access*, vol\. 12, pp\. 52 051–52 066, 2024\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/access\.2024\.3385011
- \[12\]F\. Di Marco, “Pytorch levenberg\-marquardt,” 2025\. \[Online\]\. Available:https://github\.com/fabiodimarco/torch\-levenberg\-marquardt
- \[13\]——, “Tensorflow levenberg\-marquardt,” 2025\. \[Online\]\. Available:https://github\.com/fabiodimarco/tf\-levenberg\-marquardt
- \[14\]M\. Blondel, Q\. Berthet, M\. Cuturi, R\. Frostig, S\. Hoyer, F\. Llinares\-Lopez, F\. Pedregosa, and J\.\-P\. Vert, “Efficient and modular implicit differentiation,” in*Advances in Neural Information Processing Systems 35*, ser\. NeurIPS 2022\. Neural Information Processing Systems Foundation, Inc\. \(NeurIPS\), 2022, pp\. 5230–5242\. \[Online\]\. Available:http://dx\.doi\.org/10\.52202/068431\-0378
- \[15\]S\. Kandel, S\. Maddali, Y\. S\. G\. Nashed, S\. O\. Hruszkewycz, C\. Jacobsen, and M\. Allain, “Efficient ptychographic phase retrieval via a matrix\-free levenberg\-marquardt algorithm,”*Optics Express*, vol\. 29, no\. 15, p\. 23019, July 2021\. \[Online\]\. Available:http://dx\.doi\.org/10\.1364/oe\.422768
- \[16\]L\. Höllein, A\. Božič, M\. Zollhöfer, and M\. Nießner, “3dgs\-lm: Faster gaussian\-splatting optimization with levenberg\-marquardt,” in*2025 IEEE/CVF International Conference on Computer Vision \(ICCV\)*\. IEEE, Oct\. 2025, pp\. 26 740–26 750\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/iccv51701\.2025\.02482
- \[17\]S\. Balay, S\. Abhyankar, M\. F\. Adams, S\. Benson, J\. Brown, P\. Brune, K\. Buschelman, E\. M\. Constantinescu, L\. Dalcin, A\. Dener, V\. Eijkhout, J\. Faibussowitsch, W\. D\. Gropp, V\. Hapla, T\. Isaac, P\. Jolivet, D\. Karpeev, D\. Kaushik, M\. G\. Knepley, F\. Kong, S\. Kruger, D\. A\. May, L\. C\. McInnes, R\. T\. Mills, L\. Mitchell, T\. Munson, J\. E\. Roman, K\. Rupp, P\. Sanan, J\. Sarich, B\. F\. Smith, S\. Zampini, H\. Zhang, H\. Zhang, and J\. Zhang, “PETSc Web page,” 2026\. \[Online\]\. Available:https://petsc\.org/
- \[18\]B\. Symoens, M\. Rahimi, and M\. Ahookhosh, “Inexact levenberg\-marquardt methods under hölder metric subregularity,” 2025\. \[Online\]\. Available:https://arxiv\.org/abs/2507\.16461
- \[19\]M\. R\. Hestenes and E\. Stiefel, “Methods of conjugate gradients for solving linear systems,”*Journal of research of the National Bureau of Standards*, vol\. 49, no\. 6, pp\. 409–436, 1952\. \[Online\]\. Available:https://nvlpubs\.nist\.gov/nistpubs/jres/049/6/v49\.n06\.a08\.pdf
- \[20\]S\. Agarwal, N\. Snavely, S\. M\. Seitz, and R\. Szeliski, “Bundle adjustment in the large,” in*Computer Vision \- ECCV 2010, 11th European Conference on Computer Vision, Heraklion, Crete, Greece, September 5\-11, 2010, Proceedings, Part II*, ser\. Lecture Notes in Computer Science, K\. Daniilidis, P\. Maragos, and N\. Paragios, Eds\., vol\. 6312\. Springer, 2010, pp\. 29–42\. \[Online\]\. Available:https://doi\.org/10\.1007/978\-3\-642\-15552\-9\\\_3
- \[21\]J\. Bilski, B\. Kowalczyk, A\. Marchlewska, and J\. M\. Zurada, “Local levenberg\-marquardt algorithm for learning feedforwad neural networks,”*Journal of Artificial Intelligence and Soft Computing Research*, vol\. 10, no\. 4, pp\. 299–316, June 2020\. \[Online\]\. Available:http://dx\.doi\.org/10\.2478/jaiscr\-2020\-0020
- \[22\]J\. Ansel, E\. Yang, H\. He, N\. Gimelshein, A\. Jain, M\. Voznesensky, B\. Bao, P\. Bell, D\. Berard, E\. Burovski, G\. Chauhan, A\. Chourdia, W\. Constable, A\. Desmaison, Z\. DeVito, E\. Ellison, W\. Feng, J\. Gong, M\. Gschwind, B\. Hirsh, S\. Huang, K\. Kalambarkar, L\. Kirsch, M\. Lazos, M\. Lezcano, Y\. Liang, J\. Liang, Y\. Lu, C\. Luk, B\. Maher, Y\. Pan, C\. Puhrsch, M\. Reso, M\. Saroufim, M\. Y\. Siraichi, H\. Suk, M\. Suo, P\. Tillet, E\. Wang, X\. Wang, W\. Wen, S\. Zhang, X\. Zhao, K\. Zhou, R\. Zou, A\. Mathews, G\. Chanan, P\. Wu, and S\. Chintala, “PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation,” in*29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 \(ASPLOS ’24\)*\. ACM, Apr\. 2024\. \[Online\]\. Available:https://docs\.pytorch\.org/assets/pytorch2\-2\.pdf
- \[23\]J\. Bradbury, R\. Frostig, P\. Hawkins, M\. J\. Johnson, Y\. Katariya, C\. Leary, D\. Maclaurin, G\. Necula, A\. Paszke, J\. VanderPlas, S\. Wanderman\-Milne, and Q\. Zhang, “JAX: composable transformations of Python\+NumPy programs,” 2018\. \[Online\]\. Available:http://github\.com/jax\-ml/jax
- \[24\]M\. Abadi, A\. Agarwal, P\. Barham, E\. Brevdo, Z\. Chen, C\. Citro, G\. S\. Corrado, A\. Davis, J\. Dean, M\. Devin, S\. Ghemawat, I\. Goodfellow, A\. Harp, G\. Irving, M\. Isard, Y\. Jia, R\. Jozefowicz, L\. Kaiser, M\. Kudlur, J\. Levenberg, D\. Mane, R\. Monga, S\. Moore, D\. Murray, C\. Olah, M\. Schuster, J\. Shlens, B\. Steiner, I\. Sutskever, K\. Talwar, P\. Tucker, V\. Vanhoucke, V\. Vasudevan, F\. Viegas, O\. Vinyals, P\. Warden, M\. Wattenberg, M\. Wicke, Y\. Yu, and X\. Zheng, “Tensorflow: Large\-scale machine learning on heterogeneous distributed systems,” 2016\. \[Online\]\. Available:https://arxiv\.org/abs/1603\.04467
- \[25\]F\. Darema, “SPMD computational model,” in*Encyclopedia of Parallel Computing*, D\. Padua, Ed\. Boston, MA: Springer US, 2011, pp\. 1933–1943\. \[Online\]\. Available:https://doi\.org/10\.1007/978\-0\-387\-09766\-4\_26
- \[26\]Z\. Hu, S\. Shen, T\. Bonato, S\. Jeaugey, C\. Alexander, E\. Spada, J\. Dinan, J\. Hammond, and T\. Hoefler, “Demystifying nccl: An in\-depth analysis of gpu communication protocols and algorithms,” in*2025 IEEE Symposium on High\-Performance Interconnects \(HOTI\)*\. IEEE, Aug\. 2025, pp\. 48–59\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/hoti66940\.2025\.00024
- \[27\]Gloo Contributors, “Gloo: Collective communications library with various primitives for multi\-machine training,” 2026\. \[Online\]\. Available:https://github\.com/pytorch/gloo
- \[28\]V\. Thakkar, P\. Ramani, C\. Cecka, A\. Shivam, H\. Lu, E\. Yan, J\. Kosaian, M\. Hoemmen, H\. Wu, A\. Kerr, M\. Nicely, D\. Merrill, D\. Blasig, A\. Atluri, F\. Qiao, P\. Majcher, P\. Springer, M\. Hohnerbach, J\. Wang, and M\. Gupta, “CUTLASS,” Jan\. 2023\. \[Online\]\. Available:https://github\.com/NVIDIA/cutlass
- \[29\]L\. Deng, “The mnist database of handwritten digit images for machine learning research \[best of the web\],”*IEEE Signal Processing Magazine*, vol\. 29, no\. 6, pp\. 141–142, Nov\. 2012\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/msp\.2012\.2211477
- \[30\]M\. Harris\. \(2012\) How to optimize data transfers in CUDA C/C\+\+\. NVIDIA Technical Blog\. \[Online\]\. Available:https://developer\.nvidia\.com/blog/how\-optimize\-data\-transfers\-cuda\-cc/
- \[31\]Y\. Lecun, L\. Bottou, Y\. Bengio, and P\. Haffner, “Gradient\-based learning applied to document recognition,”*Proceedings of the IEEE*, vol\. 86, no\. 11, pp\. 2278–2324, 1998\. \[Online\]\. Available:http://dx\.doi\.org/10\.1109/5\.726791
- \[32\]NVIDIA Corporation,*cuSOLVERMp*, 2026, accessed: October 16, 2026\. \[Online\]\. Available:https://docs\.nvidia\.com/cuda/cusolvermp/