# Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning

arXiv: https://arxiv.org/abs/2206.15143 · HTML: https://mlsanity.com/paper/2206.15143/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 63 references that are in the map.
- First posted on arXiv 30 Jun 2022.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Train longer, generalize better: closing the generalization gap in large batch training of neural networks](https://mlsanity.com/paper/1705.08741/) (2017)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Multi30K: Multilingual English-German Image Descriptions](https://mlsanity.com/paper/1605.00459/) (2016)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis](https://mlsanity.com/paper/1806.03884/) (2018)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)
- [Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks](https://mlsanity.com/paper/1905.10961/) (2019)
- [Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes](https://mlsanity.com/paper/1807.11205/) (2018)
- [Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters](https://mlsanity.com/paper/1706.03292/) (2017)
- [On the interplay between noise and curvature and its effect on optimization and generalization](https://mlsanity.com/paper/1906.07774/) (2019)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [When Does Preconditioning Help or Hurt Generalization?](https://mlsanity.com/paper/2006.10732/) (2020)
- [Extrapolation for Large-batch Training in Deep Learning](https://mlsanity.com/paper/2006.05720/) (2020)
- [Convolutional Neural Network Training with Distributed K-FAC](https://mlsanity.com/paper/2007.00784/) (2020)

## Built on by (most cited first, 7 of 7)

- [Compute-Efficient Deep Learning: Algorithmic Trends and Opportunities](https://mlsanity.com/paper/2210.06640/) (2022)
- [FedLPA: One-shot Federated Learning with Layer-Wise Posterior Aggregation](https://mlsanity.com/paper/2310.00339/) (2023)
- [On the Parameterization of Second-Order Optimization Effective Towards the Infinite Width](https://mlsanity.com/paper/2312.12226/) (2023)
- [Eva: A General Vectorized Approximation Framework for Second-order Optimization](https://mlsanity.com/paper/2308.02123/) (2023)
- [Efficient Bilevel Optimization with KFAC-Based Hypergradients](https://mlsanity.com/paper/2603.29108/) (2026)
- [Layerwise LQR for Geometry-Aware Optimization of Deep Networks](https://mlsanity.com/paper/2605.04230/) (2026)
- [Dual Gauss-Newton Directions for Deep Learning](https://mlsanity.com/paper/2308.08886/) (2023)

## Abstract (arXiv)

The second-order optimization methods, notably the D-KFAC (Distributed Kronecker Factored Approximate Curvature) algorithms, have gained traction on accelerating deep neural network (DNN) training on GPU clusters. However, existing D-KFAC algorithms require to compute and communicate a large volume of second-order information, i.e., Kronecker factors (KFs), before preconditioning gradients, resulting in large computation and communication overheads as well as a high memory footprint. In this paper, we propose DP-KFAC, a novel distributed preconditioning scheme that distributes the KF constructing tasks at different DNN layers to different workers. DP-KFAC not only retains the convergence property of the existing D-KFAC algorithms but also enables three benefits: reduced computation overhead in constructing KFs, no communication of KFs, and low memory footprint. Extensive experiments on a 64-GPU cluster show that DP-KFAC reduces the computation overhead by 1.55x-1.65x, the communication cost by 2.79x-3.15x, and the memory footprint by 1.14x-1.47x in each second-order update compared to the state-of-the-art D-KFAC methods.