Zero Gradients after argmax

[aniquetahir]

Hi,

I am facing this issue where I get zero gradients after using argmax in a loss function. I have created a minimal example:

import haiku as hk
import jax.numpy as jnp
import jax.random
import optax
import chex

hidden_dim = 64
input_shape = 12
num_classes = 2

class MLP(hk.Module):
    def __init__(self, name=None):
        super().__init__(name=name)


        xavier_constant_1 = jnp.sqrt(6/(input_shape + 256))
        xavier_constant_2 = jnp.sqrt(6/(256 + 256))
        xavier_constant_3 = jnp.sqrt(6/(1 + 256))
        self.seq = hk.Sequential([
            hk.Linear(hidden_dim, w_init=hk.initializers.RandomUniform(-xavier_constant_1, xavier_constant_1), b_init=hk.initializers.Constant(0.)),
            hk.Linear(hidden_dim, w_init=hk.initializers.RandomUniform(-xavier_constant_2, xavier_constant_2), b_init=hk.initializers.Constant(0.)),
            hk.Linear(num_classes, w_init=hk.initializers.RandomUniform(-xavier_constant_3, xavier_constant_3), b_init=hk.initializers.Constant(0.))
        ])

    def __call__(self, x: chex.Array):
        out = x.reshape((x.shape[0], -1))
        return self.seq(out)


def train_simulated():
    def mlp_fn(x):
        mlp = MLP('test_mlp')
        return mlp(x)
    mlp = hk.transform(mlp_fn)
    init, apply = hk.without_apply_rng(mlp)

    k1 = jax.random.PRNGKey(0)
    k2 = jax.random.PRNGKey(1)
    k3 = jax.random.PRNGKey(2)
    k4 = jax.random.PRNGKey(3)
    params = init(k1, jnp.ones((10, 12)))

    def loss_fn(parameters, x: chex.Array, y: chex.Array):
        y_hat = apply(parameters, x)
        preds = jnp.argmax(y_hat, axis=1)
        return ((preds.reshape(-1, 1) - y) ** 2).sum()

    loss_value_grad = jax.value_and_grad(loss_fn)
    v, g = loss_value_grad(params, jax.random.uniform(k2, (10, 12)), (jax.random.uniform(k3, (10, 1)) > 0.5).astype(float))
    print(g)

if __name__ == '__main__':
    train_simulated()

Also opened issue here:
google-deepmind/dm-haiku#442

[cgarciae]

argmax is not differentiable, if you want to train a classifier use optax.softmax_cross_entropy

[jakevdp]

For rank-based computations like argmax often zero is the correct gradient. For a more detailed discussion of this issue, see #8977 (comment)

For a differential version of rank-based computations, you can soften the discontinuity using softmax and similar functionality.

[jakevdp]

I've added a discussion of this to the FAQ in #11040