An iterative K-FAC algorithm for Deep Learning

TL;DR

This paper introduces CG-FAC, an iterative, matrix-free K-FAC algorithm that uses conjugate gradient to approximate the natural gradient, reducing time and memory complexity compared to standard K-FAC.

Contribution

It proposes a novel iterative K-FAC method that eliminates the need for explicit Fisher matrix and Kronecker factors, improving efficiency.

Findings

CG-FAC is matrix-free and does not require FIM or Kronecker factors.

The method has lower time and memory complexity than standard K-FAC.

CG-FAC maintains comparable accuracy in training deep neural networks.

Abstract

Kronecker-factored Approximate Curvature (K-FAC) method is a high efficiency second order optimizer for the deep learning. Its training time is less than SGD(or other first-order method) with same accuracy in many large-scale problems. The key of K-FAC is to approximates Fisher information matrix (FIM) as a block-diagonal matrix where each block is an inverse of tiny Kronecker factors. In this short note, we present CG-FAC -- an new iterative K-FAC algorithm. It uses conjugate gradient method to approximate the nature gradient. This CG-FAC method is matrix-free, that is, no need to generate the FIM matrix, also no need to generate the Kronecker factors A and G. We prove that the time and memory complexity of iterative CG-FAC is much less than that of standard K-FAC algorithm.