General and Adaptive Robust Loss Function

[chrisflesher]

Hello, I was interested in using the "General and Adaptive Robust Loss Function", it is a meta-loss that has some tunable parameters to help reject outliers. There is some existing code that implements it here:

https://github.com/google-research/google-research/tree/master/robust_loss_jax
https://github.com/google-research/google-research/tree/master/robust_loss

I was wondering if a PR for porting something like this would be feasible / encouraged?

[chrisflesher]

Think the code needed to implement would look something like this:

@jax.jit
def general_loss(x, alpha, scale):
    expm1_safe = lambda x: jnp.expm1(jnp.minimum(x, 43))  # noqa: E731
    eps = jnp.finfo(x.dtype).eps
    loss_two = 0.5 * (x / scale)**2
    scale = jnp.maximum(eps, scale)
    beta_safe = jnp.maximum(eps, jnp.abs(alpha - 2.))
    alpha_safe = jnp.where(alpha >= 0, jnp.ones_like(alpha), -jnp.ones_like(alpha)) * jnp.maximum(eps, jnp.abs(alpha))
    loss = (beta_safe / alpha_safe) * ((loss_two / (0.5 * beta_safe) + 1)**(0.5 * alpha) - 1)
    loss = jnp.where(alpha == -jnp.inf, -expm1_safe(-loss_two), loss)
    loss = jnp.where(alpha == 0., jnp.log1p(loss_two), loss)
    loss = jnp.where(alpha == 2., loss_two, loss)
    loss = jnp.where(alpha == jnp.inf, expm1_safe(loss_two), loss)
    return loss

@jax.jit
def adaptive_loss(error: jax.Array,
                  alpha: jax.Array,
                  scale: jax.Array,
                  alpha_lo: float = 0.001,
                  alpha_hi: float = 1.999):
    alpha = jax.nn.sigmoid(alpha) * (alpha_hi - alpha_lo) + alpha_lo
    scale = jnp.maximum(jnp.finfo(error.dtype).eps, scale)
    loss = general_loss(error, alpha, scale)
    x = _partition_spline_curve(alpha)
    spline = onp.load('partition_spline.npz')
    log_z = _interpolate_cubic_hermite_spline(x * spline['x_scale'], spline['values'], spline['tangents'])
    return loss + jnp.log(scale) + log_z

One issue might be the adaptive_loss function requires spline coefficients stored in an NPZ file? (currently 193kB)

[fabianp]

This sounds like a good idea to me. @jonbarron , what do you think about including these losses into optax?

[jonbarron]

Sgtm!

[fabianp]

@chrisflesher are you still interested in adding such loss? If so, please go ahead :-)

[chrisflesher]

Sure, I'll try to put together a PR.