Gradients of state variables in post-event simulations

[FFroehlich]

Sorry, took me a while to carve out some time to have a look at this again. Looks like state sensitivity in post-event simulations is now always wrong, independent of usage of ClipStepSizeController (but computed gradient values still appear to depend on it).

this is with latest jax (0.9.0.1) and diffrax (0.7.1)

I understand that for this particular problem, we could simply use the ClipStepSizeController and not add the event to diffeqsolve, but this does not work for other settings we are looking at.

import jax

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

import jax.numpy as jnp
import optimistix as optx
import diffrax

jump_time = 0.98
controller = diffrax.PIDController(rtol=1e-6, atol=1e-6)
controller = diffrax.ClipStepSizeController(controller, jump_ts=[jump_time])

def solve(event_time):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(t,event_time),True],
            [1.0, 0.0]
        )])
    )
    solver = diffrax.Heun()

    sol_event = diffrax.diffeqsolve(
        term,
        solver,
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=controller,
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: event_time - t,
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    sol = diffrax.diffeqsolve(
        term,
        solver,
        t0=sol_event.ts[-1],
        t1=2,
        dt0=None,
        y0=sol_event.ys[-1],
        stepsize_controller=controller,
        max_steps=100,
    )
    return sol_event, sol

def compute_end_state(event_time):
    sol_event, sol = solve(event_time)
    return sol.ys[-1, 0]

def compute_event_time(event_time):
    sol_event, sol = solve(event_time)
    return sol_event.ts[-1]

assert jnp.isclose(compute_end_state(jump_time),jump_time) # pass
assert jnp.isclose(compute_event_time(jump_time),jump_time) # pass

assert jnp.isclose(jax.grad(compute_event_time)(jump_time), 1.0) # pass
print(jax.grad(compute_end_state)(jump_time))
# 0.62716916609091 without `ClipStepSizeController`
# 2.220446049250313e-16 with `ClipStepSizeController`
assert jnp.isclose(jax.grad(compute_end_state)(jump_time),1.0) # fail

[patrick-kidger]

Okay, this successfully nerd-sniped for waaaay too many hours, thinking we had a bug in Diffrax itself.

Actually, it turns out that it's much simpler than that. You should replace controller = diffrax.ClipStepSizeController(controller, jump_ts=[jump_time]) with controller = diffrax.ClipStepSizeController(controller, jump_ts=[event_time]), to close over the jump time differentiably. Job done! 😁

To explain why:

1. Our first alternative option might be not to use ClipStepSizeController at all. This would definitely be a problem, and we can see why by reducing down to this MWE:

   import jax
   import jax.numpy as jnp
   import optimistix as optx
   import diffrax
   
   def solve(event_time):
       term = diffrax.ODETerm(lambda t, y, args: (t < event_time).astype(jnp.float32))
       controller = diffrax.PIDController(rtol=1e-6, atol=1e-6)
       sol = diffrax.diffeqsolve(
           term,
           diffrax.Heun(),
           t0=0,
           t1=2,
           dt0=None,
           y0=0,
           stepsize_controller=controller,
           max_steps=100,
       )
       [out] = sol.ys
       return out
   
   print(jax.value_and_grad(solve)(jnp.asarray(0.98)))

   observe that event_time is only ever consumed nondifferentiably (as part of a boolean comparison operation). So even without knowing anything about Diffrax internals, then you can be pretty sure that you will get zero gradient out (which is indeed the case).

   So the general principle to recall here is that if you want a gradient with respect to a jump time, then that jump time must be passed explicitly to ClipStepSizeController(..., jump_ts=...).

2. As for what happens when we include ClipStepSizeController, but do not pass the jump time nondifferentiably... well now hopefully it's reasonably clear that this is just another flavour of the above. Once again, the jump time is not consumed differentiably so we can reason that necessarily only a zero gradient will be obtained. Comparing to your own use-case, you can see that the expected 'mathematical' gradient of 1.0 is including the contribution from the partial derivative that corresponds to moving the jump itself.

Finally, I think the Event is a red herring. It definitely participates in the gradient computation, but in a way that ends up working out to a net zero in the end.

(I will note that I am willing to be a bit suspicious of combining jumps + events that both occur at the same t; in particular in the case that the event also has y dependence. The jump means that the dense interpolate will be very steep, whilst the event relies on that dense interpolation in order to do its solve + compute its gradient. Whilst I don't think this MWE ends up demonstrating it, it might be that your full use-case exhibits some other unresolved issue of this character.)

I hope that helps!

[FFroehlich]

Okay, this successfully nerd-sniped for waaaay too many hours, thinking we had a bug in Diffrax itself.

Actually, it turns out that it's much simpler than that. You should replace controller = diffrax.ClipStepSizeController(controller, jump_ts=[jump_time]) with controller = diffrax.ClipStepSizeController(controller, jump_ts=[event_time]), to close over the jump time differentiably. Job done! 😁

Oh right 🤦 , that explains why ClipStepSizeController no longer improved results, sorry about that.

To explain why:

1. Our first alternative option might be not to use ClipStepSizeController at all. This would definitely be a problem, and we can see why by reducing down to this MWE:
   import jax
   import jax.numpy as jnp
   import optimistix as optx
   import diffrax

   def solve(event_time):
   term = diffrax.ODETerm(lambda t, y, args: (t < event_time).astype(jnp.float32))
   controller = diffrax.PIDController(rtol=1e-6, atol=1e-6)
   sol = diffrax.diffeqsolve(
   term,
   diffrax.Heun(),
   t0=0,
   t1=2,
   dt0=None,
   y0=0,
   stepsize_controller=controller,
   max_steps=100,
   )
   [out] = sol.ys
   return out

   print(jax.value_and_grad(solve)(jnp.asarray(0.98)))

   observe that event_time is only ever consumed nondifferentiably (as part of a boolean comparison operation). So even without knowing anything about Diffrax internals, then you can be pretty sure that you will get zero gradient out (which is indeed the case).

Ah, interesting! But in my MWE above, the gradient with respect to sol.ts is nonzero (and correct!), so event_time does appear to be consumed differentiably, even when not using ClipStepSizeController. I would assume the gradient information is computed from the non-linear solve (of a here linear system) from rootfinding for the event time (possible via the implicit function theorem, same approach as you use to backprop over the non-linear solves from ODE integration). Is it possible that backpropagation for sol.ys does not (or incompletely) take the rootfinding for the event time into account?

So the general principle to recall here is that if you want a gradient with respect to a jump time, then that jump time must be passed explicitly to ClipStepSizeController(..., jump_ts=...).
2. As for what happens when we include ClipStepSizeController, but do not pass the jump time nondifferentiably... well now hopefully it's reasonably clear that this is just another flavour of the above. Once again, the jump time is not consumed differentiably so we can reason that necessarily only a zero gradient will be obtained. Comparing to your own use-case, you can see that the expected 'mathematical' gradient of 1.0 is including the contribution from the partial derivative that corresponds to moving the jump itself.

But the derivative w.r.t. sol.ts[-1] is correct even when not including the ClipStepSizeController at all, so I don't see how it would be possible to just compute it as partial derivative.

Finally, I think the Event is a red herring. It definitely participates in the gradient computation, but in a way that ends up working out to a net zero in the end.
(I will note that I am willing to be a bit suspicious of combining jumps + events that both occur at the same t; in particular in the case that the event also has y dependence. The jump means that the dense interpolate will be very steep, whilst the event relies on that dense interpolation in order to do its solve + compute its gradient. Whilst I don't think this MWE ends up demonstrating it, it might be that your full use-case exhibits some other unresolved issue of this character.)

I actually had a look what the solver steps look like for the code that you posted with and without event (always without ClipStepSizeController). The initial timepoints are exactly the same and the difference for the last time point before the event is tiny. Therefore, I struggle to see why adding the Event without ClipStepSizeController the leads to a non-zero gradient

I hope that helps!

Overall our issue is that we de-facto cannot use any events for which we cannot pre-compute jump times explicitly (e.g. non-linear root functions) to pass it to ClipStepSizeController (or at least that gradient computation is inconsistently reliable).

[patrick-kidger]

But the derivative w.r.t. sol.ts[-1] is correct even when not including the ClipStepSizeController at all

I would assume the gradient information is computed from the non-linear solve (of a here linear system) from rootfinding for the event time (possible via the implicit function theorem, same approach as you use to backprop over the non-linear solves from ODE integration)

Yup, this derivative is coming from the Event, and exactly what you speculate here is correct, including that it is via IFT. If you're curious it is this block of code that does it.

For this gradient I believe it doesn't matter whether the jump time is consumed differentiably, and that what matters is if the event time is consumed differentiably. (And the jump time and event time are conceptually two different things, even if they happen to take the same value here.)

---

Is it possible that backpropagation for sol.ys does not (or incompletely) take the rootfinding for the event time into account?

I think the event is being taken into account, although admittedly the way in which it interacts between the two solves is fairly subtle. (For example, the sol_event.ts[-1] that is found is actually slightly before the event – though within tolerance – so interestingly the second solve still has to make one very tiny step to get past that. We could change that but at least here I think it's net-harmless.)

---

Therefore, I struggle to see why adding the Event without ClipStepSizeController the leads to a non-zero gradient

Think of the jump time and the event time as being two separate things. Diffrax will compute a gradient with respect to the jump times (passed to ClipStepSizeController) and the event times (passed to Event).

If they should happen to be the same variable, then the two gradients will get added together. (Meanwhile they are different variables in your original code – you pass in the JAX tracer event_time = <JAX tracer stuff> and the raw Python float jump_time = 0.98.)

[FFroehlich]

But the derivative w.r.t. sol.ts[-1] is correct even when not including the ClipStepSizeController at all

I would assume the gradient information is computed from the non-linear solve (of a here linear system) from rootfinding for the event time (possible via the implicit function theorem, same approach as you use to backprop over the non-linear solves from ODE integration)

Yup, this derivative is coming from the Event, and exactly what you speculate here is correct, including that it is via IFT. If you're curious it is this block of code that does it.

For this gradient I believe it doesn't matter whether the jump time is consumed differentiably, and that what matters is if the event time is consumed differentiably. (And the jump time and event time are conceptually two different things, even if they happen to take the same value here.)

Is it possible that backpropagation for sol.ys does not (or incompletely) take the rootfinding for the event time into account?

I think the event is being taken into account, although admittedly the way in which it interacts between the two solves is fairly subtle. (For example, the sol_event.ts[-1] that is found is actually slightly before the event – though within tolerance – so interestingly the second solve still has to make one very tiny step to get past that. We could change that but at least here I think it's net-harmless.)

diffrax/diffrax/_integrate.py

Lines 762 to 763 in b3b3c8c

	# TODO: we might need to change the way we evaluate `_yfinal` in order to
	# get more accurate derivatives?

🤔

[patrick-kidger]

Yup! This is a note to myself that we could either evaluate the dense interpolate, or make a numerical step, in order to get the value of the solution over the last time interval. I picked the first approach but we could plausibly also pick the latter.

[FFroehlich]

Yup! This is a note to myself that we could either evaluate the dense interpolate, or make a numerical step, in order to get the value of the solution over the last time interval. I picked the first approach but we could plausibly also pick the latter.

Yeah, I've come to the conclusion that in order to get accurate gradients we probably have to do the latter. My understanding is that ClipStepSizeController does effectively the same thing with jump_time, and as gradient wrt event time are working this should be accurate.

We'll give this a shot and submit a PR if it works.

[patrick-kidger]

Yeah, I've come to the conclusion that in order to get accurate gradients we probably have to do the latter.

Why do you think that? Do you have a MWE demonstrating an issue with inaccurate gradients at the moment?

If it helps at all, this was also the rabbit hole I went down yesterday looking at this issue – before realising that at least for this MWE it something else – so I already have a branch mostly set up to do this: https://github.com/patrick-kidger/diffrax/tree/partway-event-interpolate-to-step

[FFroehlich]

On a side note, some observations that point to a more profound, but, to me, not really intelligible issue:

1. y<event_time yields correct gradients (numerically equivalent but slightly different computational graph)

def solve(et):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(y[0], et),True],
            [1.0, 0.0]
        )])
    )
    sol = diffrax.diffeqsolve(
        term,
        diffrax.Heun(),
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: et - y[0],
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    return sol.ys[-1,0]

print(jax.grad(solve)(0.98)) # == 1.0

2. removing the discontinuity in RHS yields correct gradients (solver does not require dense time steps and steps from 0.10999999940395355 to 1.1100000143051147, at which point differentiation through the interpolant just works)

import jax
import jax.numpy as jnp
import optimistix as optx
import diffrax

def solve(et):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array(1.0)
    )
    sol = diffrax.diffeqsolve(
        term,
        diffrax.Heun(),
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: et - t,
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    return sol.ys[-1,0]

print(jax.grad(solve)(0.98)) # == 1.0

3. The compute event-state gradients don't appear to be continuous

import jax
import jax.numpy as jnp
import optimistix as optx
import diffrax

def solve(et):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(t, et),True],
            [1.0, 0.0]
        )])
    )
    sol = diffrax.diffeqsolve(
        term,
        diffrax.Heun(),
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: et - t,
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    return sol.ys[-1,0]

ets = jnp.linspace(0.97, 0.99, 21)
grads = [
    jax.grad(solve)(et) for et in ets
]
import matplotlib.pyplot as plt
plt.plot(ets, grads)
plt.xlabel("Event Time")
plt.ylabel("Gradient of Final State w.r.t. Event Time")
plt.show()

1. to me suggests that jax might be messing this up with some under-the-hood optimisations?
2. & 3. rather suggest something in solver.interpolation_cls degenerates (due to small interpolation intervals from dense time-stepping?) and the computed derivatives are noise-corrupted?

[FFroehlich]

Yeah, I've come to the conclusion that in order to get accurate gradients we probably have to do the latter.

Why do you think that? Do you have a MWE demonstrating an issue with inaccurate gradients at the moment?

Yes:

import jax
import jax.numpy as jnp
import optimistix as optx
import diffrax

def solve(et):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(t, et),True],
            [1.0, 0.0]
        )])
    )
    sol = diffrax.diffeqsolve(
        term,
        diffrax.Heun(),
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: et - t,
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    return sol.ys[-1,0]

print(jax.value_and_grad(solve)(0.97))  # Should be close to (0.97, 1.0) but is (0.97, 0.5)

I've given https://github.com/patrick-kidger/diffrax/tree/partway-event-interpolate-to-step a shot and while gradients at least now appear to be stable/continuous (w.r.t. event time), they are still off by a factor 2 :(

[FFroehlich]

Yeah, I've come to the conclusion that in order to get accurate gradients we probably have to do the latter.

Why do you think that? Do you have a MWE demonstrating an issue with inaccurate gradients at the moment?

Yes:

import jax
import jax.numpy as jnp
import optimistix as optx
import diffrax

def solve(et):
    term = diffrax.ODETerm(
        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(t, et),True],
            [1.0, 0.0]
        )])
    )
    sol = diffrax.diffeqsolve(
        term,
        diffrax.Heun(),
        t0=0,
        t1=2,
        dt0=None,
        y0=jnp.array([0.0]),
        stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6),
        event=diffrax.Event(
            cond_fn=lambda t, y, args, **kw: et - t,
            root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        ),
        max_steps=100,
    )
    return sol.ys[-1,0]

print(jax.value_and_grad(solve)(0.97))  # Should be close to (0.97, 1.0) but is (0.97, 0.5)

I've given https://github.com/patrick-kidger/diffrax/tree/partway-event-interpolate-to-step a shot and while gradients at least now appear to be stable/continuous (w.r.t. event time), they are still off by a factor 2 :(

looks like gradients are actually correct (only on the partway-event-interpolate-to-step branch, not on main) when using

        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less_equal(t, et),True],
            [1.0, 0.0]
        )])

rather than

        lambda t, y, args: jnp.array([jnp.select(
            [jnp.less(t, et),True],
            [1.0, 0.0]
        )])

[patrick-kidger]

I'll try to take a look at the rest of this in a bit, but to quickly comment on a point I've bumped into before:

they are still off by a factor 2

this is probably due to the use of Heun. Change to Tsit5 or Bosh3 etc. and you'll see different results. At least for Heun's method it was clear to me why this was, since numerically that averages each endpoint.

One thing that I was thinking about here is that it might make sense to swap out the current event root find for one that, instead of returning a point estimate, returns a small interval that the root lies within. Then we could use the left point when it comes to evaluating the vector field for the gradient (as it should still lie the pre-event time interval), but the right point as our output time (e.g. to initialize any future solves, as you do). I've very definitely not chased through the details of that though. It would be easy enough to hack that together and see what happens using optx.Bisection, which maintains such a (lower, upper) state internally.

[patrick-kidger]

Oh also, for the MWE you've just posted: note that your vector field is discontinuous, so this already isn't valid input for Diffrax (unless the jump is declared via ClipStepSizeController). I'd love to remove the footgun there – and throw an explicit error instead – but I don't know of a way to detect this property of a vector field.

[FFroehlich]

I'll try to take a look at the rest of this in a bit, but to quickly comment on a point I've bumped into before:

they are still off by a factor 2

this is probably due to the use of Heun. Change to Tsit5 or Bosh3 etc. and you'll see different results. At least for Heun's method it was clear to me why this was, since numerically that averages each endpoint.

Heun: 0.5 with jnp.less, 1.0 with jnp.less_equal
Tsit5: -1.4211714 with jnp.less, 0.9035392 with jnp.less_equal
Bosh3: 0.7777778 with jnp.less, 0.7777778 with jnp.less_equal

[FFroehlich]

Oh also, for the MWE you've just posted: note that your vector field is discontinuous, so this already isn't valid input for Diffrax (unless the jump is declared via ClipStepSizeController). I'd love to remove the footgun there – and throw an explicit error instead – but I don't know of a way to detect this property of a vector field.

I was expecting you to say something like that – that's why I ended up trying jnp.less_equal since that ensures that the right hand side is continuous on the whole problem domain. I can see how the discontinuities that you get from jnp.less are something that you don't want to deal with.