Repository: kozistr/pytorch_optimizer

Issue:
VSGD


- https://openreview.net/forum?id=xu4ATNjcdy
- https://github.com/generativeai-tue/vsgd

Hints from issue comments:

Solution patch:
diff --git a/README.md b/README.md
index 0b4036cd..300c0f25 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **105 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **106 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -213,6 +213,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Simplified-Ademamix | *Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants*                     | [github](https://github.com/DepenM/Simplified-AdEMAMix/)                                                       | <https://arxiv.org/abs/2502.02431>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250202431M/exportcitation)                                                        |
 | Fira                | *Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?*                                     | [github](https://github.com/xichen-fy/Fira)                                                                    | <https://arxiv.org/abs/2410.01623>                                                          | [cite](https://github.com/xichen-fy/Fira/tree/main?tab=readme-ov-file#citation)                                                     |
 | RACS & Alice        | *Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension* |                                                                                                                | <https://arxiv.org/pdf/2502.07752>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207752G/exportcitation)                                                        |
+| VSGD                | *Variational Stochastic Gradient Descent for Deep Neural Networks*                                         | [github](https://github.com/generativeai-tue/vsgd)                                                             | <https://openreview.net/forum?id=xu4ATNjcdy>                                                | [cite](https://github.com/generativeai-tue/vsgd/tree/main?tab=readme-ov-file#cite)                                                  |
 
 ## Supported LR Scheduler
 
diff --git a/docs/changelogs/v3.5.2.md b/docs/changelogs/v3.5.2.md
index 01c692e9..22616b13 100644
--- a/docs/changelogs/v3.5.2.md
+++ b/docs/changelogs/v3.5.2.md
@@ -6,6 +6,8 @@
     * [Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?](https://arxiv.org/abs/2410.01623) 
 * Implement `RACS` and `Alice optimizer. (#376)
     * [Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension](https://arxiv.org/abs/2502.07752)
+* Implement `VSGD` optimizer. (#377, #378)
+    * [Variational Stochastic Gradient Descent for Deep Neural Networks](https://openreview.net/forum?id=xu4ATNjcdy) 
 
 ### Fix
 
diff --git a/docs/index.md b/docs/index.md
index 0b4036cd..300c0f25 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -10,7 +10,7 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **105 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Wide range of supported optimizers. Currently, **106 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
 * Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
@@ -213,6 +213,7 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Simplified-Ademamix | *Connections between Schedule-Free Optimizers, AdEMAMix, and Accelerated SGD Variants*                     | [github](https://github.com/DepenM/Simplified-AdEMAMix/)                                                       | <https://arxiv.org/abs/2502.02431>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250202431M/exportcitation)                                                        |
 | Fira                | *Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?*                                     | [github](https://github.com/xichen-fy/Fira)                                                                    | <https://arxiv.org/abs/2410.01623>                                                          | [cite](https://github.com/xichen-fy/Fira/tree/main?tab=readme-ov-file#citation)                                                     |
 | RACS & Alice        | *Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension* |                                                                                                                | <https://arxiv.org/pdf/2502.07752>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250207752G/exportcitation)                                                        |
+| VSGD                | *Variational Stochastic Gradient Descent for Deep Neural Networks*                                         | [github](https://github.com/generativeai-tue/vsgd)                                                             | <https://openreview.net/forum?id=xu4ATNjcdy>                                                | [cite](https://github.com/generativeai-tue/vsgd/tree/main?tab=readme-ov-file#cite)                                                  |
 
 ## Supported LR Scheduler
 
diff --git a/docs/optimizer.md b/docs/optimizer.md
index c62b9abd..a73ae491 100644
--- a/docs/optimizer.md
+++ b/docs/optimizer.md
@@ -436,6 +436,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.VSGD
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.WSAM
     :docstring:
     :members:
diff --git a/docs/visualization.md b/docs/visualization.md
index fd0ea788..0d039560 100644
--- a/docs/visualization.md
+++ b/docs/visualization.md
@@ -274,6 +274,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_QHM.png)
 
+### RACS
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_RACS.png)
+
 ### RAdam
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_RAdam.png)
@@ -382,6 +386,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_Tiger.png)
 
+### VSGD
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_VSGD.png)
+
 ### Yogi
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_Yogi.png)
@@ -660,6 +668,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_QHM.png)
 
+### RACS
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_RACS.png)
+
 ### RAdam
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_RAdam.png)
@@ -768,6 +780,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_Tiger.png)
 
+### VSGD
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_VSGD.png)
+
 ### Yogi
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_Yogi.png)
diff --git a/examples/visualize_optimizers.py b/examples/visualize_optimizers.py
index 247f6170..edb82a3a 100644
--- a/examples/visualize_optimizers.py
+++ b/examples/visualize_optimizers.py
@@ -16,7 +16,7 @@ from pytorch_optimizer.optimizer.alig import l2_projection
 
 filterwarnings('ignore', category=UserWarning)
 
-OPTIMIZERS_IGNORE = ('lomo', 'adalomo', 'demo', 'a2grad', 'muon')
+OPTIMIZERS_IGNORE = ('lomo', 'adalomo', 'demo', 'a2grad', 'muon', 'alice')
 OPTIMIZERS_MODEL_INPUT_NEEDED = ('lomo', 'adalomo', 'adammini')
 OPTIMIZERS_GRAPH_NEEDED = ('adahessian', 'sophiah')
 OPTIMIZERS_CLOSURE_NEEDED = ('alig', 'bsam')
diff --git a/pyproject.toml b/pyproject.toml
index f72a0c54..fc52920d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,8 +20,9 @@ keywords = [
     "PNM", "Prodigy", "PSGD", "QHAdam", "QHM", "RACS", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM",
     "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo",
     "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM", "StableSPAM", "SRMM", "StableAdamW",
-    "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice",
-    "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "SWATS", "TAM", "Tiger", "TRAC", "VSGD", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1",
+    "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD",
+    "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
diff --git a/pytorch_optimizer/__init__.py b/pytorch_optimizer/__init__.py
index 1fdc8a69..d51b0efb 100644
--- a/pytorch_optimizer/__init__.py
+++ b/pytorch_optimizer/__init__.py
@@ -69,6 +69,7 @@ from pytorch_optimizer.optimizer import (
     SWATS,
     TAM,
     TRAC,
+    VSGD,
     WSAM,
     A2Grad,
     AccSGD,
diff --git a/pytorch_optimizer/optimizer/__init__.py b/pytorch_optimizer/optimizer/__init__.py
index a4e60483..9551d80e 100644
--- a/pytorch_optimizer/optimizer/__init__.py
+++ b/pytorch_optimizer/optimizer/__init__.py
@@ -90,7 +90,7 @@ from pytorch_optimizer.optimizer.schedulefree import (
     ScheduleFreeWrapper,
 )
 from pytorch_optimizer.optimizer.scion import SCION, SCIONLight
-from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, AccSGD, SGDSaI, SignSGD
+from pytorch_optimizer.optimizer.sgd import ASGD, SGDW, VSGD, AccSGD, SGDSaI, SignSGD
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
 from pytorch_optimizer.optimizer.sm3 import SM3
@@ -318,6 +318,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     Fira,
     RACS,
     Alice,
+    VSGD,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
diff --git a/pytorch_optimizer/optimizer/sgd.py b/pytorch_optimizer/optimizer/sgd.py
index e29f5ec0..ba78315f 100644
--- a/pytorch_optimizer/optimizer/sgd.py
+++ b/pytorch_optimizer/optimizer/sgd.py
@@ -519,3 +519,129 @@ class SGDSaI(BaseOptimizer):
                 p.add_(buf, alpha=-group['lr'] * state['gsnr'])
 
         return loss
+
+
+class VSGD(BaseOptimizer):
+    r"""Variational Stochastic Gradient Descent for Deep Neural Networks.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param ghattg: float. prior variance ratio between ghat and g, Var(ghat_t-g_t)/Var(g_t-g_{t-1}).
+    :param ps: float. prior strength.
+    :param tau1: float. remember rate for the gamma parameters of g.
+    :param tau2: float. remember rate for the gamma parameter of ghat.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-1,
+        ghattg: float = 30.0,
+        ps: float = 1e-8,
+        tau1: float = 0.81,
+        tau2: float = 0.9,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        eps: float = 1e-8,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_non_negative(ghattg, 'ghattg')
+        self.validate_non_negative(ps, 'ps')
+        self.validate_non_negative(tau1, 'tau1')
+        self.validate_non_negative(tau2, 'tau2')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'tau1': tau1,
+            'tau2': tau2,
+            'pa2': 2.0 * ps + 1.0 + 1e-4,
+            'pbg2': 2.0 * ps,
+            'pbhg2': 2.0 * ghattg * ps,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'VSGD'
+
+    @torch.no_grad()
+    def reset(self):
+        pass
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' not in group:
+                group['step'] = 1
+            else:
+                group['step'] += 1
+
+            pa2, pbg2, pbhg2 = group['pa2'], group['pbg2'], group['pbhg2']
+
+            rho1: float = math.pow(group['step'], -group['tau1'])
+            rho2: float = math.pow(group['step'], -group['tau2'])
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['mug'] = torch.zeros_like(p)
+                    state['bg'] = torch.zeros_like(p)
+                    state['bhg'] = torch.zeros_like(p)
+
+                self.apply_weight_decay(
+                    p,
+                    grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=False,
+                )
+
+                bg, bhg = state['bg'], state['bhg']
+
+                if group['step'] == 1:
+                    sg = pbg2 / (pa2 - 1.0)
+                    shg = pbhg2 / (pa2 - 1.0)
+                else:
+                    sg = bg / pa2
+                    shg = bhg / pa2
+
+                mug = state['mug']
+                mug_prev = mug.clone()
+
+                mug.mul_(shg).add_(grad * sg).div_(sg + shg)
+
+                sigg = (sg * shg) / (sg + shg)
+                mug_sq = mug.pow(2).add_(sigg)
+
+                bg2 = pbg2 + mug_sq - 2.0 * mug * mug_prev + mug_prev.pow(2)
+                bhg2 = pbhg2 + mug_sq - 2.0 * grad * mug + grad.pow(2)
+
+                bg.mul_(1.0 - rho1).add_(bg2, alpha=rho1)
+                bhg.mul_(1.0 - rho2).add_(bhg2, alpha=rho2)
+
+                p.add_(group['lr'] / mug_sq.sqrt().add_(group['eps']) * mug, alpha=-1.0)
+
+        return loss


Test patch:
diff --git a/tests/constants.py b/tests/constants.py
index 4a2dfc07..28a1bee4 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -25,6 +25,7 @@ from pytorch_optimizer.optimizer import (
     SRMM,
     SWATS,
     TAM,
+    VSGD,
     A2Grad,
     AccSGD,
     AdaBelief,
@@ -584,6 +585,7 @@ OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaGC, {'lr': 1e-1, 'warmup_steps': 2, 'weight_decay': 1e-3}, 5),
     (RACS, {'lr': 1e0}, 5),
     (Alice, {'lr': 1e0, 'rank': 2, 'leading_basis': 1, 'update_interval': 2}, 5),
+    (VSGD, {'lr': 1e0}, 5),
     (Ranger25, {'lr': 1e-1}, 3),
     (Ranger25, {'lr': 1e-1, 't_alpha_beta3': 5}, 3),
     (Ranger25, {'lr': 5e-2, 'stable_adamw': False, 'orthograd': False, 'eps': None, 'lookahead_merge_time': 2}, 3),
diff --git a/tests/test_load_modules.py b/tests/test_load_modules.py
index 30a44bd1..0bc071a1 100644
--- a/tests/test_load_modules.py
+++ b/tests/test_load_modules.py
@@ -34,7 +34,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 103
+    assert len(get_supported_optimizers()) == 104
     assert len(get_supported_optimizers('adam*')) == 8
     assert len(get_supported_optimizers(['adam*', 'ranger*'])) == 11