Schedules for Non-Learning Rate Hyper-parameters

[n2cholas]

Is it possible to have schedules for hyper-parameters such as momentum?

For example, onecycle momentum from this paper: https://arxiv.org/abs/1708.07120. There are also benefits to using schedules for parameters like epsilon (as a trust-region / damping parameter), highlighted by this blog post: http://zna.do/epsilon.

[n2cholas]

Most of the current optimizers keep a step count in their state, so adding a schedule to their hyper-parameters would be straightforward. The signature for Adam, for example, could be changed from:

def scale_by_adam(b1: float = 0.9,
                  b2: float = 0.999,
                  eps: float = 1e-8,
                  eps_root: float = 0.0) -> GradientTransformation:

to

Schedule = Callable[[int], Union[float, jnp.ndarray]]

def scale_by_adam(b1: Union[float, Schedule] = 0.9,
                  b2: Union[float, Schedule] = 0.999,
                  eps: Union[float, Schedule] = 1e-8,
                  eps_root: Union[float, Schedule] = 0.0) -> GradientTransformation:

Then, within update, we would just need to do something like:

def update_fn(updates, state, params=None):
  count_inc = _safe_int32_increment(state.count)
  if callable(b1): b1 = b1(count_inc)
  if callable(b2): b2 = b2(count_inc)
  if callable(eps): eps = eps(count_inc)
  if callable(eps_root): eps_root = eps_root(count_inc)
  ...

I'd be happy to try contributing this if it seems reasonable.

[n2cholas]

Another option is to store the hyperparameters in the optimizer state, similar to this new proposed optimizer API for Flax, so users can explicitly change them as needed in the training loop. This has the added benefit of making logging these hyperpameters (e.g. to tensorboard) much easier.

[n2cholas]

Another possibility is to create a wrapper, like below:

class ScheduledTransformState(NamedTuple):
    count: jnp.ndarray  # shape=(), dtype=jnp.int32
    hparams: Dict[str, Any]
    inner_state: Any

def scheduled_transform(inner, **schedule):
    def schedule_fn(step):
        return {k: f(step) for k, f in schedule.items()}
    
    def init_fn(params):
        count = jnp.zeros([], jnp.int32)
        hparams = schedule_fn(count)
        return ScheduledTransformState(
            count=count,
            hparams=hparams,
            inner_state=inner(**hparams).init(params))
        
    def update_fn(updates, state, params=None):
        count_inc = _safe_int32_increment(state.count)
        hparams = schedule_fn(count_inc)
        updates, inner_state = inner(**hparams).update(
            updates, state.inner_state, params)
        return updates, ScheduledTransformState(
            count=count_inc,
            hparams=hparams,
            inner_state=inner_state)
    
    return GradientTransformation(init_fn, update_fn)

Example usage:

import optax
from functools import partial

optimizer = optax.scheduled_transform(
    partial(optax.sgd, nesterov=True), 
    learning_rate=optax.exponential_decay(...),
    momentum=optax.piecewise_constant_schedule(...))

This is much less intrusive (don't need to change existing transformations), easier to test, and more general than the other ideas above.

[n2cholas]

Any thoughts on this wrapper @mtthss?

[mtthss]

I find your last proposal very interesting!
Apologies for the delay in replying.
I'll get back to you tomorrow

[mtthss]

Off the top of my head comments:

1. I love that it allows to support scheduling of arbitrary hypers of arbitrary transforms without complicating the code of individual optimizers/transforms.

2. I think the scheduled_transform wrapper probably should accept also fixed hypers. This way the entire thing actually just becomes a very special partial that supports both fixed values and schedules with the same syntax of a standard Python partial. E.g. we could easily change it so that your example above would just become:

optimizer = optax.scheduled_transform(
    optax.sgd,
    nesterov=True,
    learning_rate=optax.exponential_decay(...),
    momentum=optax.piecewise_constant_schedule(...))

3. We would probably want to find a better name than scheduled_transform though :) Maybe it should just be optax.partial (once we support both fixed hypers and schedules).

[mtthss]

@avital

[n2cholas]

I agree with 2. I think supporting this behaviour is as simple as checking whether the argument is callable in schedule_fn:

def schedule_fn(step):
    return {k: (f(step) if callable(f) else f) 
            for k, f in schedule.items()}

I agree with 3 as well, names are not my forté 😅 (the alternative name i was thinking of was optax.schedule, though this is not a particularly good name either since we now accept fixed hyper-parameters as well).

I think the name partial is interesting. My only reservation is that normal partial application accepts a Callable[..., ReturnType] and returns a Callable[..., ReturnType], whereas here we accept a Callable[..., GradientTransformation] and return a GradientTransformation. I think if optax.partial returned a Callable[..., GradientTransformation], partial would be a more fitting name. What do you think?

[mtthss]

Naming is not my forte either tbh :)

But I agree on your comment about the name partial
If we did call it partial then it should not return an already instantiated gradient transformation;

Do you think it would be possible to adapt the wrapper to actually behave as a partial? So that for instance you would do:

scheduled_sgd = optax.partial(
    optax.sgd,
    learning_rate=optax.exponential_decay(...),
    momentum=optax.piecewise_constant_schedule(...))

optim = scheduled_sgd(nesterov=True)

[mtthss]

The downside of the partial signature is that this way it would not be immediately chainable because it doesn't give you back directly a GradientTransformation, so user would need to remember to separately instantiate it.

Overall I think your approach is actually best,
so the only concrete changes I'd suggest are

• could you rename it to optax.schedule_hyperparams?
• could you support the fixed hyperparam case?

It really is quite nice, and will allow us to cleanly disentangle the support for scheduling hyper-parameters from the individual gradient transformations, thereby keeping the latter nice and clean.

Could you put together an actual PR and then we circulate to gather some more feedback?

[n2cholas]

@mtthss sounds good, I'll incorporate those two changes and create a PR!