Title: KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
Venue: NeurIPS 2025
Accepted: yes
Keywords: Optimization

Kalman Filter

Deep Learning
Abstract: 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 second-order optimizers while maintaining the efficiency of first-order methods.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a Kalman-based optimization algorithm designed for training neural networks. Rather than relying on expensive second-order gradient calculations or simplistic diagonal covariance assumptions, proposed method explicitly models structured gradient uncertainty through recursive updates of compact gradient-covariance products. \n\nThe goals of the work is to achieve the accuracy of second-order methods while maintaining the efficiency of first-order optimizers. Empirical rsults are rpesented on image classification and language modeling tasks.\n\nSTRENGTHS AND WEAKNESSES:\nI find this to be a interesting paper solving an important problem in optimization. Some specific points,\n\nStrengths:\n- The work addresses an important limitatoin existing in the prior work (KOALA where there are simplistic assumptions made). This is a good improvement.\n- The algorithm maintains computational complexity comparable to first-order methods, which is a significant advantage over many traditional second-order optimizers that suffer from high computational costs. This makes the proposed method appealing to practitioners in my opinion. Second order methods are back in fashion more recently.\n- The paper produces strong empirical results across diverse tasks and architectures which I greatly appreciate. \n- Detailed comparison study across various other key optimizers is provided (e.g. Table 4).\n- Motivation and writing in the paper is clear in my opinion.\n\nWeaknesses:\n- My main concern is the lack of direct comparison with all major second-order optimizers: While it compares against AdaFisher, the more recent popular ones like K-FAC, Shampoo are only briefly mentioned in tables without extensive discussion of their performance characteristics. Probably other reviewers noted this too. Comparison of training time, efficiency and convergence with these more popular methods will further strenghten the paper.\n\nQUESTIONS:\n- The paper mentions that KOALA++ retains scalability similar to first-order methods. Can the authors provide a more direct comparison of the computational overhead (e.g., FLOPs, memory usage) of KOALA++ versus Adam or SGD for a single training step on a large model?\n\n- Given the acknowledged limitation regarding the positive semi-definite (PSD) nature of key quantities in the optimizer, how much impact does this have on the algorithm's performance in practice, particularly in very deep or sensitive models?\n\nLIMITATIONS:\nN/A (do not think this applies).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a modification of the Koala algorithm that considers a different posterior covariance update to the Kalman filter based on a directional covariance projection.\n\nSTRENGTHS AND WEAKNESSES:\n- **Strengths**: The use of a KF-inspired optimisation algorithm to tackle the learning of a high-dimensional model is commendable.  The empirical performance looks good.\n\n- **Weaknesses**: The paper feels like a quick extension of the original Koala paper. The introduction of the directional covariance projection and its optimisation objective is simply introduced without any major motivation. Various approximations are made, but no error analysis is presented as to the choice of approximate covariance. Furthermore, some of the results seem counterintuitive to what one would expect in using proxy second-order optimisation methods as opposed to first-order optimisation methods (see question below).\n\nThe paper is also missing references that have also tackled the problem of learning with a KF-type optimisation algorithm that *do not* assume a diagonal covariance matrix. I believe these would've been better choices of benchmarks, as the proposed Koala++ method is not a first-order optimisation algorithm.\n\nFor example:\n1.⁠ ⁠Chang, Peter G., et al. \"Low-rank extended Kalman filtering for online learning of neural networks from streaming data.\" Conference on Lifelong Learning Agents. PMLR, 2023.\n2.⁠ ⁠Mishkin, Aaron, et al. \"Slang: Fast structured covariance approximations for bayesian deep learning with natural gradient.\" Advances in neural information processing systems 31 (2018).\n3.⁠ ⁠Duran-Martin, Gerardo, Aleyna Kara, and Kevin Murphy. \"Efficient online bayesian inference for neural bandits.\" International conference on artificial intelligence and statistics. PMLR, 2022.\n\nA more thorough literature review, comparison to prior work, and the points mentioned above **must** be addressed by the authors.\n\nQUESTIONS:\nKoala++, being a quasi-second-order optimiser, I would've expected the test loss to be much lower than first-order optimisers because they can navigate the loss landscape much more efficiently than their first-order counterparts. However, Figure 1 shows the exact opposite: it is only at the very end of training that we see Koala++ obtain a lower test loss and error. I find this to be counterintuitive. Could you explain why we see this behaviour, and is this the case for all other experiments?\n\n(typo?): I believe all $QI$ in Table 1 should be only $Q$ as you treat $Q$ as a matrix and not a scalar.\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces KOALA++, an extension of the Kalman-inspired optimizer KOALA, designed to improve efficiency and robustness in deep neural network training. KOALA++ innovatively tracks a low-rank surrogate of the gradient-covariance product to approximate directional uncertainties in optimization without storing or inverting full covariance matrices. It bridges the gap between first-order optimizers (e.g., SGD, Adam) and second-order optimizers (e.g., K-FAC, Shampoo) by maintaining efficiency while capturing richer curvature information. Extensive experiments on image classification (CIFAR-10/100, ImageNet) and language modeling (WikiText-2) show that KOALA++ performs competitively with state-of-the-art optimizers in terms of both accuracy and convergence stability, with manageable computational overhead.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper is well-written, mathematically solid, and offers promising empirical evidence.\n\nThe paper provides a significant methodological advance by designing a Kalman-inspired optimizer that efficiently models structured uncertainty without requiring full covariance matrices or expensive matrix inversions.\n\nDerivations are carefully detailed, with two variants of covariance estimation derived rigorously.\n\nWeaknesses:\nWhile the method is claimed to be efficient, the paper would benefit from more detailed profiling of compute time or FLOP comparisons versus competitors like AdaFisher or K-FAC, especially on large-scale models.\n\nThe method introduces several hyperparameters (e.g., Q, R, initial σ²), but guidance on selecting these in practice is limited, and sensitivity analyses are sparse.\n\nThe language modeling results are encouraging but are reported on a relatively small dataset (WikiText-2). Testing KOALA++ on larger LLM pretraining settings would solidify claims of scalability.\n\nQUESTIONS:\nCould the authors provide quantitative measurements (e.g., time per epoch, memory footprint) comparing KOALA++ to other second-order optimizers like AdaFisher, Shampoo and NG+ on larger models like Swin Transformers or GPT-based architectures?\n\nIn practice, how much wall-clock speedup is retained relative to plain SGD?\n\nHave the authors considered hybrid methods that might store low-rank factorizations (e.g., via sketching techniques) instead of purely vector products?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new kalman filter based gradient descent methods that extends a current SOTA method called KOALA. KOALA++ i using a low-rank (rank = 1) estimate of the covariance matrix in the Kalman filter as opposed to a diagonal estimate in the original KOALA paper. In order to obtain the covariance update of previous time steps, the authors formulate a constraint optimization problem based on the minimum frobenius norm and symmetry constraints. The authors show the competitiveness of their approach on several image-based (CIFAR-10 and CIFAR 100) and language based tasks (training a GPT-1 model).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Sound approach to gradient descent \n- The idea of estimating a low-dimensional covariance in the sub-space of the direction of the current gradient is interesting\n- The evaluations are exhaustive and convincing. solid improvement over SOTA in most tasks\n\nWeaknesses:\n- Some motivations are unclear (see questions)\n- Many results do not contain confidence intervals\n- for the experiments with the learning rate schedule, it seems the benefit only comes after the last reduction of the learning rate. Even then, the benefit is rather small. \n- It is unclear how much the reported benefits can be attributed to better chosen hyper-parameters. \n- Computation time per epoch should be reported for all experiments\n\nQUESTIONS:\n- Why is the Frobenius norm a good measure for the optimization problem? Minimizing the Frobernius norm without constraints would result in 0 covariance. Would an objective that would measure the distance to a prior covariance (> 0) not make more sense?\n- Its unclear to me wether the hyperparameters are fairly choosen. Do you use individual hyperparameters per task/dataset for your method and fixed hyper-parameters for the baselines (or where the hyperparameters there also adapted per task). Did you try to optimize hyperparameters of the baseline?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}