How to have different optimizer hyperparameters for different type of parameters (weights) in keras with jax backend?

[VachanVY]

• I'm using a fully custom training loop (NOT using model.fit) with LossScaleOptimizer
• I want to do something like the below in keras that the below pytorch code does

torch code from

# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
            {'params': decay_params, 'weight_decay': weight_decay},
            {'params': nodecay_params, 'weight_decay': 0.0}
        ]
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)

my jax code (in the Training section of the notebook link)

Please help. Thanks!

[fchollet]

Just use separate optimizers for each group of variables. If you have your own custom training loop that should be easy.

[fchollet]

this seems a bit complicated. Is it how it's done?

It sounds like you are using JAX. The stateless nature of JAX makes such variable manipulations always a little bit involved.

You do need a "control group" for each set of variables (to be used with each optimizer). Whether that's complicated or not depends on your factoring. I recommend using model.variables directly to track variable order (this is the order accepted by stateless_call), and to rely on variable paths to make your control groups (lists?).

def filter_group(all_variables, path_prefix):
    return [v if v.path.startswith(path_prefix)]

var_group = filter_group(model.variables, path_prefix="something")

If you want to use such a filter function with arrays (which aren't Keras Variables, but plain JAX arrays) you will want to keep a map between Keras variable instances and corresponding arrays. Or just... assign your variables before doing the filter, e.g.

my_arrays = ...
for v, a in zip(model.variables, my_arrays):
   v.assign(a)

var_group = filter_group(model.variables)

[VachanVY]

Can you also clarify the use of gradient scaling in this situation?

[fchollet]

In your training loop, just scale the loss after computing it. You can take a look at how the JAX train_step does it.

[VachanVY]

@fchollet

decay_opt = LossScaleOptimizer(AdamW(weight_decay=1e-1))
nodecay_opt = LossScaleOptimizer(AdamW(weight_decay=0.0))

What LossScaleOptimizer does (correct me if I'm wrong..):

• Scales the loss by a factor using optimizer.scale_loss. Now all gradients are scaled by that factor
• optimizer.apply takes care of unscaling the gradients and then updates the trainable vars

Question:
Let's say I use the decay_opt grad scaler to scale the loss, it'll scale the loss by a factor f
As all grads are scaled by that factor f, but the loss was scaled using decay_opt, so when applying nodecay_grads how will nodecay_opt know the factor f (as it's known only by decay_opt as it scaled the loss) by which it should unscale and apply nodecay_grads.

[VachanVY]

@fchollet Could please clarify how the gradients will be unscaled and updated?