diff --git a/pytensor_ml/optim/rules.py b/pytensor_ml/optim/rules.py index 63d94a0..3ed499d 100644 --- a/pytensor_ml/optim/rules.py +++ b/pytensor_ml/optim/rules.py @@ -91,9 +91,47 @@ 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 @@ -101,8 +139,8 @@ def adam_updates( 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( diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py index 31b9698..74e91be 100644 --- a/tests/optim/test_rules.py +++ b/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."""