KFAC explained

Let's consider a linear layer (without bias for simplicity)

\begin{align*} \vz^{(l)} = \mW^{(l)} \vz^{(l-1)} \end{align*}

with parameters \(\vtheta^{(l)} = \vec \mW^{(l)}\). The output-parameter Jacobian is

\begin{align*} \jac_{\vtheta^{(l)}}\vz^{(l)} = {\vz^{(l-1)}}^{\top} \otimes \mI_{\dim(\vz^{(l)})} \end{align*}

and has Kronecker structure. This Kronecker structure carries through to the parameter Hessian: insertion into the Hessian backpropagation equation \eqref{orglatexenvironment1} yields

\begin{align*} \nabla_{\vtheta^{(l)}}^{2}\ell &= \left[ {\vz^{(l-1)}}^{\top} \otimes \mI_{\dim(\vz^{(l)})} \right]^{\top} \nabla_{\vz^{(l)}}^{2}\ell \left[ {\vz^{(l-1)}}^{\top} \otimes \mI_{\dim(\vz^{(l)})} \right] \\ &= \vz^{(l-1)} {\vz^{(l-1)}}^{\top} \otimes \nabla_{\vz^{(l)}}^{2}\ell\,. \end{align*}

In KFAC, the second Kronecker factor is approximated through an outer product of a backpropagated vector \(\vg^{(l)}\) with \(\dim(\vg^{(l)}) = \dim(\vz^{(l)})\). The details of this vector don't matter here; just consider it "one way" to approximate \(\nabla_{\vz^{(l)}}^{2}\ell \approx \vg^{(l)} {\vg^{(l)}}^{\top}\).

We have arrived at the Kronecker-factored curvature approximation for linear layers:

\begin{align} \label{orglatexenvironment2} \mathrm{KFAC}(\nabla_{\vtheta^{(l)}}^{2}\ell) = \vz^{(l-1)} {\vz^{(l-1)}}^{\top} \otimes \vg^{(l)}{\vg^{(l)}}^{\top}\,. \end{align}

Next up, we consider a convolution layer (without bias for simplicity). It takes an image-shaped tensor \(\tZ^{(l-1)} \in\sR^{C_{\text{in}} \times H_{\text{in}} \times W_{\text{in}}}\) and maps it to an image-shaped output \(\tZ^{(l)} \in\sR^{C_{\text{out}} \times H_{\text{out}} \times W_{\text{out}}}\). The kernel is a rank-4 tensor \(\tW^{(l)} \in \sR^{C_{\text{out}} \times C_{\text{in}} \times K_{H} \times K_{W}}\), and we can write

\begin{align*} \tZ^{(l)} = \tZ^{(l-1)} \star \tW^{(l)} \end{align*}

with \(\vtheta^{(l)} = \vec \tW^{(l)}\). Instead of tensors, we will view the forward pass in terms of the matrices \(\mW^{(l)} \in \sR^{C_{\text{out}} \times C_{\text{in}} K_{H} K_{W}}\) and \(\mZ^{(l)} \in\sR^{C_{\text{out}} \times H_{\text{out}} W_{\text{out}}}\), which are just reshapes of the above tensors. With these matrices, convolution can be expressed as matrix multiplication

\begin{align*} \mZ^{(l)} = \mW^{(l)} \llbracket\tZ^{(l-1)}\rrbracket \end{align*}

where \(\llbracket\tZ^{(l-1)}\rrbracket \in \sR^{C_{\text{in}} K_{H} K_{W} \times H_{\text{out}} W_{\text{out}}}\) is the unfolded input, sometimes also referred to as im2col. We don't have to worry how to obtain the unfolded input; deep learning libraries provide such functionality. Now, the forward pass looks already very similar to the linear layer.

To use the Hessian backpropagation equation \eqref{orglatexenvironment1}, we need to flatten the output into a vector, \(\vz^{(l)} = \vec \mZ^{(l)} \in \sR^{C_{\text{out}} H_{\text{out}} W_{\text{out}}}\). For this vector, the Jacobian is given by

\begin{align*} \jac_{\vtheta^{(l)}}\vz^{(l)} = \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \mI_{C_{\text{out}}}\,. \end{align*}

Inserting this, we get

\begin{align*} \nabla_{\vtheta^{(l)}}^{2} \ell &= \left( \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \mI_{C_{\text{out}}} \right)^{\top} \nabla_{\vz^{(l)}}^{2} \ell \left( \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \mI_{C_{\text{out}}} \right) \\ &= \left( \llbracket \tZ^{(l-1)} \rrbracket \otimes \mI_{C_{\text{out}}} \right) \nabla_{\vz^{(l)}}^{2} \ell \left( \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \mI_{C_{\text{out}}} \right)\,. \end{align*}

Unlike for the linear layer, this cannot be simplified to a Kronecker product without further approximations! (Grosse & Martens, 2016) present the assumptions under which the above expression simplifies to a Kronecker product.

I want to take a different approach here. It is relatively easy to see that if the \(C_{\text{out}} H_{\text{out}} W_{\text{out}} \times C_{\text{out}} H_{\text{out}} W_{\text{out}}\) matrix \(\nabla_{\vz^{(l)}}^{2} \ell\) could be expressed as a Kronecker product,

\begin{align*} \nabla_{\vz^{(l)}}^{2}\ell = \mA \otimes \mB\,, \qquad \mA \in \sR^{H_{\text{out}} W_{\text{out}} \times H_{\text{out}} W_{\text{out}}}, \mB \in \sR^{C_{\text{out}} \times C_{\text{out}}}\,, \end{align*}

this would imply a Kronecker structure for the weight Hessian,

\begin{align*} \implies \nabla_{\vtheta^{(l)}}^{2}\ell = \llbracket \tZ^{(l-1)} \rrbracket \mA \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \mB\,. \end{align*}

If we now fix \(\mA = \mI_{H_{\text{out}} W_{\text{out}}}\) (see Figure 2), and determine \(\mB\) by finding the best possible approximation to \(\nabla_{\vtheta^{(l)}}^{2}\ell\) (see Figure 3), we recover the Kronecker factors for convolutions as presented in (Grosse & Martens, 2016).

Concretely, we will minimize the squared Frobenius norm between \(\nabla_{\vz^{(l)}}^{2}\ell\) and \(\mI_{H_{\text{out} W_{\text{out}}}} \otimes \mB\) to determine the best choice for \(\mB\). From a high-level perspective, \(\mB\) is the mean over the \(C_{\text{out}} \times C_{\text{out}}\) diagonal blocks of \(\nabla_{\vz^{(l)}}^{2} \ell\), see Figure 3. To see this in more detail, let's write down the elements of \(\mI_{H_{\text{out} W_{\text{out}}}} \otimes \mB\). To address elements, we require 4 indices to distinguish channels from spatial dimensions. This yields

\begin{align*} \left[ \mI_{H_{\text{out}} W_{\text{out}}} \otimes \mB \right]_{(x_1, c_1),(x_2, c_2)} = \delta_{x_1,x_2} \left[ \mB \right]_{c_1,c_2}\,, \qquad x_{1,2} = 1, \dots, H_{\text{out}} W_{\text{out}}\,, \quad c_{1,2} = 1, \dots, C_{\text{out}}\,. \end{align*}

The squared Frobenius norm of the residual between \(\nabla_{\vz^{(l)}}^{2}\ell\) and \(\mI_{H_{\text{out} W_{\text{out}}}} \otimes \mB\) is

\begin{align*} \left\lVert\nabla_{\vz^{(l)}}^{2}\ell - \mI_{H_{\text{out}} W_{\text{out}}} \otimes \mB \right\rvert_{F}^{2} = \sum_{x_1, x_2}\sum_{c_1, c_2} \left( \left[ \nabla_{\vz^{(l)}}^{2}\ell \right]_{(x_1, c_1), (x_2, c_2)} - \delta_{x_1,x_2} \left[ \mB \right]_{c_1, c_2} \right)^2 \end{align*}

Taking the gradient w.r.t. an element of \(\mB\) we get

\begin{align*} \frac{\partial \left\lVert\nabla_{\vz^{(l)}}^{2}\ell - \mI_{H_{\text{out}} W_{\text{out}}} \otimes \mB \right\rVert_{F}^{2}}{ \partial \left[ \mB \right]_{\tilde{c}_1, \tilde{c}_2}} &= \sum_{x_1, x_2}\sum_{c_1, c_2} -2 \delta_{c_1,\tilde{c}_1} \delta_{c_2,\tilde{c}_2} \delta_{x_1,x_2} \left( \left[ \nabla_{\vz^{(l)}}^{2}\ell \right]_{(x_1, c_1), (x_2, c_2)} - \delta_{x_1,x_2} \left[ \mB \right]_{c_1, c_2} \right) \\ &= -2 \sum_{x} \left( \left[ \nabla_{\vz^{(l)}}^{2}\ell \right]_{(x, \tilde{c}_1), (x, \tilde{c}_2)} - \left[ \mB \right]_{\tilde{c}_1, \tilde{c}_2} \right) \\ &= 2 H_{\text{out}} W_{\text{out}} \left[ \mB \right]_{\tilde{c}_1, \tilde{c}_2} - 2 \sum_{x} \left[ \nabla_{\vz^{(l)}}^{2}\ell \right]_{(x, \tilde{c}_1), (x, \tilde{c}_2)}\,, \end{align*}

and setting this gradient to zero, we arrive at the expression for an element of \(\mB\)

\begin{align*} \left[ \mB \right]_{c_1, c_2} = \frac{1}{H_{\text{out}} W_{\text{out}}} \sum_{x} \left[ \nabla_{\vz^{(l)}}^{2}\ell \right]_{(x, c_1), (x, c_2)}\,, \end{align*}

i.e. we obtain \(\mB\) by average-tracing over the spatial dimension of the Hessian w.r.t. the output, see Figure 3 for an illustration.

As a last step, we will insert the approximation that KFAC uses for the backpropagated Hessian w.r.t. the convolution output. Just like for the linear layer, this is an outer product of vectors \(\vg^{(l)} \in \sR^{C_{\text{out}} H_{\text{out}} W_{\text{out}}}\) of same dimension as the output (again, it does not matter here how these vectors are obtained in detail). This leads to

\begin{align*} \left[ \mB \right]_{c_1, c_2} &= \frac{1}{H_{\text{out}} W_{\text{out}}} \sum_{x} \left[ \vg^{(l)} {\vg^{(l)}}^{\top} \right]_{(x, c_1), (x, c_2)} \\ &= \frac{1}{H_{\text{out}} W_{\text{out}}} \sum_{x} \left[ \vg^{(l)}\right]_{(x, c_1)} \left[\vg^{(l)} \right]_{(x, c_2)} \end{align*}

which can be expressed as

\begin{align*} \mB = \frac{1}{H_{\text{out}} W_{\text{out}}} {\mG^{(l)}}^{\top} \mG^{(l)} \end{align*}

were \(\mG^{(l)} \in \sR^{H_{\text{out}} W_{\text{out}} \times C_{\text{out}}}\) is a matrix-view of \(\vg^{(l)}\) that separates the spatial dimensions into rows, and the channel dimensions into columns.

We have now derived the Kronecker approximation for weights in convolution layers from (Grosse & Martens, 2016):

\begin{align} \label{orglatexenvironment3} \mathrm{KFAC}(\nabla_{\vtheta^{(l)}}^{2} \ell) = \llbracket \tZ^{(l-1)} \rrbracket \llbracket \tZ^{(l-1)} \rrbracket^{\top} \otimes \frac{1}{H_{\text{out}} W_{\text{out}}} {\mG^{(l)}}^{\top} \mG^{(l)}\,. \end{align}

Note that, in contrast to the linear layer in Equation \eqref{orglatexenvironment2}, we had to make an additional approximation to the backpropagated Hessian to impose a Kronecker structure.

Transpose convolution is structurally quite similar to convolution. We will use this similarity to define a Kronecker-factorized Hessian approximation for its weight.

Just like a convolution layer, a transpose convolution layer maps an image-shaped tensor \(\tZ^{(l-1)} \in\sR^{C_{\text{in}} \times H_{\text{in}} \times W_{\text{in}}}\) to an image-shaped output \(\tZ^{(l)} \in\sR^{C_{\text{out}} \times H_{\text{out}} \times W_{\text{out}}}\) using a kernel, a rank-4 tensor, \(\tW^{(l)} \in \sR^{C_{\text{out}} \times C_{\text{in}} \times K_{H} \times K_{W}}\). We can write

\begin{align*} \tZ^{(l)} = \tZ^{(l-1)} \star_{\top} \tW^{(l)} \end{align*}

with \(\vtheta^{(l)} = \vec \tW^{(l)}\), and using \(\star_{\top}\) to denote transpose convolution. Instead of tensors, we will view the forward pass in terms of the matrices \(\mW^{(l)} \in \sR^{C_{\text{out}} \times C_{\text{in}} K_{H} K_{W}}\) and \(\mZ^{(l)} \in\sR^{C_{\text{out}} \times H_{\text{out}} W_{\text{out}}}\), which are just reshapes of the above tensors. In complete analogy to convolution, transpose convolution can be expressed as matrix multiplication

\begin{align*} \mZ^{(l)} = \mW^{(l)} \llbracket\tZ^{(l-1)}\rrbracket_{\top} \end{align*}

where \(\llbracket\tZ^{(l-1)}\rrbracket_{\top} \in \sR^{C_{\text{in}} K_{H} K_{W} \times H_{\text{out}} W_{\text{out}}}\) is the unfolded input for a transpose convolution. This unfolding operation \(\llbracket \cdot \rrbracket_{\top}\) only differs in low-level details from its counterpart \(\llbracket \cdot \rrbracket\) for convolution. We don't have to worry about how to obtain this unfolded input; there are libraries that provide such functionality.

Since the we have expressed the forward pass of transpose convolution in the same form as for convolution, we can fast-forward through all the intermediate steps and write down the approximation of KFAC for the kernel of a transpose convolution:

\begin{align} \label{orglatexenvironment4} \mathrm{KFAC}(\nabla_{\vtheta^{(l)}}^{2} \ell) = \llbracket \tZ^{(l-1)} \rrbracket_{\top} \llbracket \tZ^{(l-1)} \rrbracket^{\top}_{\top} \otimes \frac{1}{H_{\text{out}} W_{\text{out}}} {\mG^{(l)}}^{\top} \mG^{(l)}\,. \end{align}

There is one caveat that you may have noticed already if you have ever worked with transpose convolutions in deep learning libraries like PyTorch: while the kernel for convolution is commonly stored as a \(C_{\text{out}} \times C_{\text{in}} \times K_{H} \times K_{W}\) tensor, the kernel for transpose convolution is usually stored in the permuted format \(C_{\text{in}} \times C_{\text{out}} \times K_{H} \times K_{W}\). To work with the Kronecker representation, we therefore need to properly pre- and post-process quantities to and from \(C_{\text{in}} \times C_{\text{out}} \times K_{H} \times K_{W}\) format.