LAprop as an optimizer along with frobenius weight normalization + associated tests

emerging_optimizers/scalar_optimizers/__init__.py

@@ -12,4 +12,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from emerging_optimizers.scalar_optimizers.laprop import LaProp
 from emerging_optimizers.scalar_optimizers.lion import Lion

emerging_optimizers/scalar_optimizers/laprop.py

@@ -0,0 +1,186 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections.abc import Callable
+from typing import TYPE_CHECKING, override
+
+
+if TYPE_CHECKING:
+    from typing import overload
+
+import torch
+from absl import logging
+from torch.optim.optimizer import ParamsT
+
+from emerging_optimizers import registry
+from emerging_optimizers.mixin import WeightDecayMixin, WeightDecayT
+from emerging_optimizers.scalar_optimizers.update_functions import calculate_laprop_update
+
+
+__all__ = [
+    "LaProp",
+]
+
+
+@registry.register_optimizer("laprop")
+class LaProp(WeightDecayMixin, torch.optim.Optimizer):
+    """LaProp optimizer.
+
+    LAProp can be seen as RMSProp with a momentum term, or normalized SGD with momentum.
+    This optimizer tracks Adam-style first and second moments, but normalizes the gradient
+    before the first-moment update.
+
+    The update rule below assumes ``weight_decay_method="decoupled"`` (the default).
+    See :class:`~emerging_optimizers.mixin.WeightDecayMixin` for the other modes.
+
+    .. math::
+        p = p \\cdot (1 - \\text{lr} \\cdot \\lambda) \\\\
+        v_t = \\beta_2 v_{t-1} + (1 - \\beta_2) g_t^2 \\\\
+        \\hat{v}_t = \\frac{v_t}{1 - \\beta_2^t} \\\\
+        g'_t = \\frac{g_t}{\\sqrt{\\hat{v}_t} + \\epsilon} \\\\
+        m_t = \\beta_1 m_{t-1} + (1 - \\beta_1) g'_t \\\\
+        \\hat{m}_t = \\frac{m_t}{1 - \\beta_1^t} \\\\
+        p = p - \\text{lr} \\cdot \\hat{m}_t
+
+    References:
+        - Ziyin, L., Wang, Z. T., & Ueda, M. *LaProp: Separating Momentum and
+          Adaptivity in Adam.* arXiv:2002.04839 (2020).
+          [`arXiv:2002.04839 <https://arxiv.org/abs/2002.04839>`_]
+
+    Args:
+        params: Iterable of parameters to optimize or dicts defining parameter groups.
+        lr: Learning rate.
+        betas: Coefficients (beta1, beta2) for first and second moment EMAs.
+        eps: Term added to the denominator for numerical stability.
+        weight_decay: Weight decay coefficient.
+        correct_bias: Whether to apply bias correction to the first and second moment EMAs.
+        frob_normalize: Whether to normalize each updated parameter back to its pre-update Frobenius norm.
+        weight_decay_method: Method to apply weight decay, see
+            :class:`~emerging_optimizers.mixin.WeightDecayMixin` for more details.
+    """
+
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: float = 1e-3,
+        betas: tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        correct_bias: bool = True,
+        frob_normalize: bool = False,
+        weight_decay_method: WeightDecayT = "decoupled",
+    ) -> None:
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta at index 1: {betas[1]}")
+        if not 0.0 <= eps:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if frob_normalize and weight_decay != 0.0:
+            logging.warning("LaProp with frob_normalize=True is intended to be used with weight_decay=0.0.")
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias)
+        self.weight_decay_method = weight_decay_method
+        self.frob_normalize = frob_normalize
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def _init_group(
+        self,
+        group: dict,
+        skip_non_grad_params: bool = True,
+    ) -> None:
+        """Performs lazy state initialization for parameters.
+
+        Args:
+            group: Parameter group dictionary.
+            skip_non_grad_params: If True, skip parameters without gradients.
+        """
+        for p in group["params"]:
+            if skip_non_grad_params and p.grad is None:
+                continue
+            state = self.state[p]
+
+            if len(state) == 0:
+                state["step"] = 0
+                state["exp_avg"] = torch.zeros_like(p.data)
+                state["exp_avg_sq"] = torch.zeros_like(p.data)
+
+    if TYPE_CHECKING:
+
+        @overload
+        def step(self, closure: None = ...) -> None: ...
+
+        @overload
+        def step(self, closure: Callable[[], float]) -> float: ...
+
+    @torch.no_grad()  # type: ignore[misc]
+    @override
+    def step(self, closure: Callable[[], float] | None = None) -> float | None:
+        """Perform a single optimization step.
+
+        Note:
+            When ``weight_decay_method="l2"``, ``p.grad`` is modified in-place
+            (the L2 penalty ``weight_decay * p`` is added to the gradient).
+            If you need the original gradient after this call, clone it beforehand.
+
+        Args:
+            closure: A closure that reevaluates the model and returns the loss (optional).
+
+        Returns:
+            The loss from the closure, if provided.
+        """
+        if closure is None:
+            loss = None
+        else:
+            loss = closure()
+
+        for group in self.param_groups:
+            self._init_group(group)
+
+            lr = group["lr"]
+            betas = group["betas"]
+            eps = group["eps"]
+            weight_decay = group["weight_decay"]
+            correct_bias = group["correct_bias"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue  # pragma: no cover
+
+                grad = p.grad
+                state = self.state[p]
+                state["step"] += 1
+                pre_norm = p.data.norm() if self.frob_normalize else None
+
+                self._apply_weight_decay_inplace(p.data, grad, lr, weight_decay)
+
+                update = calculate_laprop_update(
+                    grad,
+                    state["exp_avg"],
+                    state["exp_avg_sq"],
+                    correct_bias,
+                    betas,
+                    state["step"],
+                    eps,
+                )
+                p.data.add_(update, alpha=-lr)
+                if self.frob_normalize:
+                    assert pre_norm is not None
+                    p.data.mul_(pre_norm / p.data.norm().clamp_min(eps))
+
+        return loss

tests/test_laprop.py

@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from absl import flags, logging
+from absl.testing import absltest, parameterized
+
+from emerging_optimizers.scalar_optimizers import LaProp
+from emerging_optimizers.scalar_optimizers.update_functions import calculate_laprop_update
+
+
+flags.DEFINE_enum("device", "cpu", ["cpu", "cuda"], "Device to run tests on")
+flags.DEFINE_integer("seed", None, "Random seed for reproducible tests")
+FLAGS = flags.FLAGS
+
+
+def setUpModule() -> None:
+    if FLAGS.seed is not None:
+        logging.info("Setting random seed to %d", FLAGS.seed)
+        torch.manual_seed(FLAGS.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(FLAGS.seed)
+
+
+class LaPropOptimizerTest(parameterized.TestCase):
+    def setUp(self):
+        self.device = FLAGS.device
+
+    @parameterized.parameters(
+        {"shape": (3, 3)},
+        {"shape": (15, 31)},
+        {"shape": (127, 255)},
+    )
+    def test_smoke(self, shape) -> None:
+        """LaProp optimizer can be instantiated and stepped."""
+        param = torch.nn.Parameter(torch.randn(*shape, device=self.device))
+        optimizer = LaProp([param], lr=1e-4)
+        param.grad = torch.randn_like(param)
+        optimizer.step()
+
+    @parameterized.parameters(
+        {"shape": (3, 3)},
+        {"shape": (15, 31)},
+        {"shape": (127, 255)},
+    )
+    def test_state_initialization(self, shape) -> None:
+        """LaProp initializes first moment, second moment, and step state."""
+        beta1, beta2 = 0.5, 0.75
+        param = torch.nn.Parameter(torch.randint(1, 5, shape, device=self.device, dtype=torch.float32))
+        optimizer = LaProp([param], lr=0.25, betas=(beta1, beta2), weight_decay=0.0, correct_bias=True)
+        grad = torch.randint_like(param, 1, 5)
+        param.grad = grad.clone()
+        optimizer.step()
+
+        self.assertEqual(optimizer.state[param]["step"], 1)
+        self.assertIn("exp_avg", optimizer.state[param])
+        self.assertIn("exp_avg_sq", optimizer.state[param])
+
+        expected_exp_avg_sq = (1 - beta2) * grad.square()
+        normalized_grad = grad / (grad.abs() + optimizer.param_groups[0]["eps"])
+        expected_exp_avg = (1 - beta1) * normalized_grad
+        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], expected_exp_avg_sq, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg"], expected_exp_avg, atol=0, rtol=0)
+
+    @parameterized.parameters(
+        {"shape": (3, 3)},
+        {"shape": (15, 31)},
+        {"shape": (127, 255)},
+    )
+    def test_optimizer_step_matches_update_function(self, shape) -> None:
+        """LaProp optimizer delegates update math to calculate_laprop_update."""
+        lr = 0.25
+        betas = (0.5, 0.75)
+        eps = 1e-8
+        param = torch.nn.Parameter(torch.randint(-5, 5, shape, device=self.device, dtype=torch.float32))
+        grad = torch.randint(-5, 5, shape, device=self.device, dtype=torch.float32)
+        optimizer = LaProp([param], lr=lr, betas=betas, eps=eps, weight_decay=0.0)
+
+        old_param = param.detach().clone()
+        exp_avg = torch.zeros_like(param)
+        exp_avg_sq = torch.zeros_like(param)
+        expected_update = calculate_laprop_update(grad, exp_avg, exp_avg_sq, True, betas, 1, eps)
+
+        param.grad = grad.clone()
+        optimizer.step()
+
+        torch.testing.assert_close(param, old_param - lr * expected_update, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg"], exp_avg, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], exp_avg_sq, atol=0, rtol=0)
+
+    @parameterized.parameters(
+        {"shape": (3, 3)},
+        {"shape": (15, 31)},
+        {"shape": (127, 255)},
+    )
+    def test_no_grad_no_update_params_unchanged(self, shape) -> None:
+        """Parameters without gradients are not updated."""
+        param = torch.nn.Parameter(torch.randn(*shape, device=self.device))
+        original = param.detach().clone()
+        optimizer = LaProp([param], lr=1e-4)
+        optimizer.step()
+        torch.testing.assert_close(param, original, atol=0, rtol=0)
+
+    @parameterized.parameters(
+        {"shape": (3, 3)},
+        {"shape": (15, 31)},
+        {"shape": (127, 255)},
+    )
+    def test_frob_normalize_preserves_parameter_norm(self, shape) -> None:
+        """LaProp can normalize updated parameters back to their pre-update norm."""
+        param = torch.nn.Parameter(torch.randint(1, 5, shape, device=self.device, dtype=torch.float32))
+        optimizer = LaProp([param], lr=0.25, weight_decay=0.0, frob_normalize=True)
+        param.grad = torch.randint(1, 5, shape, device=self.device, dtype=torch.float32)
+        original_norm = param.norm()
+
+        optimizer.step()
+
+        torch.testing.assert_close(param.norm(), original_norm)
+
+    @parameterized.parameters(True, False)
+    def test_init_group_skip_non_grad_params(self, skip_non_grad_params) -> None:
+        """Test _init_group with skip_non_grad_params flag."""
+        param_with_grad = torch.nn.Parameter(torch.randn(5, 7, device=self.device))
+        param_without_grad = torch.nn.Parameter(torch.randn(5, 7, device=self.device))
+        param_with_grad.grad = torch.randn_like(param_with_grad)
+
+        opt = LaProp([param_with_grad, param_without_grad], lr=1e-4)
+        opt._init_group(opt.param_groups[0], skip_non_grad_params=skip_non_grad_params)
+
+        self.assertIn("exp_avg", opt.state[param_with_grad])
+        self.assertIn("exp_avg_sq", opt.state[param_with_grad])
+        self.assertEqual(opt.state[param_with_grad]["step"], 0)
+        self.assertEqual("exp_avg" in opt.state[param_without_grad], not skip_non_grad_params)
+
+    def test_negative_lr_raises_value_error(self) -> None:
+        """Test that LaProp raises ValueError for negative learning rate."""
+        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
+        with self.assertRaisesRegex(ValueError, "Invalid learning rate"):
+            LaProp([param], lr=-1.0)
+
+    def test_beta0_out_of_range_raises_value_error(self) -> None:
+        """Test that LaProp raises ValueError for invalid beta at index 0."""
+        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
+        with self.assertRaisesRegex(ValueError, "Invalid beta at index 0"):
+            LaProp([param], betas=(1.0, 0.999))
+
+    def test_beta1_out_of_range_raises_value_error(self) -> None:
+        """Test that LaProp raises ValueError for invalid beta at index 1."""
+        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
+        with self.assertRaisesRegex(ValueError, "Invalid beta at index 1"):
+            LaProp([param], betas=(0.9, 1.0))
+
+    def test_negative_eps_raises_value_error(self) -> None:
+        """Test that LaProp raises ValueError for negative eps."""
+        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
+        with self.assertRaisesRegex(ValueError, "Invalid epsilon"):
+            LaProp([param], eps=-1e-8)
+
+    def test_negative_weight_decay_raises_value_error(self) -> None:
+        """Test that LaProp raises ValueError for negative weight_decay."""
+        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
+        with self.assertRaisesRegex(ValueError, "Invalid weight_decay"):
+            LaProp([param], weight_decay=-0.1)
+
+
+if __name__ == "__main__":
+    absltest.main()

tests/test_registry.py

@@ -61,6 +61,7 @@ def __init__(self, params, lr=0.01):
         ("scion", scion.Scion),
         ("soap", soap.SOAP),
         ("lion", scalar_optimizers.Lion),
+        ("laprop", scalar_optimizers.LaProp),
     )
     def test_get_optimizer(self, opt_name, expected_opt_cls):
         opt_cls = registry.get_optimizer_cls(opt_name)