KOALA++: Efficient Kalman-Based Optimization of Neural Networks with Gradient-Covariance Products
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.
Code (0)
등록된 구현이 없습니다.
Tasks
image-classificationImage ClassificationLanguage ModelingLanguage ModellingSecond-order methodsSimilar Papers 제목 키워드 기반
KOALA: A Kalman Optimization Algorithm with Loss Adaptivity
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 (iterati…
Language ModelingLanguage ModellingStochastic OptimizationDuality-Based Stochastic Policy Optimization for Estimation with Unknown Noise Covariances
Duality of control and estimation allows mapping recent advances in data-guided control to the estimation setup. This paper formalizes and utilizes such a mapping to consider learning the optimal (steady-state) Kalman ga…
Riemannian Trust-Region based Adaptive Kalman filter with unknown noise Covariance matrices
The problem of adaptive Kalman filtering for a discrete observable linear time-varying system with unknown noise covariance matrices is addressed in this paper. The measurement difference autocovariance method is used to…
Supervisory Measurement-Guided Noise Covariance Estimation
Reliable state estimation hinges on accurate specification of sensor noise covariances, which weigh heterogeneous measurements. In practice, these covariances are difficult to identify due to environmental variability, f…
Computational EfficiencyBilevel OptimizationThe Extended Kalman Filter is a Natural Gradient Descent in Trajectory Space
The extended Kalman filter is perhaps the most standard tool to estimate in real time the state of a dynamical system from noisy measurements of some function of the system, with extensive practical applications (such as…