PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Levenberg-Marquardt Training Framework for Deep Learning

Xianghao Wang, Johan Barthélemy, John Taylor

arXiv:2610.22925v1cs.LGcs.DC

TL;DR

LMA’s Jacobian and parameter-sized linear system make exact training difficult to scale, while PCG-based alternatives approximate the solve. LeMA combines an exact dual formulation with Jacobian partitioning for distributed multi-GPU training. On 8 GPUs, it reports 7.97× speedup and 99.7% strong scaling efficiency for a billion-parameter model.

  • Problem

    LMA’s Jacobian and parameter-sized coefficient matrix create prohibitive memory and computational demands for large-scale training.

    Source: §I · p. 1

  • Method

    LeMA uses an exact dual formulation and partitions Jacobian computation to support distributed multi-GPU LMA training.

    Source: §I · pp. 1–2

  • Results

    7.97× speedup and 99.7% strong scaling efficiency on 8 GPUs for CNN-1B, which has 1,000,022,632 parameters.

    Source: §I · p. 2, E. Strong Scaling · p. 9

  • Takeaways & Limitations

    LeMA scales exact LMA steps to models with billions of parameters while reducing peak memory relative to the full-Jacobian dual form.

    Source: C. Memory Footprint · p. 8, §VI · pp. 9–10

    The current design assumes samples are processed independently, excluding intra-batch operations such as Batch Normalization.

    Source: §VI · p. 10

Abstract

from arXiv · show

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.

I. INTRODUCTION

LeMA makes exact LMA training practical at large scale through a dual formulation and partitioning strategy, with strong multi-GPU scaling and reduced memory use in reported tests. The section also distinguishes LMA’s curvature-scaled damping from the earlier Levenberg algorithm and reviews existing scaling limits.

A. The Levenberg-Marquardt Algorithm

  • LMA minimizes squared residuals using a damped Gauss-Newton system; its diagonal damping scales parameters by curvature, unlike the earlier identity-damped Levenberg algorithm.Lower damping moves toward Gauss-Newton, while higher damping moves toward gradient descent; the described implementation accepts steps that reduce squared residuals and adjusts damping accordingly.Source: §I · p. 1, A. The Levenberg-Marquardt Algorithm · p. 2

I. INTRODUCTION

  • LMA’s O(MN + M^2) memory footprint and costly Jacobian, matrix construction, and solve operations hinder scaling to large neural networks.The linear system uses an M × M coefficient matrix and an N × M Jacobian, where M is the parameter count and N the batch size.Source: §I · p. 1
  • PCG avoids explicitly constructing large matrices but approximates the solution, requires many network propagations, and depends on a task-specific preconditioner.An exact neuron-wise alternative assumes one output per neuron, making it incompatible with parameter-sharing convolutional layers.Source: §I · p. 1
  • The exact dual formulation shifts LMA’s linear system from parameter dimension to batch dimension, while partitioning avoids storing the full Jacobian and supports multi-GPU parallelism.The paper also identifies robust residual functions tailored for LMA as a contribution.Source: §I · p. 2
  • 7.97× speedup and 99.7% strong scaling efficiency on 8 GPUs; LeMA scales to models with 1,000,022,632 parameters.For LeNet-5, LeMA reports an 8.3× speedup and 33× memory reduction on one GPU, and a 37.8× speedup and 41× memory reduction over a standard multi-GPU LMA baseline on 8 GPUs.Source: §I · p. 2

B. Jacobian Matrix Computation

The framework computes Jacobians row-wise with VJPs because samples are assumed independent; this requires fewer evaluations than column-wise JVP construction under that assumption.

B. Jacobian Matrix Computation

  • VJPs use reverse-mode automatic differentiation with a forward pass that stores activations and a backward pass; JVPs use forward mode in a single forward pass without storing activations.The two primitives compute JT v and Jv, respectively, without explicitly materializing the full Jacobian.Source: B. Jacobian Matrix Computation · p. 3
  • Under sample independence, one single-sample VJP yields one Jacobian row, so assembling the full Jacobian requires N VJPs.The framework therefore constructs the Jacobian row-wise.Source: B. Jacobian Matrix Computation · p. 3
  • Column-wise construction would require MN single-sample JVPs, because each Jacobian column involves all N samples.Equivalently, extracting a column requires one N-sample JVP.Source: B. Jacobian Matrix Computation · p. 3
  • The independence assumption excludes intra-batch operations such as Batch Normalization.Under this assumption, each residual depends only on its corresponding sample.Source: B. Jacobian Matrix Computation · p. 3

III. ALGORITHM DESIGN

For M > N, the paper transforms LMA into an exact dual form that uses an N × N system instead of the parameter-sized M × M Gram matrix. It proves the update is equivalent to the standard LMA form and distinguishes this derivation from earlier dual formulations.

A. Dual Transformation

  • When M > N, the dual system uses the N × N matrix JD−1JT, shifting the linear system from parameter dimension to batch dimension.The transformation targets the M × M Gram matrix JT J that dominates memory in this setting.Source: A. Dual Transformation · p. 3
  • The paper establishes that the dual update yields the same update as the standard LMA equation, with invertibility relying on λ > 0 and positive diagonal D.It assumes D = diag(JT J) is invertible.Source: A. Dual Transformation · p. 3
  • The exact dual formulation retains Marquardt’s curvature-based damping, unlike recent dual implementations corresponding to the Levenberg algorithm.The paper derives and proves an exact dual formulation of complete LMA.Source: A. Dual Transformation · p. 3

B. Residual Formulation

LeMA represents common losses with per-sample residuals whose squared sum matches the LMA objective. Residual functions must also avoid unstable derivatives near zero error.

B. Residual Formulation

  • MSE, cross-entropy, and MAE are expressed as per-sample residuals for LMA’s sum-of-squared-residuals objective.For MSE, the residual is direct; cross-entropy and MAE use square-root-based residuals, with the MAE form also including the error sign.Source: B. Residual Formulation · p. 3
  • Adding ε inside the cross-entropy square root prevents a zero-loss sample from causing division by zero in the Jacobian.The derivative of the unmodified square root is singular at zero loss.Source: B. Residual Formulation · p. 3
  • For MAE, ε and the error sign keep the residual derivative continuous and positive as the error crosses zero.Without the sign correction, the derivative jumps between large values when ε is small, making the Jacobian discontinuous and the optimizer oscillate around the minimum.Source: B. Residual Formulation · p. 3
  • LeMA accepts a step only when the squared residual norm decreases; otherwise, it increases λ and recomputes the step.An accepted step decreases λ.Source: B. Residual Formulation · p. 4

C. The LeMA Algorithm

LeMA chooses the standard or dual form according to the relative sizes of the parameter count and batch size, reducing the Gram matrix to the smaller dimension. Its distributed implementation runs across GPU processes with local mini-batches.

C. The LeMA Algorithm

  • LeMA selects the dual form when M > N and the standard form otherwise, so the Gram matrix has min(M, N)^2 elements.This reduces the algorithm’s memory footprint from O(MN + M^2) to O(MN + min(M, N)^2); Jacobian sharding further eliminates the O(MN) term.Source: C. The LeMA Algorithm · p. 4
  • LeMA forms the Gram matrix once per iteration and reuses it across rejected steps, which change only λ.Residuals from an accepted step carry over to the next iteration.Source: C. The LeMA Algorithm · p. 4

IV. LEMA IMPLEMENTATION

  • LeMA executes the algorithm on P GPUs using one SPMD process per GPU, with large tensors communicated through NCCL and small scalars through Gloo.Keeping scalar communication on the host avoids small host-device copies that would interrupt GPU streams.Source: §IV · p. 4, A. Distributed Setup · p. 4

A. Distributed Setup

  • Each process loads a disjoint local mini-batch, and the P mini-batches together form the global batch required by the dual form.The parameters are stored in one flat tensor with each layer’s weights mapped to their positions.Source: A. Distributed Setup · p. 4

B. Partition and Parallelization

LeMA partitions the Jacobian by rows and computes shards on demand, avoiding storage of the full matrix. It distributes shard calculations across GPUs and uses distinct procedures for the standard and dual forms.

B. Partition and Parallelization

  • Computing and freeing row shards as they are added to the Gram matrix reduces Jacobian storage from O(NM) to O(SM).Each shard contains S sample rows and is computed on the fly; row partitioning requires S single-sample VJPs per shard.Source: B. Partition and Parallelization · p. 4

1) Standard Form:

  • In the standard form, processes sum their local shard contributions and all-reduce J^T J and J^T r before solving the M × M system.The diagonal entries of J^T J are clamped to at least 1% of their mean to avoid singular D entries.Source: 1) Standard Form · p. 4

2) Dual Form:

  • The dual form solves an N-dimensional system using JD^-1J^T, then computes the parameter update as D^-1J^T v without forming an M × M matrix.It saves memory at the cost of re-evaluating every global shard for each local shard; forming the Gram matrix uses an all-gather.Source: 2) Dual Form · pp. 4–5
  • For small-row shard products, split-K divides the inner dimension into chunks and sums partial blocks to keep more GPU streaming multiprocessors active.LeMA uses split-K only when shards have few rows, because larger shards already provide enough parallelism.Source: 2) Dual Form · p. 5

C. Jacobian Shard Computation

LeMA evaluates Jacobian shards by vectorizing per-sample VJPs in slices, balancing execution speed against activation memory. It computes residuals alongside the Jacobian and slices dual-form products to limit activation storage.

C. Jacobian Shard Computation

  • JAC vectorizes VJPs for L samples at a time, producing an L × M Jacobian slice while trading execution speed against activation memory.The per-sample function returns residuals as an auxiliary output, avoiding an extra forward pass.Source: C. Jacobian Shard Computation · pp. 5–6
  • The dual-form VJP processes L samples per call and sums the results to avoid retaining activations for all N/P local samples at once.Unlike JAC, a VJP computes J^T v directly from a backward pass on the scalar v^T r, without building a Jacobian slice.Source: C. Jacobian Shard Computation · pp. 5–6

D. Memory Footprint and Performance Analysis

The cost analysis measures per-device memory and computation while accounting for vectorized passes over independent samples. It omits the input mini-batch and drops asymptotic terms only when a retained term strictly dominates.

D. Memory Footprint and Performance Analysis

1) Memory Footprint:

LeMA’s memory depends quadratically on the smaller of model size and batch size, with additional linear shard and slice costs. The dual form trades extra Jacobian evaluation for reduced memory when M > N.

1) Memory Footprint:

  • Memory comprises the Gram matrix, Jacobian shards, and slice activations; LeMA’s per-device bound is min(M, N)^2 + SM + LA.The dual form’s nested loop holds an additional shard, raising Jacobian-shard storage to 3SM; the standard form uses 2SM.Source: 1) Memory Footprint · p. 6

2) Computation:

  • LU factorization costs O(M^3) for the standard system or O(N^3) for the dual system, and this solve cost does not shrink with GPU count.Source: 2) Computation · p. 6
  • When M > N, dual-form block products cost a factor of M/N less than standard-form Gram accumulation, but require Jacobian re-evaluation N/S times.Source: 2) Computation · p. 6

3) Communication:

The evaluation examines correctness and performance rather than training accuracy, using specified hardware, datasets, models, and measurements. Communication cost and shard and slice sizes are analyzed as system-design considerations.

3) Communication:

  • The standard form all-reduces JT J and JT r, while the dual form all-gathers JD−1JT and all-reduces d and JT v.The Gram matrix dominates communication volume, which is O(min(M, N)^2 + M) elements.Source: 3) Communication · p. 6
  • Larger shard size S reduces dual-form Jacobian re-evaluation at O(SM) memory cost; larger slice size L improves vectorized VJP utilization at O(LA) memory cost.S and L are tuned together to use memory left free by the Gram matrix.Source: 3) Communication · pp. 6–7

V. EVALUATION

  • LeMA changes how an LMA step is computed, not which update it takes, so the evaluation tests correctness and performance rather than training accuracy.The evaluation plan covers dual-form equivalence, single-GPU model capacity, form comparison, strong scaling, shard and slice sizes, and split-K.Source: §V · p. 7

1) Hardware and Software:

  • The experiments use two 80 GB NVIDIA A100-SXM4 GPUs? No; the server has eight 80 GB NVIDIA A100-SXM4 GPUs, connected by NVLink 3.0.The implementation uses Python 3.14, PyTorch 2.14, CUDA 12.6, and NCCL 2.29, with FP32 computations.Source: 1) Hardware and Software · p. 7

2) Datasets and Models:

  • MNIST experiments include CNN-100M and CNN-1B with 100,002,598 and 1,000,022,632 parameters, respectively.LeNet-5 and a 784-16-10 MLP are used for experiments requiring standard LMA.Source: 2) Datasets and Models · p. 7

3) Measurements:

  • Step time is divided into four phases, and peak memory is analyzed by classifying live tensors at the moment of maximum use.The phases are Jacobian shard evaluation, Gram block products, collectives, and the remaining work.Source: 3) Measurements · p. 7

B. Exactness of the Dual Form

LeMA’s dual form supports near-ideal strong scaling for billion-parameter CNN training, while sharding enables much larger models to fit in memory. Its FP32 dual-form update also remains closer to an FP64 reference than the standard form under low damping.

B. Exactness of the Dual Form

  • The exactness check compares single-step updates against an FP64 reference across LeNet-5 and MLP batch sizes and three damping factors.The reference explicitly forms J and JT J and solves the standard form in FP64.Source: B. Exactness of the Dual Form · p. 7
  • At λ = 10^3, both forms’ relative errors are at most 4.4 × 10^-6; at λ = 10^-3, standard-form error reaches 21% to 280%, while dual-form error stays within 1.3×10^-2.The FP64 dual form differs from the FP64 reference by at most 5.3 × 10^-9.Source: B. Exactness of the Dual Form · p. 7

C. Memory Footprint

  • The standard forms hit the M^2 memory bottleneck first; the full-Jacobian standard implementation exhausts memory at M ≈ 3×10^5.The figure marks the first model to run out of memory and breaks down each variant at its largest model.Source: C. Memory Footprint · pp. 7–8, Figure 1 · p. 8
  • At 1,000,022,632 parameters, LeMA trains CNN-1B on one GPU with 60.0 GB peak memory, while full-Jacobian dual form reaches only 3.0 × 10^7 parameters.The 60.0 GB consists of 48.0 GB for shards and 12.0 GB for vectors; Figure 1 shows peak memory and its breakdown.Source: C. Memory Footprint · p. 8, Figure 1 · p. 8

D. Standard Form vs. Dual Form

  • For LeNet-5, the dual form is faster and uses less memory than the standard form across GPU counts, with 8.3× speedup on one GPU and 37.8× on 8 GPUs.The comparison uses N = 8,192 < M = 61,706; the figure includes ideal-scaling reference lines.Source: D. Standard Form vs. Dual Form · p. 8, Figure 2 · p. 8
  • On 8 GPUs, dual-form LeNet-5 reaches 5.84× speedup, while the standard form reaches 1.28× because its M×M solve is redundantly performed on each GPU.The standard solve accounts for 73% of single-GPU step time; dual-form memory is 33× to 41× lower.Source: D. Standard Form vs. Dual Form · p. 8

E. Strong Scaling

  • 7.97× speedup and 99.7% parallel efficiency: CNN-1B step time falls from 317.2 s to 39.8 s on 8 GPUs.The result is shown in the fixed-global-batch strong-scaling figure.Source: Figure 3 · p. 8, E. Strong Scaling · p. 9
  • For CNN-1B, Jacobian evaluation and Gram block products take 68% and 29% of step time, while communication is 0.16% on 8 GPUs.The two dominant phases are split evenly across GPUs; the redundant solve takes less than 0.03% of the step.Source: Figure 3 · p. 8, E. Strong Scaling · p. 9

F. Shard and Slice Sizes

The section tunes shard size and split-K for throughput, then concludes that exact dual-form LeMA scales to billion-parameter models while retaining design limitations.

F. Shard and Slice Sizes

  • At S = 32, increasing slice size L from 1 to 32 speeds the step by only 1.5×, while larger shards substantially reduce step time.With L = S, step time drops from 77.1 s at S = 2 to 4.75 s at S = 32; S = 64 with L = 8 reaches 3.70 s. Memory scales with S, and S = L = 64 runs out of memory because reserved memory is fragmented.Source: F. Shard and Slice Sizes · p. 9

G. Split-K of the Gram Block Products

VI. CONCLUSION AND FUTURE WORK

  • LeMA combines an exact dual formulation with partitioning and split-K optimizations to scale LMA to models with billions of parameters.The dual formulation shifts the linear system from the parameter dimension to the batch dimension.Source: §VI · pp. 9–10
  • The current design assumes independent sample processing, excludes intra-batch operations such as Batch Normalization, and leaves convergence evaluation on practical deep-learning tasks for future work.It also re-evaluates global Jacobian shards and redundantly gathers and solves the full coefficient matrix on each GPU.Source: §VI · p. 10
Loading 2610.22925v1…