Eigenvalue-Corrected Kronecker-Factored Approximate Curvature

Aug 15, 202610 minute read

This is meant to be a readable explanation of KFAC and EKFAC.

KFAC

Kronecker-Factored Approximate Curvature (KFAC), introduced by Martens and Grosse, is a method to approximate the Fisher Information Matrix (FIM). This is often useful for computing influence functions, where we need to invert a large Hessian. This is costly, so we seek to approximate the Hessian via the Gauss-Newton Hessian $J^T H_y J$ .

Fact. The FIM equals the expected Hessian of the negative log-likelihood (where the expectation is taken over $y \sim p(y \mid x; \theta)$ ). This expected Hessian coincides with the Gauss-Newton Hessian $J^T H_y J$ .

So approximating the FIM $\mathbb{E_{x \sim p_{data}(x), \; y \sim p(y \mid x; \theta)}}[(\nabla_{\theta} \log p(y \mid x; \theta))(\nabla_{\theta} \log p(y \mid x; \theta))^{T}]$ as an efficiently invertible matrix allows us to also efficiently invert the Hessian.

Approximating the FIM

Consider the gradient of the logprobs w.r.t. a certain layer $l$ 's weights $\nabla_{W_{l}} \log p(y \mid x, \theta)$ . By the chain rule:

\nabla_{W_{l}} \log p(y \mid x, \theta) = \nabla_{W_{l} a_{l-1}} \log p(y \mid x, \theta) \cdot \nabla_{W_{l}} W_{l} a_{l-1} \overset{*}{=} \nabla_{W_{l} a_{l-1}} \log p(y \mid x, \theta) \cdot a_{l-1}^{T}

where $*$ follows from the fact that $\nabla_{W_{l}} a_{l-1} = 0$

Now define $\theta_{l} = \operatorname{vec}(W_l)$ and $z_{l} = W_{l} a_{l-1}$ . So we have:

\nabla_{\theta_{l}} \log p (y \mid x, \theta) = \operatorname{vec}(\nabla _{W_{l}} \log p(y \mid x; \theta))

Substituting,

\nabla_{\theta_{l}} \log p (y \mid x, \theta) = \operatorname{vec}(\nabla_{z_l} \log p(y \mid x, \theta) \cdot a_{l-1}^{T})

Lemma. $\operatorname{vec}(u \cdot v^{T}) = v \otimes u$ when $u \in \mathbb{R}^m$ and $v \in \mathbb{R}^n$ .

This is easy to verify yourself! Applying the lemma,

\nabla_{\theta_{l}} \log p(y \mid x; \theta) = a_{l-1} \otimes \nabla_{z_{l}} \log p(y \mid x; \theta)

As I mentioned previously, the Fisher Information Matrix for layer $l$ is given by:

F_l = \mathbb{E}[(\nabla_{\theta_{l}} \log p(y \mid x; \theta))(\nabla_{\theta_{l}} \log p(y \mid x; \theta))^{T}]

Substituting our expression for the gradient,

F_l = \mathbb{E}[(a_{l-1} \otimes \nabla_{z_{l}} \log p(y \mid x; \theta)) \cdot (a_{l-1} \otimes \nabla_{z_{l}} \log p(y \mid x; \theta))^T]

We can write:

F_l = \mathbb{E}[a_{l-1}a_{l-1}^T \otimes (\nabla_{z_{l}} \log p(y \mid x; \theta)) (\nabla_{z_{l}}\log p(y \mid x; \theta))^T]

Here we used the mixed product property $(A \otimes B) \cdot (C \otimes D) = AC \otimes BD$ and that $(u \otimes v)^T = u^T \otimes v^T$ .

To simplify this expression, we make the following assumption: $a_{l-1}$ and $\nabla_{z_{l}} \log p(y \mid x; \theta)$ are independent, i.e. when taking the expected value over many examples, $a_{l-1}$ doesn't provide much information about $\nabla_{z_{l}} \log p(y \mid x; \theta)$ .

Remark. This assumption is not accurate in practice, and can cause trouble. This is why we need eigenvalue correction!

This assumption allows us to separate the expected value:

F_l = \mathbb{E}[a_{l-1}a_{l-1}^T] \otimes \mathbb{E}[(\nabla_{z_{l}} \log p(y \mid x; \theta)) (\nabla_{z_{l}}\log p(y \mid x; \theta))^T]

Designating the first matrix $A_{l-1}$ and the second $B_{l}$ , we have:

F_l = A_{l-1} \otimes B_{l}

Note that we have obtained the FIM only for a single layer $l$ . The FIM for the entire network is a block matrix:

F = \begin{bmatrix} F_{11} & F_{12} & F_{13} \\ F_{21} & F_{22} & F_{23} \\ F_{31} & F_{32} & F_{33}\end{bmatrix}

Here, we make another assumption: that $F$ is block-diagonal, i.e. $F_{ij} = 0$ when $i \neq j$ . This means that for any two layers $l$ and $k$ , we are assuming the following:

\mathbb{E}[(\nabla_{\theta_{l}} \log p(y \mid x; \theta)) (\nabla_{\theta_{k}}\log p(y \mid x; \theta))^T] = 0

Then we can assemble $F$ as:

F = \operatorname{blockdiag}(F_1, F_2, ..., F_L)

Eigenvalue-Corrected KFAC

(Spectral Theorem for Real Symmetric Matrices). Let $V$ be a real inner product space with $\dim(V) < \infty$ , and $S:V \to V$ be a symmetric linear map so that $S=S^T$ . Then $S$ is normal ( $SS^T = S^2 = S^T S$ ) and admits an orthonormal basis of eigenvectors.

Remark. In this section, I write $A = A_{l-1}$ and $B = B_l$ for simplicity's sake.

Notice that $A$ and $B$ are real symmetric matrices. Therefore, by the spectral theorem, they admit orthonormal bases of eigenvectors, which means we can diagonalize them:

A = P_A \Lambda_A P_A^{T} \qquad B = P_B \Lambda_B P_B^{T}

where $P_A$ and $P_B$ are orthogonal, i.e. $P_A \cdot P_A^T = P_A^T \cdot P_A = I$ i.e. $P_A^T = P_A^{-1}$

Substituting,

F_l = (P_A \Lambda_A P_A^{T}) \otimes (P_B \Lambda_B P_B^{T})

Applying the same identities as before,

F_l = (P_A \otimes P_B) (\Lambda_A \otimes \Lambda_B) (P_A^{T} \otimes P_B^{T}) = (P_A \otimes P_B) (\Lambda_A \otimes \Lambda_B) (P_A \otimes P_B)^{T}

This means that the matrix $\Lambda_A \otimes \Lambda_B$ contains the eigenvalues of $F_l$ .

Problem. What happens if these eigenvalues are wrong because the independence assumption we made earlier was imperfect?

Solution. We compute the eigenvalues of $F_l$ directly, i.e. a matrix $\Lambda$ which is not necessarily equal to $\Lambda_A \otimes \Lambda_B$ .

To do this, we project the gradient onto the columns of $P_A \otimes P_B$ , i.e. our (assumed) eigenbasis for $F_l$ and take the expected value of the squared projection.

First, recall the expression for the general quadratic form of a column vector $x$ : $q(x) = x^T M x$ . Now with $M = F_l$ , we have:

x^T F_lx = \mathbb{E}[x^T \nabla_{\theta_l} \log p(y \mid x; \theta) \nabla_{\theta_l} \log p(y \mid x; \theta)^T x] = \mathbb{E}[(x^T \nabla_{\theta_l} \log p(y \mid x; \theta))^2]

Note that for an eigenvector $x$ of $M$ such that $\lVert x \rVert = 1$ , $x^T M x = \lambda x^T x = \lambda$ . Since the columns of the matrix $P_A \otimes P_B$ are orthonormal eigenvectors, we can assemble $\Lambda$ as follows:

\Lambda_{ii} = \mathbb{E} \left[ \left ([(P_A \otimes P_B)^{T} \nabla_{\theta_l} \log p(y \mid x; \theta )]_{i} \right)^2 \right]

This allows us to keep our convenient assumption about the eigenbasis while correcting the inaccurate eigenvalues.