Added double adjoint

[patrick-kidger]

Just creating a draft PR to let you know I'm working on double-backward through the adjoint.

In informal tests it seems to work as expected. The only thing left to do is write some formal tests. I'll put them in test_adjoint.py and switch that file over to pytest while I'm at it.

[lxuechen]

Thanks for pulling off a draft!

Some general comments

• Do we need to include t and v in _unpack_y_aug. My understanding is that we shouldn't unless 1) there's an explicit use-case where we need to take derivative w.r.t. t, and 2) there's an explicit use-case where v is not from the Brownian motion but some other tensor that requires grad. I haven't seen cases for 1), but I'm not entirely sure there isn't such a case. I don't think we'll ever see cases for 2).

• It seems to me that since _SdeintAdjointMethod.apply is called in backward, the context returned by torch.is_grad_enabled in the adjoint functions (e.g. f_uncorr) should be False due to the implicit context setting with forward of torch.autograd.Function.

• The current paradigm is adjoints-through-adjoints, and doesn't support naive backprop through adjoints. We should probably document this, even though this may seem trivial to us.

In informal tests it seems to work as expected.

What are the informal tests? Are there finite-difference tests? Or did we directly use torch.autograd.gradgradcheck? Does the old tests still pass? If you could also push the companion tests, then I think I'd be able to help more from running and tweaking the code.

Hopefully asking all of these question at once doesn't seem too annoying, but I just want to gain a better understanding of what's happening here. Also, we don't have to do everything at once or address all of the questions above at once, but it's good to know what works and what doesn't so that we can place meaningful TODOs.

[patrick-kidger]

• Including t and v: AFAIK both t and v should never require gradients in the current set-up.

But that's a detail that's beyond AdjointSDE's scope. AdjointSDE requires that certain conditions are satisfied, so it should enforce them. Practically: if that detail ever changes down the line then I'd prefer to know about it without things silently going wrong. The logic here can get quite hairy when you start getting into adjoints-of-adjoints.

• implicit contexts: this isn't true in the double-adjoint case, when f_uncorrected of the double adjoint sets enable_grad and then calls f_uncorrected of the first adjoint.

• Naive backprop: yeah, it should probably be added to the docstring for sdeint_adjoint. I'll do that.

Old single-adjoint tests still pass.

[lxuechen]

Including t and v: AFAIK both t and v should never require gradients in the current set-up.

Per current usage, I can't really come up with a case where t and v are not leaf variables, so I would prefer to not let _get_state check these just to reduce redundancy. The idea is that this PR adds the requires grad checks for ts upfront, and that v should always be generated by the Brownian motion. I think I would leave some comments documenting the situation here if we're really feeling unsure.

implicit contexts: this isn't true in the double-adjoint case, when f_uncorrected of the double adjoint sets enable_grad and then calls f_uncorrected of the first adjoint.

You're right, thanks for explaining.

Naive backprop: yeah, it should probably be added to the docstring for sdeint_adjoint. I'll do that.

Thanks for the fix for documentation!

Additionally, could we do a numerical test? Ideally, just replacing the torch.autograd.gradcheck here with torch.autograd.gradgradcheck (and slightly modifying the surrounding) should tell us something.

As a side note, I'll be quite busy Tues and Wed this week. But if you're willing to wait a little, then I can do the tests on Thursdays and Friday. In any case, I think testing is something that must be accomplished. Obviously you're a pretty good coder and software engineer, but I wouldn't feel the code is complete or absolutely trustworthy before we get those tests.

[lxuechen]

I'm a little worried about this line. Consider the scenario where we set y0 to not require gradients. By current standards, neither t nor v should require gradients. So in the backprop-through-solver setting, requires_grad evaluates to False.

The problematic scenario is then solving an SDE with parameters that require gradients. Since requires_grad is False, we don't create the extra graph backpropping through vg_dg_vjp towards the parameters. However, in the ideal situation, we should still backprop through this node so that our models get better gradients.

Overall, I think this possible failure case indicates that we should have a numerical test for backprop-through-solver comparing against finite differences.

[patrick-kidger]

I think it's worth noting that we currently don't (a) backprop-through-the-solver with the adjoint, nor do we (b) differentiate wrt parameters without also differentiating wrt y, so we at least shouldn't be any running into any issues with the current codebase.

That said, this is a subtle point that we'd probably never catch in the future if this scenario does come up, so I agree that I'd prefer to do something about it now.

I'm not completely sure about a fix.
The current requires_grad is meant to correspond to "if I hadn't forced y to have a gradient, would some_func(t, y, v) require gradient?"
I'm -1 on any approach that just fixes it to True, as else we'll be leaking graphs from the function.
I think my proposed fix would be to:

• If y requires gradient, set requires_grad to True. (As before.)
• If y doesn't require gradient, then we're going to be forcing it to require gradient. Then we call some_func(t, y, v) as before. Prior to calling vjp we manually traverse its autograd history (via .grad_fn, .next_functions) and set requires_grad = not (is y the only leaf).

Not a huge fan of that but I don't have any better ideas. Thoughts?

[lxuechen]

I think it's worth noting that we currently don't (a) backprop-through-the-solver with the adjoint, nor do we (b) differentiate wrt parameters without also differentiating wrt y, so we at least shouldn't be any running into any issues with the current codebase.

I don't think the line I'm specifically referring to has much to do with adjoints. gdg_prod is part of ForwardSDE, so in the backprop through solver case, it would get called when we choose, say, the Milstein method.

More specifically, I think this will cause a problem in the first step of sdeint (w/o adjoints), where y0 doesn't require gradient. The later steps shouldn't be affected by too much, since the typical y beyond y0 should require gradients due to being sums of tensors that require gradient with tensors that don't.

I'm not completely sure about a fix.

I'll think about this in more detail tomorrow, when I actually may have a chunk of sparse time.

[patrick-kidger]

Oh sorry, you're commenting on base_sde rather than adjoint_sde. You're right, this would break Milstein.

I think the analysis for the forward case is the same as the backward case that I actually discussed, though, so I think I'd end up proposing the same fix. I'll see what you say.

Unrelatedly, I've updated test_adjoint.py to pytest. Scalar noise+Ito SDE seems to be failing the tests btw. (It wasn't tested previously.) I've not dug into why that is - I know this isn't a supported case on master but I thought it was in now?

[lxuechen]

I think the issue I mentioned is mostly about when sde.parameters() is non-empty. So I would perhaps add a literal that checks whether ForwardSDE has any parameters that require gradients.

I'll fix the Scalar noise + Ito case.

[patrick-kidger]

So that relies on the implicit assumption that everything going into the function is a parameter. That needn't be the case, c.f. our discussion a bit back on contextualisation.

[patrick-kidger]

I agree that every current case involves t and v being leaf variables already. (Moreover non-gradient-requiring variables.)
I'm quite strongly in favour of maintaining the checks though, for what I think is the same reason that you brought up the discussion on requires_grad: we might not be doing it now, but it's an implicit assumption that will silently do the wrong thing if we ever change that assumption in the future.

Regarding tests - I feel like you think I'm a bit cavalier about tests! Don't worry, I appreciate their importance for any software project; I just tend to put them in a bit later than you.
For this PR specifically, it sounds like you're offering to write the tests? If you get time and are happy to do it, that would be great.

[lxuechen]

I agree that every current case involves t and v being leaf variables already. (Moreover non-gradient-requiring variables.)
I'm quite strongly in favour of maintaining the checks though, for what I think is the same reason that you brought up the discussion on requires_grad: we might not be doing it now, but it's an implicit assumption that will silently do the wrong thing if we ever change that assumption in the future.

Regarding tests - I feel like you think I'm a bit cavalier about tests! Don't worry, I appreciate their importance for any software project; I just tend to put them in a bit later than you.
For this PR specifically, it sounds like you're offering to write the tests? If you get time and are happy to do it, that would be great.

I appreciate the detailed thoughts, though I still believe there's a concrete difference between the issue I mentioned and the argument about t and v. I think we currently do have a use case of backprop through solver, and the issue I mentioned could affect this behavior. On the other hand, I couldn't really come up with a case where gradients wrt t and v need to be taken.

[patrick-kidger]

Hmm I don't quite follow. Gradients wrt t and v are a separate issue to this breaking Milstein. (If that's what you're saying then I agree!)
I agree that we never need gradients wrt t or v, I'm just quite assert-happy as a way of making sure things aren't silently misbehaving.
I agree that this PR introduces an issue wrt Milstein (whilst fixing the issue wrt leaking graphs).

Reflecting on the main (gradient) issue, I am actually inclined to just set requires_grad = torch.is_grad_enabled() and leave it at that. This does leak graphs, but I don't think the overhead should be too large. I think the alternate fix, or manually checking the graph, is probably a bit too magic.

[lxuechen]

Hmm I don't quite follow. Gradients wrt t and v are a separate issue to this breaking Milstein. (If that's what you're saying then I agree!)

Yes, I'm trying to argue that these are separate situations.

I agree that we never need gradients wrt t or v, I'm just quite assert-happy as a way of making sure things aren't silently misbehaving.

I still don't think it's necessary to check t or v, as I'm unable to come up with any likely scenario where these variables would require grad. Given that this check also adds complexity to the code and a tiny overhead, I'm not sure of the advantage.

I agree that this PR introduces an issue wrt Milstein (whilst fixing the issue wrt leaking graphs).

Reflecting on the main (gradient) issue, I am actually inclined to just set requires_grad = torch.is_grad_enabled() and leave it at that. This does leak graphs, but I don't think the overhead should be too large. I think the alternate fix, or manually checking the graph, is probably a bit too magic.

I agree with this fix, i.e. just leaving requires_grad = torch.is_grad_enabled().

Side note, I think this PR might be ready to be converted to review-mode from draft mode modulo the t v grad check discussion.

[patrick-kidger]

Changed requires_grad.

Regarding the asserts: as I understand it, you want to remove the assert t.is_leaf and assert v.is_leaf lines from adjoint_sde.py? I think it's important that AdjointSDE (or indeed any other library component) is self-contained; explicitly checked assumptions are better than implicit assumptions, especially in cases like this where a failure of that assumption will fail silently rather than fail loudly.

I can't think of a use case for t or v needing gradients either - but that could change not just because of some unforeseen feature, but because of a bug. (And the cost/complexity of two asserts is very low.)

[lxuechen]

Ok, let's leave the checks as they are for now. Happy to have this merged and thanks again for the great work in getting the adjoint-adjoints working!