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

We sincerely thank the reviewer for the thoughtful and constructive feedback. We are especially grateful for your recognition of our key contributions and your clear understanding of the motivation and design behind KOALA++. Your positive comments, particularly acknowledging KOALA++ as a valuable improvement over the original KOALA, its empirical robustness across diverse settings, and the clarity of our writing, are very encouraging to us.

Below, we address the specific concerns raised:

---

FLOPs and Memory Usage Comparison

We thank the reviewer for this insightful question. To address it comprehensively, below we have included an efficiency profile summarizing the optimizer-side FLOPs, per-epoch wall-clock time, and peak memory usage across larger architectures: ResNet‑101 (CNN), Swin Transformer Base (ViT) with larger batch size(1024). For the results of GPT-1, see Table 5 of the main paper. All experiments use the same dataloader, hardware, and exclude warm-up steps.

Here are our observations:

• KOALA++ maintains $O(N)$ optimizer-side complexity via structured vector operations and avoids matrix inversions entirely. This contrasts with methods like Shampoo that involve more matrix operations (e.g., Cholesky, blockwise inverses).

• KOALA++ is consistently close to AdamW in epoch time and significantly faster than AdaHessian/Shampoo. As batch size grows, KOALA++ approaches near parity with AdamW due to forward/backward FLOPs dominating total runtime.

Efficiency Profile

ResNet‑101 (Batch Size = 1024)

Swin Transformer Base(Batch Size = 1024)

---

PSD Impact

We thank the reviewer for this insightful question regarding the PSD (positive semi-definite) property of the covariance-like quantities in KOALA++.

As mentioned in the Limitations section, we point out that the approximated $P_{k - 2}$ is not necessarily guaranteed to be PSD. To better understand the impact of this missing constraint, we carry out analysis and experiments in the Supplementary Material in section B.3. This investigation shows that the algorithm is quite stable in terms of the $P_{k-2}$ matrix eigenvalues despite lacking the PSD property. This initial analysis is further confirmed by the fact that in all our empirical results across various datasets and model architectures we observe that this does not degrade convergence or stability of the optimizer.

---

Minor typos and grammar issue

We thank the reviewer for pointing out the minor typos and grammar issues. We will carefully revise the manuscript to correct all such issues.

We sincerely thank the reviewer for their time and thoughtful evaluation of our work. We greatly appreciate the reviewer that they notice the value of using Kalman Filter as a deep neural network optimizer.

Below, we address the specific concerns raised:

---

Thinking of KOALA++

We appreciate the reviewer’s summary and understand the impression that KOALA++ may appear as a quick extension of the original KOALA method. However, we respectfully clarify that while KOALA++ is indeed based on KOALA, we do not view it as a mere incremental extension. Instead, KOALA++ addresses core limitations in KOALA’s original design and introduces a substantially different treatment of uncertainty modeling in optimization.

As noted in our paper, KOALA treats neural network optimization as estimating the state in the Kalman Filter's dynamic system. However, it makes several simplifying assumptions that we believe limit its expressiveness（see lines 77-79 in the main paper).

1. Covariance Approximation: Although the covariance $P_{k-1}$ is formally a matrix, KOALA approximates it as a scaled identity to reduce computational complexity. We found this to be limiting, especially in high-dimensional neural network optimization settings.

2. Gradient Product Approximation: KOALA further approximates the rank-1 matrix $H_k^\top H_k$ as another scaled identity:

   $$H_k^\top H_k \approx \|H_k\|^2 \cdot I$$

   This approximation discards the interactions between different weights of the network.

To address these issues, KOALA++ introduces a core methodological shift: instead of using scaled identity matrices, we do not make any structural assumptions on $P_k$ and only update it implicitly by tracking vector–matrix products. While expressions like $H_k P_{k-2}$ still appear in the update, we design the algorithm to never explicitly store or construct the covariance $P_{k-2}$ by using the relationship, $H_{k-1} P_{k-2} = v_{k-1}$ and recovering $P_{k-2}$ via a simple least-squares fit, which preserves the efficiency of the original design while offering significantly more expressive updates.

---

Motivation of directional covariance optimization objective

We appreciate the request for a clearer motivation. Below we explain why we project the covariance onto a directional subspace and why we solve a Frobenius minimum‑norm problem to approximate it.

Our decision to minimize the Frobenius norm of $P_{k-2}$ was guided by two main aspects: (1) it yields a convex optimization problem with a unique solution, ensuring an unambiguous covariance estimate; and (2) it helps keep the estimate bounded, which is important for maintaining numerical stability in our method. We will revise the paper to make this design decision and its motivation clearer in the main text, and we thank the reviewer again for bringing this to our attention.

---

Error Analysis

We thank the reviewer for raising this point. In this work, our primary focus has been on the empirical evaluation of the method. Indeed, we plan to work on a formal theoretical analysis of the approximation and other aspects of the method in future work.

---

Related Work

We thank the reviewer for pointing to the related work in Kalman-inspired optimization methods with structured (non-diagonal) covariance approximations. We acknowledge the relevance of these methods and will include a discussion of them in our revised manuscript.

While KOALA++ is indeed inspired by Kalman filtering principles and incorporates uncertainty modeling, it differs from the cited works in design choices and application settings:

• Chang et al. (2023) relies on EKF-like updates with diagonal plus low-rank covariance approximations to manage computational cost. Mishkin et al. (2018) also propose a diagonal-plus-low-rank approxmation of the inverse covariance for efficient posterior approximation in the variational Bayesian inference. In contrast to these methods, KOALA++ does not make explicit assumptions on the structure of the covariance and only tracks its directional projections. Moreover, Chang et al. (2023) formulates the training as online Bayesian inference, while KOALA++ stays within classical deep learning optimization settings. Nonetheless, we agree that the underlying methodology is related and will discuss this in the revised paper.

• Duran-Martin et al. (2022) reduces the dimensionality of the state space and performs inference in a lower-dimensional subspace using full EKF updates, thereby avoiding structural assumptions on the covariance in the original weight space. In contrast, KOALA++ operates directly in the original parameter space, benefiting from flexibile navigation in the high-dimensional space, while maintaining linear computational complexity without imposing explicit structural constraints on the covariance.

---

Counterintuition about the training dynamics

We would like to thank the reviewer for the insightful comment. We recognize that the most significant improvements with KOALA++ come closer to the end of the training where the learning rate is small and the method navigates the loss landscape at the smallest scale. Our reasoning is that during the early stages of training, the optimization landscape is typically more non-convex and noisy. The aggressive updates from first-order methods are particularly effective at these stages. As training advances and the loss surface becomes smoother at the smaller scale, leveraging the curvature of the loss, as (quasi-)second-order methods do (e.g. KOALA++ modelling the uncertainty of the weights' updates), becomes more beneficial and enhances generalization. We would like to thank the reviewer again for pointing our attention to this interesting behavior and leave its further investigation for the future work.

---

Potential typo in $QI$

While $Q$ is a matrix in the standard EKF, in KOALA++, we adopt the design choice from KOALA and use a scaled identity matrix $QI$.

---

Thanks for your response. I am not totally convinced about some of your justifications. First, while you highlight several aspects that are novel to KOALA++, there is no actual analysis (only linear algebra) whereas it would be valuable to include some insight or theoretical understanding of the optimizer itself. Second, the “simple least squares fit” appears to be a rather basic component, but I am willing to acknowledge it as a practical advantage of the proposed approach.

However, a point that concerns me is your explanation that KOALA++ only gains from its directional covariance projections late in training conflicts. This seems to be contrary to what has been the reported behavior of other quasi–second-order optimizers in the literature:

1. Figure 2 in [1] shows that Muon (despite using a similar curvature-aware update) outperforms Adam from the very first step, reaching a validation loss of $\approx 3.95$ after only 2s, compared to Adam’s higher loss at the same wall-clock time. If curvature information were only useful once the learning rate is small, Muon would not hold such an advantage immediately.

2. Figure 4 in [2] illustrates Shampoo’s log-perplexity on LM1B, where it already beats Adam at 50k steps and maintains that gap through 500k steps. A preconditioned method that only yielded late-stage improvements would not show a steady, early reduction in perplexity like Shampoo does.

These two cases directly contradict the claim that “aggressive first-order updates dominate initial training and quasi-second-order gains emerge only at the end.” To reconcile this, please either:

• Provide a clear theoretical argument for why KOALA++’s directional least-squares projection fails to leverage curvature information during the early phase, or
• Include Muon and Shampoo in your empirical comparisons to demonstrate that KOALA++ truly differs in its convergence profile.

Absent such an explanation or data, the hypothesis about delayed curvature benefits is unsupported. I appreciate that the time to include new experiments at this point is limited, but without proof of this point, I find it hard to raise my score. I would appreciate at least a convincing justification.

References:

[1] https://kellerjordan.github.io/posts/muon/

[2] Gupta, Vineet, Tomer Koren, and Yoram Singer. "Shampoo: Preconditioned stochastic tensor optimization." International Conference on Machine Learning. PMLR, 2018.

We sincerely thank the reviewer for their time and thoughtful evaluation of our work. We greatly appreciate the reviewer’s positive feedback on our theoretical deriviation of KOALA++, and we are pleased that they found it solid and rigorous.

Below, we address the specific concerns raised:

---

Efficiency profiling

To illustrate the compotational efficiency of KOALA++, we reported the average training time per epoch for the NLP experiment on WikiText-2 (see Table 5). As our method's computational complexity is task-independent, we considered this signle example to be sufficiently representative. Nonetheless, as requested by the reviewer, we have run additional experiments across different tasks/architectures to better represent the efficiency profile of KOALA++. We provide the results here below:

ResNet‑50 (Batch Size = 256)

Swin Transformer (Batch Size = 256)

As the time for the rebuttal is limited, we were not able to run the experiments with larger models. However, we will include them in the revised paper.

---

Hyperparameter guidance & sensitivity

Notice that in the experiments on CIFAR‑10/100 (Tables 2, 3 and 4, CNN-based models), we used the same hyperparameter settings across different models. The results demonstrate KOALA++’s low‑tuning usability (it achieves better or on par results with the baselines), which we consider the practical strength of our method. When moving to the Transformer-based architectures, we adopted a small grid search (see Supplementary Material lines 195-198, 220-222). We will better clarify our evaluation scheme in the revised paper.

Below we provide some insights on choosing the hyperparameters of KOALA++, which we will also include in the revised manuscript to better support practical deployment of our method:

• As for the measurement noise $R$, we estimate it online as an exponential moving average of the mini-batch loss variance, using a smoothing factor $\alpha = 0.9$ as implemented in our code. This procedure effectively removes $R$ from the set of hyperparameters that require tuning.
• $\sigma_0$ determines only the scale of the initial directional state $v_1 = \sigma_0 H_1$. We have done some ablation studies for this parameter (see Supplementary Material Table 3), and we found that the convergence of KOALA++ is robust to the choice of $\sigma_0$, so for simplicity and symmetry, we suggest to set the prior variance $\sigma_{0}$ and the variance of the process noise $Q$ to the same value.
• Given above choices, $Q$ is the only hyperparameter except for the learning rate that requires some careful tuning. Empirically we found that $Q \in [0.1, 0.4]$ works relatively well across different datasets.

---

Speed vs SGD

Below we provide the throughput comparisons vs SGD on CIFAR-10 with ResNet50 using different batch sizes. All the runs are done on a single H100 GPU.

The trend matches our design intuition: KOALA++ performs a small set of additional $O(N)$ operations per iteration compared to plain SGD. These optimizer costs are independent of the batch size, while forward/backward cost grows with the batch size. As a result, the relative overhead shrinks from 1.62× (bs=128) to 1.00× (bs=1024), approaching parity with SGD.

Additionally, while this already demonstrates the efficiency of the method, the throughput of KOALA++ can be further improved by reusing repeated dot products across consecutive iterations. We plan to include updated results reflecting these improvements in the revision.

---

Scalability in LM

We thank the reviewer for the valuable suggestion. We agree that demonstrating the effectiveness of KOALA++ in larger-scale training would further strengthen the paper. However, due to the limited scope and time constraints of the rebuttal period, we are unable to run such experiments at this stage. That said, we plan to evaluate KOALA++ on larger architectures from the GPT family in future work.

---

Consideration of hybrid methods

We thank the reviewer for this insightful suggestion.

We did consider hybrid approximations of the covariance matrix, such as restricting its structure to a diagonal plus rank-1 form (i.e., $P_k = \sigma I + v_k^\top v_k$), with corresponding updates to the vector $v_k$. While this approach captures slightly richer curvature information, it has been explored in prior work (e.g., Lo-Fi by Chang et al., 2023) and is closely related to KOALA’s original scaled identity formulation, which potentially limits its novelty.

In contrast, KOALA++ takes a fundamentally different approach by operating directly on vector–matrix products, without ever explicitly constructing $P_k$. This avoids imposing structural assumptions and opens up greater flexibility. However, we acknowledge that there may be more effective hybrid approximations or sketching-based factorizations that could offer better trade-offs. We appreciate the reviewer for pointing out this direction and will reflect on it more carefully in future work.

---

We appreciate your initial review, and we are looking forward to your feedback on our rebuttal. We would greatly appreciate the opportunity to engage in further discussion to clarify any remaining concerns or questions you may have.

Thank you again for your time and for reviewing our work.

We sincerely thank the reviewer for their time and thoughtful evaluation of our work. We greatly appreciate the reviewer’s positive feedback on our experimental evaluation of KOALA++, and we are pleased that they found it extensive and convincing.

Below, we address the specific concerns raised:

---

Clarification on Covariance Approximation

The reviewer mentions in the paper summary that KOALA++ employs a rank-1 approximation of the covariance matrix. We would like to kindly clarify that the approximation of $P_{k-2}$ in equation (9) is, in fact, rank-2, as it is based on both $v_{k - 1}$ and $H_{k - 1}$, which are generally non-collinear vectors (see equation (15)). This implies that the rank of the implicitly tracked covariance is at least 2.

---

Missing Confidence Intervals

Thank you for bringing this to our attention. As noted in lines 181–182 and 221–223, we ran each experiment multiple times and reported average performance metrics. In Table 4, we also reported standard deviations. However, confidence intervals are indeed missing in Tables 2 and 3. In some cases, the values were sourced from the respective original papers, which did not provide confidence intervals. We will include them wherever possible in the revised manuscript.

---

The Benefit Comes Only After the Final Learning Rate Reduction

Thank you for the thoughtful comment. We acknowledge that the strongest gains from KOALA++ occur after the final learning rate drop. Our intuition behind this is based on the fact that in the earlier phases of training, the optimization landscape tends to be more non-convex and noisy, where agressive updates from first-order methods are often sufficient or even preferrable. However, as training progresses and the loss surface becomes smoother (after the last learning rate drop), exploiting the curvature of the loss through modelling the uncertainty of the updates in KOALA++ becomes most effective and leads to better generalization. This behavior deserves further investigation.

---

Fairness of Hyperparameter Tuning

In most cases, we reported performance numbers directly from prior work, assuming that hyperparameters had been appropriately tuned in those original studies. For baselines where we conducted our own experiments, we performed hyperparameter tuning via grid search. Notice however that in Table 2, 3 and 4 we used the same hyperparameter settings for all the different CNN-based model architectures and datasets. Nonetheless, KOALA++ is consistently the top or the second top performer in all cases. We thus attribute this result more to the method than to the tuning. We will clarify this in the revised the manuscript to make our evaluation more explicit.

---

Computation Time per Epoch

We reported the average training time per epoch for the NLP experiment on WikiText-2 to demonstrate the computational efficiency of KOALA++, particularly in comparison with second-order methods such as AdaHessian. Since the computational complexity of our method is task-independent, we felt one representative example would suffice. However, we include the training times for other tasks here below.

ResNet‑50 (Batch Size = 256)

Swin Transformer Tiny(Batch Size = 256)

---

Frobenius Norm Justification

Thank you for the insightful suggestion. Our choice to minimize the Frobenius norm of $P_{k-2}$ was motivated by two considerations: (1) it leads to a convex optimization problem with a unique solution, thereby avoiding ambiguity in the covariance estimate; and (2) it helps keep the estimate bounded, which is important for the numerical stability of our method. That said, we appreciate the reviewer’s proposal to exploit alternative priors, such as minimizing distance to a scaled identity matrix. We will explore this direction in future work.

---

Please let us know if further clarification is needed. We once again thank the reviewer for their valuable feedback and constructive suggestions.