Repository: stockeh/mlx-optimizers (language: python)

Issue:
Add ADOPT
ADOPT: Modified Adam Can Converge with Any β2 with the Optimal Rate

https://arxiv.org/abs/2411.02853
https://github.com/iShohei220/adopt


Discussion / hints:

Fix patch:
diff --git a/docs/media/rosenbrock_ADOPT.png b/docs/media/rosenbrock_ADOPT.png
new file mode 100644
index 0000000..9bba623
Binary files /dev/null and b/docs/media/rosenbrock_ADOPT.png differ
diff --git a/docs/media/rosenbrock_all.png b/docs/media/rosenbrock_all.png
index 162eff5..bc0dd99 100644
Binary files a/docs/media/rosenbrock_all.png and b/docs/media/rosenbrock_all.png differ
diff --git a/examples/rosenbrock.py b/examples/rosenbrock.py
index 5611ddd..31dd58e 100644
--- a/examples/rosenbrock.py
+++ b/examples/rosenbrock.py
@@ -116,6 +116,7 @@ def execute_experiments(optimizers, objective, func, plot_func, initial_state):
         (optim.QHAdam, 0, 0.5, {}),
         (optim.DiffGrad, 0, 0.5, {}),
         (optim.MADGRAD, 0, 0.5, {}),
+        (optim.ADOPT, 0, 0.25, {}),
     ]
     execute_experiments(
         optimizers, objective_rosenbrock, rosenbrock, plot_rosenbrock, ROSENBROCK_INITIAL
diff --git a/mlx_optimizers/__init__.py b/mlx_optimizers/__init__.py
index 926a91b..6dc32b6 100644
--- a/mlx_optimizers/__init__.py
+++ b/mlx_optimizers/__init__.py
@@ -2,9 +2,10 @@
 
 __version__ = VERSION
 
+from .adopt import ADOPT
 from .diffgrad import DiffGrad
 from .madgrad import MADGRAD
 from .muon import Muon
 from .qhadam import QHAdam
 
-__all__ = ["DiffGrad", "Muon", "QHAdam", "MADGRAD", "__version__"]
+__all__ = ["DiffGrad", "Muon", "QHAdam", "MADGRAD", "ADOPT", "__version__"]
diff --git a/mlx_optimizers/adopt.py b/mlx_optimizers/adopt.py
new file mode 100644
index 0000000..59b7af8
--- /dev/null
+++ b/mlx_optimizers/adopt.py
@@ -0,0 +1,70 @@
+from typing import Callable, List, Union
+
+import mlx.core as mx
+from mlx.optimizers import Optimizer
+
+
+class ADOPT(Optimizer):
+    r"""The ADOPT optimizer [1].
+
+    [1]: Taniguchi, Shohei, et al., 2024. ADOPT: Modified Adam Can
+    Converge with Any :math:`\beta_2` with the Optimal Rate. NeurIPS 2024.
+
+    .. math::
+
+        v_0 &= g_0^2,\quad m_1 = g_1 / \max{\sqrt{v_0}, \epsilon} \\
+        
+        \theta_{t} &= \theta_{t-1} - \alpha m_{t-1} \\
+        v_{t} &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2
+        m_{t+1} &= \beta_1 m_{t} + (1 - \beta_1) (g_{t+1} / \max{\sqrt{v_t}, \epsilon})
+
+    Args:
+    learning_rate (float or callable): The learning rate :math:`\lambda`.
+    betas (Tuple[float, float], optional): The coefficients
+        :math:`(\beta_1, \beta_2)` used for computing running averages of the
+        gradient and its square. Default: ``(0.9, 0.9999)``
+    eps (float, optional): The term :math:`\epsilon` added to the
+        denominator to improve numerical stability. Default: ``1e-6`
+    """
+
+    def __init__(
+        self,
+        learning_rate: Union[float, Callable[[mx.array], mx.array]],
+        betas: List[float] = [0.9, 0.9999],
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+
+        self._maybe_schedule("learning_rate", learning_rate)
+        self.betas = betas
+        self.eps = eps
+
+    def init_single(self, parameter: mx.array, state: dict):
+        """Initialize optimizer state"""
+        state["m"] = mx.zeros_like(parameter)
+        state["v"] = mx.zeros_like(parameter)
+        state["first"] = True
+
+    def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict):
+        """Performs a single optimization step, updating :math:`m` and :math:`v`"""
+        lr = self.learning_rate.astype(gradient.dtype)
+        b1, b2 = self.betas
+        eps = self.eps
+
+        m = state["m"]
+        v = state["v"]
+
+        if state["first"]:  # TODO: eval of self.step not allowed with mx.compile
+            state["first"] = False
+            v = mx.square(gradient)
+            m = gradient / mx.maximum(mx.sqrt(v), eps)
+        else:
+            m = b1 * m + (1 - b1) * gradient / mx.maximum(mx.sqrt(v), eps)
+
+        parameter = parameter - lr * m
+        v = b2 * v + (1 - b2) * mx.square(parameter)
+
+        state["m"] = m
+        state["v"] = v
+
+        return parameter


Test patch:
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 727fec6..99e7148 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -37,6 +37,7 @@ def ids(v):
     (optim.QHAdam, {"learning_rate": 0.25}, 300),
     (optim.DiffGrad, {"learning_rate": 0.3}, 300),
     (optim.MADGRAD, {"learning_rate": 0.03}, 300),
+    (optim.ADOPT, {"learning_rate": 0.8}, 1200),
 ]
 
 
diff --git a/tests/test_neuralnet.py b/tests/test_neuralnet.py
index 053a16a..3ef9e85 100644
--- a/tests/test_neuralnet.py
+++ b/tests/test_neuralnet.py
@@ -58,10 +58,11 @@ def ids(v):
 
 
 optimizers = [
-    (optim.QHAdam, {"learning_rate": 0.01}, 500),
-    (optim.DiffGrad, {"learning_rate": 0.01}, 500),
-    (optim.Muon, {"learning_rate": 0.01}, 500),
-    (optim.MADGRAD, {"learning_rate": 0.01}, 500),
+    (optim.QHAdam, {"learning_rate": 0.01}, 50),
+    (optim.DiffGrad, {"learning_rate": 0.01}, 100),
+    (optim.Muon, {"learning_rate": 0.01}, 100),
+    (optim.MADGRAD, {"learning_rate": 0.01}, 50),
+    (optim.ADOPT, {"learning_rate": 0.01}, 50),
 ]