Randomized K-FACs: Speeding up K-FAC with Randomized Numerical Linear Algebra
K-FAC is a successful tractable implementation of Natural Gradient for Deep
Learning, which nevertheless suffers from the requirement to compute the
inverse of the Kronecker factors (through an eigen-decomposition). This can be
very time-consuming (or even prohibitive) when these factors are large. In this
paper, we theoretically show that, owing to the exponential-average
construction paradigm of the Kronecker factors that is typically used, their
eigen-spectrum must decay. We show numerically that in practice this decay is
very rapid, leading to the idea that we could save substantial computation by
only focusing on the first few eigen-modes when inverting the
Kronecker-factors. Randomized Numerical Linear Algebra provides us with the
necessary tools to do so. Numerical results show we obtain $\approx2.5\times$
reduction in per-epoch time and $\approx3.3\times$ reduction in time to target
accuracy. We compare our proposed K-FAC sped-up versions with a more
computationally efficient NG implementation, SENG, and observe we perform on
par with it.