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.
