# KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks

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

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 57 references that are in the map.
- First posted on arXiv 4 Jul 2021.

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

- [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)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [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)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Clipper: A Low-Latency Online Prediction Serving System](https://mlsanity.com/paper/1612.03079/) (2016)
- [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)
- [Noisy Natural Gradient as Variational Inference](https://mlsanity.com/paper/1712.02390/) (2017)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [TensorFlow-Serving: Flexible, High-Performance ML Serving](https://mlsanity.com/paper/1712.06139/) (2017)
- [Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks](https://mlsanity.com/paper/1905.10961/) (2019)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads](https://mlsanity.com/paper/2008.09213/) (2020)
- [SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks](https://mlsanity.com/paper/1801.04380/) (2018)
- [Large-Scale Distributed Second-Order Optimization Using Kronecker-Factored Approximate Curvature for Deep Convolutional Neural Networks](https://mlsanity.com/paper/1811.12019/) (2018)
- [Convolutional Neural Network Training with Distributed K-FAC](https://mlsanity.com/paper/2007.00784/) (2020)
- [DLHub: Model and Data Serving for Science](https://mlsanity.com/paper/1811.11213/) (2018)
- [Scale MLPerf-0.6 models on Google TPU-v3 Pods](https://mlsanity.com/paper/1909.09756/) (2019)

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

- [Compute-Efficient Deep Learning: Algorithmic Trends and Opportunities](https://mlsanity.com/paper/2210.06640/) (2022)
- [Merging by Matching Models in Task Parameter Subspaces](https://mlsanity.com/paper/2312.04339/) (2023)
- [Influence Functions for Scalable Data Attribution in Diffusion Models](https://mlsanity.com/paper/2410.13850/) (2024)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matrices](https://mlsanity.com/paper/2211.14133/) (2022)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [SLoPe: Double-Pruned Sparse Plus Lazy Low-Rank Adapter Pretraining of LLMs](https://mlsanity.com/paper/2405.16325/) (2024)
- [Efficient Quantized Sparse Matrix Operations on Tensor Cores](https://mlsanity.com/paper/2209.06979/) (2022)
- [On the Parameterization of Second-Order Optimization Effective Towards the Infinite Width](https://mlsanity.com/paper/2312.12226/) (2023)
- [Scalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning](https://mlsanity.com/paper/2206.15143/) (2022)
- [MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates](https://mlsanity.com/paper/2306.01685/) (2023)
- [Scale-invariant Learning by Physics Inversion](https://mlsanity.com/paper/2109.15048/) (2021)
- [PETScML: Second-order solvers for training regression problems in Scientific Machine Learning](https://mlsanity.com/paper/2403.12188/) (2024)
- [PoF: Post-Training of Feature Extractor for Improving Generalization](https://mlsanity.com/paper/2207.01847/) (2022)
- [FOSI: Hybrid First and Second Order Optimization](https://mlsanity.com/paper/2302.08484/) (2023)
- [Eva: A General Vectorized Approximation Framework for Second-order Optimization](https://mlsanity.com/paper/2308.02123/) (2023)
- [Gradient Descent with Provably Tuned Learning-rate Schedules](https://mlsanity.com/paper/2512.05084/) (2025)
- [Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvature](https://mlsanity.com/paper/2608.09763/) (2026)
- [Minibatching Offers Improved Generalization Performance for Second Order Optimizers](https://mlsanity.com/paper/2307.11684/) (2023)
- [Adversarial Vulnerability as a Consequence of On-Manifold Inseparibility](https://mlsanity.com/paper/2410.06921/) (2024)

## Abstract (arXiv)

Kronecker-factored Approximate Curvature (K-FAC) has recently been shown to converge faster in deep neural network (DNN) training than stochastic gradient descent (SGD); however, K-FAC's larger memory footprint hinders its applicability to large models. We present KAISA, a K-FAC-enabled, Adaptable, Improved, and ScAlable second-order optimizer framework that adapts the memory footprint, communication, and computation given specific models and hardware to improve performance and increase scalability. We quantify the tradeoffs between memory and communication cost and evaluate KAISA on large models, including ResNet-50, Mask R-CNN, U-Net, and BERT, on up to 128 NVIDIA A100 GPUs. Compared to the original optimizers, KAISA converges 18.1-36.3% faster across applications with the same global batch size. Under a fixed memory budget, KAISA converges 32.5% and 41.6% faster in ResNet-50 and BERT-Large, respectively. KAISA can balance memory and communication to achieve scaling efficiency equal to or better than the baseline optimizers. KAISA is open source and available at https://github.com/gpauloski/kfac_pytorch.