[Paper Review] Training (Overparametrized) Neural Networks in Near-Linear Time

Jan van den Brand, Binghui Peng+2 more|arXiv (Cornell University)|Jun 20, 2020

Stochastic Gradient Optimization Techniques89 references23 citations

TL;DR

This paper presents a near-linear time algorithm for training overparametrized ReLU neural networks using a novel reformulation of the Gauss-Newton method as an ℓ2-regression problem, accelerated via Fast-Johnson-Lindenstrauss (Fast-JL) dimension reduction to precondition the Gram matrix. The method achieves eO(mnd + n³) per-iteration cost and eO((mnd + n³) log(1/ǫ)) total runtime for training loss reduction to ǫ, marking a quadratic speedup over prior second-order methods and demonstrating the feasibility of advanced randomized linear algebra in deep learning optimization.

ABSTRACT

The slow convergence rate and pathological curvature issues of first-order gradient methods for training deep neural networks, initiated an ongoing effort for developing faster $\mathit{second}$-$\mathit{order}$ optimization algorithms beyond SGD, without compromising the generalization error. Despite their remarkable convergence rate ($\mathit{independent}$ of the training batch size $n$), second-order algorithms incur a daunting slowdown in the $\mathit{cost}$ $\mathit{per}$ $\mathit{iteration}$ (inverting the Hessian matrix of the loss function), which renders them impractical. Very recently, this computational overhead was mitigated by the works of [ZMG19,CGH+19}, yielding an $O(mn^2)$-time second-order algorithm for training two-layer overparametrized neural networks of polynomial width $m$. We show how to speed up the algorithm of [CGH+19], achieving an $ ilde{O}(mn)$-time backpropagation algorithm for training (mildly overparametrized) ReLU networks, which is near-linear in the dimension ($mn$) of the full gradient (Jacobian) matrix. The centerpiece of our algorithm is to reformulate the Gauss-Newton iteration as an $\ell_2$-regression problem, and then use a Fast-JL type dimension reduction to $\mathit{precondition}$ the underlying Gram matrix in time independent of $M$, allowing to find a sufficiently good approximate solution via $\mathit{first}$-$\mathit{order}$ conjugate gradient. Our result provides a proof-of-concept that advanced machinery from randomized linear algebra -- which led to recent breakthroughs in $\mathit{convex}$ $\mathit{optimization}$ (ERM, LPs, Regression) -- can be carried over to the realm of deep learning as well.

Motivation & Objective

• To address the high computational cost of second-order optimization in deep learning, particularly the O(mn²) per-iteration cost of prior Gauss-Newton methods for overparametrized networks.
• To extend second-order optimization to ReLU networks, which are more complex and realistic than smooth-activation networks.
• To achieve near-linear time training in the dimension of the full gradient (mn), overcoming the bottleneck of Hessian inversion in traditional second-order methods.
• To demonstrate that advanced randomized linear algebra techniques—previously used in convex optimization—can be effectively transferred to non-convex deep learning settings.

Proposed method

• Reformulate the Gauss-Newton update as an ℓ2-regression problem on the Jacobian matrix to enable efficient solution via conjugate gradient.
• Apply Fast-Johnson-Lindenstrauss (Fast-JL) type dimension reduction to precondition the Gram matrix J_t J_t^T, reducing its size while preserving solution quality.
• Use first-order conjugate gradient to solve the preconditioned regression problem in time independent of the original matrix size M.
• Leverage Neural Tangent Kernel (NTK) theory to justify local linearization of the overparametrized network, enabling reduction to kernel regression.
• Ensure convergence by bounding the error in the Jacobian approximation and the solution to the regression subproblem using concentration inequalities.
• Integrate the regression solver into a backpropagation framework with per-iteration cost eO(mnd + n³), dominated by Jacobian evaluation and regression solve.

Experimental results

Research questions

• RQ1 Can second-order optimization for overparametrized ReLU networks be made practical by reducing the per-iteration cost below O(mn²)?
• RQ2 Can Fast-JL dimension reduction be used to precondition the Gauss-Newton Hessian approximation in a way that preserves convergence guarantees?
• RQ3 Is it possible to achieve near-linear time training in the dimension of the full gradient (mn) for two-layer ReLU networks?
• RQ4 Can the machinery of randomized linear algebra—successful in convex optimization—be adapted to non-convex deep learning training?
• RQ5 Does the proposed algorithm maintain generalization performance while achieving significantly faster convergence than SGD?

Key findings

• The algorithm achieves a quadratic speedup over the O(mn²) method of [CGH+19], reducing per-iteration cost to eO(mnd + n³).
• The method is the first second-order algorithm for ReLU networks with near-linear time complexity in the gradient dimension (mn), under mild overparametrization (m = Ω(max{λ⁻⁴n⁴, λ⁻²n²d log(n/δ)})).
• The algorithm ensures ∥ft+1 − y∥² ≤ ½∥ft − y∥² per iteration, guaranteeing linear convergence to the target loss.
• The total running time to reduce training loss below ǫ is eO((mnd + n³) log(1/ǫ)), which becomes eO((mnd + n^ω) log(1/ǫ)) with fast matrix multiplication.
• The algorithm maintains bounded weight updates, ensuring that network weights do not deviate significantly from initialization, supporting generalization.
• The method is successfully applied to convex optimization, improving Newton method runtime to eO((nd log(κ) + d³) log(1/ǫ)) for γ-strongly convex, β-smooth functions with L-Lipschitz Hessian.