Source-linked AI summary
Training Neural Networks Without Gradients: A Scalable ADMM Approach
Gavin Taylor, Ryan Burmeister, Zheng Xu, Bharat Singh, Ankit Patel, Tom Goldstein
TL;DR
Large-scale neural-network training is limited by gradient methods’ parallelization and optimization difficulties. The paper uses ADMM and Bregman iteration to replace gradient steps with globally solved minimization sub-problems. Experiments report strong scaling and faster benchmark attainment on a large problem, while the study remains a proof of concept over a narrow classification range.
Problem
Gradient-based training methods lack strong scaling across large numbers of cores and suffer from vanishing gradients, poor conditioning, saturation, and saddle points.
Method
ADMM and Bregman iteration reduce network training to minimization sub-problems that are solved globally in closed form without gradient steps.
Results
7.8 seconds on 7200 cores reached 64% test accuracy with ADMM, compared with 181 seconds for L-BFGS and 44 minutes for conjugate gradients; SGD did not reach 64% in seven hours.
Takeaways & Limitations
The method scales linearly to thousands of cores and can outperform other methods on problems involving extremely large datasets.
Takeaways & Limitations
The experiments cover a narrow range of classification problems and are intended as a proof of concept rather than evidence of absolute superiority.
Abstract
from arXiv · showhide
With the growing importance of large network models and enormous training datasets, GPUs have become increasingly necessary to train neural networks. This is largely because conventional optimization algorithms rely on stochastic gradient methods that don't scale well to large numbers of cores in a cluster setting. Furthermore, the convergence of all gradient methods, including batch methods, suffers from common problems like saturation effects, poor conditioning, and saddle points. This paper explores an unconventional training method that uses alternating direction methods and Bregman iteration to train networks without gradient descent steps. The proposed method reduces the network training problem to a sequence of minimization sub-steps that can each be solved globally in closed form. The proposed method is advantageous because it avoids many of the caveats that make gradient methods slow on highly non-convex problems. The method exhibits strong scaling in the distributed setting, yielding linear speedups even when split over thousands of cores.
1. Introduction
The paper motivates an alternative to gradient-based neural-network training because large-scale parallelization is difficult and non-convex optimization introduces several convergence problems. It proposes ADMM and Bregman iteration, decomposing training into globally solvable sub-steps for scalable distributed computation.
- Large datasets and many model parameters motivate specialized hardware and substantial training and hyperparameter-tuning investments.
- Gradient-based methods scale poorly when parallelized and are slowed by saddle points, poor conditioning, and vanishing gradients.
- ADMM and Bregman iteration decompose neural-network training into sub-steps solved to global optimality.
- The proposed method is designed to scale linearly when data is distributed across cores while remaining robust to gradient saturation and poor conditioning.
2. Background and notation
The paper formulates feed-forward neural-network training over weights and activations, then explains why gradient methods face scaling and optimization difficulties. Its alternative separates layerwise objectives so weights and nonlinear activations can be updated through globally solvable sub-problems, with layerwise parallelism.
- Background and notation: A feed-forward network applies linear operators and nonlinear activations across layers, mapping input activations to output activations.
- Background and notation: Training tunes weight matrices so the final activations match targets under a loss function, using all training samples simultaneously in the batch formulation.
- Background and notation: The formulation permits convolutional operators and assumes feed-forward architectures, while allowing modification to more complex topologies such as recurrent networks.
- What’s wrong with backprop?: SGD scales poorly across many cores because its inexpensive mini-batch steps provide too little computation to split efficiently and require frequent communication.
- What’s wrong with backprop?: Gradient methods also face vanishing gradients, saddle points, local-minimum concerns, and limited parallelization across layers.
- Alternative optimization: The proposed method separates weight-input relations from nonlinear activations, enabling independent layer updates and globally closed-form sub-problem solutions.
- Related work: Unlike MAC, the proposed method is designed with closed-form sub-problems and uses Lagrange multipliers to enforce equality constraints exactly.
- Related work: The paper identifies alternating least squares as a way to distribute computation across a cluster, whereas related work considered multi-threading on one machine.
3. Alternating minimization for neural networks
The method splits neural-network training into constrained variables for weights, pre-activations, and activations, then alternately minimizes each block. Its design makes weight and activation updates least-squares problems and decomposes nonlinear output updates into independent one-dimensional problems with closed-form solutions.
- Problem formulation: The formulation introduces z_l = W_l a_{l−1} and a_l = h_l(z_l) to decouple linear weights from nonlinear link functions.The split problem enforces these relationships through constraints and penalty terms, with Lagrange multipliers added for exact enforcement.
- Algorithm: The algorithm repeatedly minimizes the variable blocks and then updates the Lagrange multipliers.Algorithm 1 takes training features and labels as input and initializes the activation variables and multiplier before repeating the layer updates.
- Alternating scheme: Alternating minimization updates one variable block at a time: {W_l}, {a_l}, or {z_l}, while holding the others fixed.The split formulation is designed so every sub-step has a simple closed-form solution.
- Minimization sub-steps: Weight updates solve linear least-squares problems, with W_l updated using the pseudoinverse of the preceding activation matrix.Activation updates are also least-squares problems, balancing the next-layer linear relation and the nonlinear activation penalty.
- Minimization sub-steps: Because h acts entry-wise, the non-convex output update decomposes into independent one-dimensional problems that can be solved globally in closed form for piecewise-linear activations.For smoother sigmoid curves, precomputed lookup tables can solve these one-dimensional problems quickly.
4. Lagrange multiplier updates via method of multipliers and Bregman iteration
The paper interprets its multiplier update as Bregman iteration closely related to the method of multipliers, combined with alternating minimization. This formulation improves stability, but the resulting multi-block non-convex scheme lies outside established ADMM convergence results.
- Convergence and stability: The multi-block, non-smooth, non-convex formulation lies outside known ADMM convergence results, and conventional separate-multiplier ADMM is highly unstable here.The destabilization is attributed to the large number of coupled non-smooth, non-convex terms.
- Convergence and stability: The Bregman Lagrange update avoids non-smooth constraint terms, and the resulting method is reported to be extremely stable.This stability contrasts with the conventional ADMM application described for the same constrained problem.
- Bregman interpretation: Bregman iteration repeatedly minimizes an objective under an inexact constraint penalty and adds a linear correction to make the quadratic penalty more influential in the next iteration.In this paper, the objective is the loss and the operator contains the network constraints.
- Bregman interpretation: The proposed scheme combines a Bregman Lagrange update with alternating minimization, making it an instance of the split Bregman method.The multiplier update can be interpreted as updating the sub-gradient in the Bregman iterative method.
- Method of multipliers: The paper also views the update as an approximation to the method of multipliers, whose linear-constraint update adds the penalty gradient through A^T(Au − b).The network formulation assigns the loss to J(u) and the constraints to A.
5. Distributed implementation using data parallelism
The distributed implementation uses data parallelism: workers store different training-data columns and independently perform most activation and output updates. Weight updates aggregate reduced matrix products centrally, with transpose reduction lowering communication volume.
- Data parallelism: Data parallelization distributes different subsets of training data across N workers, splitting activation, output, and multiplier matrices by columns.Each worker stores local activations and outputs corresponding to its assigned training samples.
- Communication-free updates: Activation and output optimization sub-steps require no communication and parallelize across workers using local data.Output entries are decoupled, allowing each worker to solve one-dimensional problems in closed form.
- Communication reduction: Transpose reduction reduces transmitted matrix dimensionality before products are sent to the central node.This is especially useful when the number of features is smaller than the number of training samples; the resulting reduction substantially lowers transmitted data.
- Weight update: Weight updates use the pseudoinverse relation W_l ← z_l a_l† and aggregate products computed separately on each worker.The worker-level products are summed through a single reduce operation before the central weight update.
- Weight update: After reduction, the central server computes the inverse-related quantity, updates W_l, and broadcasts the result to worker nodes.Each server maintains a full representation of the weight matrix and can form its local matrix inverse.
6. Implementation details
The implementation uses warm starts, fixed penalty parameters, binary labels with a hinge loss, and Gaussian initialization for activations and outputs. The selected loss yields closed-form minimization sub-problems.
- Initialization and convergence: Warm-starting ADMM with several iterations before updating Lagrange multipliers addresses poor initial iterates from random initialization.The motivation is that convergence theory requires a good minimizer before multiplier updates.
- Penalty parameters: The experiments use γ_i = 10 and β_i = 1 for all reported trials, and this choice worked reliably across a wide range of problems and architectures.The classical ADMM formulation guarantees convergence for any quadratic penalty parameters, according to the passage.
- Loss function: The implementation uses binary class labels, with each output target a_L equal to 1 or 0, and a separable hinge-penalty loss.The hinge terms are defined piecewise according to the target label.
- Loss function: The hinge loss yields minimization sub-problems that are easily solved in closed form.This supports the alternating minimization design used by the training method.
- Initialization and convergence: Activation and output matrices are initialized with i.i.d. Gaussian random variables, while weights require no initialization because they are updated first.Reported results use unit-variance Gaussian variables, and performance appears fairly insensitive to the variance.
7. Experiments
Experiments compare ADMM with SGD, conjugate gradients, and L-BFGS on SVHN and Higgs, measuring scaling and time to accuracy. ADMM scales linearly with cores and is especially competitive on the larger Higgs task.
- Experiments compare optimization time, test accuracy over time, and scaling across cores on SVHN and Higgs classification tasks.The datasets differ substantially in size and difficulty.
- SVHN: On SVHN, ADMM exhibits linear core scaling and competes with GPU-based gradient methods despite CPU-to-CPU communication.The network uses two hidden layers with 100 and 50 ReLU nodes.
- SVHN: 13.3 seconds: ADMM reaches 95% test accuracy on SVHN with 1,024 cores, versus 28.3 seconds for SGD, 3.3 seconds for L-BFGS, and 10.1 seconds for conjugate gradients.These values average 10 runs and follow hyperparameter tuning for the competing methods.
- Higgs: On Higgs, parallelizing over additional cores dramatically decreases optimization time and again exhibits linear scaling.The benchmark is 64% test-set accuracy for a network with one 300-node hidden layer.
- Higgs: 7.8 seconds: ADMM reaches 64% accuracy on Higgs with 7,200 cores, compared with 181 seconds for L-BFGS and 44 minutes for conjugate gradients.SGD never reached 64% accuracy during seven hours of training.
8. Discussion & Conclusion
The paper presents gradient-free neural-network training with linear scaling to thousands of cores. It reports strong performance on extremely large datasets while framing the experiments as a proof of concept rather than evidence of universal superiority.
- The proposed method trains neural networks without gradient steps and scales linearly up to thousands of cores.
- Strong scaling enables the approach to outperform other methods on problems involving extremely large datasets.
- L-BFGS reached 75% test accuracy on Higgs, producing the only superior classifier despite requiring more time than ADMM to reach 64%.
- The experiments cover a fairly narrow range of classification problems and are intended as a proof of concept, not a demonstration of absolute superiority.
- Looking forward: Future work targets recurrent and convolutional networks, while momentum and alternative initialization schemes may improve convergence speed.