Repository: google-deepmind/optax

Issue:
Support specifying end_value for exponential_decay
It would be convenient to support specifying an `end_value` for exponential decay. So, the `exponential_decay` would require `init_value` and `transition_steps` and one of (`end_value` or `decay_rate`). [Here](https://github.com/tensorflow/kfac/blob/master/kfac/python/keras/callbacks.py#L116) is an example of what I'm proposing. 

I would be happy to contribute this, if you think it is suitable.


Discussion:
I like the idea of adding an `end_value` argument, to support these.

When a `decay_rate` and `transtion_steps` is provided:
```
step = min(global_step - num_delay_steps, num_decay_steps)
    decayed_value = init_value * decay_rate^step
```

When a `decay_rate` and `end_value` are provided:
```
step = global_step - num_delay_steps
decayed_value = max(init_value * decay_rate^step, final_value)
```

But I would be inclined to keep `decay_rate` compulsory,
and have a check that only of `end_value` and `transtion_steps` is provided.

For the third use case in the docs you linked,
we could instead provide a utility function that given a `end_value` and `transtion_steps`
computes for you the appropriate `decay_rate`

Does that seem reasonable to you?

That sounds good @mtthss, I'll make a PR for that! 

The utility function would be pretty trivial:

```python
def exponential_decay_rate(init_value: float, end_value: float, transition_steps: int):
    return (final_value / init_value) ** (1.0 / transition_steps)
```

To be honest, calling this function is about as much code as the computation itself, so I don't think it makes sense to provide it. What do you think? 


EDIT: Sorry, I misunderstood the current exponential decay schedule. I thought the computation was `decayed_value = init_value * decay_rate ^ count`, but actually,  the computation is `decayed_value = init_value * decay_rate ^ (count / transition_steps)`).

I initially thought `transition_steps` represented the total number of decay steps (like in `polynomial_schedule`), but it's really used to control the staircase lengths (if `staircase=True`). 

In this case, it would make sense to allow the user to set both the `transition_steps` and `end_value`, right? Since they may want to control the staircase width as well as the `end_value`. 
Good point, and apologies for the confusion, you are absolutely right.
Maybe worth clarifying better in the doc string.
Great, I'll work on a PR for this. I think the doc string is fine (it shows the update expression), I just didn't read it closely 😅.

Resolving patch:
diff --git a/optax/_src/schedule.py b/optax/_src/schedule.py
--- a/optax/_src/schedule.py
+++ b/optax/_src/schedule.py
@@ -120,8 +120,9 @@ def exponential_decay(
     transition_steps: int,
     decay_rate: float,
     transition_begin: int = 0,
-    staircase: bool = False) -> Schedule:
-  """Constructs a schedule with either cointinuous or discrete exponential decay.
+    staircase: bool = False,
+    end_value: Optional[float] = None) -> Schedule:
+  """Constructs a schedule with either continuous or discrete exponential decay.
 
   This function applies an exponential decay function to a provided initial
   value. The function returns the decayed value as follows:
@@ -140,6 +141,9 @@ def exponential_decay(
     transition_begin: must be positive. After how many steps to start annealing
       (before this many steps the scalar value is held fixed at `init_value`).
     staircase: if `True`, decay the values at discrete intervals.
+    end_value: the value at which the exponential decay stops. When
+      `decay_rate` < 1, `end_value` is treated as a lower bound, otherwise as
+      an upper bound. Has no effect when `decay_rate' = 0.
 
   Returns:
     schedule: A function that maps step counts to values.
@@ -147,8 +151,8 @@ def exponential_decay(
 
   if transition_steps <= 0:
     logging.info(
-        'An exponential schedule was set with a non-positive `transition_steps` '
-        'value; this will result in a constant schedule with value '
+        'An exponential schedule was set with a non-positive `transition_steps`'
+        ' value; this will result in a constant schedule with value '
         '`init_value`.')
     return lambda count: init_value
 
@@ -164,13 +168,19 @@ def exponential_decay(
         'value; this will result in `transition_begin` falling back to `0`.')
     transition_begin = 0
 
+  if end_value is not None:
+    clip_fn = jnp.maximum if decay_rate < 1.0 else jnp.minimum
+
   def schedule(count):
     count -= transition_begin
     p = count / transition_steps
     if staircase:
       p = jnp.floor(p)
-    return jnp.where(
+    decayed_value = jnp.where(
         count <= 0, init_value, init_value * jnp.power(decay_rate, p))
+    if end_value is not None:
+      decayed_value = clip_fn(decayed_value, end_value)
+    return decayed_value
 
   return schedule
 


Test patch:
diff --git a/optax/_src/schedule_test.py b/optax/_src/schedule_test.py
--- a/optax/_src/schedule_test.py
+++ b/optax/_src/schedule_test.py
@@ -234,6 +234,47 @@ def _staircased(count):
     np.testing.assert_allclose(
         expected_vals, np.array(generated_vals), atol=1e-3)
 
+  @chex.all_variants()
+  @parameterized.parameters(
+      (0.2, 0.1, False), (1.0, 0.1, False), (2.0, 3.0, False),
+      (0.2, 0.1, True), (1.0, 0.1, True), (2.0, 3.0, True))
+  def test_end_value_with_staircase(self, decay_rate, end_value, staircase):
+    # Get schedule function.
+    init_value = 1.
+    num_steps = 11
+    transition_steps = 2
+    transition_begin = 3
+    schedule_fn = self.variant(
+        schedule.exponential_decay(
+            init_value=init_value, transition_steps=transition_steps,
+            decay_rate=decay_rate, transition_begin=transition_begin,
+            staircase=staircase, end_value=end_value))
+
+    # Test that generated values equal the expected schedule values.
+    def _staircased(count):
+      p = count / transition_steps
+      if staircase:
+        p = np.floor(p)
+      return p
+
+    generated_vals = []
+    for count in range(num_steps + transition_begin):
+      generated_vals.append(schedule_fn(count))
+    expected_vals = np.array(
+        [init_value] * transition_begin + [
+            init_value * np.power(decay_rate, _staircased(count))
+            for count in range(num_steps)
+        ],
+        dtype=np.float32)
+
+    if decay_rate < 1.0:
+      expected_vals = np.maximum(expected_vals, end_value)
+    else:
+      expected_vals = np.minimum(expected_vals, end_value)
+
+    np.testing.assert_allclose(
+        expected_vals, np.array(generated_vals), atol=1e-3)
+
 
 class CosineDecayTest(chex.TestCase):