KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866918168687017984 |
|---|---|
| author | Xia, Zixuan Davtyan, Aram Favaro, Paolo |
| author_facet | Xia, Zixuan Davtyan, Aram Favaro, Paolo |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_04432 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products Xia, Zixuan Davtyan, Aram Favaro, Paolo Machine Learning Optimization and Control 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. |
| title | KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products |
| topic | Machine Learning Optimization and Control |
| url | https://arxiv.org/abs/2506.04432 |