NaN occurs when backprop through the kernel including ReLu layer

[ZuowenWang0000]

Hello,
I have encountered a problem when I try to get the gradient of some loss function with respect to some input variable x, I got a NaN after several iterations. And this only appears when the NTK is including ReLu layer. I've tried Erf or Sigmoid both don't have this problem.

The kernel function I am getting from:

  self.init_fn, self.f, self.kernel_fn = ntstax.serial(
      ntstax.Dense(1, parameterization='ntk'),
      ntstax.Relu(do_backprop=True, do_stabilize=True),
      ntstax.Dense(1, parameterization='ntk')
  )

And I try to grad via:
grads = grad(model_loss, argnums=1)(params, (x, y))[0]

and model_loss = lambda params, (x,y) : loss_func(pred(params, x), y)

btw,
Is the _safe_sqrt function here:

neural-tangents/neural_tangents/stax.py

Line 3847 in c6f759d

def _safe_sqrt(x):

is not back_prop safe? We might need np.where to np.maximum, just like in _sqrt.

Thanks!

[romanngg]

Thanks for the report - could you check if this^ commit helps? It should remove the do_backprop argument and fix nans, as well as improve numerical stability of differentiating nonlinearities - lmk if this works!

[ZuowenWang0000]

Thanks for the report - could you check if this^ commit helps? It should remove the do_backprop argument and fix nans, as well as improve numerical stability of differentiating nonlinearities - lmk if this works!

ReLU works now. thanks for the timely update!