What is stop_gradient used for?

[cool-RR]

I'm a little confused about jax.lax.stop_gradient.

I'm debugging the behavior of the POLA algorithm. In a critical part of the code, it passes the results through this expression: jnp.exp(x - jax.lax.stop_gradient(x). This causes the resulting expression to always be 0, but still have a non-zero gradient. I want to understand why this is done.

Putting aside my specific algorithm, why does stop_gradient even exist? I understand what it does, i.e. it causes JAX to not differentiate by that expression. I've toyed around with examples using stop_gradient and the results make sense to me. But I'm not understanding why anyone would want this. Could you please explain it in a simple way?

[cool-RR]

Ah, I found an explanation in the TensorFlow docs: https://www.tensorflow.org/api_docs/python/tf/stop_gradient

I think that will be enough for me to understand. Thank you anyway.

[jakevdp]

Thanks for the question: we do have some documentation of jax.lax.stop_gradient here: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.stop_gradient.html

Is there anything in particular we could add there that would have answered your question?

[cool-RR]

The documentation explains what stop_gradient does but not what problem it solves.

I'm now diving into https://www.tensorflow.org/api_docs/python/tf/stop_gradient and while I get what problem it's trying to solve with the softmax example, I'm getting correct results even when I remove the stop_gradient, so I'm still trying to understand how stop_gradient is contributing to solving the problem.

[jakevdp]

Thanks - we'll update the JAX stop_gradient docs with a similar example.

[jakevdp]

Doc update in #23707

[cool-RR]

Thank you. I'll try to understand the example at https://jax.readthedocs.io/en/latest/_tutorials/advanced-autodiff.html#stopping-gradients

[cool-RR]

I'm confused about that explanation. It says "This update is not the gradient of any loss function. However, it can be written as the gradient of the pseudo loss function." There's no explanation for why we'd want to write it as a gradient of a loss function, rather than the existing expression.