Gradient Descent on Neurons and its Link to Approximate Second-Order Optimization
This work addresses the optimization bottleneck in neural network training for researchers and practitioners by revealing misconceptions about KFAC and proposing a more efficient alternative.
The paper challenges the belief that Kronecker-Factored Approximate Curvature (KFAC) closely approximates second-order optimization, showing it significantly outperforms true second-order updates and instead approximates gradient descent on neurons, leading to improved computational cost and data-efficiency.
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.