Share one implementation between adam and adamw

pytensor_ml/optim/rules.py

@@ -91,18 +91,56 @@ def adam_updates(
     Updates
         Mapping from each parameter and its moment buffers to their next values.
     """
+    return _adam_family_updates(
+        loss_or_gradients,
+        parameters,
+        learning_rate=learning_rate,
+        beta1=beta1,
+        beta2=beta2,
+        epsilon=epsilon,
+        amsgrad=amsgrad,
+        namespace="adam",
+    )
+
+
+def _adam_family_updates(
+    loss_or_gradients: LossOrGradients,
+    parameters: Sequence[Parameter],
+    *,
+    learning_rate: Rate,
+    beta1: float,
+    beta2: float,
+    epsilon: float,
+    amsgrad: bool,
+    namespace: str,
+    weight_decay: float = 0.0,
+    mask: Callable[[Parameter], bool] | None = None,
+) -> Updates:
+    """
+    Adam's update, shared by :func:`adam_updates` and :func:`adamw_updates`.
+
+    Parameters
+    ----------
+    namespace : str
+        Prefix for every state slot this rule allocates, so two rules in one graph keep separate moments
+        and separate step counts rather than reusing each other's.
+    weight_decay : float
+        Coefficient of the decoupled decay term added to the update. Zero adds no term at all. Default 0.0.
+    mask : callable, optional
+        Predicate selecting which parameters the decay reaches. Every parameter when omitted.
+    """
     gradients = get_gradients(loss_or_gradients, parameters)
 
-    step_count = counter("adam/step_count")
+    step_count = counter(f"{namespace}/step_count")
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
     first_moment_bias_correction = 1 - beta1**new_step_count_float
     second_moment_bias_correction = 1 - beta2**new_step_count_float
 
     updates: Updates = {step_count: new_step_count}
     for parameter, gradient in zip(parameters, gradients):
-        first_moment = state_for(parameter, "adam/first_moment")
-        second_moment = state_for(parameter, "adam/second_moment")
+        first_moment = state_for(parameter, f"{namespace}/first_moment")
+        second_moment = state_for(parameter, f"{namespace}/second_moment")
 
         new_first_moment = beta1 * first_moment + (1 - beta1) * gradient
         new_second_moment = beta2 * second_moment + (1 - beta2) * gradient**2
@@ -111,17 +149,25 @@ def adam_updates(
 
         second_moment_for_denominator = (
             _running_max_second_moment(
-                parameter, new_second_moment, updates, "adam/max_second_moment"
+                parameter, new_second_moment, updates, f"{namespace}/max_second_moment"
             )
             if amsgrad
             else new_second_moment
         )
         corrected_first_moment = new_first_moment / first_moment_bias_correction
         corrected_second_moment = second_moment_for_denominator / second_moment_bias_correction
+        denominator = pt.sqrt(corrected_second_moment) + epsilon
+
+        # Decoupled decay is added inside the rate multiplication -- that is what decouples it from the
+        # moments -- and without decay there is nothing to add. The two groupings agree in exact arithmetic
+        # and differ in the last bits, so folding them into one moves a third of float32 adam results.
+        if weight_decay:
+            decay_term = weight_decay * parameter if (mask is None or mask(parameter)) else 0.0
+            step = learning_rate * (corrected_first_moment / denominator + decay_term)
+        else:
+            step = learning_rate * corrected_first_moment / denominator
 
-        updates[parameter] = parameter - learning_rate * corrected_first_moment / (
-            pt.sqrt(corrected_second_moment) + epsilon
-        )
+        updates[parameter] = parameter - step
 
     return updates
 
@@ -193,38 +239,18 @@ def adamw_updates(
     Updates
         Mapping from each parameter and its moment buffers to their next values.
     """
-    gradients = get_gradients(loss_or_gradients, parameters)
-
-    step_count = counter("adamw/step_count")
-    new_step_count = step_count + 1
-    new_step_count_float = new_step_count.astype(config.floatX)
-    first_moment_bias_correction = 1 - beta1**new_step_count_float
-    second_moment_bias_correction = 1 - beta2**new_step_count_float
-
-    updates: Updates = {step_count: new_step_count}
-    for parameter, gradient in zip(parameters, gradients):
-        first_moment = state_for(parameter, "adamw/first_moment")
-        second_moment = state_for(parameter, "adamw/second_moment")
-
-        new_first_moment = beta1 * first_moment + (1 - beta1) * gradient
-        new_second_moment = beta2 * second_moment + (1 - beta2) * gradient**2
-        updates[first_moment] = new_first_moment
-        updates[second_moment] = new_second_moment
-
-        second_moment_for_denominator = (
-            _running_max_second_moment(
-                parameter, new_second_moment, updates, "adamw/max_second_moment"
-            )
-            if amsgrad
-            else new_second_moment
-        )
-        adam_update = (new_first_moment / first_moment_bias_correction) / (
-            pt.sqrt(second_moment_for_denominator / second_moment_bias_correction) + epsilon
-        )
-        decay_term = weight_decay * parameter if (mask is None or mask(parameter)) else 0.0
-        updates[parameter] = parameter - learning_rate * (adam_update + decay_term)
-
-    return updates
+    return _adam_family_updates(
+        loss_or_gradients,
+        parameters,
+        learning_rate=learning_rate,
+        beta1=beta1,
+        beta2=beta2,
+        epsilon=epsilon,
+        amsgrad=amsgrad,
+        namespace="adamw",
+        weight_decay=weight_decay,
+        mask=mask,
+    )
 
 
 def nadam_updates(

tests/optim/test_rules.py

@@ -201,6 +201,30 @@ def test_adamw_names_its_own_state():
     }
 
 
+def test_adam_and_adamw_in_one_step_keep_separate_state():
+    """The collision the name test above describes in prose, actually run. Both rules derive every slot from
+    one namespace argument, so a namespace that defaulted or went missing would hand them the same memoized
+    buffers -- and sharing a step counter is invisible, since both write the same increment to it."""
+    by_adam = trainable(np.array([1.0, -2.0]), name="by_adam")
+    by_adamw = trainable(np.array([1.0, -2.0]), name="by_adamw")
+    loss = 0.5 * (by_adam**2).sum() + 0.5 * (by_adamw**2).sum()
+
+    adam_state = adam_updates(loss, [by_adam], learning_rate=0.1, amsgrad=True)
+    adamw_state = adamw_updates(loss, [by_adamw], learning_rate=0.1, amsgrad=True)
+
+    assert not set(adam_state) & set(
+        adamw_state
+    )  # not one buffer between them, including the counter
+    function([], loss, updates={**adam_state, **adamw_state})()
+
+    counters = {
+        key.name: key.get_value()
+        for key in (*adam_state, *adamw_state)
+        if key.name.endswith("step_count")
+    }
+    assert counters == {"adam/step_count": 1, "adamw/step_count": 1}
+
+
 @pytest.mark.parametrize(
     "make_rule",
     [lambda: adam(learning_rate=1e-2), lambda: sgd(learning_rate=1e-2, momentum=0.9)],
@@ -292,6 +316,26 @@ def test_adamw_mask_excludes_parameters_from_decay():
     np.testing.assert_allclose(b.get_value() - 2.0, -lr * 1.0, rtol=RTOL)
 
 
+def test_adamw_without_decay_is_adam():
+    """The decoupled decay term is the *only* difference between the two rules, which is what lets them share
+    an implementation. Run over enough steps that the moments, the bias corrections and the accumulated
+    trajectory all have to agree, and compare exactly: with no decay these are meant to be the same
+    computation, not merely close, so the two drifting apart at all means they stopped sharing it."""
+    start = np.array([1.5, -0.5, 2.0])
+    by_adam = trainable(start.copy(), name="by_adam")
+    by_adamw = trainable(start.copy(), name="by_adamw")
+
+    adam_step = function([], [], updates=adam_updates(0.5 * (by_adam**2).sum(), [by_adam]))
+    adamw_step = function(
+        [], [], updates=adamw_updates(0.5 * (by_adamw**2).sum(), [by_adamw], weight_decay=0.0)
+    )
+    for _ in range(20):
+        adam_step()
+        adamw_step()
+
+    np.testing.assert_array_equal(by_adamw.get_value(), by_adam.get_value())
+
+
 def test_adagrad_step_decays_as_inverse_sqrt_t():
     """Under a constant gradient the accumulator grows as ``t * g**2``, so AdaGrad's step magnitude decays as
     ``lr / sqrt(t)`` for every coordinate — independent of the gradient magnitude itself."""