KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
Authors: Zixuan Xia
Aram Davtyan
Paolo Favaro
Venue: NeurIPS 2025 poster ()
Keywords: Optimization
Kalman Filter
Deep Learning
Primary area: optimization
TL;DR: A quasi-second order optimizer with first-order efficiency by using Kalman Filter

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.

KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
Zixuan Xia* zixuan xia@students, wnwee. ch
Aram Davtyan aram. dww.daaaaannuiibb..... ch
Paolo Favaro
faolo.favarooumibecch
Abstract
We propose KOALA++, 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 calcula- tion, method directly estimates the parameter covariance matrix by recur- sively updating compact gradient covariance products. This design improves upon the original KOALA framework that assumed diagonal covariance by im- plicitly 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 better than state-of-the-ard first- and second-order optimizers while main- taining the efficiency of first-order methods. The code publicly available at https:///gttuue
Introduction
Optimization lies at the heart of deep learning. As neural networks continue to scale in depth, width, and data complexity, the efficiency, stability, and scalability optimization algorithms have become increasingly important. First-order methods such as stochastic gradient descent (SGD) [28 and its
adaptive variants like Adam [17 and RMSProp (00 offer practical efficiency and have driven much
of the progress in deep learning. However, they often suffer from instability, slow convergence, and sensitivity to hyperparameter tuning particularly in large-scale or noisy training regimes [16] (32.
To mitigate these issues, second-order optimizers such as K-FAC 20 and Shampoo (8) introduce curvature-aware updates by approximating the Hessian Fisher information matrix. These methods often improve convergence and generalization, but require costly matrix operations, limiting their scalability. Recently, AdaFisher [ has emerged as a more scalable alternative, using block -diaoona Kronecker approximations of the Fisher information matrix to reduce overhead while preserving second-order benefits. Despite its effectiveness, AdaFisher fundamentally depends on the quality of the Fisher matrix estimation, which remains sensitive to noise and mini-batch variability.
An alternative and theoretically grounded direction estimation theory. the posterior covariance matrix of an unbiased estimator defines the tightest achievable lower bound on parameter uncertainty, formalized by the Cramér-Rao bound [24] 3) Although the Fisher information matrix serves as an inverse bound proxy, it does not directly capture the true uncertainty in the parameters. Kalman filters, on the other hand, propagate the full covariance matrix of the estimate through time, providing a more faithful representation of uncertainty particularly in noisy or nonstationary environments.
from Kalman filtering. In Bayesian
Work done during Master studies at University of Bem Computer Vision Group, University of Bern wtttps://ww............t unibe.ch
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Motivated by this insight, KOALA proposed to frame stochastic optimization Kalman filtering problem, which models parameter updates as recursive state estimation, using covariance matrix to adjust step sizes adaptively. However, KOALA's original formulation relies on overly simplistic approximations, such as diagonal covariances and idealized gradient assumptions, which limit its practical effectiveness.
In this paper, we introduce KOALA++, a computationally efficient extension of the KOALA frame- work that explicitly maintains structured gradient covariance information. KOALA+ tracks the evolution of uncertainty directly through a recursive update at each training iteration ic of the gradient- covariance product ty = - which we call surrogate, where H₃ is the gradient of the loss function and P is the prior estimate of the parameter covariance. The key idea our method is to approximate the recursive 1- By doing so, KOALA++ effectively avoids expensive matrix inversions and storing the entire covariance matrix P while benefiting from the lifted constraints on its structure. Moreover, KOALA+ retains scalability akin to first-order methods, while achieving optimization dynamics and final performance competitive with state-of-the-art second-order methods.
sondiient-ovvriance update so that it never uses the covariance matrix
Our main contributions are summarized as follows:
We propose KOALA++, a Kalman-ine Kalman-inspired optimizer that explicitly propagates structured gradient uncertainty via covariance modeling.
We derive an efficient recursive update for the surrogate va based on minimum-norm solution to the covariance estimation problem.
We implement KOALA+ with minimal overhead and validate its effectiveness across diverse architectures and tasks, including the ResNet family [9. vision trans- formers [1], and language models.
Extensive experiments show that KOALA+ consistently matches exceeds the accuracy of first- and second-order baselines while offering improved stability and a better efficiency/performmnce trade-off.
Related Work and Motivation
2.1
Optimization Methods: A Brief Overview
First-order optimizers such as SGD 28 and Adam 17 are widely adopted for their scalability and simplicity, but often ignore the structure of gradient noise and require manually tuned learning rate schedules.
Second-order optimization methods, including K-FAC 20 Shampoo and more recently AdaFisher (7]. seek to address these shortcomings by explicitly incorporating curvature information. Nevertheless, practical implementations rely heavily fixed approximations, such block dock-diagonal or Kronecker-factored representations of the Hessian or Fisher matrices, to maintain computational feasibility. However, such approximations either still show limitations in terms of efficiency, when compared to first order methods, or adopt drastic simplifications that make the estimation process less stable.
2.2 Optimization via Kalman Filtering
The idea of using Kalman filtering for optimization has long history. Early works include scalar Kalman updates for stochastic optimization [13. applications in reinforcement learning [29], and connections between the Extended Kalman Filter (EKF) and natural gradient methods [22]. Although these approaches demonstrate the versatility of Kalman-based updates, they are often task-specific limited to small-scale models, and computing full covariance updates remains prohibitively expensive.
More recent research has revisited the Kalman filtering perspective in the context of large-scale neural optimization. Duran-Martin al. (6] and Chang et al. 2 interpret learning as a Bayesian state estimation problem, where the network parameters are treated as latent states, and the likelihood of the data provides an explicit probabilistic observation model. These methods apply low-rank or subspace- strutuurd versions of the Extended Kalman Filter to maintain tractable posterior updates.

---

Their formulation operates in the full parameter or feature space (i.e., with matrix-valued state and observation models), which implies a computational cost of ( O(n²) per update, where is the number of model parameters, even under low-rank approximations. Consequently, they are mainly applicable to some relatively small scale settings.
In contrast, KOALA 4 reinterprets gradient-based optimization as a scalar state estimation problem, where the loss function itself serves as an implicit observation. This design removes the need for an explicit probabilistic target and allows Kalman-style uncertainty propagation to be applied directly to general stochastic objectives. The formulation therefore scales linearly with the number of parameters, achieving (n) complexity in practice, while remaining compatible with arbitrary differentiable loss functions. This implicit observation model makes KOALA more general and applicable wider range of optimization problems than Kalman filters formulated on explicit probabilistic observations. Nevertheless, the original KOALA formulation still involves certain simplifying assumptions that limit its expressiveness and stability in complex, anisotropic optimization landscapes.
For comparison, SLANG [21 focuses purely on efficient covariance estimation through the diagonal plus low-rank structure without invoking any Kalman filtering or Bayesian state-update interpretation. Although SLANG achieves improved posterior approximations compared diagonal methods, its updates are not dynamically informed by temporal prediction-correction as in Kalman filtering, and its computational cost remains least On?²).
In general, these methods differ in their dimensionality, observation modeling, and scalability. A gap remains between the expressiveness of structured covariance estimation and the efficiency of scalar Kalman filtering. This motivates the design of KOALA++ which preserves KOALA's first-order scalability while introducing directionally aware covariance updates for large-scale optimization.
2.3
Motivation for KOALA++: Towards Structured Covariance Updates
The original KOALA formulation addresses efficiency by interpreting optimization as a scalar Kalman filtering problem, where the loss function acts as implicit observation. However, its covariance update assumes Ps Rs o²1 and substitutes H Hg with gradient norms, which may reduce some information to single scalar, limiting its ability to adapt to anisotropic curvature or correlated gradient directions, particularly in deep or noisy regimes.
On the other hand, as discussed in subsection 2.2 better approximations of the covariance matrix can provide richer uncertainty representations, but they are computationally prohibitive. This trade-off leaves a gap between expressiveness and scalability that current methods have yet to bridge.
Thus, KOALA+ is designed to overcome this limitation while retaining the first-order scalability of KOALA. Rather than maintaining or constraining a full covariance matrix, KOALA++ explicitly tracks directional covariance projection
2k H₂ Ph-11
which quantifies how uncertainty evolves along current gradient directions. This projection is updated via an approximate recursive least-squares step that implicitly captures curvature anisotropy without requiring explicit matrix storage inversion. As result, KOALA+ provides lightweight mecha- nism for propagating structured uncertainty, offering an effective balance between expressiveness and computational efficiency.
3 KOALA++
3.1 Problem Formulation and Kalman Filtering Notation
Neural network training is fundamentally an optimization problem, where, given a dataset D {(r,, y:) y the objective of the learning task is to adjust the model's parameters 0 E IR. so that the network's predictions f( 0) minimize the empirical risk
1 min L()), with L(0) = N [((())))..) 9
(1)
where is chosen loss function. However, in modern deep learning applications, N is often large, which makes full-batch optimization computationally infeasible. As result, practitioners rely
3

---

Table 1: EKF equations from
Sp HA(PA-1 Q + R
innovation covariance, where HA VL (0x-1)
KE (PR-1 + +Q)H 5K
Kalman gain
P: (I -KHHH))PPPT Q1)
posterior of the state update state estimate update
O2 = 0k-1 K L (0k-))
stochastic optimization methods that work minibatches of data and minimize estimates of the loss (+) = Σ 0),
BB
where By is set of indices in the k-th minibatch and [Bx] the cardinality of the minibatch. Although minibatch-basee approaches such as SGD [28] and its adaptive variants are effective and scalable, they introduce noise due to the stochastic nature of loss and gradient estimates.
To mitigate this issue, recent approach, KOALA casts neural network training as a recursive state estimation problem using Kalman filtering 15 estimating the hidden state of a system from noisy observations. prior estimates with new measurements update the state 02 and its associated uncertainty in the form of a covariance matrix Px. Importantly, by keeping track of an estimate the system's state's covariance, Kalman filtering is essentially implicitly using second order information about the state variables. The standard Extended Kalman Filter (EKF) generalizes this to non-linear systems using local linearization. In KOALA, model parameters 0₂ are treated latent states with stationary evolution under process noise, and the minibatch loss function L3 serves as a noisy observation of target loss value L In these settings, the EKF equations can be summarized as in Table To remain tractable in high-dimensional settings, KOALA simplifies the EKF update by:
principled Bayesian framework for recursively
each iteration k. it combines
Assuming an identity covariance matrix Py ofT:
Approximating H H3 with ||
Reducing all matrix operations scalar computations.
The resulting update closely resembles SGD, but with an adaptively scaled learning rate governed by the uncertainty estimate Pr.
While KOALA demonstrates promising stability and simplicity, its reliance on scalar approximations limits capacity capture directional or structured noise high-dimensional loss landscapes. To address this, we propose KOALA++, which enhances the original framework with an efficient estimation of the directional state covariance while preserving the Kalman- update algorithm.
3.2 Kalman Update with us -ased Reparameteriizttoon
To overcome the limitations of isotropic covariance assumptions KOALA, KOALA+ introduces more expressive formulation that captures directional uncertainty without incurring the full cost of matrix-valued updates. Instead of updating the full posterior covariance P& E guxn, we propagate a low-rank surrogate
UA = H P E IR 1xx
which represents the projection of the parameter covariance along the current measurement direction. This allows KOALA+ to maintain richer understanding of uncertainty and curvature structure while keeping the computational complexity comparable to first-order methods.
Low-Rank Reparameterization of the Kalman Update: By assuming constant process noise
Qk = Q for all k, the innovation covariance S₁ becomes
S (P + H + HaN + H&QH, R.
Using this expression, the Kalman gain becomes
K₁ (P 1 + QH
of + QHH
(5)
S
Sk

---

and the corresponding state update is
2/ +QH
O2 s- 02 + (I.E Larget (xx)))
(6)
Sk
Notice that, suggested in 4., the learning rate of the update can be controlled by selecting specific l* larget For example, 1 T. (1 m) (0x-)) ensures that the update is scaled by f]--
Recursive Update of Ub: To avoid storing P explicitly, we recursively update by substituting the Kalman covariance update P (I - Q)
Vk = H1Pk-11 -1 H&ll
ck- Q 1 HK-1
HI - -1
(Px-2 +
Sk-1
Lk k-1 + QHH
UR + +QH+-1 R-1
HHPPa-2+ H&Q
H
Sk-1
UK
Hs Sk-1
Q
For simplicity, define the recurring term
+ +QH
A = HR
(8)
Sk-1
which allows to rewrite the update for Up in more compact way:
U2 Ps- P + (H₂ x H
In the last equation, all terms except the first can be efficiently calculated in iteration k. The first term cannot be easily evaluated due to the term Pk-2. The objective is to never store it explicitly to avoid could approximate H3 with H and replace
the quadratic storage demand. As first attempt, the problematic term with Uk-1- However, we observed experimentally that this does not work well because of the high variability of the minibatch gradients.
As an alternative, we aim to retrospectively approximate the prior covariance matrix Pa-2 instead. We do this by taking advantage of the relation f H However, since this equation is generally underdetermined in high-dimensional settings, we formulate regularized inverse problem to obtain unique and tractable solution. Specifically, adopt the minimum Frobenius-nom formulation
min 2
subject to H HR-1PR-22 = Uh
(10)
which yields closed-form least-squares solution. This construction allows approximate Pk-2 efficiently without explicitly storing or inverting full covariance matrices.
For computational and memory efficiency, we use only one linear constraint, Pk-1 H Pa 2- although more could be used in principle. However, introducing more constraints, e.g., projections of other previous gradients HR-2, HR-3. etc., would also require memorizing more gradient terms and using more calculations to update Uk- The use of the minimum Frobenius-norm criterion is to ensure stable update for U2 by limiting the norm of Pk-2-
In the next subsection, we discuss two variants of our algorithm, depending on the nature of the solution to the optimization problem 10
problee
3.3 Least-Squares Estimate of P
Variant I: Vanilla Least-Squares Estimation In the first variant of our approach, we solve the lquates problem without imposing any structural constraints on the covariance estimate. Given the relation fk-1 = H Pt-2, we compute Pe-2 RRax by introducing the Lagrange multipliers J 1xe and write the Lagrangian
C(Px-2,)) 1 (Ha-1Ph-2 vx-1)p.
(1I)
5

---

Algorithm +
Initialize 0o, 21, Q, R. and fix the learning rate schedule 7 for 2 to do
For simplicity, denote H₂ VL 11)
Calculate Cge, Tk 7k respectively from Equations Update:
13 and (16)
Uk = (ox A + (HA
A HE 1 H
(19)
mn (0
0₂ = 0k-1 HaU H&QH +
(Np + x
(20)
end for
return 0r
This formulation admits the closed-form solution
PR-2 = 1 vk-1
(12)
H
By substituting this solution into H&Pk-2 gives
H k-1
= = where a =
(13)
1|
Finally, we can obtain the following vanilla uk update
02 = (arss A (HE- H1
(14)
Variant II: Symmetric Covariance Estimation To improve stability and theoretical consistency with classical Kalman filtering, we introduce a second variant that enforces symmetry in the ap- proximation Pk-2. 2- In particular, the additional regularization term Px-2 P is added the objective. the previous derivation, we obtain closed-form solution for Pa-
H Uh
+ Uk-1 Ha-1
H H
PR-2
H Un-1
(15)
H
When substituting this into H4. Pr.PP-2. the only difference between the variants is the term
(Hxv (H+ H (Ha-10 1
(16)
|Ha-1|1
To simplify notation and implementation, we adopt unified update rule for both variants. In the unconstrained case, we simply set r = O, effectively reducing the update the asymmetric version. This allows the general form of the update to be written compactly as
(ax A + (H₃ > Ha- 1)Q H
(17)
To maintain consistency with the original KOALA formulation and to simplify the calculation in the early stages of training, we initialize the directional covariance product i1 H,Po using scaled identity assumption on the initial covariance matrix Po. Specifically, we set:
H1,
(18)
where of is scalar hyperparameter representing the initial uncertainty. This corresponds to assuming Po = T, which aligns with the isotropic prior used KOALA and allows avoid explicit matrix computations at initialization. The general procedure of KOALA COALA++ ++ is summarized in Algorithm
'See Appendix B. for detailed derivation.
"See Appendix B for detailed derivation.

---

Table 2: Experimental results on CIFAR-I0 and CIFAR-100 with different optimizers for 100 and 200 epochs. Best (blue) and second best (orange) performances are in boldface and with shading.
Dataset
Architecture
Method
Error
100-cpochs
200-epochs
Top-1 1 Err
Top-5 Err
Top-1 Err
Top-5 Err
SGD
5.699.19
0.18u.03
4.83u.20
0.15u.ua
Adam
6.960.06
0.270.05
6.35a.13
0.27o.
ResNet-18
KOALA-M
6.29g.04
0.37o.02
6.06o.14
0.280 08
KOAIA++(Ours)
5.61o.10
0.260.04
5,43a.19
0.24
SGD
6.6la.12
0.17u.02
5.38a.10
0.14u.
CIFAR-10
Adam
6.38a.12
0.266.02
5.77a.20
0.21o.02
ResNet-50
KOALA-M
8.090.16
0.290.04
7.780.39
0.32o 06
KOALA++OUrs)
5.28
0.21o.04
4.90g.21
0.180.06
SGD
6.19g.a6
0.170.03
4.91a.14
0.13u.0a
W-ResNee-50-2
Adam
6.03g.07
0.24n.03
5.530.21
0.27
KOALA-M
7.62a.13
0.300.04
7.670.43
0.32a.
KOALA++(Ours)
4.94
0.160.05
4.70g.17
0.170.01
SGD
25.070.88
6.890.39
22.280.31
5.55o
Adam
24.25n.34
7.06n.
24.14
7.470.34
ResNet-50
KOALA-M
23.71d.33
6.58p.04
22.910.38
6.09a.33
KOALA+H(Ors)
22.249.44
5,970.32
21.86g.06
5.790.00
CIFAR-100
SGD
23.030.29
5.970.11
23.500.15
6.470.16
W-ResNet-50-2
Adam
23.996.17
6.690.07
23.740.19
7.10n.17
KOALA-M
21.59
5.71n.12
22.220.34
6.15n.17
KOALA+<Ours)
21.23
5,33o.13
20.330.21
5.02a.19
Experiments
We evaluate KOALA+ on range of vision and language tasks to demonstrate its generality, stability, and convergence behavior compared to existing optimizers. Our experiments are organized into three parts: image classification, language modeling and ablation studies. Due space constraints, we report representative subset of results the main text. Additional evaluations, including ablations, hyperparameter tuning, and extended visualizations, are provided in Appendix
4.1 Image Classification
CIFAR-10/100 Classification: We follow the experimental setup of the original KOALA paper for CIFAR-10 10 and CIFAR-100 classification tasks, including data augmentation, model architectures, and optimization settings. We introduce two modifications for CIFAR-10: (i) training is repeated with three random seeds (42, 3407, and 2025) to ensure statistical robustness, and (ii) we adopt two-stage learning rate decay strategy the 100th and 150th epochs for 200-epoch training (inspired by successful multi-step decay strategies reported in prior works such as (27)) which improves
generalization compared the single-step decay used the original paper. The result are shown in Table 2
Comparison with Recent Advances To comprehensively evaluate the effectiveness and general- ization capability of KOALA++, we design two experimental settings that progressively incorporate more recent optimizers, modern network architectures, stronger learning rate schedulers. and richer data augmentations. This design mirrors recent benchmarking practices in the literature, particularly those from KOALA and AdaFisher 77.
(1) CIFAR-100 + ResNet-50 + Advanced Optimizers. To evaluate KOALA+ in standardized yet competitive setting, we follow the benchmarking protocol introduced in KOALA (4, and compare our method with wide range of recent optimizers on the CIFAR-100 dataset using ResNet-50, trained for 100 epochs. All optimizers are tested under the same learning rate schedule (StepL.) and training configuration. For each optimizer, we adopt hyperparameters as suggested in the original papers or official implementations. The settings and results are shown in Table 3
"Some of earlier reported results were obtained from existing public benchmarks, and therefore mean and standard deviation values are not available for all entries. To ensure reproducibility, we reran the experiments under identical settings; the full code and configurations are available at: sttps://gither com/ Sumxiaa/ KOALA_ Plus__luPPPP PPuss

---

Table 3: CIFAR-100 (ResNet-50,, 100 epochs, StepLR) results and optimizer hyperparameters.
Optimizer
Top-1 Err.
Top-5 Err.
Hyperparameters Source
Yogi 25
33.99
10.90
Ir=10 31099, 32=0.999, c=10-
Adamax
32.42
10.74
Same as Adam
AdamW
27.23
7.98
Same as Adam
AdamP 10
26.62
7.61
Same Adam
Amsgrad 26
25.27
6.78
Same as Adam
Adan 31
24.92
6.86
Paper and repo default
Fromage
24.65
6,71
Ir=10 (recommended in GitHut
AdaFisher 7
23.38
6.05
Original AdaFisher hyperparams
Adabelief 33
23.07
6.05
Official config
KOALA-M -M
23.71
6,58
Redid the experiments with our implementation settings
KOALA++ (Ours)
22.24
5.97
Based our implementation settings
GXLARS KXLAR Alm
100
KINLANN
2.0
(Validation)
3.5
60
0.0
20
40
60
80 100
20
40
60
80
100
Epochs
Epocha
(a) Test Loss (log)
(b) Test Top-1 Error
Figure I: Comparison of training loss, validation loss, and validation error for different optimizers on CIFAR-100 using ResNet-50. KOALA+ demonstrates the strongest performance drop at scheduled learning rate decays (epochs 30, 60, and 90), highlighting its superior scheduler responsiveness.
Figure I presents the training and validation dynamics on CIFAR-100 with ResNet-50 under a multi-step learning rate schedule epochs 30, 60, and 90. Compared to SGD, Adam, AdaFisher, and KOALA-M, our proposed KOALA+ demonstrates consistently lower validation loss on the logarithmic scale, along with improved generalization as shown by the lower Top-1 error.
In particular, the benefits of KOALA+ are most evident as scheduled learning rate drops: the optimizer exhibits sharp improvements in both loss and error immediately following the milestones, indicating more effective adaptation to the changing optimization dynamics. This suggests that KOALA++ better integrates structural information from past gradients and uncertainty estimates, allowing to respond more effectively to scheduler-inducee, shifts.
Moreover, KOALA+ maintains competitive or superior stability throughout all training phases, without exhibiting the overfitting spikes or stagnation observed in other optimizers. These obser- vations collectively validate K as a robust and efficient optimizer for large-scale vision tasks.
(2) CIFAR-10 + Modern Models Cosine Annealing + Strong Augmentation To align with the stronger benchmark protocol proposed by AdaFisher [7], we conduct a second set of experiments that adopt more modern practices deep learning optimization:

---

Table 4: Top-1 accuracy (%) CIFAR-ID across various architectures. We report the average over seeds. The large number is the mean and the subscript is the standard deviation.
Optimizer
ResNet50
ResNet101
DenseNet121
MobileNetV3
Tiny Swin
SGD
95.7
95.980.2
96.09u.1
94.430.2
82.34o.2
Adam
94.45
94.570.1
94.86o.1
93.32g.1
87.370.6
AdaHessian
95.54o.1
95.290.6
96.11u.1
92.86u.1
84.15u.2
K-FAC
95.666.1
96.01o.1
96.12..1
94.34..1
64.79p.5
Shampoo
94.599.1
94.630.1
95.66p.1
93.81p.2
63.91g.4
(daFisher(W)
96.340.2
96.39..1
96.72..1
95.28
88.74o.4
KOALA+ (Ours)
96.449.1
96.660.1
96.60g.1
94.52u.1
87.64g.2
MobileNetV3 [.] and the Swin Transformer (Tiny) [18), reflecting diversity of architectures including convolutional and attention-based designs.
Advanced models: ResNet-101, DenseNet-121 [12],
Modern learning rate scheduler: We use CosineAnnealingll. for all optimizers, which has been shown to improve convergence stability and final accuracy in modern setups.
Stronger data augmentation: In addition to standard cropping and flipping, we apply Cutout 5 during training to improve generalization.
Larger batch size: Following current best practices, we increase the batch size to 256, which can benefit optimizers that sensitive noisy gradients.
More Random Seeds: Following the AdaFisher benchmark protocol 77 we report top-1 accuracy averaged over five random seeds: 42, 3407, 9331, 7, and 2025, The standard deviation across these runs is shown in the bottom-right corner of each table entry.
As shown in Table 4 KOALA+ performs par with AdaFisher across range of architectures, while exhibiting consistently lower standard deviations, indicating greater stability among random obtained without extensive hyperparameter adjustment, highlight-
seeds. All KOALA++ results ing the robustness of the method in the default settings. is also worth noting that on Swin-Tiny, KOALA+ reaches strong performance with significantly smaller variance, demonstrating better adaptability to Transformer-based architectures.
Efficiency profiling. Beyond accuracy, also compare the computational efficiency of KOALA++ with strong baselines. KOALA+* achieves runtime and memory comparable to Adam(W), while being substantially more efficient than Shampoo and AdaHessian. A detailed comparison, including per-epoch time, FLOPs, and peak memory usage on ResNet-50 and Swin-Tiny, provided in the Appendix E.3
4.2
Language Model
We evaluated KOALA++ on the Wikitext-2 dataset, following the same language modeling bench- mark setting introduced in the AdaFisher paper [7]. The model is scaled-down version of GPT-1 with four masked self-attention layers and approximately 28 million parameters. We adopt the same architecture, tokenizer, and 50-epoch training schedule as used in the original benchmark.
In line with the original AdaFisher benchmark (see Githut corresponding to the best-performing model on the validation set, along with the total training time, provide holistic view performance and efficiency. The results are summarized in Table 5 In addition to reporting test perplexity and total training time, we further provide visualizations to better illustrate the training dynamics (see Appendix E).
report the perplexity (PPL)
4.3 Ablation Study
Due to space constraints, we report only the core ablation comparing KOALA-M and two variants of KOALA+ CIFAR-100 with ResNet-50. The default KOALA+ uses symmetric gain

---

6/1 Am0
1.25
Top-1 Einor (Velldetoon)
100
0.50
90
0.25
90
80
100
Epochs
Epochs
(a) Test Loss (log)
(b) Test Top-1 Error
Figure 2: Ablation study comparing KOALA++ its asymmetric variant KOALA+ (NS), and KOALA-M on CIFAR-100 with ResNet-50. Left: Log-scale test loss curves. Right: Validation Top-1 error rates.
Table 5:
Language modeling performance on Wikitext-2. We report the test perplexity (lower is better) corresponding to the best validation model, well the total training time over 50 epochs.
Optimizer
Test PPL 4
Average Training Time per Epoch (s)
AdamW
173.98
23.11
AdaHessian
407.69
62.91
AdaFisherW
152.72
23.39
KOALA+ (Ours)
151.75
25.56
approximation, while KOALA+ (NS) adopts asymmetric covariance. As shown in Figure 2 KOALA+ achieves faster convergence and better final performance. KOALA++ (NS) is competitive but less stable in later stages, underscoring the benefit of structural symmetry. Full ablations provided in the Appendix D
5 Conclusions, Limitations and Future Work
We proposed KOALA++, an efficient optimization method that extends the Kalman- KOALA algorithm by incorporating lightweight directional covariance approximation. KOALA+ retains the stability and simplicity of its predecessor, while substantially improving capacity capture general structured uncertainty in digh-dimensional loss landscapes. Through extensive experiments across vision and language tasks, we demonstrate that KOALA+ achieves competitive supe- rior performance compared to both first- and second-order optimizers, while maintaining training efficiency on par with standard first-order methods.
Limitations. A notable limitation of our method is that the approximated covariance matrix Ps is not explicitly constrained be positive semi-definite (PSD). Although we empirically observe that the eigenvalues of P& remain stable and bounded throughout the training (see the Appendix B.3). the lack of formal PSD guarantee may pose risks certain settings or applications involving sensitive numerical stability. Addressing this limitation efficiently is left to future work.
Future Work. To address this issue, future work may explore enforcing PSD constraints via projec- tion, spectral regularization, or reparameterization techniques. Additionally, we plan to investigate the integration of KOALA+ with adaptive gradient clipping and apply large-scale pretraining scenarios such as vision-language models, diffusion models, or instruction- LLMs.
Acknowledgements
Calculations were performed on UBELIX (https://www.id) id. unibe. ch/hpe), the HPC cluster at the University of Bern. Aram Davtyan has been supported by SNSF Grant 10001278.
10

---

References
[1] Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems, 53:21370-21381, 2020.
[2] Peter G Chang, Gerardo Durán-Martin Alexander Y Shestopaloff, Matt Jones, and Kevin Murphy. Low-rank extended kalman filtering for online learning of neural networks from streaming data. arXiv preprint ar0iv:2305. 190555555 2023.
[3] Harald Cramér. Mathematical methods of statistics, volume 9. Princeton university press, 1999.
[4] Aram Davtyan, Sepehr Sameni, Llukman Cerkezi, Givi Meishvili, Adam Bielski, and Paolo Favaro. Koala: A kalman optimization algorithm with loss adaptivity. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6471-6479 2022.
[5] Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint 0rXiv:1708.0455:: 2017.
[6] Gerardo Duran-Martin, Aleyna Kara, and Kevin Murphy. Efficient online bayesian inference for neural bandits. In International conference on artificial intelligence and statistics, pages 6002-6021. PMLR, 2022.
[7] Damien Martins Gomes, Yanlei Zhang, Eugene Belilovsky, Guy Wolf, and Mahdi S Hos- seini. Adafisher: Adaptive second order optimization via fisher information. arXiv preprint arXiv:2405. 16397, 2024.
[8] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018
[9] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition. pages 770-778, 2016.
[10] Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and Jung-Woo Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights. arXiv preprint 2rXiv:2006.082 2020....
[11] Andrew Howard, Mark Sandler, Grace Chu, Liang- Chinh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, al. Searching for mobilenetv3. In Proceedings of the IEFF/CVF international conference on computer vision, pages 1314-1324, 2019.
[12] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 4700-4708, 2017.
[13] Mahmoud Ismail, Mina Attari, Saeid Habibi, and Samir Ziada. Estimation theory and neural networks revisited: Rekf and rsvsf as optimization techniques for deep-learning. Neural Networks, 108:509-526, 2018.
[14] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.
[15] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of Basic Engineering, 82(1):35 45, 1960.
[16] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations (ICLR), 2017.
[17] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference Learning Representations (ICLR), 2015.

---

[18] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[19] Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[20] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approx- imate curvature. In International conference machine learning, pages 2015.
PMLR,
[21] Aaron Mishkin, Frederik Kunstner, Didrik Nielsen, Mark Schmidt, and Mohammad Emtiyaz Khan. Slang: Fast structured covariance approximations for bayesian deep learning with natural gradient. Advances in neural information processing systems, 31, 2018.
[22] Yann Ollivier. The extended kalman filter is natural gradient descent in trajectory space. arXiv preprint arXiv: 10t0:/901.0000 000996, 2019.
[23] Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
[24] C.Radhakrishna Rao. Information and the accuracy attainable in the estimation of statistical parameters. Bulletin of the Calculta Mathematical Society, 37:81-91, 1945. Reprinted in: Kotz and Johnson (Eds.), Breakthroughs in Statistics, Springer, 1992.
[25] S Reddi, Manzil Zaheer, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Proceeding of 32nd conference on neural information processing systems (NIPS 2018), volume 2018.
profrmationprocessiing
[26] Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv: 192709977 2019.
[27] Leslic Rice, Eric Wong, and Zico Kolter. Overfitting in adversarially robust deep learning, In International Conference on Machine Learning (ICML), pages 8093-8104. PMLR, 2020.
[28] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, 22(3)4400-477 1951.
[29] Shirli Di-Castro Shashua and Shie Mannor. Trust region value optimization using kalman filtering. arXiv preprint arXiv:1901.07866 2019.
[30] Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5 rmsprop: Divide the gradient by running average its recent magnitude. Coursera: Neural Networks for Machine Learning. 2012.
[31] Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions Pattern Analysis and Machine Intelligence, 2024.
[32] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations (ICLR), 2017.
[33] Juntang Zhuang. Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvoreek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33:1878 189918951111 18806, 2020.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly state our main contributions and support them with theoretical and experimental results (see Sections and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in Section including assumptions covariance symmetry and potential instability.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
13

---

Answer: [Yes]
Justification: The assumptions are made in Section and the derivations are provided the Appendix B..
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All experimental settings are described in Section and Appendix including models, datasets, and hyperparameters.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
14

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We plan release the code upon publication, and all datasets used (CIFAR- 10/100, WikiTeet-2) are publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnis.cc
aublic/guidees/Coeeebbbmsss
caapssioooolyyyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. wttps/cccpbbccggiui.ceeee wtw.eeeeeeoSSSeeesSSSeeeeee.
www.ssssiicggoooooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Details of model architectures, optimizer configurations, learning rates, and schedules are provided in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We report averages over 3-5 seeds and include standard deviations in Table and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
15

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We describe compute resources in Appendix including training time on an H100 GPU.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't itdn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips... http:/riipppppppppce/ee//eeeaaaa casetteeeeeessc
Answer: [Yes]
Justification:
We adhere to the NeurIPS Code of Ethics. No private sensitive data is used.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper does not discuss societal impacts, as it focuses solely on the design and evaluation of optimization algorithms, which are technical contributions without direct societal applications.
Guidelines:
16

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not describe such safeguards, as it does not involve the release high-risk models or datasets such as pretrained language models, image generators, or scraped data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All external assets used in the paper (e.g., the AdaFisher benchmark) are prop- erly credited, and their licenses and terms of use (e.g., MIT License) have been respected.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
17

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pa se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's ' creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The reimplementation of the KOALA+ optimizer is new asset introduced in this paper. well documented, and usage instructions will provided alongside the code repository.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The study does not involve human participants, so IRB approval and risk disclosures are not applicable.
Guidelines:
18

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [Yes]
Justification: We used LLMs for proofreading and structuring (e.g.. ChatGPT), but they were not used as part of the core methodology in this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy for what should or should not be described.
19

---

Appendix Contents
A
KOALA
20
A.1
Risk Minimization as Loss Adaptivity
20
A.2
KOALA-V KOALA-V(Yanilla
21
A.3
KOALA-M (Incorporating Momentum Dynamics)
21
B
Mathematical Foundations of KOALA++
22
B.1
Vanilla Optimal Solution of Covariance Matrix P
22
B.2
Symmetric Optimal Solution of Covariance Matrix
22
B.3
Eigenvalues Analysis of the Optimal Matrix 1
23
C
Convergence Behavior KOALA++
25
C.I
On the Stability and Implicit Positive Semi Definitenes of KOALA++
25
C.2
Convergence Characteristics of KOALA+ Quasi-Second-Order Optimizers
26
Ablations
27
D.1
Effect of Batch Size on KOALA+ Performance
27
D.2
Different Schedulers
27
D.3
Different Initializations of KOALA+
28
E
Experiments Details
29
E.1
Hardware
29
E.2
Image Classification
29
E.3
Efficiency Analysis
E.4
Language Modeling
A
KOALA
In this section, going revisit the original KOALA framework, which only briefly introduced in the main text. KOALA 4 is a Calman-inspire optimization algorithm that in- terprets gradient-based updates recursive state estimation process. To ensure scalability in high-dimensional deep learning settings, the original work proposed two tractable variants: KOALA- and KOALA-M. These variants differ in how they approximate the parameter covariance and in whether they incorporate momentum dynamics. In what follows, we provide detailed and self-contained exposition of both KOALA-V and KOALA-M. including their derivations, update rules, and practical considerations. This not only offers a clearer understanding of KOALA as foundational method, but also sets the stage for the improvements introduced KOALA++
A.1 Risk Minimization Loss Adaptivity
The core idea behind KOALA is to cast the optimization deep neural networks as a state estimation problem under uncertainty. Specifically, neural network training is considered to minimize noisy risk function La, which is stochastic estimate of the empirical risk L. Motivated by the central limit theorem, KOALA assumes that the mini-batch risk Lg can be modeled Gaussian random variable with scalar noise
(0) ~ L(0) - Uks where Up (O,,R+).
(21)
0 denotes the model parameters and Rs. is the variance of the observation noise.
Rather than optimizing minş L()) directly, KOALA reformulates the problem by adapting the model parameters that the stochastic risk (Ox) matches desired target risk L jugel
i.e..
La (0x)
target
Uk.
(22)
20

---

This formulation enables the use of Kalman filtering, which provides a principled way to recursively update the estimate of 0g by combining prior knowledge with new and noisy observations. The underlying system can be described as standard dynamical model:
O2 = 0x-1 + wk-1,
Wh ~ (0, Qu)
(23)
Larget L
(0n) + Uk: Uk ~ (0, Rs).
(24)
A.2 KOALA-V (Vanilla)
The idea behind KOALA-V is remarkably natural, which is directly apply the standard Kalman filtering update equations to the system in Equations (23) and (24). However, directly computing the posterior covariance Pt R=*=" is infeasible for high-dimensional models. To address this, KOÅLA-V introduces the following approximations:
The posterior covariance Ps. is approximated by a scaled identity matrix:
Pa La
(25)
.g Hy is approximated by the squared norm of the gradient:
H H& ((.) In-
(26)
The resulting parameter update rule is given in Algorithm 2
This update closely resembles the SGD update:
0₂ 1 L.
(27)
but with key difference: in KOALA-V, the scaling of the gradient is automatically adapted based on the current loss, gradient norm, and covariance estimate. This adaptivity makes KOALA-V more robust noisy training dynamics than static first-order methods like SGD or Adam.
Algorithm 2 KOALA-V (Vanilla)
Initialize Bu. Po. Q and R for i in range(1, 7) do Predict:
02 = Ok-13 P
Update:
PLLL - L target
0, = 02
vLu((k)
(28)
PAII (0x)) 2 + R
R Ps
Ps PRI (0,))1 +R
(29)
end for return 0r
A.3 KOALA-M (Incorporating Momentum Dynamics)
KOALA-M extends the original KOALA framework by explicitly integrating momentum-based dynamics into the Kalman filtering formulation. Instead of modeling the parameter update as simple random walk, KOALA-M introduces an auxiliary momentum state Pk. resulting in the augmented state dynamics:
0g 0x-1 + Pk-1 + Wk-1,
Wk-1 ~ N/O,, Q).
(30)
Pk Pk-1 1 2
My-1 N(O,U),
(31)
21

---

where и controls the momentum persistence. This momentum augmentation enables KOALA-M M to better capture and leverage correlations across sequential updates, effectively enhancing stability.
However, KOALA-M inherits computational complexity due to its augmented state. To mitigate this, KOALA-M approximates the posterior covariance with structured form, typically assuming he a Kronecker product of a 2 x 2 covariance matrix with an identity matrix 4, Similar to KOALA-V, KOALA-M employs layer-wise covariance updates to obtain better approximation of parameter uncertainty, partially alleviating the limitation of overly simplistic covariance structures.
In contrast, our proposed KOALA++ algorithm does not require explicit momentum augmentation or restrictive covariance approximations due to its inherently richer and adaptive covariance repre- sentation. KOALA+ directly maintains structured gradient-covariance products without imposing additional structural assumptions on the covariance matrix P, thus naturally capturing both directional information and momentum-like dynamics implicitly.
B
Mathematical Foundations of KOALA++
In the main paper, we provided both the non-symmetric and symmetric solutions for Ps-2. this section, provide the detailed derivations of those. For notational simplicity, where possible, drop the subscripts in k and we can restate the optimization problem as follows:
min PI P
subject to =P=0
(32)
B.1 Vanilla Optimal Solution of Covariance Matrix
To solve this constrained optimization problem, we introduce Lagrange multiplier μ 1x and construct the corresponding Lagrangian function:
((P.)) |PI| + H(HP - T)
(33)
We then find the optimal solution by differentiating the Lagrangian with respect to P, and set this derivative to zero to find critical points:
OC(P.P) p) = -
1 HT 2
P
(34)
OP
Now, substituting this optimality condition into the constraint equation HP v, have:
H
1 (-1HHT
H 1
-リ=リ= HHH
(35)
Thus, the optimal solution for P is obtained by substituting back the expression for Jc
H T
(36)
H
This solution is referred to as the vanilla optimal solution (asymmetric) because it does not explicitly impose symmetry constraints on P. In the next subsection, we derive a symmetric solution variant that respects the symmetry requirement typically imposed on covariance matrices.
B.2 Symmetric Optimal Solution of Covariance Matrix P
For derivation of the symmetric solution, we use the same Lagrangian framework, but now explicitly incorporate the symmetry constraint into the formulation. The Lagrangian function, considering the symmetry constraint, remains the same as before (see Equation (33)).
22

---

We use the symmetric matrix derivative form from the Matrix Cookbook 23 Taking the derivative with respect to P and setting it to zero yields
OC
0C
0L 0P
OC 0P
0P Symmetric) = OP
+
diag
=
(37)
By plugging in Equation 33 we have
2P + p + +P + H diag(2P + p H) = 0
(38)
Since P is symmetric, we can rewrite Equation 38 as:
4P + + 11 diag(2P + p H) = 0,
(39)
which gives the following equations
4P + p T + H == (off diagonal) (on diagonal).
(40)
2P + H 1=
We observe that the off- off-diagonal solution satisfies the on-diagonal equation the diagonal. Therefore, we only need to solve the off-diagonal case. From the off-diagonal constraint 4P + H H =
have P - (1² H H p). By substituting in HP U, obtain
1 (H H + HH p1) =. HP = - 4
(41)
We solve y as
/ == - H H HH T)-,,
(42)
where denotes the identity matrix. To compute the inverse, we recall the Woodbury matrix identity (23)
(A+ (-+UV--12 - --
(43)
By letting A HH HT and V H, find the following expression
1
(H 1 H) HHT 1 ( H -
(44)
2
HH
This provides expression for μ without matrix inverses, i.e.
HHT (
H 2 HH
(45)
Finally, we can plug this expression have a closed-form optimal solution P that is,
D 1
(I [
1H HH v T H 2HHT)
1 HHT 2 HHT
1H T H
I
HHT
v H H HuT HHT (HHTJ2 H H.
(46)
B.3 Eigenvalues Analysis of the Optimal Matrix P Recall the closed-form expression for the optimal matrix
1 ( b
HUT
P H (H
H H. (H
(47)
To analyze the definiteness of P*, we consider its eigenvalues. Let w E gax1 be eigenvector of P and λ its corresponding cigenvector, then,
Aw
P w.
(48)
We can immediately see that w lies the span of span. Then, can write
T and T because the range of P lies in their
aHl T + BU
(49)
23

---

By substituting it into the eigenvalue equation we obtain
Xall + +B0 HHT (H + H)laH + BeT) HE (HHT)Z² H' T H(aft + BBo (50)
We now compute both terms, For simplicity, let us name the following three scalar values as
x =HH",, y = HoT, vo
(51)
Then, we have
ww== (H" vaH T + H uBo" + +o Hall + HBb
(52)
1
x2 (Hu" H HaH + Hu H H3v
(53)
1 (ayll T + +B2l + taxv" Byr )
1
1
ayall
312111""
11
(54)
x2
12
By 3z HT + aov
By
0
(55)
I
x²
v²) f(zz
at
By
(56)
x²
By matching both sides coefficient-wisee we derive a 2D eigenvalue system as follows
a1 B(xz -2)
(57)
x² 0. + By
8A
(58)
x
If a 0, then from Equation (57) we know that 33 = 0 or 2z y² = 0. However, 3 = 0 implies that = aHH + Nr 0. which contradicts the assumption that w is non-zero eigenvector. On the other hand, by the Cauchy-Schward v become linearly dependent. In this case, the eigenvalues of 1* are 0 and a, where H T av T due to the linear dependence. The linear dependence is, however, very unlikely in practice, as, in general H and span two-dimensional subspace.
inequality, æz y² = implies that the two directions H and
Therefore, we very likely have a + O. In this case, the eigenvalue must satisfy
8
F2 y²
(59)
x
Now, we need to solve the following equation to examine the sign of #
(xz - 1/2) (g)
2
ry
x²
(60)
a
We get
B ry + V y22y2 + 4(xz x²)
(61)
a
2(xz 82)
Finally, the eigenvalues of the matrix P are
zy -
22 + 4(12 y2)
\12
(62)
2x
Although the matrix P obtained through least-squares minimization does not guarantee positive definiteness, we conducted an empirical analysis by visualizing the evolution of its eigenvalues at different training stages.
Specifically, performed eigenvalue tracking the ResNet-50 model trained CIFAR-100 at 100 epochs. The learning rate was decayed at epochs 30, 60, and 90, which correspond to the early, middle, and late stages of training, respectively.
To better understand the behavior of the matrix P. we analyzed the distribution of its eigenvalues at four critical epochs: 0, 30, 60, and 90. The results are shown in Figure 3
across all

---

Purticiio of Large and Smrvll Rigenselioss
Panticioon of Lange and Small figensaiom
Lum
Oper
-02
306 th
200
159
150
25.3
(a) Epoch O
(b) Epoch 30
Figure 3: Evolution of the large and small eigenvalues of the matrix at epoch and epoch 30 across all
We found the cases for 60 and 90 epochs to give identical trends as for 30 epochs case. Thus, avoid unnecessary repetition, we are not showing the plots for these cases.
At epoch 0 (Figure 3 left), the eigenvalues exhibit a wide dynamic range with both large and small eigenvalues increasing significantly over time. The large eigenvalue (orange) and small eigenvalue (blue) rise rapidly and display sharp oscillations, indicating an evolving optimization landscape in the early stage of training-
However, by epoch 30 (Figure 3 right), the eigenvalue trajectories stabilize, and the magnitude of fluctuations reduces dramatically. From epoch 30 onwards, the eigenvalue distributions become increasingly similar. In these later epochs, both the large and small eigenvalues remain relatively flat and nearly constant throughout training, reflecting stabilized matrix structure and learning dynamics.
C Convergence Behavior of KOALA++
C.1 On the Stability and Implicit Positive Semi-Definiteness of KOALA++
While the current formulation of KOALA++ does not explicitly enforce the positive semi-definite (PSD) constraint on the covariance- like matrix, we empirically observe that the optimizer remains numerically stable and convergent across diverse architectures and datasets. This suggests that, in practice, the lack an explicit PSD guarantee does not adversely affect convergence behavior.
This is probably because, when we are calculating the eigenvalues in Equation (62). the quantities λ1 are obtained from the least-squares solution of the auxiliary matrix Pk-2, which represents the optimal estimate under the minimal-norm fitting objective. In other words, the eigenvalues in Equation (62) correspond to an analytically fitted surrogate of the covariance update, rather than the actual covariance-like matrix used in the algorithmic recursion. In practice, the true update referred to as the ground truth matrix-is defined by integrating the solution from Equation 12 into (9,, then have:
H TH (Hu T H)
+QI (OT QQT)(( -QH)
Uk
2
H(wT + +QHT)+
(63)
We call the ground-truie truth matrix Mt and express it as follows
HTD To+ (HoT)(H H)
(o + QQ ( QH) H(vT QH
(
M
+QI
(64)
|
[
Based on this formulation, we can further examine the directional definiteness of the update. Specifi- cally, by evaluating the product Hau we can infer whether the update direction induced by M₁ corresponds to a positive-definit step along the current gradient direction. In particular, when
HaU = HrM,HH 0,
(65)
25

---

Angle between V, and H. per Epoch
Angle between Vi and H: per Epoch
Angle between V, and H2 per Epoch
Averago: Angle
Avenge Angle
Average Angle
24
Seorogg Angle
18
2.5
25
12
1.5
20
40
120
0
20
40
63
100
20
40
60
50
100
Epoch
Epoch
Esoch
(a) CIFAR-10
(b) CIFAR-100 100
(c) ImageNet32
Figure 4: Empirical verification of the positive semi-definiteness of Ms in KOALA++. The plots show the evolution of the angle between the curvature vector Hg and the update direction eg across training for three datasets. The consistently acute angles (< 90%) indicate that Hş and uk remain directionally aligned, confirming the effective PSD behavior of M2 in practice.
the update matrix M2 can be considered positive definite with respect to the local curvature defined by Hę. This provides a practical criterion for assessing the stability and consistency of the covariance update in KOALA++.
Empirically, we verify this property on three benchmark datasets: CIFAR-10. 10, CIFAR- 100, and ImageNet32. For each case, we monitor the angle between the curvature vector H₃ and the update direction va throughout the training. As shown in Figure 4 the angle remains consistently acute (< 90%) in all training stages, which implies that Hx 0, thus confirming the numerical positive semi-definiteness of Mg and the empirical stability of the update rule.
However, from a theoretical standpoint, ensuring postdive-eemideefiiienss could further strengthen the stability guarantees of the method, In particular, one promising direction is to modify the optimization objective so that the covariance estimate minimizes its distance-for example, in the Frobenius norm norm-to a structured PSD matrix such as a scaled identity. Such formulation naturally leads to convex optimization problem with a unique and bounded solution, thereby eliminating ambiguity and improving theoretical soundness.
Although this PSID-enforced variant is beyond the current scope of our study, we view as an impor- tant extension future work aimed at establishing stronger convergence guarantees for KOALA++
C.2 Convergence Characteristics of KOALA+ vs. Quasi-Second-Order Optimizers
To better understand how KOALA+ differs from quasi-second-order optimizers in convergence behavior, we conduct a controlled comparison on CIFAR-10 10 using a ResNet-50 backbone trained for 100 epochs under cosine learning rate scheduler.
Although the Shampoo optimizer g is a well-established quasi- ceassiiseoond--orrr -order method, its computa- tional overhead makes direct comparison impractical under identical epoch budgets. To ensure fair and efficient evaluation, we instead adopt Muon [14, a more recent and scalable curvature-aware optimizer that follows similar quasi-second-order philosophy but operates with significantly lower wall-clock cost. This substitution allows us to better examine the convergence profile of KOALA+ against comparable quasi-seconde sordd approaches.
As shown in Figure 5 Muon achieves faster convergence in the early phases, while KOALA+ initially exhibits larger fluctuations but gradually stabilizes and reaches comparable or better final accuracy. We attribute this difference to the distinct curvature modeling mechanisms of the two optimizers. Muon performs explicit matrix preconditioning with guaranteed positive semi-definite (PSD) curvature estimates, whereas KOALA++ relies on an implicit least-squares projection derived from scalar observations. Since the resulting covariance-like matrix is not explicitly constrained to be PSD, KOALA++ may require several epochs to implicitly correct unstable directions before converging to well-conditioned update regime. This interpretation remains hypothesis, and
26

---

(allidatoon
Top- Emror (Velidatton
40
100
100
Epacho
Epochs
(a) Log-scale Test Loss Epochs
(b) Top-1 Validation Error vs Epochs
Figure 5: CIFAR-10 with ResNet18 (Left) and ResNer50 (Right) trained for 100 epochs under cosine learning-rate scheduler. Validation Top-1 errors (the inset highlights the late-epoch region) are reported.
rigorous theoretical analysis of the PSD property and its impact convergence stability is left to future work.
D Ablations
D.1 Effect of Batch Size on KOALA+ Performance
We evaluated the impact of batch size on the performance KOALA+ using ResNet50 on the CIFAR100 dataset. As shown in Table increasing the batch size leads noticeable degradation in both Top-1 and Top-5 accuracy. The results suggest that KOALA+ is more effective smaller-bate batch training regimes, possibly due to better gradient variance estimation and more frequent updates.
Table 6: Impact of batch size on KOALA+ performance (ResNet50, CIFAR100).
Batch Size
Top-1 Error (%)
Top-5 Error (%)
64
20.27
4.91
128
20.38
4.76
256
21.78
5,84
512
23.63
6.59
D.2 Different Schedulers
In this experiment, we evaluate the impact of different learning rate schedulers on the performance of KOALA+ using the CIFAR-100 dataset. We consider two representative architectures: ResNet-50 as a typical LNN-based model and Swin-Tiny as typical ViT-based model. All models are trained for 200 epochs using the same optimizer (KOALA++ and hyperparameters.
The partially recorded results are summarized in Table 7
Table 7: Top-1 and Top-5 error (%) on CIFAR-100 using KOALA+ with different schedulers.
ResNet-50 (CNN)
Swin-Tiny (ViT)
Schedulers
Top-1 I Error Top-5 Error
Top-1 Error Top-5 Error
Multi-step CosineAmnealiggLR
21.80
5.89 5.07
20.57
35.05
12.07
Warmup Cosine
32.18
10.64
As Table 7 shows, different schedulers can lead to notable differences final top-1 and top-5 error. We now take closer look at how these schedulers influence the test loss dynamics during training.
27

---

12-
0.
Test
0.-
0.4
0.
0.2
02-
0.4
100
150
200
00
80
100
Epochs
(a) ResNet-50: Multi-step vs. CosinnAnnealingl.R
(b) Swin-Tiny: CosineAmnallingL.R vs. Warmup Cosine
Figure 6: Log-scale test loss curves of KOALA+ on CIFAR-100 under different learning rate schedulers. Left: CosineAmealigLLR provides a smoother convergence than Multi-step on ResNet- 50. Right: Warmup + Cosine improves stability and convergence for Swin- Tiny compared to ConialAnne alone.
Table 8: Effect of varying on KOALA+ performance. Dataset: CIFAR-100, Model: ResNet-50,
Epochs: 200, q = 0.2, Ir 2.0
Top-1 Error (%)
Top-5 Error (%)
0.05
22.95
6.60
0.10
22.04
6.11
0.20
21.80
5.89
For ResNet-50, we compare the Multi-step scheduler and the CosineAnnealingL... scheduler. As shown in Figure 6 Cosine results in smoother and more stable descent in test loss, especially during the late training stages. This stable convergence leads to slightly better final performance compared to the traditional multi-step decay.
For the ViT-based Swin-Tiny model, we evaluate both CosineAnnealingL..t and Warmup + Cosine schedule. We find that adding a warm-up stage significantly enhances convergence speed and test loss. This consistent with the behavior of transformer-bas based architectures, which are more sensitive to large initial learning rates. The Multi-step scheduler is not included in this setting, as it is generally less effective for transformers.
Note: We observed that the Swin-Tiny model tends to overfit in the later stages when trained for 200 epochs. While we report the final Top-1 and Top-5 error at epoch 200 Table for consistency, the plotted curves in Figure onlly show the first 100 epochs. This decision was made to focus on the informative convergence behavior and improve the visual clarity of the comparison.
D.3 Different Initializations of KOALA+
The initialization of the directional covariance product va in KOALA+ depends directly on the initial covariance matrix Po. To study the impact of on the performance of the model, we conducted an ablation experiment in which we vary E {0.05, 0.1, 0.2] while keeping other hyperparameters fixed.
We perform this experiment on the CIFAR-100 dataset using the ResNet-50 architecture, training for 200 epochs. The learning rate is set to 2.0, and the Kalman process noise parameter q is fixed at 0.2. The results are summarized in Table
We observe that increasing leads to improved performance, suggesting that initializing en with stronger magnitude provides better early-stage directionality, which helps escape suboptimal regions in the optimization landscape. The results are shown in Figure 7
28

---

10020 KIHLA--10.2)
12-
23.0
10-
Top-1 Entor (Velidetoon
70-
0.
193
04
02
00
20
90
100
150
200
25
125
200
Epochs
(a) Log-scale Test Loss
(b) Top-1 Validation Error (%)
Figure 7: Effect of varying KOALA++ performance 200 epochs CIFAR-100. Larger values yield more stable convergence and lower final error.
E
Experiments Details
E.1
Hardware
All experiments reported in this work were conducted on server equipped with single NVIDIA H100 GPU with 80 gigabytes of VRAM and 128 gigabytes of RAM. Unless otherwise stated, all training and evaluation tasks were executed using this configuration.
E,2 Image Classification
E.2.1
HP Settings for CIFAR-10/100 and Hyperparameter Sensitivity
We evaluate KOALA+ on both CIFAR-10 and CIFAR-100 using a variety of architectures, including CNN-based models (e.g., ResNet (9] and vision transformers (e.g., Swin-Tiny (18). Across all experiments, we used the same hyperparameter settings for CNN-based models to highlight KOALA++'s low-tuning usability, as it achieves competitive or better results than baselines without dataset-specific tuning. When moving Transformer based architectures, we conducted a small grid search following the procedure described in the Supplementary Material.
Below provide further clarification how the key hyperparameters selected, consistent with the analysis included in the Supplementary Material to support practical deployment of KOALA++:
Measurement noise R: estimated online as an exponential moving average of the mini- batch loss variance, using smoothing factor 0.9. This removes R from the set of hyperparameters requiring manual tuning.
Initial variance dg: controls only the scale of the initial directional state Ablation studies indicate that KOALA++'s convergence is robust to this value, so fix do to match the process noise variance Q for simplicity and symmetry.
H1.
Process noise Q: the main hyperparameter that requires tuning besides the learning rate. Empirically, we found that Q E [0.1, 11,0..1 0.4) performs well across datasets.
For CIFAR-10, we initialize both og and Q to 0.1, with an initial learning rate of 1.0. For CIFAR-100, which has more classes and richer data distribution, adopt slightly larger values d Q 0.2 and increase the initial learning rate to 2.0. A weight decay of 5 x 10 applied to all ResNet and other CNN models.
When applying KOALA++ to Swin- Tiny, we perform coarse grid search over the weight decay parameter, evaluating values {1 x 10-1,1 x 1 2,5 x 10-3, 10 10-4,5 x 10-4]. Interestingly, KOALA+ performs best with a smaller weight decay (1 10-4), whereas other optimizers typically favor higher decay values (e.g., 1 10-2). We hypothesize that this stems from KOALA++'s implicit regularization through dynamic covariance updates, which already stabilize parameter updates. Thus, combining KOALA+ with large explicit regularizer (c.g., strong weight decay) can lead to
29

---

opochs)
70
Adafater
5GD
5.75
503
Adam
FroorValidatten)
5.50
0.23
00-
5.00
4.73
00
4.50
4.25
125
1.5
200
100
170
180
130
200
Epachs
(a) Validation Error (Epoch 1-200)
Validation Error (Epoch 160-200)
Figure 8: Validation Top-1 Error (%) CIFAR-10 using ResNet-50. Left: Full 200-epoch curve.
Right: Zoomed view of epochs 160-200.
1.4
cOILAee
KINLAN+
Ataliter
12
0.8
0.0
0.4
0.2
40-
30-
0
50
300
125
150
175
200
25
50
100
127
150
175
200
Epochs
Epochs
(a) Lcglscale Test Loss
(b) Top-1 Validation Error (%)
Figure 9: Results on CIFAR-100 with ResNet-50 over 200 epochs. Left: Log-scale test loss. Right:
T validation error.
underfiting- in transformer-based models a smaller weight decay more suitable such cases.
E.2.2 More Results for CIFAR10/100
In the main body of the paper, we only visualized 100-epoch results on CIFAR-100 using ResNet-50 to compare KOALA+ with several baselines. Here, we provide more comprehensive results by extending the training to 200 epochs and including CIFAR-10 10 as well.
CIFAR10 Figure (Left) presents the Top-1 validation error on CIFAR-10 with ResNet-50 across 200 training epochs. To highlight the behavior in the final stage, we also plot zoomed-in view of validation error during epochs 160-200 in Figure g (Right).
Nhie:This figure shows one random seed; KOALA+ ++ outperforms SGD on average across three runs.
CIFARI00
We additionally report the 200-epoch results on CIFAR-100 using ResNet-50, compar- ing KOALA++ with several standard optimizers. The log-scale test loss and Top-1 validation error are shown Figure
E.2.3
Experiments for ImageNet32
We train both ResNet18 and ResNet50 models on the ImageNet32 dataset. Following stan- dard augmentation practices for low-resolution image classification, we apply RandomCrop (32, padding=4) and BandomHorizzotallFiig converted tensors and normalized using the standard ImageNet statistics. For validation, we directly normalize the images without applying any cropping or resizing. All models trained for 100
during training. The input images are subsequently
30

---

Table 9: Final Top-I and Top-5 validation error (%) on ImageNet32.
Model
Optimizer
Top-1 Error
Top-5 Error
SGD
45.81
22.02
ResNet18
KOALA+
45.71
21.94
SGD
40.17
17.57
ResNet50
KOALA+
39.03
17.12
KOALA
90
HOWLA
56D
SDD
Top-1 Emnr (Valldaatoo))
Log-scalo Test Loss
3.0
60-
b0
60
Epochs
Epochs
Figure 10: Left: Log-scale test loss; Right: Top-1 validation error for ResNet18 on ImageNet32.
epochs with a learning rate scheduler that decays the learning rate by factor of 0.1 at epochs 30, 60, and 90, and a weight decay of 0.0001.
Note: This setup differs slightly from the data augmentation strategy used in KOALA's ImageNet experiments [4] where validation images center-cropped prior evaluation.
Figure TO shows the log-scale test loss and Top-1 validation error ResNet18. We observe that SGD reaches a lower test error faster in carly training but tends to overfit slightly in the later stages, In contrast, KOALA++ shows improved generalization in the long run, despite being less stable early
Table summarizes the final Top-1 and Top-5 errors for both ResNet18 and ResNet50. The results show that KOALA++ achieves comparable or better performance to SGD without requiring any task-specific hyperparameter tuning.
E.3 Efficiency Analysis
We provide additional results the computational efficiency of KOALA++ Besides accuracy, it important that an optimizer remains efficient in terms of runtime and memory, particularly when scaling larger models and batch sizes.
ResNet-50 and Swin-Tiny (batch size 256). Table 10 reports average training time per epoch, asymptotic FLOPs, and peak memory usage ResNet-50 and Swin-Tiny. KOALA+ achieves runtime and memory comparable to Adam(W) (W), while being substantially more efficient than Shampoo and AdaHess.... For K-FAC, note that the more precise per-iteration complexity is 0(6d2 although it is commonly reported as O(n²) ) in the literature.
ResNet- 101 and Swin-Base (batch size 1024). To further scalability, report results ResNet- 101 and Swin Transformer Base with batch size 1024 in Table [I] KOALA++ remains competitive with Adam(W), while scaling better than AdaHessian and Shampoo both runtime and memory. This confirms that KOALA++ maintains near- far--fist-orderr efficiency even at larger scales,
Here is the number of layers. d the typical hidden dimension, and the batch size. "ddHHessian uses Hutchinson's trick Hessian diagonal estimation, which requires additional Hessian- vector products, leading to O(n²) complexity.
12 More precisely Ofddlm) m)), where l is the number of layers, d the hidden dimension, and the batch size.
31

---

Table 10: Efficiency profiling on ResNet-50 and Swin-Tiny (batch size 256). We report average training time per epoch, asymptotic FLOPs, and peak memory usage.
ResNet-50 (Batch size 256)
Optimizer
Time/epoch (s)
FLOPs (class)
Peak Mem (GB)
Adam
11.08
((n)
10.51
AdaFisher
11.64
O(n)
11.91
AdaHessiar
63.43
O(n²)
22.79
K-FAC
17.16
O
12.34
Shampoo
134.63
O
10.41
KOALA+ (ours)
13.23
O(n)
10.27
Swin-Tiny (Batch size 256)
Adam(W)
16.35
((n)
3.46
AdaFisher
12.00
O(n)
4.38
AdaHessian
42.00
O
5.97
KOALA+ (ours)
16.10
((n)
4.50
Table 11: Efficiency profiling on ResNet-100 and Swin Transformer Base with batch size 1024.
ResNet-101 (Batch size 1024)
Optimizer
Time/epoch (s)
FLOPs (class)
Peak Mem (GB)
Adam
16.79
O(n)
48. 18.16
AdaFisher
16.62
O(n)
48.11
K-FAC
23,79
O(n²)
52.09
Shampoo
64.29
O(n (5)
48.90
KOALA+ (ours)
17.61
O(n)
48.25
Swin Transformer Base (Batch size 1024)
Optimizer
Time/epoch (s)
FLOPs (class)
Peak Mem (GB)
Adam(W) (
19.07
((n)
23.23
AdaFisherW
18.95
O(n)
22.28
AdaHessian
57.72
O(n²)
64.08
KOALA+ (ours)
19.52
O(n)
22.51
Comparison with SGD across batch sizes. We also compare throughput with SGD on CIFAR-10 using ResNet-50 under varying batch sizes (Table KOALA++ incurs a small constant overhead per iteration due (n) covariance-ve vector products, but this cost is independent of batch size. As the batch size grows, forward/backward cost dominates, and the relative overhead of KOALA+ shrinks, approaching parity with SGD batch size 1024.
Table 12: Throughput comparison of KOALA++ vs SGD across different batch sizes on CIFAR-I0 with ResNet-50.
Batch Size
SGD (s/epoch)
KOALA+ (s/epoch)
Ratio (KOALA++/SGD)
Overhead
128
11.08
17.92
1.62
+61.7%
256
10.60
13.23
1.25x
+24.8%
512
9.70
11.41
1.18% x
+17.6%
1024
10.56
10.60
1.00x
+0.4%
This trend aligns with our design intuition: within each batch, the effective covariance dimension becomes smaller stochastic noise is averaged out, the extra operations of KOALA+ contribute
32

---

less to the total runtime. Consequently, KOALA++ scales gracefully with batch size and approaches the efficiency of SGD large-batch regimes.
E.4 Language Modeling
We adopt the experimental benchmark setup provided by AdaFisher 77 available at https; //github. //github.com/Attll htt.a//asaaaaalllllaaaaaaaaaaaaa
For fair comparison, perimental pipeline and only test on the WikiText-2 dataset using the small GPTI model.
the same ex-
To tune the hyperparameters of KOALA++, we performed grid search over the following ranges:
Learning Rate (Ir): (1e-1, 5e-2, 2e-2, le-2, 5e-3, 2e-3, le-3, le-3,5e-4) 5e-4)
sigma and q (set equal): [0.01, 0.02, 0.05, 0.1}
Weight Decay: (le-1, 5e-2, 2e-2, 1e-2, 5e-3, 2e-3, 1e-3, 5e-4, 2e-4, le-4]
The optimal configuration we found was: Learning Rate: 2c-3, sigma = q: 0.1, Weight Decay: le-4, which achieved the best perplexity among all tested settings. Figure II displays the validation loss and test perplexity curve of KOALA+ on the WikiText-2 dataset.
Valication PiLiLog Scaler: Smilll GPTL on Willbut
Volidanion Lossilog Scalol: Small GPT) WilTfett
Validation FPI, (Logl: Smat GPTL WirTen
60ALA11
KOALA
KORLA1
Adam/
Welddaton Lovallog Scale)
Adanti
Adanty
Valiation PPLI.og
2.0
6.5
6.0-
1.8
Validation
16-
5.5
5.0
5.0
200
002
802
1000
1200
200
430
000
800
1000
1200
10
20
30
Wall Clock Time 1s)
Wall Clock Time (3)
Epoch
(a) Validation PPL vs Time (log) (b) Validation Loss vs Time (log)
(c) Validation PPL vs Epoch (log)
Figure 11: Performance of KOALA+ on WikiText-2 with the small GPTI model. Left: Log-scale
validation perplexity wall clock time. Right: Log-scale validation perplexity over training epochs. Following (7), A was trained for 55 epochs, while KOALA++ were trained for 50 epochs.
33


Reviews:
[{"review_id": "HQi80Jvhhw", "reviewer": "Reviewer_csNz", "summary": "The 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).", "questions": "- 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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", "was_revised": true, "final_justification": "Insightful new algorithm that outperforms ADAM and is based on Kalman filtering. The performance improvement is rather small for some tasks, thats why I won't rate it as a 5.", "rebuttal": "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.\n\nBelow, we address the specific concerns raised:\n\n---\n\n#### **Clarification on Covariance Approximation**\n\nThe 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.\n\n---\n\n#### **Missing Confidence Intervals**\n\nThank 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.\n\n---\n\n#### **The Benefit Comes Only After the Final Learning Rate Reduction**\n\nThank 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.\n\n___\n\n#### **Fairness of Hyperparameter Tuning**\n\nIn 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.\n\n---\n\n#### **Computation Time per Epoch**\n\nWe 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.\n\n##### **ResNet‑50 (Batch Size = 256)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| Adam              | 11.08            | O(N)          | 10.51         |\n| AdaFisher         | 11.64            | O(N)          | 11.91         |\n| AdaHessian        | 63.43            |          | 22.79         |\n| K‑FAC             | 17.16            |       | 12.34         |\n| Shampoo           | 134.63           | O(N^1.5)      | 10.41         |\n| **KOALA++ (ours)**| **13.23**        | **O(N)**      | **10.27**     |\n\n\n\n##### **Swin Transformer Tiny(Batch Size = 256)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| AdamW          | 16.35            | O(N)          | 3.46          |\n| AdaFisher         | 12.00            | O(N)          | 4.38          |\n| AdaHessian        | 42.00            |           | 5.97          |\n| **KOALA++ (ours)**| **16.10**        | **O(N)**      | **4.50**      |\n\n\n---\n\n#### **Frobenius Norm Justification**\n\nThank 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.\n\n---\n\nPlease let us know if further clarification is needed. We once again thank the reviewer for their valuable feedback and constructive suggestions."}, {"review_id": "BIgdi216vP", "reviewer": "Reviewer_esMW", "summary": "This 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.", "questions": "Could 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nBelow, we address the specific concerns raised:\n\n---\n\n#### **Efficiency profiling**\n\n\nTo 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:\n\n##### **ResNet‑50 (Batch Size = 256)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| Adam              | 11.08            | O(N)          | 10.51         |\n| AdaFisher         | 11.64            | O(N)          | 11.91         |\n| AdaHessian        | 63.43            |           | 22.79         |\n| K‑FAC             | 17.16            |       | 12.34         |\n| Shampoo           | 134.63           | O(N^1.5)      | 10.41         |\n| **KOALA++ (ours)**| **13.23**        | **O(N)**      | **10.27**     |\n\n\n\n##### **Swin Transformer (Batch Size = 256)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| Adam(W)           | 16.35            | O(N)          | 3.46          |\n| AdaFisher         | 12.00            | O(N)          | 4.38          |\n| AdaHessian        | 42.00            |           | 5.97          |\n| **KOALA++ (ours)**| **16.10**        | **O(N)**      | **4.50**      |\n\n\nAs 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.\n\n---\n\n#### **Hyperparameter guidance & sensitivity**\n\nNotice 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.\n\nBelow 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:\n\n- 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.\n- $\\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.\n- 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.\n\n---\n\n\n#### **Speed vs SGD**\n\nBelow 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. \n\n| Batch Size | SGD (s/epoch) | KOALA++ (s/epoch) | KOALA++ / SGD | Overhead     |\n|-----------:|--------------:|------------------:|---------------:|--------------:|\n| 128        | 11.08         | 17.92             | **1.62×**      | **+61.7%**    |\n| 256        | 10.60         | 13.23             | **1.25×**      | **+24.8%**    |\n| 512        | 9.70          | 11.41             | **1.18×**      | **+17.6%**    |\n| 1024       | 10.56         | 10.60             | **1.00×**      | **+0.4%**     |\n\n\nThe 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. \n\nAdditionally, 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.\n\n---\n\n#### **Scalability in LM**\n\nWe 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.\n\n---\n#### **Consideration of hybrid methods**\n\nWe thank the reviewer for this insightful suggestion.\n\nWe 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.\n\nIn 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.\n\n---"}, {"review_id": "TRML7hjEg8", "reviewer": "Reviewer_RxbJ", "summary": "The 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.", "questions": "Koala++, 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.", "limitations": "See above.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "- **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.", "was_revised": true, "final_justification": "I appreciate the authors’ follow-up, but my key concerns remain unaddressed. The central issue was the lack of empirical or theoretical support for the claim that KOALA++’s curvature-aware component only provides benefits late in training. While the authors now attribute differences in convergence profiles to entanglement with learning rate schedulers, this remains a hypothesis without direct experimental validation or comparison to methods like Muon and Shampoo, as I originally requested. Their explanation still does not reconcile the documented early-stage advantages of these quasi–second-order optimizers in the literature. In addition, the paper provides no substantive theoretical analysis of KOALA++’s capabilities, relying instead on linear algebraic formulation without deeper insight, and it oversells the significance of the simple least-squares step, which is a rather basic component. Without (1) a principled theoretical grounding and (2) empirical evidence demonstrating a genuine difference in early-phase behavior, the core claims remain speculative. Consequently, I do not find sufficient basis to raise my score.", "rebuttal": "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.\n\nBelow, we address the specific concerns raised:\n\n---\n#### **Thinking of KOALA++**\nWe 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.\n\nAs 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).\n\n1. **Covariance Approximation**: \n   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.\n\n2. **Gradient Product Approximation**: \n   KOALA further approximates the rank-1 matrix $H_k^\\top H_k$ as another scaled identity: \n   $$\n   H_k^\\top H_k \\approx \\|H_k\\|^2 \\cdot I\n   $$  \n   This approximation discards the interactions between different weights of the network.\n\nTo 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.\n\n\n---\n#### **Motivation of directional covariance optimization objective**\n\nWe 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.\n\nOur 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.\n\n---\n\n#### **Error Analysis**\n\nWe 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.\n\n---\n\n#### **Related Work**\n\nWe 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.\n\nWhile KOALA++ is indeed inspired by Kalman filtering principles and incorporates uncertainty modeling, it differs from the cited works in design choices and application settings:\n\n- **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.\n\n- **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.\n\n\n\n---\n#### **Counterintuition about the training dynamics**\n\nWe 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.\n\n---\n\n\n#### Potential typo in $QI$ \n\nWhile $Q$ is a matrix in the standard EKF, in KOALA++, we adopt the design choice from KOALA and use a scaled identity matrix $QI$.\n\n---"}, {"review_id": "dibSv8VfYQ", "reviewer": "Reviewer_oqam", "summary": "This 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.", "questions": "- 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?", "limitations": "N/A (do not think this applies).", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "I 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.", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nBelow, we address the specific concerns raised:\n\n---\n#### **FLOPs and Memory Usage Comparison**\nWe 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.\n\nHere are our observations:\n\n- 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). \n\n\n\n\n- 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.\n\n\n\n\n#### **Efficiency Profile**\n\n\n##### **ResNet‑101 (Batch Size = 1024)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| Adam              | 16.79            | O(N)          | 48.16         |\n| AdaFisher         | 16.62            | O(N)          | 48.11         |\n| K‑FAC             | 23.79            |       | 52.09         |\n| Shampoo           | 64.29          | O(N^1.5)      | 48.9         |\n| **KOALA++ (ours)**| **17.61**        | **O(N)**      | **48.25**     |\n\n\n\n##### **Swin Transformer Base(Batch Size = 1024)**\n\n| Optimizer         | Time per epoch (s) | FLOPs (class) | Peak Mem (GB) |\n|------------------:|-----------------:|:--------------|--------------:|\n| AdamW           | 19.07            | O(N)          | 23.23          |\n| AdaFisherW         | 18.95            | O(N)          | 22.28          |\n| AdaHessian        | 57.72            |           | 64.08          |\n| **KOALA++ (ours)**| **19.52**        | **O(N)**      | **22.51**      |\n\n\n\n---\n\n#### **PSD Impact**\nWe thank the reviewer for this insightful question regarding the PSD (positive semi-definite) property of the covariance-like quantities in KOALA++.\n\n\nAs 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. \n\n\n\n---\n\n#### **Minor typos and grammar issue**\n\nWe thank the reviewer for pointing out the minor typos and grammar issues. We will carefully revise the manuscript to correct all such issues."}]

Author rebuttal:

Meta-review: The paper introduces a new Kalman filter based gradient descent methods that extends a current SOTA method called KOALA. The reviewers highlight the strong empirical results and practical usefulness. But a reviewer argues that the performance improvement is rather small for some tasks.  The central issue for a reviewer was the lack of empirical or theoretical support for the claim that KOALA++’s curvature-aware component only provides benefits late in training. While the authors now attribute differences in convergence profiles to entanglement with learning rate schedulers, this remains a hypothesis without direct experimental validation or comparison to methods like Muon and Shampoo. In summary, this is a borderline paper.
Decision: Accept (poster)