Add cyclical_step flag to TransformerSchedule.

Cyclic LR escape the current local minimum to explore other minima. [1]
[1] https://arxiv.org/abs/1704.00109

PiperOrigin-RevId: 488968190

Author: ds-hwang

lingvo/core/schedule.py

@@ -319,12 +319,19 @@ def Params(cls):
              'decay_end-th step.')
     p.Define('decay_factor', -0.5, 'Decay factor after warmup.')
     p.Define('start_step', 0, 'Translate the function left in the step axis.')
+    p.Define(
+        'cyclical_step', None, 'Int, if set, at the step, the cycle restarts. '
+        'Cyclic LR escapes a bad local minimum. arxiv.org/abs/1704.00109')
     return p
 
   def Value(self, step=None):
     """Returns the current learning rate decay."""
     p = self.params
-    current_step = tf.cast(self.GetStep(step), tf.float32)
+    current_step = self.GetStep(step)
+    if p.cyclical_step is not None:
+      assert isinstance(p.cyclical_step, int)
+      current_step = tf.math.mod(current_step, p.cyclical_step)
+    current_step = tf.cast(current_step, tf.float32)
     start_step = tf.cast(p.start_step, tf.float32)
     warmup_steps = tf.cast(p.warmup_steps * p.worker_replicas, tf.float32)
     if p.decay_end is not None:

lingvo/core/schedule_test.py

@@ -307,6 +307,22 @@ def testTransformerScheduleNoWarmUp(self):
           self.assertAllClose(base_lrs.Value().eval(), lrs.Value().eval())
           self.assertAllClose(base_lrs.Value().eval(), lrs.Value().eval())
 
+  def testTransformerScheduleCycle(self):
+    ref_p = schedule.TransformerSchedule.Params().Set(warmup_steps=0)
+    cycle = 3000
+    cyc_p = ref_p.Copy().Set(cyclical_step=cycle)
+    with self.session():
+      ref = ref_p.Instantiate()
+      cyc = cyc_p.Instantiate()
+      ref_pts = []
+      cyc_pts = []
+      for step in range(0, 1000, 10_000):
+        with py_utils.GlobalStepContext(step % cycle):
+          ref_pts.append(ref.Value().eval())
+        with py_utils.GlobalStepContext(step):
+          cyc_pts.append(cyc.Value().eval())
+      self.assertAllClose(ref_pts, cyc_pts)
+
   def testTransformerMLPerfSchedule(self):
     params = schedule.TransformerMLPerfSchedule.Params().Set(
         warmup_steps=4000, warmup_init_fraction=.3, model_dim=512)