Kronecker-factored Approximate Curvature (KFAC) From Scratch

TL;DR

This paper offers a comprehensive, step-by-step tutorial on implementing Kronecker-factored Approximate Curvature (KFAC) in deep learning, including mathematical explanations, code, and test cases to facilitate understanding and correct implementation.

Contribution

It provides the first detailed, reliable implementation guide for KFAC, combining math, code, and testing strategies based on recent insights to lower barriers for practitioners.

Findings

Provides test cases for KFAC implementation

Clarifies common pitfalls and solutions in coding KFAC

Enhances understanding of KFAC's mathematical foundations

Abstract

Kronecker-factored approximate curvature (KFAC) is arguably one of the most prominent curvature approximations in deep learning. Its applications range from optimization to Bayesian deep learning, training data attribution with influence functions, and model compression or merging. While the intuition behind KFAC is easy to understand, its implementation is tedious: It comes in many flavours, has common pitfalls when translating the math to code, and is challenging to test, which complicates ensuring a properly functioning implementation. Some of the authors themselves have dealt with these challenges and experienced the discomfort of not being able to fully test their code. Thanks to recent advances in understanding KFAC, we are now able to provide test cases and a recipe for a reliable KFAC implementation. This tutorial is meant as a ground-up introduction to KFAC. In contrast to the existing work, our focus lies on providing both math and code side-by-side and providing test cases based on the latest insights into KFAC that are scattered throughout the literature. We hope this tutorial provides a contemporary view of KFAC that allows beginners to gain a deeper understanding of this curvature approximation while lowering the barrier to its implementation, extension, and usage in practice.