Add KFAC optimizer

[n-gao]

🚀 Feature

Add the Kronecker-factored Approximate Curvature (KFAC) optimizer and/or E-KFAC.

Motivation

These are both approximate second order methods with promising results in terms of convergence rate. KFAC uses an block-diagonal approximation to the Fisher information matrix to approximate natural gradients. This allows for larger learning rates and thus improved convergence rate. Having a fully features built-in optimizer allows researchers and deep learning engineers to quickly iterate and test.

Pitch

A KFAC implementation that works out of the box for all supported layers should be added to the built-in optimizers.

Additional context

An unmaintained Pytorch implementation
The standard TF implementation
KFAC for CNN
KFAC for RNN

cc @vincentqb

[gchanan]

@vincentqb can you make a call on if we want this in PyTorch core?

[n-gao]

I want to add two figures to illustrate why KFAC is, imo, such an important algorithm. The figures show the optimization landscape and how KFAC and SGD reach the optimum.

[n-gao]

In case someone is looking for a KFAC implementation in KFAC: https://github.com/n-gao/pytorch-kfac