Fiddling with Kron (PSGD) optimizer

README.md

@@ -12,6 +12,11 @@
 
 ## What's New
 
+## Jan 27, 2025
+* Add Kron Optimizer (PSGD w/ Kronecker-factored preconditioner) 
+  * Code from https://github.com/evanatyourservice/kron_torch
+  * See also https://sites.google.com/site/lixilinx/home/psgd
+
 ## Jan 19, 2025
 * Fix loading of LeViT safetensor weights, remove conversion code which should have been deactivated
 * Add 'SO150M' ViT weights trained with SBB recipes, decent results, but not optimal shape for ImageNet-12k/1k pretrain/ft
@@ -461,6 +466,7 @@ Included optimizers available via `timm.optim.create_optimizer_v2` factory metho
 * `adamp` and `sgdp` by [Naver ClovAI](https://github.com/clovaai) - https://arxiv.org/abs/2006.08217
 * `adan` an implementation of Adan adapted from https://github.com/sail-sg/Adan - https://arxiv.org/abs/2208.06677
 * `adopt` ADOPT adapted from https://github.com/iShohei220/adopt - https://arxiv.org/abs/2411.02853
+* `kron` PSGD w/ Kronecker-factored preconditioner from https://github.com/evanatyourservice/kron_torch - https://sites.google.com/site/lixilinx/home/psgd
 * `lamb` an implementation of Lamb and LambC (w/ trust-clipping) cleaned up and modified to support use with XLA - https://arxiv.org/abs/1904.00962
 * `laprop` optimizer from https://github.com/Z-T-WANG/LaProp-Optimizer - https://arxiv.org/abs/2002.04839
 * `lars` an implementation of LARS and LARC (w/ trust-clipping) - https://arxiv.org/abs/1708.03888

tests/test_optim.py

@@ -290,7 +290,7 @@ def _build_params_dict_single(weight, bias, **kwargs):
     return [dict(params=bias, **kwargs)]
 
 
-@pytest.mark.parametrize('optimizer', list_optimizers(exclude_filters=('fused*', 'bnb*')))
+@pytest.mark.parametrize('optimizer', list_optimizers(exclude_filters=('fused*', 'bnb*', 'kron*')))
 def test_optim_factory(optimizer):
     assert issubclass(get_optimizer_class(optimizer, bind_defaults=False), torch.optim.Optimizer)
 
@@ -386,6 +386,14 @@ def test_adam(optimizer):
     _test_model(optimizer, dict(lr=5e-2))
 
 
+@pytest.mark.parametrize('optimizer',  ['kron'])
+def test_kron(optimizer):
+    _test_rosenbrock(
+        lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
+    )
+    _test_model(optimizer, dict(lr=1e-3))
+
+
 @pytest.mark.parametrize('optimizer',  ['adopt', 'adoptw'])
 def test_adopt(optimizer):
     _test_rosenbrock(

timm/optim/_optim_factory.py

@@ -23,6 +23,7 @@
 from .adamw import AdamWLegacy
 from .adan import Adan
 from .adopt import Adopt
+from .kron import Kron
 from .lamb import Lamb
 from .laprop import LaProp
 from .lars import Lars
@@ -693,6 +694,19 @@ def _register_other_optimizers(registry: OptimizerRegistry) -> None:
             has_betas=True,
             second_order=True,
         ),
+        OptimInfo(
+            name='kron',
+            opt_class=Kron,
+            description='PSGD optimizer with Kronecker-factored preconditioner',
+            has_momentum=True,
+        ),
+        OptimInfo(
+            name='kronw',
+            opt_class=Kron,
+            description='PSGD optimizer with Kronecker-factored preconditioner and decoupled weight decay',
+            has_momentum=True,
+            defaults={'decoupled_decay': True}
+        ),
         OptimInfo(
             name='laprop',
             opt_class=LaProp,