如何在 Jax 中获得包含 argmax 的损失函数的梯度？

Question

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()

The output of the code is the gradients of the loss function for the parameters. However, all of the gradients are zero. This is not expected because the labels and the inputs are generated randomly.

Answer 1

When you're using a sorting-based computation like argmax , often zero is the correct gradient. For more discussion of this, see FAQ: Why are gradients zero for functions based on sort order? in the JAX documentation.