Add LaProp to JAX SC optimizer specs and modify base optimizer to return general hyperparameter array instead of learning_rate only.

PiperOrigin-RevId: 746093821

Author: nicolagp

jax_tpu_embedding/sparsecore/lib/nn/BUILD

@@ -30,6 +30,7 @@ pytype_strict_library(
     ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_adagrad",
+        "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_laprop",
         "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_grad_with_sgd",
         pypi_requirement("jax"),
         pypi_requirement("jax/extend"),

jax_tpu_embedding/sparsecore/lib/nn/embedding.py

@@ -660,8 +660,7 @@ def tpu_sparse_dense_matmul_grad(
     embedding_variable = embedding_variables[stacked_table_name]
     activation_gradient = gradients[stacked_table_name]
     stack_table_spec = stacked_table_specs[stacked_table_name]
-    learning_rate = stack_table_spec.optimizer.get_learning_rate(step)
-    hyper_params = [learning_rate]
+    hyper_params = stack_table_spec.optimizer.get_hyperparameters(step)
     # The MLIR computation symbol names need to be different. We attach the
     # table name to the symbol name to ensure that.
     symbol_name = "{}-{}{}".format(

jax_tpu_embedding/sparsecore/lib/nn/embedding_spec.py

@@ -26,6 +26,7 @@
 import jax.extend as jex
 import jax.numpy as jnp
 from jax_tpu_embedding.sparsecore.lib.core.primitives import sparse_dense_matmul_grad_with_adagrad
+from jax_tpu_embedding.sparsecore.lib.core.primitives import sparse_dense_matmul_grad_with_laprop
 from jax_tpu_embedding.sparsecore.lib.core.primitives import sparse_dense_matmul_grad_with_sgd
 
 HyperParameterType: TypeAlias = Callable[[], jax.Array] | float
@@ -45,6 +46,10 @@
     "AdagradSlotVariables", ["accumulator"]
 )
 
+LaPropSlotVariables = collections.namedtuple(
+    "LaPropSlotVariables", ["mu", "nu"]
+)
+
 
 # TODO(b/365975374): Create helper functions for generating OptimizerSpecs.
 @dataclasses.dataclass(frozen=True, order=True)
@@ -81,6 +86,12 @@ def get_learning_rate(self, step: int | None = None) -> jax.Array:
     else:
       return jnp.array(self.learning_rate, dtype=jnp.float32)
 
+  def get_hyperparameters(
+      self, step: int | None = None
+  ) -> tuple[jax.Array, ...]:
+    """Returns the hyperparameters for the optimizer."""
+    return (self.get_learning_rate(step),)
+
   def slot_variables_initializers(self) -> tuple[CallableTableInitializer, ...]:
     """Slot variables initializers for the optimizer.
 
@@ -199,6 +210,90 @@ def get_optimizer_primitive(self) -> jex.core.Primitive:
     )
 
 
+class LaPropOptimizerSpec(OptimizerSpec):
+  """Spec for the LaProp optimizer.
+
+  Laprop decouples momentum and adaptivity in the Adam-style methods, leading to
+  improved speed and stability compare to Adam.
+  https://arxiv.org/abs/2002.04839
+
+  Attributes:
+    learning_rate: The learning rate for the training variables or embeddings.
+    b1: decay rate for the exponentially weighted average of grads.
+    b2: decay rate for the exponentially weighted average of squared grads.
+    eps: term added to the squared gradient to improve numerical stability.
+    rms_clip_threshold: Clipping threshold for RMS.
+    initial_slot_value: Initial value for the slot variables.
+  """
+
+  def __init__(
+      self,
+      learning_rate=0.001,
+      b1: float = 0.9,
+      b2: float = 0.95,
+      eps: float = 1e-30,
+      rms_clip_threshold: float = 1.0,
+      initial_slot_value: float = 0.0,
+  ):
+    super().__init__(
+        learning_rate=learning_rate,
+    )
+    self.b1 = b1
+    self.b2 = b2
+    self.eps = eps
+    self.rms_clip_threshold = rms_clip_threshold
+    self.initial_slot_value = initial_slot_value
+
+  def slot_variables_initializers(self) -> tuple[CallableTableInitializer, ...]:
+    return LaPropSlotVariables(
+        mu=jax.nn.initializers.constant(self.initial_slot_value),
+        nu=jax.nn.initializers.constant(self.initial_slot_value),
+    )
+
+  def get_decay_rate(self, step: int | None = None) -> jax.Array:
+    """Returns the decay rate for the optimizer."""
+
+    if step is None:
+      return jnp.array(self.b2, dtype=jnp.float32)
+
+    decay_rate = (
+        self.b2
+        * (1.0 - jnp.power(self.b2, step))
+        / ((1.0 - jnp.power(self.b2, step+1.0)))
+    )
+
+    return jnp.array(decay_rate, dtype=jnp.float32)
+
+  def get_hyperparameters(
+      self, step: int | None = None
+  ) -> tuple[jax.Array, ...]:
+    """Returns the LaProp hyperparameters: (learning_rate, b1, decay_rate, eps)."""
+    return (
+        self.get_learning_rate(step),
+        jnp.array(self.b1, dtype=jnp.float32),
+        self.get_decay_rate(step),
+        jnp.array(self.eps, dtype=jnp.float32),
+    )
+
+  def __hash__(self) -> int:
+    return hash((
+        self.learning_rate,
+        self.b1,
+        self.b2,
+        self.eps,
+        self.rms_clip_threshold,
+        self.initial_slot_value,
+    ))
+
+  def short_name(self) -> str:
+    return "laprop"
+
+  def get_optimizer_primitive(self) -> jex.core.Primitive:
+    return (
+        sparse_dense_matmul_grad_with_laprop.tpu_sparse_dense_matmul_grad_with_laprop_primitive
+    )
+
+
 @dataclasses.dataclass(eq=True, frozen=True, order=True)
 class FeatureIdTransformation:
   """Transformation to apply to the input feature ids."""

jax_tpu_embedding/sparsecore/lib/nn/tests/BUILD

@@ -184,6 +184,7 @@ pytype_strict_contrib_test(
     deps = [
         "//jax_tpu_embedding/sparsecore/lib/nn:embedding_spec",
         pypi_requirement("absl/testing:absltest"),
+        pypi_requirement("jax"),
         pypi_requirement("optax/schedules"),
     ],
 )

jax_tpu_embedding/sparsecore/lib/nn/tests/embedding_spec_test.py

@@ -14,6 +14,7 @@
 """Tests for embedding spec."""
 
 from absl.testing import absltest
+import jax.numpy as jnp
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 from optax import schedules
 
@@ -72,6 +73,44 @@ def test_compare_adagrad(self):
     self.assertEqual(op.learning_rate, 0.1)
     self.assertEqual(op.initial_accumulator_value, 0.1)
 
+  def test_compare_laprop(self):
+    self.assertEqual(
+        embedding_spec.LaPropOptimizerSpec(
+            learning_rate=0.1,
+            b1=0.9,
+            b2=0.95,
+            eps=1e-30,
+            rms_clip_threshold=1.0,
+            initial_slot_value=0.0,
+        ),
+        embedding_spec.LaPropOptimizerSpec(
+            learning_rate=0.1,
+            b1=0.9,
+            b2=0.95,
+            eps=1e-30,
+            rms_clip_threshold=1.0,
+            initial_slot_value=0.0,
+        ),
+    )
+    self.assertNotEqual(
+        embedding_spec.LaPropOptimizerSpec(
+            learning_rate=0.1,
+            b1=0.8,
+            b2=0.95,
+            eps=1e-30,
+            rms_clip_threshold=1.0,
+            initial_slot_value=0.0,
+        ),
+        embedding_spec.LaPropOptimizerSpec(
+            learning_rate=0.1,
+            b1=0.9,
+            b2=0.95,
+            eps=1e-30,
+            rms_clip_threshold=1.0,
+            initial_slot_value=0.0,
+        ),
+    )
+
   def test_learning_rate_callable(self):
     def lr():
       return 0.1
@@ -90,6 +129,32 @@ def test_learning_rate_schedule(self):
     self.assertEqual(op.get_learning_rate(50), 0.55)
     self.assertEqual(op.get_learning_rate(100), 0.1)
 
+  def test_hyperparameters(self):
+    op = embedding_spec.AdagradOptimizerSpec(
+        learning_rate=schedules.linear_schedule(
+            init_value=1.0, end_value=0.1, transition_steps=100
+        )
+    )
+    self.assertEqual(op.get_hyperparameters(0), (1.0,))
+
+    op = embedding_spec.LaPropOptimizerSpec(
+        learning_rate=schedules.linear_schedule(
+            init_value=1.0, end_value=0.1, transition_steps=100
+        ),
+        b1=0.9,
+        b2=0.95,
+        eps=1e-30,
+        rms_clip_threshold=1.0,
+        initial_slot_value=0.0,
+    )
+    expected_hyperparameters = (
+        jnp.array(1.0, dtype=jnp.float32),
+        jnp.array(0.9, dtype=jnp.float32),
+        jnp.array(0.0, dtype=jnp.float32),
+        jnp.array(1e-30, dtype=jnp.float32),
+    )
+    self.assertEqual(op.get_hyperparameters(0), expected_hyperparameters)
+
 
 if __name__ == "__main__":
   absltest.main()