Exploiting Subgradient Sparsity in Max-Plus Neural Networks

Ikhlas Enaieh; Olivier Fercoq

stat.MLMar 4, 2026

Exploiting Subgradient Sparsity in Max-Plus Neural Networks

Authors: Ikhlas Enaieh, Olivier Fercoq

Organizations: LTCI, S2A · Image, Data, Signal Department (IDS), Telecom Paris, Institut Polytechnique de Paris, France. · S2A, LTCI

Abstract

Deep Neural Networks are powerful tools for solving machine learning problems, but their training often involves dense and costly parameter updates. In this work, we use a novel Max-Plus neural architecture in which classical addition and multiplication are replaced with maximum and summation operations respectively. This is a promising architecture in terms of interpretability, but its training is challenging. A particular feature is that this algebraic structure naturally induces sparsity in the subgradients, as only neurons that contribute to the maximum affect the loss. However, standard backpropagation fails to exploit this sparsity, leading to unnecessary computations. In this work, we focus on the minimization of the worst sample loss which transfers this sparsity to the optimization loss. To address this, we propose a sparse subgradient algorithm that explicitly exploits the algebraic sparsity. By tailoring the optimization procedure to the non-smooth nature of Max-Plus models, our method achieves more efficient updates while retaining theoretical guarantees. This highlights a principled path toward bridging algebraic structure and scalable learning.

Figures & tables

Operation	Target	Standard	SCT
Compute	$\underset{1\leq i\leq N}{\max}\;x_{i}$	$\mathcal{O}(N)$	$\mathcal{O}(N)$
Update	$\underset{1\leq i\leq N}{\max}\;\tilde{x}_{i}$	$\mathcal{O}(N)$	$\bm{\mathcal{O}(\log N)}$

Table 1: Complexity of maximum computation: Standard Maximum vs. SCT representation

Model	MLP ( $\frac{1}{N}$ )	LMM ( $\frac{1}{N}$ )	LMM ( $\max_{n}$ )
max loss	1.839	1.025	0.426
avg loss	0.025	0.16	0.245
accuracy	99%	98.1%	100%

Table 2: Training performance on the Iris dataset.

Model	MLP ( $\frac{1}{N}$ )	LMM ( $\frac{1}{N}$ )	LMM ( $\max_{n}$ )
max loss	8.744	1.253	1.03
avg loss	0.047	0.227	0.311
accuracy	88%	91.11%	93.33%

Table 3: Testing performance on the Iris dataset.

Figure 1: Final Max-SCCE values on the IRIS dataset obtained with three different weight initialization strategies, evaluated over 10 independent runs.

Figure 2: Convergence of the LMM model on MNIST

	0	1	2	3	4	5	6	7	8	9
0	922	0	8	5	3	26	5	2	5	4
1	0	1089	12	2	8	2	3	4	14	1
2	9	1	923	30	11	44	21	25	4	1
3	7	3	31	849	0	64	0	9	30	17
4	4	2	3	4	820	1	5	8	14	121
5	8	0	2	42	3	784	14	8	20	11

Table 4: Confusion matrix of the LMM model on the MNIST test set (

88.6\%

accuracy).

Update mode	Time (s)	s / iter	Mem. (GB)
Sparse + skip $W^{0}$	36	0.12	513
Sparse	1045	3.48	513
Dense	5688	18.96	10

Table 5: Average computational time per iteration on MNIST.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Distribution of predicted probabilities for the true labels on MNIST using the LMM model.

Model	MLP ( $\frac{1}{N}$ )	LMM ( $\frac{1}{N}$ )	LMM ( $\max_{n}$ )
max loss	27.63	5.02	1.64
avg loss	0.05	1.23	1.21
accuracy	99.15%	69%	91.54%

Appendix

Table 6: Training performance on the MNIST dataset.

Model	MLP ( $\frac{1}{N}$ )	LMM ( $\frac{1}{N}$ )	LMM ( $\max_{n}$ )
max loss	27.63	4.55	3.59
avg loss	0.047	1.22	1.23
accuracy	88%	69%	88.75%

Appendix

Table 7: Testing performance on the MNIST dataset.

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 332 papers

Scalable Levenberg-Marquardt Training Framework for Deep Learning
Oct 16, 2026Xianghao Wang, Johan Barthélemy, John TaylorNeural Network Optimization Second-Order Optimization
School of Computing Australian National University Canberra, Australia · NVIDIA CA, United States
Training Non-Differentiable Networks via Optimal Transport
May 3, 2026An T. LeBlack-Box Optimization Neural Network Optimization
Center for AI Research, VinUniversity, Vietnam · Intelligent Autonomous Systems, TU Darmstadt, Germany
Subspace Levenberg Marquardt Algorithms in Training Neural Networks
Sep 1, 2026M. Duc HoangNeural Network Optimization Krylov Subspace Methods
Department of Mathematics, University of California, Davis, Davis, CA, USA

Oct 16, 2026J/K move · Enter open · S save

Scalable Levenberg-Marquardt Training Framework for Deep Learning

Xianghao Wang, Johan Barthélemy, John Taylor

School of Computing Australian National University Canberra, Australia · NVIDIA CA, United States

Neural Network Optimization Second-Order Optimization Memory-Efficient Training

Open paper AskPDF arXiv

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.

Abstract

Neural Network Optimization

Explore similar work

Scalable Levenberg-Marquardt Training Framework for Deep Learning

Training Non-Differentiable Networks via Optimal Transport

Subspace Levenberg Marquardt Algorithms in Training Neural Networks