[Paper Review] Fast Approximation of the Gauss-Newton Hessian Matrix for the Multilayer Perceptron

Chao Chen, Severin Reiz+3 more|arXiv (Cornell University)|Oct 27, 2019

Stochastic Gradient Optimization Techniques46 references4 citations

TL;DR

This paper presents a fast Monte Carlo sampling algorithm for entry-wise approximation of the Gauss-Newton Hessian (GNH) matrix in multilayer perceptrons, reducing computational cost from 𝒪(Nn) to 𝒪(n + d/ε²) per entry. The method enables efficient hierarchical-matrix (ℋ-).

ABSTRACT

We introduce a fast algorithm for entry-wise evaluation of the Gauss-Newton Hessian (GNH) matrix for the fully-connected feed-forward neural network. The algorithm has a precomputation step and a sampling step. While it generally requires $O(Nn)$ work to compute an entry (and the entire column) in the GNH matrix for a neural network with $N$ parameters and $n$ data points, our fast sampling algorithm reduces the cost to $O(n+d/ε^2)$ work, where $d$ is the output dimension of the network and $ε$ is a prescribed accuracy (independent of $N$). One application of our algorithm is constructing the hierarchical-matrix (H-matrix) approximation of the GNH matrix for solving linear systems and eigenvalue problems. It generally requires $O(N^2)$ memory and $O(N^3)$ work to store and factorize the GNH matrix, respectively. The H-matrix approximation requires only $O(N r_o)$ memory footprint and $O(N r_o^2)$ work to be factorized, where $r_o \ll N$ is the maximum rank of off-diagonal blocks in the GNH matrix. We demonstrate the performance of our fast algorithm and the H-matrix approximation on classification and autoencoder neural networks.

Motivation & Objective

• Develop a fast, scalable method for entry-wise evaluation of the Gauss-Newton Hessian (GNH) matrix in fully connected feed-forward neural networks.
• Reduce the computational cost of GNH matrix evaluation from 𝒪(Nn) to 𝒪(n + d/ε²) per entry, independent of the number of parameters N.
• Enable efficient construction of hierarchical-matrix (ℋ-matrix) approximations of the GNH for solving linear systems and eigenvalue problems.
• Demonstrate the method's effectiveness on classification and autoencoder networks using MNIST and CIFAR-10 datasets.
• Provide theoretical guarantees on sampling accuracy using McDiarmid’s inequality and validate convergence in numerical experiments.

Proposed method

• Introduce a two-stage algorithm: a precomputation phase to store intermediate Jacobian and Hessian terms, and a sampling phase to estimate GNH entries via randomized sampling.
• Leverage the structure of the GNH matrix as H = JᵀQJ, where J is the Jacobian of network outputs w.r.t. weights and Q is the Hessian of the loss function.
• Use Monte Carlo sampling to approximate individual entries of H by drawing random data points and computing stochastic estimates of JᵢᵀQᵢJᵢ.
• Apply McDiarmid’s inequality to bound the sampling error, ensuring convergence to the true GNH entry with high probability.
• Construct ℋ-matrix approximations of the GNH using sampled entries, reducing memory from 𝒪(N²) to 𝒪(Nrₒ) and factorization cost from 𝒪(N³) to 𝒪(Nrₒ²), where rₒ ≪ N is the off-diagonal block rank.
• Integrate the sampling algorithm into the ℋ-matrix construction pipeline, enabling fast and accurate low-rank approximations of the GNH matrix.

Experimental results

Research questions

• RQ1 Can we significantly reduce the computational cost of evaluating individual entries in the Gauss-Newton Hessian matrix for deep neural networks?
• RQ2 Does a randomized sampling approach based on precomputed terms achieve high accuracy with sublinear dependence on the number of parameters N?
• RQ3 Can the proposed sampling method enable efficient and accurate hierarchical-matrix (ℋ-matrix) approximation of the GNH matrix?
• RQ4 How does the sampling error scale with the number of samples and desired accuracy ε, and can theoretical bounds be established?
• RQ5 Can the resulting ℋ-matrix approximation be used effectively for solving linear systems and eigenvalue problems involving the GNH matrix in practice?

Key findings

• The proposed algorithm reduces the cost of evaluating a single GNH matrix entry from 𝒪(Nn) to 𝒪(n + d/ε²), achieving a significant speedup for large networks.
• For the MNIST autoencoder, the method achieved a 1.68% compression rate with a Frobenius error of 6.1×10⁻² using 10,000 samples, and the error decreased with more samples.
• On CIFAR-10, the ℋ-matrix approximation achieved a 4.83% compression rate with a Frobenius error of 7.3×10⁻² at high accuracy using 10,000 samples.
• The sampling error converged to zero as the number of samples increased, and the ℋ-matrix approximation error also decreased, validating the overall approach.
• The method outperformed uniform sampling by a factor of 100 in terms of required samples for equivalent accuracy, demonstrating superior sampling efficiency.
• Preliminary experiments on VGG networks suggest the method generalizes to convolutional architectures, indicating potential for broader applicability.