PaperlayerView exact revision on arXiv

Source-linked AI summary

The Hessian of tall-skinny networks is easy to invert

Ali Rahimi

arXiv:2601.06096v3cs.LG

TL;DR

Explicit Hessian inversion is impractical for large deep nets, and approximations or iterative inverse-vector methods have limitations. The paper gives an exact layerwise method for Hessian-inverse-vector products without forming the Hessian. Its storage and computation scale linearly with depth, unlike the naive approach, though the proposed LDU solve is prone to numerical instability.

  • Problem

    For large deep nets, storing or inverting the Hessian is impractical, and Krylov-based inverse-vector computation depends on the Hessian’s poor conditioning.

    Source: §1 · p. 1

  • Method

    The algorithm exploits layerwise structure to compute Hessian-inverse-vector products without forming or storing the Hessian.

    Source: §1 · pp. 1–2

  • Results

    Storage scales quadratically and computation cubically in the largest per-layer dimension, but both scale linearly with layer count, unlike the naive method’s cubic depth cost.

    Source: §1 · p. 2, §7 · p. 6

  • Takeaways & Limitations

    The depth-scaling advantage is greatest for tall, skinny networks, and the authors hope to use the method as a preconditioner for stochastic gradient descent.

    Source: §1 · p. 2, §7 · p. 6

    The proposed LDU factorization can be numerically unstable for large linear systems; the paper suggests an off-the-shelf banded solver as an improvement.

    Source: §7 · p. 6

Abstract

from arXiv · show

We describe an exact algorithm to solve linear systems of the form $Hx=b$ where $H$ is the Hessian of a deep net. The method computes Hessian-inverse-vector products without storing the Hessian or its inverse. It requires time and storage that scale linearly in the number of layers. This is in contrast to the naive approach of first computing the Hessian, then solving the linear system, which takes storage and time that are respectively quadratic and cubic in the number of layers. The Hessian-inverse-vector product method scales roughly like Pearlmutter's algorithm for computing Hessian-vector products.

1 Introduction

For modern deep nets, explicitly forming and inverting the Hessian is impractical, while common low-rank or diagonal approximations are inexact. The paper presents an exact Hessian-inverse-vector method whose depth dependence is linear, making it especially efficient for tall, skinny networks.

1 Introduction

  • A billion-parameter model’s Hessian would have a quintillion entries, making it impractical to store, multiply, or invert.Low-rank and diagonal approximations are common workarounds.Source: §1 · p. 1
  • Krylov iterations can use Hessian-vector products to obtain inverse-vector products, but their result depends on the Hessian’s notoriously poor conditioning.Pearlmutter’s Hessian-vector product scales linearly with network depth, but no direct inverse-vector variant had been identified.Source: §1 · p. 1
  • The proposed exact algorithm exploits layerwise structure, representing the Hessian through layer derivatives and an inverse block bi-diagonal backpropagation operator.Multiplication by this structure computes Hessian-vector products using the same operations as Pearlmutter’s algorithm.Source: §1 · p. 1
  • Storage is O(L max(a, p)^2) and computation is O(L max(a, p)^3), with only linear dependence on layer count L.Here, each layer has at most p parameters and a activations; the per-layer dependence remains cubic for computation.Source: §1 · pp. 1–2
  • The method ameliorates the naive algorithm’s cubic depth dependence to linear dependence, so it runs faster on tall, skinny networks.The paper notes that modern networks are typically short and wide and proposes tall-skinny networks as a potential preconditioning application.Source: §1 · p. 2

2 Overview

The overview’s goal is to solve Hx = b without explicitly forming the Hessian. It outlines a lifted, permuted block-system approach that reduces the solve to block-LDU factorization and substitutions.

2 Overview

  • The objective is to solve Hx = b, with H the neural-network Hessian, without explicitly forming H.The algorithm’s stated input is b, with damping ϵ, and its output is x = (H + ϵI)^−1b.Source: §2 · p. 2
  • The method writes backpropagation as a block bi-diagonal system, differentiates it to express the Hessian, then lifts that expression into a linear system.The lifted system is pivoted into block tri-diagonal form.Source: §2 · p. 2
  • After pivoting, the block tri-diagonal system is factorized into lower, diagonal, and upper block matrices using LDU.Forward and backward substitution solve the resulting system.Source: §2 · p. 2
  • The solution is un-pivoted and un-lifted to recover x.The system can be represented as a permuted block-tri-diagonal system before solving.Source: §2 · p. 2
  • The algorithm is described as straightforward to implement with a suitable library of block-matrix operations.Algorithm 2 summarizes its steps.Source: §2 · p. 2

3 Notation

The notation represents a deep net as a pipeline of layer functions, with parameters and intermediate activations carried through to a scalar loss. The derivations assume differentiable layer operations and specify dimensions for complexity accounting.

3 Notation

  • The network is written as a pipeline of functions fℓ, with z0 as input and zℓ as the output of layer ℓ.The vectors x1 through xL are the layer parameters.Source: §3 · p. 3
  • Training labels are subsumed in z0 and propagated through the layers, while the final layer fL combines the last layer with the training loss.The scalar zL is the model’s loss on input z0.Source: §3 · p. 3
  • The derivations assume first and partial derivatives exist, largely setting aside nondifferentiable operations such as ReLU by assuming differentiable approximations.The text notes that symmetry of the Hessian requires differentiability everywhere.Source: §3 · p. 3
  • For complexity accounting, all but the last layer are assumed to have at most a-dimensional activations and p-dimensional parameters.The layer shapes are otherwise unrestricted by the derivations.Source: §3 · p. 3

4 Backpropagation, the matrix way

Backpropagation is expressed as a backward recursion for the stochastic gradient, then rewritten as a block bi-diagonal linear system. Its triangular structure supports solving by back-substitution.

4 Backpropagation, the matrix way

  • The parameter gradient for stochastic input z0 is computed by the chain rule using a backward recursion.The recursion has base case bL = 1.Source: §4 · p. 3
  • The recursion is written in vector form as a block bi-diagonal system, with identity diagonal blocks and layer-derivative off-diagonal blocks.M is invertible because its triangular spectrum consists of ones on the diagonal.Source: §4 · p. 4
  • Solving the system for b and substituting back yields the gradient expression.The system’s solution uses back-substitution, coinciding with backpropagation computations.Source: §4 · p. 4

5 The Hessian

The section defines the Hessian through changes in the gradient and describes computing Hessian-vector products without forming the matrix. It contrasts repeated Hessian-vector products for solving systems with a direct inverse-vector approach.

5 The Hessian

  • The Hessian is the matrix mapping infinitesimal parameter changes to changes in the gradient.It is defined as the unique matrix H(x) such that dg approaches H(x) dx as dx approaches zero.Source: §5 · p. 4
  • The Hessian-vector product can be computed without forming H, using the same operations as Pearlmutter’s trick.The cost is dominated by multiplying by the Dxx, Dzx, and Dzz matrices.Source: §5 · p. 5
  • Krylov methods solve Hx = b through repeated Hessian multiplications, with the number of applications depending on H’s condition number.The next section presents a way to solve the inverse system with only max(a, p) times more operations than computing Hx.Source: §5 · p. 5

6 Applying the inverse of the Hessian

Although the Hessian is dense, its polynomial structure can be lifted into a larger system that is easier to solve. Reordering blocks makes that system block-tridiagonal, enabling layerwise substitution.

6 Applying the inverse of the Hessian

  • The Hessian is a second-order matrix polynomial in M^-1, which is dense even though M is block-bidiagonal.The polynomial can be lifted into a higher-order object whose inverse is easy to compute.Source: §6 · p. 5
  • Auxiliary variables convert the damped system (H + ϵI)x = g into a larger block-linear system whose solution yields x.The block-linear system can be pivoted into block-tridiagonal form, which is more efficient to solve than Gaussian elimination on K or H.Source: §6 · p. 5
  • A permutation that transposes blocks with sub-blocks turns K into a block-tridiagonal system.The reordering groups x, y, and z variables layer by layer; the permutation is its own inverse.Source: §6 · pp. 5–6
  • Block LDU decomposition followed by back and forward substitution solves the reordered system in time linear in network depth and cubic in block dimension.Block dimensions depend on the layer’s parameter or activation count.Source: §6 · p. 6

7 Conclusion

The method computes inverse-Hessian-vector products without storing the Hessian, reducing depth dependence relative to the naive method. Its final block-tridiagonal solve can be numerically unstable for large systems; the authors suggest a banded solver and hope to use the technique as a preconditioner.

7 Conclusion

  • The method avoids storing the Hessian and reduces computation’s depth dependence from cubic for the naive method to linear.It computes the product of the inverse Hessian with a vector.Source: §7 · p. 6
  • Large linear systems solved with LDU factorization are prone to numerical instability; an off-the-shelf banded solver is proposed as an improvement.The final stage uses forward and backward substitution to solve a block-tridiagonal system.Source: §7 · p. 6
  • The authors hope to use the technique as a preconditioner for stochastic gradient descent, with possible renewed interest in extremely deep architectures.They expect the greatest speedup for tall, skinny networks.Source: §7 · p. 6

A The Hessian

The appendix derives a gradient expression using the derivative of M^-1 and defines a convention for gradients of tensor-valued functions. That convention also accommodates Hessians of vector-valued functions and the chain rule.

A The Hessian

  • The derivative identity for M^-1 is used to express the differential of Dx in terms of dDx and dM.The derivation also uses b = eL M^-1.Source: A The Hessian · p. 7
  • For tensor-valued functions, the gradient is defined as the matrix satisfying the stated limiting relation between changes in g and dx.The convention defines the Hessian of a vector-valued function and accommodates the chain rule.Source: A The Hessian · p. 7
  • For restrictions y and z of x to coordinate subsets, the mixed derivative ∇yzg is defined as ∇y∇zg.Its stated shape is R^(|g||z|)×|y|.Source: A The Hessian · p. 7

A.1 The term involving dDx

The section derives how changes in layer derivative blocks depend on parameter and activation changes, then combines these terms into the Hessian form.

A.1 The term involving dDx

  • Each dDℓ varies with parameter changes through both xℓ and the preceding activation zℓ−1, whose differential follows a layerwise recurrence.The recurrence starts from dz0 = 0 and combines the direct parameter change with the propagated activation change.Source: §A.1 · p. 8

A.2 The term involving dM

  • Each dMℓ likewise varies through changes in xℓ and zℓ−1, extending the derivative calculation to the blocks involving ∇zfℓ.The passage states d vec(Mℓ) = (∇xzfℓ) dxℓ + (∇zzfℓ) dzℓ−1.Source: §A.2 · p. 9

A.3 Putting it all together

  • The preceding derivative terms are assembled into the Hessian form for the deep net.Source: §A.3 · p. 9

B Equation (9) is Pearlmutter’s Hessian-vector multiplication algorithm

Equation (9) specializes Pearlmutter’s Hessian-vector multiplication machinery to pipeline networks, with operations that coincide exactly. The pivoting examples show how permutations change a matrix’s block structure.

B Equation (9) is Pearlmutter’s Hessian-vector multiplication algorithm

  • Equation (9) computes Hessian-vector products in time linear in L and its operations coincide exactly with Pearlmutter’s algorithm for pipeline networks.The section identifies this as a specialization of Pearlmutter’s framework for networks with arbitrary topologies.Source: B Equation (9) is Pearlmutter’s… · pp. 9–10

C Examples of Pivoting Matrices to Block-Banded Form

Loading 2601.06096v3…