Marking a function as nondifferentiable with jax.scipy.optimize

[jdongg]

I have a problem in the context of neural networks in which my loss function (to be optimized over params) has an extra parameter which is essentially computed by using the current value of params. The way this is structured is that the fun argument to jax.scipy.optimize.minimize contains a call to a helper function helperFunc which takes as input params. But helperFun should not be differentiable wrt params.

To give some more technical context, the hidden parameters (weights and biases) of the NN are treated as optimizable variables while the parameters of the activation layer are computed using the hidden parameters and are not viewed as part of the gradient or hessian and not an optimization variable. It's somewhat similar to this work: Robust Training and Initialization of Deep Neural Networks where they solve a linear least squares system involving the hidden parameters in order to obtain the activation parameters.

In my current implementation, jax would view the helper function as differentiable and I believe there is numerical instability due to trying to differentiable through a linear least squares solver. Either way, my intention like I said it to have the helper function not be differentiable. I tried doing something like passing a copy of params into the helper function but jax still of course views the copy as differentiable.

To present an MWE, I tried to simplify things as much as possible. I know this example is probably nonsensical, but it computes the minimum of the function x^2 + a(x) where a(x) = (x-1)^4 using BFGS. The current example returns a solution of x=0.410245, which is the correct location of the minimum for x^2 + (x-1)^4. I would like for the optimizer to not view helperFunc/a(x) as differentiable. In this case, I believe it should return x=0 as the minimum. What would be the easiest way for me to accomplish this with jax? I tried looking for similar questions but couldn't find anything on here. TIA!

import jax.numpy as jnp
import jax.scipy.optimize

from jax.config import config
config.update("jax_enable_x64", True)

def helperFunc(x):
	return jnp.power(x-1., 4)

def f(x):
	xcopy = x
	a = helperFunc(xcopy)
	return jnp.squeeze(jnp.power(x, 2) + a)

jOptimRes_x0 = jax.scipy.optimize.minimize(fun=f, x0=0.6*np.ones([1,]), method='BFGS')

[jakevdp]

I suspect what you're looking for is jax.lax.stop_gradient:

def f(x):
  a = lax.stop_gradient(helperFunc(x))
  return jnp.squeeze(jnp.power(x, 2) + a)