Repository navigation
sample method returns zero gradients #216
Description
Hi, I'm curious if it could be possible to differentiate through sampled properties in flowjax. For example, this is possible when sampling other distributions, like multivariate_normal within jax:
def test_normal(std):
randkey = jax.random.key(1)
mean = jnp.array([0.0])
cov = jnp.array([[std**2]])
return jax.random.multivariate_normal(randkey, mean=mean, cov=cov)[0]
std = 1.0
jax.grad(test_normal)(std)
>>> Array(-0.15443718, dtype=float32, weak_type=True)
However, when I attempt to do something similar with a flowjax model, the gradients always seem to be zero. The gradients are zero both with respect to the condition variable and the flowjax model parameters, but since it is simpler I will just use the condition variable for this example:
def test_flowjax(condition):
randkey = jax.random.key(1)
flow = flowjax.flows.block_neural_autoregressive_flow(
key=jax.random.key(0),
base_dist=flowjax.distributions.Normal(jnp.zeros(1)),
cond_dim=1
)
condition = jnp.array([condition])
return flow.sample(randkey, condition=condition)[0]
condition = 1.0
jax.grad(test_flowjax)(condition)
>>> Array(0., dtype=float32, weak_type=True)
For all sets of flow parameters, all random seeds, and all concrete values of condition that I have tried, this gradient is always zero.
Activity
danielward27 commentedon Apr 10, 2025
If you set
invert=Falseinblock_neural_autoregressive_flow, the gradient should not be zero. The problem is that eitherlog_proborsample{_and_log_prob}require a numerical method. In the setup you have here, sampling occurs through a numerical method, which results in a stepwise function with zero gradient.Technically the computed gradient of the numerical method isn't wrong, but it's probably usually unintentional to differentiate through the numerical method, so it probably would be better to error! The same problem as #176.
Alternatively, if you use a different model (e.g.
masked_autoregressive_floworcoupling_flow), this problem shouldn't exist.AlanPearl commentedon Apr 10, 2025
Ah, I see. Thanks for the quick and clear response! My test passes just by adding
invert=Falsedanielward27 commentedon Feb 18, 2026
This should be fixed by #234. Thanks @noahewolfe!
Care should still be taken to know the orientation of the flow though (
invert=Falseorinvert=True), because of efficiency reasons and more risk for numerical precision issues when using the numerical method.