# KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 33 references that are in the map.
- First posted on arXiv 4 Jun 2025.

## Builds on (most cited first, 21 of 21)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Searching for MobileNetV3](https://mlsanity.com/paper/1905.02244/) (2019)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [Overfitting in adversarially robust deep learning](https://mlsanity.com/paper/2002.11569/) (2020)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://mlsanity.com/paper/2208.06677/) (2022)
- [On the distance between two neural networks and the stability of learning](https://mlsanity.com/paper/2002.03432/) (2020)
- [SLANG: Fast Structured Covariance Approximations for Bayesian Deep Learning with Natural Gradient](https://mlsanity.com/paper/1811.04504/) (2018)
- [Low-rank extended Kalman filtering for online learning of neural networks from streaming data](https://mlsanity.com/paper/2305.19535/) (2023)
- [Efficient Online Bayesian Inference for Neural Bandits](https://mlsanity.com/paper/2112.00195/) (2021)
- [AdaFisher: Adaptive Second Order Optimization via Fisher Information](https://mlsanity.com/paper/2405.16397/) (2024)
- [The Extended Kalman Filter is a Natural Gradient Descent in Trajectory Space](https://mlsanity.com/paper/1901.00696/) (2019)
- [Trust Region Value Optimization using Kalman Filtering](https://mlsanity.com/paper/1901.07860/) (2019)
- [KOALA: A Kalman Optimization Algorithm with Loss Adaptivity](https://mlsanity.com/paper/2107.03331/) (2021)

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

- [Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models](https://mlsanity.com/paper/2606.25086/) (2026)
- [Infinite-Parameter LLMs: Generating and Adapting Weights from Live Data](https://mlsanity.com/paper/2609.18842/) (2026)

## Abstract (arXiv)

We propose KOALA++, a scalable Kalman-based optimization algorithm that explicitly models structured gradient uncertainty in neural network training. Unlike second-order methods, which rely on expensive second order gradient calculation, our method directly estimates the parameter covariance matrix by recursively updating compact gradient covariance products. This design improves upon the original KOALA framework that assumed diagonal covariance by implicitly capturing richer uncertainty structure without storing the full covariance matrix and avoiding large matrix inversions. Across diverse tasks, including image classification and language modeling, KOALA++ achieves accuracy on par or better than state-of-the-art first- and second-order optimizers while maintaining the efficiency of first-order methods.