# KOALA: A Kalman Optimization Algorithm with Loss Adaptivity

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 65 references that are in the map.
- First posted on arXiv 7 Jul 2021.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [Adaptive Gradient Methods with Dynamic Bound of Learning Rate](https://mlsanity.com/paper/1902.09843/) (2019)
- [On the Convergence of A Class of Adam-Type Algorithms for Non-Convex Optimization](https://mlsanity.com/paper/1808.02941/) (2018)
- [Practical Gauss-Newton Optimisation for Deep Learning](https://mlsanity.com/paper/1706.03662/) (2017)
- [Painless Stochastic Gradient: Interpolation, Line-Search, and Convergence Rates](https://mlsanity.com/paper/1905.09997/) (2019)
- [Closing the Generalization Gap of Adaptive Gradient Methods in Training Deep Neural Networks](https://mlsanity.com/paper/1806.06763/) (2018)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Variants of RMSProp and Adagrad with Logarithmic Regret Bounds](https://mlsanity.com/paper/1706.05507/) (2017)
- [Biased Stochastic First-Order Methods for Conditional Stochastic Optimization and Applications in Meta Learning](https://mlsanity.com/paper/2002.10790/) (2020)
- [On the distance between two neural networks and the stability of learning](https://mlsanity.com/paper/2002.03432/) (2020)
- [Probabilistic Line Searches for Stochastic Optimization](https://mlsanity.com/paper/1502.02846/) (2015)
- [diffGrad: An Optimization Method for Convolutional Neural Networks](https://mlsanity.com/paper/1909.11015/) (2019)
- [L4: Practical loss-based stepsize adaptation for deep learning](https://mlsanity.com/paper/1802.05074/) (2018)
- [WNGrad: Learn the Learning Rate in Gradient Descent](https://mlsanity.com/paper/1803.02865/) (2018)
- [Parabolic Approximation Line Search for DNNs](https://mlsanity.com/paper/1903.11991/) (2019)
- [Adam with Bandit Sampling for Deep Learning](https://mlsanity.com/paper/2010.12986/) (2020)
- [The Extended Kalman Filter is a Natural Gradient Descent in Trajectory Space](https://mlsanity.com/paper/1901.00696/) (2019)

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

- [KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products](https://mlsanity.com/paper/2506.04432/) (2025)
- [Signal Processing Meets SGD: From Momentum to Filter](https://mlsanity.com/paper/2311.02818/) (2023)
- [LoKO: Low-Rank Kalman Optimizer for Online Fine-Tuning of Large Models](https://mlsanity.com/paper/2410.11551/) (2024)
- [Bayesian Natural Gradient Fine-Tuning of CLIP Models via Kalman Filtering](https://mlsanity.com/paper/2511.01694/) (2025)
- [Dynamic Momentum Recalibration in Online Gradient Learning](https://mlsanity.com/paper/2603.06120/) (2026)
- [Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models](https://mlsanity.com/paper/2606.25086/) (2026)

## Abstract (arXiv)

Optimization is often cast as a deterministic problem, where the solution is found through some iterative procedure such as gradient descent. However, when training neural networks the loss function changes over (iteration) time due to the randomized selection of a subset of the samples. This randomization turns the optimization problem into a stochastic one. We propose to consider the loss as a noisy observation with respect to some reference optimum. This interpretation of the loss allows us to adopt Kalman filtering as an optimizer, as its recursive formulation is designed to estimate unknown parameters from noisy measurements. Moreover, we show that the Kalman Filter dynamical model for the evolution of the unknown parameters can be used to capture the gradient dynamics of advanced methods such as Momentum and Adam. We call this stochastic optimization method KOALA, which is short for Kalman Optimization Algorithm with Loss Adaptivity. KOALA is an easy to implement, scalable, and efficient method to train neural networks. We provide convergence analysis and show experimentally that it yields parameter estimates that are on par with or better than existing state of the art optimization algorithms across several neural network architectures and machine learning tasks, such as computer vision and language modeling.