Gradient Descent on Neurons and its Link to Approximate Second-Order Optimization

TL;DR

This paper reveals that KFAC, a popular second-order optimizer approximation, actually behaves more like a first-order method and outperforms true second-order updates in neural network training.

Contribution

The study demonstrates that KFAC is not a true second-order method but approximates a first-order neuron-level gradient descent, challenging previous assumptions.

Findings

KFAC significantly outperforms true second-order updates.

KFAC effectively approximates a first-order neuron-level algorithm.

The neuron-level gradient descent often improves over KFAC in efficiency.

Abstract

Second-order optimizers are thought to hold the potential to speed up neural network training, but due to the enormous size of the curvature matrix, they typically require approximations to be computationally tractable. The most successful family of approximations are Kronecker-Factored, block-diagonal curvature estimates (KFAC). Here, we combine tools from prior work to evaluate exact second-order updates with careful ablations to establish a surprising result: Due to its approximations, KFAC is not closely related to second-order updates, and in particular, it significantly outperforms true second-order updates. This challenges widely held believes and immediately raises the question why KFAC performs so well. Towards answering this question we present evidence strongly suggesting that KFAC approximates a first-order algorithm, which performs gradient descent on neurons rather than weights. Finally, we show that this optimizer often improves over KFAC in terms of computational cost and data-efficiency.