Milstein (Strat), Milstein grad-free (Ito + Strat)

[mtsokol]

@patrick-kidger Hi!
I'm sharing a WIP draft with three methods: Milstein for Strat, Milstein grad-free for Ito and Strat - done as described - grad-free variants are inside original methods where Milstein Strat is a separate one (as agreed for code duplication, but whole code differs in - dt - should I unify it into one method with or leave that duplicaton?)
Grad-free mode can be used by passing {'grad_free': True} in options: False or lack of that option results in default sde.gdg_prod usage.

Issues/questions:

1. I think that method sde.g_prod shouldn't compute diffusion g and perform mul with I_k -> in Milstein grad-free I need to compute g, g * I_k and g * v so I end up with g being calculated three times instead of one. What do you think about extending sde API by adding def prod(self, g, v) method that will apply correct seq_mul according to the noise type? (that will introduce a bit of code duplication I guess)
2. I will finish writing diagnostics when fix for gdg_prod bug is merged. But stratanovich_diagonal.py with grad-free Milstein works fine. Here's one plot and rate:
   
   
   I think that usage of grad-free variant changes order to 0.5, right? (I can if it in code)
3. sqrt_dt = torch.sqrt(dt) if isinstance(dt, torch.Tensor) else math.sqrt(dt) appears in few places now - how about introducing misc.sqrt?

[patrick-kidger]

Heyup.

Overall, looks good!

Re: code duplication: I'd suggest creating an abstract base class BaseMilstein with nearly everything factored into that, with an abstract method that's called where their functionality differs. Then have two subclasses MilsteinIto and MilsteinStratonovich with each defining the correct method. Keep everything in one file.

1. I agree that's an issue; your proposed fix sounds good. (Additionally then adjust g_prod to just be a composition of g and prod. Eventually we may then remove the simple g_prod functionality but I don't want to break the adjoint stuff Chen is working on atm, which uses this.)

2. With luck, the fix for gdg_prod will be merged very shortly! Just been merged.

I'm surprised that the grad-free is looking like 0.5; it should still be an order-1 scheme for commutative noise.

3. Ideally dt should always be of just a single dtype (and we just use the appropriate sqrt function) - do you know where these different types are coming from?

Other comments:

• Can the 'euler' example be removed from the Strat diagnostics? Simplifies things slightly and it's not a strat method anyway.
• Can we go for MilsteinStratonovich rather than MilsteinStrat?
• Let's keep a single 'milstein' method from the user's point of view: the select function should additionally be made aware of the sde type and choose between them appropriately.

[patrick-kidger]

(I'll do a proper code review once these points are organised.)

[mtsokol]

@patrick-kidger Ok, applied all above comments.
I only haven't looked at 3. - the reason for that dt casting, but it existed in codebase in methods/skr.py - but will look at it.

Current rates for modified ito_diagonal.py:

...and stratonovich_diagonal.py:

For Ito grad-free keeps the same rate as expected but for Stratonovich it differs from 1.0.

side note: right now running diagnostics prints one additional warning but I haven't looked at it yet:

UserWarning: BrownianInterval is optimised for interval-based queries, not point evaluation. Consider using BrownianPath or BrownianTree instead.

[patrick-kidger]

I'm perturbed by the fact that the Strat derivative-free Milstein doesn't seem to be order 1. But I've been through your implementation and it looks correct to me.

Can you try it on other problems / other noise types / with smaller steps, and see what happens? Hopefully that will help to diagnose the issue.

Regarding the warning: this can safely be ignored. (Something I put in to warn about an easy-to-make inefficiency, but we know what we're doing. ;) ) If you like you can suppress it.

Overall this PR looks in good shape to me.

[mtsokol]

@patrick-kidger All comments addressed.

I've run stratonovich_diagonal also for Ex1 and Ex3 (considering correction term) and the plots looked correct (same as for Ex2) but slope is still 0.5.
Same for step size dt = 1e-2. When ito_scalar is fixed I can run it also against it later to check.
Few plots:

Ex1:

Ex3:

Smaller step size:

[patrick-kidger]

Other than the order for Derivative-free Stratonovich Milstein, this all looks good to me.

Pinging @lxuechen - maybe another set of eyes will help. Any idea why this solver seems to be getting order 0.5 rather than 1.0?

[lxuechen]

What is the reference on the current impl of derivative-free Milstein?

I think there might be a couple of issues in the math of the code I see, but I'm not completely certain given that I only went over this quickly.

[patrick-kidger]

What is the reference on the current impl of derivative-free Milstein?

I think there might be a couple of issues in the math of the code I see, but I'm not completely certain given that I only went over this quickly.

Equation (24): https://infoscience.epfl.ch/record/143450/files/sde_tutorial.pdf

Alternate source - equation (9) in: https://core.ac.uk/download/pdf/82209565.pdf

[mtsokol]

I think I've spotted one divergence in pdfs:
In https://infoscience.epfl.ch/record/143450/files/sde_tutorial.pdf there's no difference in Y_tilde that is being computed as: Y_tilde = Y_n + fn * h + g_n * sqrt(h) for both Ito and Stratonovich versions (25).

Where in Julia DiffEq they differ in fn * h term (it's removed for Stratonovich): https://github.com/SciML/StochasticDiffEq.jl/blob/master/src/perform_step/low_order.jl#L168 (line 175 vs 179)

And also in second pdf you linked those two different definitions of Z_aux are present (end of page 4.) but grad-free milstein for Stratonovich formula is not present (only: Again SDELab includes versions that do not requireuser-supplied derivatives under (13))

I tried it with that different Y_tilde but it's still 0.6 so there's still something else there.

[mtsokol]

And according to DiffEq here's t0 being used (But with those fixes slope is still 0.6).

[patrick-kidger]

Fair comment that Stratonovich specifically isn't mentioned in the second pdf, but that shouldn't affect things: the difference is just in what zeta means.

Yeah, I'd spotted the difference in t, but (a) our test problem is independent of t, and (b) if we treat t as part of the state (with unit drift and zero diffusion) then we recover the version you've implemented, so I don't think it should matter.

[patrick-kidger]

Maybe double check what convergence DiffEq.jl gets? In principle it should be possible (if perhaps a bit laborious) to get the two libraries to compute the exact same thing, and see where things diverge.

[mtsokol]

Hmm.. and running ito_diagonal for DiffEq Stratonovich form for Y_tilde doesn't change slope (still 1.0) so maybe that's completely unrelated.

[lxuechen]

I probably won't have time to look into this carefully this week, and I think it's likely only safe for us to merge this when we see the correct rate or figure out why otherwise.

One comment I'd have in general for the derivative-free scheme is to look into its motivation. The whole argument behind the derivative-free scheme is that it approximates the gdg term with the difference between function evaluations. The way to derive these schemes is via Taylor expansion (this is classical calculus and doesn't require any stochastic analysis). Concretely, you could do an expansion of g at y0 evaluated at y_tilde, and the overall finite diff term (after div root(h) and other ops) should match the overall gdg term in the vanilla Milstein scheme. This explains why that removing the f h term is valid as you described in Julia diffeq's impl and that you should eval at t0 (and not at t1). To double check that the math is correct, I'd also recommend starting from the Taylor expansion.

More generally, it makes sense to learn about why the scheme works that way while coding up the scheme. This will also make it clear why different people implement it differently and validly.

[mtsokol]

@patrick-kidger Hi! I've tried to recreate Ex2 in Julia here: https://github.com/mtsokol/torchsde/blob/dev-notebook/examples/Milstein%20order%20check.ipynb
Within few executions Ito version was mostly close to 0.9 but Stratonovich is also rather 0.6-0.7
WDYT?

[patrick-kidger]

Fascinating. Perhaps reassuring that this seems to be happening. Btw, do you know what the difference is between RKMil and RKMilCommute? The former says it handles scalar and diagonal noise, which are both commutative; but that sounds like the remit of the latter.

My next guess is that we're observing an intermediate regime, where we might only get 0.5 convergence for the timescales we're looking at, and would need to go much smaller. (Would need to think a bit harder about whether this is possible in this case, but I do know of examples going the other way, when you can observe 1.0 convergence on intermediate scales despite 0.5 being the true rate.)
To test this: can you try making the noise term larger/smaller and seeing what happens? e.g. just multiplying the output of g by 0.1 and seeing if that changes things.

[mtsokol]

Ok, I've tried it few times and that's results:

I couldn't come up with analytical solution for Ex2 with g multiplied by 0.1 so I decided to do it with Ex1 where solution is obvious. But before that I've tried to run Ex1 without modifying g.

Here's updated notebook: https://github.com/mtsokol/torchsde/blob/dev-notebook/examples/Milstein%20order%20check%202.ipynb

For Ito slope is 1.0 (but in few random runs it sometimes ended up as 0.6 lowest)

But for Stratonovich with proper correction it's 1.0 (in few runs it's in range 0.8 - 1.1)

And now fun part: I changed stratonovich_diagonal.py to be using Ex1 (also added proper correction) and modified Y_tilde just like in Julia implementation (when I checked it for Ex1 earlier I didn't change Y_tilde and the result was again 0.5 (also I checked it now once again)) and now that grad-free Milstein Stratonovich is sharp 1.0 🤔 (also checked it in few runs, here's branch: https://github.com/mtsokol/torchsde/tree/dev-notebook)

My questions:

1. What do you think about it?
2. Could it be due to the choice of Ex2 content and it's analytical solution formula?

[lxuechen]

prod will break the adjoint code. I don't think we should optimize this for now, and I'd rather just use g_prod at the cost of a few extra function evaluations. Feel free to add this to #23.

[mtsokol]

Sure! Will reverse it.

[lxuechen]

Thanks! We can brainstorm about ways to speed things up after this PR gets merged. On a high-level, the reason we made g_prod in the first place is that it's a vjp for the adjoint, and vjps are faster than directly computing the Jacobian. Evaluating g directly for the adjoint would mean computing the whole Jacobian, which scales poorly w.r.t. the dimension.

[mtsokol]

Sure thing!
For now I just pushed changes.

[patrick-kidger]

Nicely figured out! I can completely believe that it's down to the choice of SDE, and that we're just observing funny phenomena from this one. This is particularly backed up by the fact that the same behaviour occurs in a completely different package.

I am surprised that the choice of Y_tilde affects things. I thought either one was acceptable. Still, happy to go with what works.

(And yeah, as Chen says we're finding that things are a bit complicated in adjoint-land so the prod changes will have to come later.)

[mtsokol]

Ok, so current version is with all the latest remarks:
g_prod reversed,
diagnostics get correct orders (plots above) after changing to Ex1 for stratonovich_diagonal and impl tweak from Julia,
t1 changed to t0

[lxuechen]

LGTM. Feel free to pull the squash and merge trigger after taking another look @patrick-kidger.