Research map / Optimization / Optimizers
KOALA++: Efficient Kalman-Based Optimization with Gradient-Covariance Products
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes KOALA++, a Kalman-based optimizer that recursively tracks a gradient-covariance product v_k = H_k P_{k-1} via a minimum-Frobenius-norm approximation, avoiding storing the full covariance, and reports results on CIFAR-10/100 and WikiText-2 against first- and second-order optimizers.
Problem
Efficient, stable optimization of deep networks that captures structured gradient uncertainty without the cost of full covariance matrices or second-order gradient computation.
Approach
Extends the scalar/diagonal KOALA Kalman filter by recursively updating a compact gradient-covariance product v_k = H_k P_{k-1}, approximating the unavailable P_{k-2} with a minimum-Frobenius-norm least-squares solution that uses only the last gradient and v_{k-1}, yielding an O(n) per-step update.
Key results
- On CIFAR-100 with ResNet-50 (100 epochs, StepLR), KOALA++ reaches 22.24% top-1 error versus 23.38% for AdaFisher and 23.71% for KOALA-M.
- On CIFAR-10 with five modern architectures averaged over 5 seeds, KOALA++ performs on par with AdaFisher (e.g., 96.44% vs 96.34% on ResNet50, 88.74% AdaFisher vs 87.64% KOALA++ on Tiny Swin).
- On WikiText-2 with a scaled-down GPT-1, KOALA++ achieves 151.75 test PPL versus 152.72 for AdaFisherW and 173.98 for AdamW, at 25.56 s/epoch versus 23.11 s for AdamW.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace KOALA-M with KOALA++ | KOALA · CIFAR-100, ResNet-50, 100 epochs | Top-1 error: 23.71 -> 22.24 | Helped |
| replace KOALA-M with KOALA++ | KOALA · CIFAR-100, ResNet-50, 200 epochs | Top-1 error | Helped |
| use KOALA++ instead of the strongest listed baseline | CIFAR-100 ResNet-50 training · CIFAR-100, ResNet-50, 100 epochs | Top-1 error: 23.07 -> 22.24 | Helped |
| use KOALA++ instead of the strongest listed baseline | CIFAR-10 ResNet-101 training · CIFAR-10, ResNet-101 | Top-1 accuracy: 96.39 -> 96.66 | Helped |
| use KOALA++ instead of the strongest listed baseline | CIFAR-10 DenseNet-121 training · CIFAR-10, DenseNet-121 | Top-1 accuracy: 96.72 -> 96.60 | Hurt |
| use KOALA++ instead of the strongest listed baseline | CIFAR-10 MobileNetV3 training · CIFAR-10, MobileNetV3 | Top-1 accuracy: 95.28 -> 94.52 | Hurt |
| use KOALA++ instead of the strongest listed baseline | CIFAR-10 Tiny Swin training · CIFAR-10, Tiny Swin | Top-1 accuracy: 88.74 -> 87.64 | Hurt |
| use KOALA++ instead of AdaFisherW | WikiText-2 language modeling · WikiText-2 | test perplexity: 152.72 -> 151.75 | Helped |
| increase batch size from 64 to 128 (ablation) | KOALA++ · CIFAR-100, ResNet-50 | Top-1 error: 20.27 -> 20.38 | Hurt |
| increase batch size from 64 to 256 (ablation) | KOALA++ · CIFAR-100, ResNet-50 | Top-1 error: 20.27 -> 21.78 | Hurt |
| increase batch size from 64 to 512 (ablation) | KOALA++ · CIFAR-100, ResNet-50 | Top-1 error: 20.27 -> 23.63 | Hurt |
| use CosineAnnealingLR instead of Multi-step (ablation) | KOALA++ · CIFAR-100, ResNet-50 | Top-1 error: 21.80 -> 20.57 | Helped |
Surprising findings
- The approximated covariance is not guaranteed positive semi-definite, and the authors rely on an empirical angle test rather than a proof.
- On Tiny Swin, KOALA++ is slightly below AdaFisher(W) (87.64% vs 88.74%) and on CIFAR-10 ResNet-18 SGD outperforms KOALA++ at 100 epochs.
Limitations
- The approximated covariance matrix is not explicitly constrained to be positive semi-definite, so there is no formal stability guarantee.
- Only one linear constraint (the most recent gradient projection) is used, so the covariance approximation is coarse.
- Performance degrades with larger batch sizes (e.g., 512 gives 23.63% top-1 error on CIFAR-100 ResNet-50), and the method is weaker on some settings (Tiny Swin vs AdaFisher).
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet32 WikiText-2
Benchmarks: CIFAR-10 CIFAR-100 WikiText-2
Metrics: top-1 error top-5 error top-1 accuracy test perplexity training time per epoch
Baselines: SGD Adam AdamW KOALA-M AdaFisher AdaHessian K-FAC Shampoo Muon
Builds on: KOALA Extended Kalman Filter AdaFisher
Closest papers by meaning
- KOALA: A Kalman Optimization Algorithm with Loss Adaptivity (AAAI 2022 · similarity 0.93)
- Accelerating Neural Network Training with Augmented Koopman Dynamics (NeurIPS 2026 · similarity 0.86)
- Fast Approximate Natural Gradient Descent in a Kronecker Factored Eigenbasis (NeurIPS 2018 · similarity 0.85)
- An Implicit Form of Krasulina's k-PCA Update without the Orthonormality Constraint (AAAI 2020 · similarity 0.85)
- Deep-Koopman-KANDy: Dictionary Discovery for Deep-Koopman Operators with Kolmogorov-Arnold Networks for Dynamics (NeurIPS 2026 · similarity 0.85)
- SDS-LoRA: Overcoming Anisotropic Gradient Scaling in Low-Rank Adaptation (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.