Understanding K-FAC Optimization in Deep Learning

This article explains how Kronecker-Factored Approximate Curvature (K-FAC) makes second-order optimization practical for deep artificial intelligence models. You will learn why traditional Hessian matrices are too large to compute, how K-FAC substitutes the Hessian with the Fisher Information Matrix, and how it uses Kronecker products to break down massive calculations into small, fast operations.

The Problem with Second-Order Optimization

In deep learning, training a model requires minimizing a loss function. First-order methods, like Stochastic Gradient Descent (SGD), only look at the slope (gradient) of the loss. Second-order methods also look at the curvature of the loss using the Hessian matrix. By understanding the curvature, an optimizer can take much smarter, larger steps toward the minimum, drastically reducing the number of training iterations needed.

However, the Hessian matrix requires computing the second derivative for every pair of weights in a neural network. If a network has 10 million parameters, the Hessian matrix will have 100 trillion elements. Storing and inverting a matrix of this size is computationally impossible on modern hardware.

Step 1: Replacing the Hessian with the Fisher Matrix

Instead of directly approximating the massive Hessian matrix, K-FAC approximates the Fisher Information Matrix (FIM). In probabilistic models—which include almost all deep neural networks—the Fisher Information Matrix acts as a natural stand-in for the Hessian. It captures the exact same local curvature information of the loss landscape while guaranteeing that the matrix is positive semi-definite, a property that prevents the optimizer from stepping in the wrong direction.

Step 2: Block-Diagonal Approximation

Even with the Fisher matrix, the size remains a problem. To shrink it, K-FAC assumes that the weights in different layers of the neural network are statistically independent of one another.

While this is not entirely true in practice, it is a good enough approximation for optimization. By ignoring the correlations between different layers, K-FAC zeroes out most of the Fisher matrix, turning it into a "block-diagonal" matrix. Now, instead of inverting one massive matrix for the whole network, the algorithm only needs to invert one smaller matrix per layer.

Step 3: The Kronecker Factorization

Even a single layer's Fisher block can be too large to invert. If a layer has 1,000 inputs and 1,000 outputs, it has 1 million weights, resulting in a Fisher block of 1 trillion elements. This is where the core math of K-FAC comes in.

The gradient of a weight in a specific layer is simply the product of two things:

  1. The activation coming into the layer (forward pass).
  2. The error gradient coming back into the layer (backward pass).

K-FAC makes a second major assumption: it assumes that the incoming activations and the backward error gradients are statistically independent.

Because of this independence, the massive Fisher block for a single layer can be rewritten as the Kronecker product of two much smaller matrices:

  • Matrix A: The covariance matrix of the layer's inputs (size: 1,000 x 1,000).
  • Matrix G: The covariance matrix of the layer's output gradients (size: 1,000 x 1,000).

Why This Solves the Computation Bottleneck

The magic of the Kronecker product lies in a specific mathematical identity: the inverse of a Kronecker product is equal to the Kronecker product of the inverses.

Instead of inverting the 1,000,000 x 1,000,000 Fisher block, K-FAC simply inverts Matrix A (1,000 x 1,000) and Matrix G (1,000 x 1,000) separately. Inverting these two small matrices takes a fraction of a millisecond on a modern GPU. Once inverted, K-FAC uses them to scale the standard gradients, effectively applying second-order curvature information to update the network's weights with minimal memory and compute overhead.