Gradient clipping based on histogram ?

[mchagneux]

Hi,

I'm looking for a way to clip gradients based on their distributions (in a minibatch) at each time step, lowering the norm of only the extreme ones. Basically I want to max out those gradients that lie outside some percentile at a given value (https://en.wikipedia.org/wiki/Winsorizing). I'm surprised this doesn't exist in most frameworks ? Is there a principled way to do this or should I write this myself ? This function https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.mstats.winsorize.html has not been implemented in Jax yet.

Thanks in advance,

Mathis

[mtthss]

It does seem an interesting idea in principle, I havn't read about any canonical way of doing it though.

[quantshah]

Is it very complicated to just clip gradients in the standard update function, e.g.,

def update(params, opt_state, x, y):
    """Single update step.

    Args:
        params (dict): Parameters of the model
        opt_state (tuple): Tuple of optimizer and params state.
        x (array): Array of input data.
        y (array): Target values.
    """
    grads = jax.grad(loss)(params, x, y)

    ####################################################################
    # Gradient clipping/normalizing for all elements in the dictionary
    ####################################################################
    grads['encoder/linear']["b"] = grads['encoder/linear']["b"].clip(0.0001)
    grads['encoder/linear']["b"] = grads['encoder/linear']["b"]/jnp.linalg.norm(grads['encoder/linear']["b"])

    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state

[mtthss]

You can use optax.chain() when defining the optimiser to add gradient clipping as part of your optimizer.update call