Gradient clipping

[JuanFMontesinos]

How to carry out gradient clipping in JAX's ecosystem? I havent seen a clear cannonical method.

[mattjj]

Thanks for the question! Do you mean like this gradient clipping, or did you have something else in mind?

[JuanFMontesinos]

Dear @mattjj
Thanks for you input! That actually seem a way of overriding how gradients are computed for a given function, right?
I was more asking about a general way to clip gradients in a generic way when training deep learning models.

For example, taking a segment of my code:

        (loss, vars), grads = jax.value_and_grad(_loss, has_aux=True)(state.params)

        if self.cfg.model.norm in ["batch"]:
            state = state.apply_gradients(grads=grads, batch_stats=vars["batch_stats"])
        else:
            state = state.apply_gradients(grads=grads)

I imagine I iterate in a self recursive manner through grads and clip? But are there any functionality already predefined to do this? Namely, a cannonical, recommended way?

[dfm]

@mattjj might be able to say more, but I think the typical approach would be to use something like the clip transformation from optax (as documented here), especially if you're already using optax for training. Even if you're not using optax, I expect that you could repurpose those functions.