How to sample fresh noise at each step for a non-linear noise-perturbed ODE ``dX = f(X, z)dt``?

[vadmbertr]

Hi,
Thanks for the great work in developing diffrax!

I have a use case involving the introduction of some noise at every integration time-step, but then apply a non-linear transformation to it (so it's not a "classical" SDE, more something like a perturbed ODE if that makes sense).
It would broadly write as $dX = f(X, z) dt$, where $z \sim \mathcal{N}(0, 1)$ and $f$ is an arbitrary function (a neural network for example).

I'm not very sure about how to implement a vector field corresponding to $f$, as a fresh sample $z$ is needed at every integration time-step (and that we can not pass around a "fresh" jax.Key, as args must be fixed during the whole integration).
The way I see this would be to pass a VirtualBrownianTree to the vector field, through args, and interpolate this path at every integration time $t$ to sample $z$. Is this correct? If so, would it be "safe" to backpropagate through the solve if using UnsafeBrownianPath with a fixed step-size?

Do you think about another approach to this?

Thank you for the feedback!

[jpbrodrick89]

Hi @vadmbertr , thanks for your question. It would be good to understand your use case a bit better. If you literally want to take a N(0,1) draw at every timestep and f is square integrable then as $dt \rightarrow 0$ your accumulated variance will also go to zero. If what you really want is for $z$ to be a Brownian motion (as suggested by your second paragraph) then instead of pre-simulating W I would recommend storing z in your state so something like

$$y = (X, z)$$

$$ dz = dW $$ $$ dX = f(X, z)dt$$

However, there are a few gotcha's with this case. First it introduces correlations between the effective draws at each time step, and more importantly the variance of $z$ will grow unbounded with time $t$. To partially handle this you could just replace with an Ornstein-Uhlenbeck process that saturates at (for example) variance 1 and reduces auto-corrleation:

$$dz = -\theta z dt + \sigma dW$$

For some $\theta$ and $\sigma$ (to saturate at a variance of 1 you will want $\sigma = \sqrt{2\theta}$).

If what you actually want is genuine white noise (zero correlation time), that's a different model — a classical SDE. Note you can't feed white noise through a nonlinear f directly (only the linear coupling survives in the limit), so you'd approximate it by matching moments:

$$dX = \text{E}[f(X,z)|x] dt + \sqrt{\text{Var}}[f(X,z)|x]dW$$

Would either of these approaches help? Happy to discuss further.

[vadmbertr]

Hi @jpbrodrick89, thanks for the answer and leads!
Indeed, I should have given more context...

What I want is closer to a learned perturbation than to a real SDE: perturbing the parameters or the latent space of a model (a NN) with iid noise at every step, as in AIFS-CRPS (arXiv:2412.15832) and FGN (arXiv:2506.10772). There $dt$ is fixed and large (the model step, not a discretization), the rollout is autoregressive, and there is no $dt \to 0$ limit. So the vanishing variance, the issue with passing the white noise through $f$, and the need for an OU process probably don't apply in this case.

One of the main reasons for looking at diffrax was to reuse its checkpointing mechanism rather than reimplement one on top of equinox internals — not really because I need the SDE or adaptive-stepping machinery.

On standard normal vs Brownian path: passing a VirtualBrownianTree / UnsafeBrownianPath and evaluating it at $t$ was mostly a way to avoid computing a step index into a pre-sampled noise array. But it seems to create more problems than it solves (increment vs level scaling, the consistency/correlation differences between UBP and VBT, and re-evaluation at several stage times for non-Euler solvers).

In the end the cleanest approach looks like pre-sampling the noise, passing it through args, and indexing by step (idx = round((t - t0) / dt)) inside the term: one fixed sample per step, independent of where the solver evaluates. But since this is really a discrete per-step perturbation rather than something the term/control abstraction is built for, it may be cleaner to drop diffrax and use equinox checkpointing directly.

Does this seem reasonable, or am I missing a way diffrax abstractions could still fit?

[patrick-kidger]

In the end the cleanest approach looks like pre-sampling the noise, passing it through args, and indexing by step (idx = round((t - t0) / dt)) inside the term: one fixed sample per step, independent of where the solver evaluates. But since this is really a discrete per-step perturbation rather than something the term/control abstraction is built for, it may be cleaner to drop diffrax and use equinox checkpointing directly.

Does this seem reasonable, or am I missing a way diffrax abstractions could still fit?

I think you've hit the nail on the head. Because Diffrax may perform adaptive stepping, then the general model is that a control term (such as Brownian motion) should prespecify all its randomness 'up front', and is then queried for its value at particular times; there's no fixed-step-size abstraction at that level. In such a model then your round((t - t0)/dt) approach is the right one.

Alternatively, if you wanted to use eqx.internal.while_loop directly, then go for it! That would also be totally fine.

---

Footnote, the state of affairs described here has been non-ideal for a while. It's theoretically elegant, but in practice it means that fixed step size SDE solvers aren't as efficient as they could be, as e.g. one needs to perform this round business on each step.

I'm wondering if we should extend the Diffrax model for controls to also pass an index: None | Int[Array, ""]. When using adaptive stepping then we pass a None. When using fixed stepping (to be precise, without rejections), then we could pass an integer. This would allow controls to handle this case and do something more efficient.

@lockwo WDYT of this? It's not quite 'stateful controls' – our previous proposal for solving issues of this sort – but would this be the right abstraction to solve some of our historical issues around efficiency of SDE solves?

[vadmbertr]

In the end the cleanest approach looks like pre-sampling the noise, passing it through args, and indexing by step (idx = round((t - t0) / dt)) inside the term: one fixed sample per step, independent of where the solver evaluates. But since this is really a discrete per-step perturbation rather than something the term/control abstraction is built for, it may be cleaner to drop diffrax and use equinox checkpointing directly.
Does this seem reasonable, or am I missing a way diffrax abstractions could still fit?

I think you've hit the nail on the head. Because Diffrax may perform adaptive stepping, then the general model is that a control term (such as Brownian motion) should prespecify all its randomness 'up front', and is then queried for its value at particular times; there's no fixed-step-size abstraction at that level. In such a model then your round((t - t0)/dt) approach is the right one.

Alternatively, if you wanted to use eqx.internal.while_loop directly, then go for it! That would also be totally fine.

Thanks for the feedback! I was slightly reluctant to the use of eqx.internal (because it's "internal" and not documented) but I guess that as it is used extensively in Diffrax (and maybe elsewhere) I should be safe.

Footnote, the state of affairs described here has been non-ideal for a while. It's theoretically elegant, but in practice it means that fixed step size SDE solvers aren't as efficient as they could be, as e.g. one needs to perform this round business on each step.

I'm wondering if we should extend the Diffrax model for controls to also pass an index: None | Int[Array, ""]. When using adaptive stepping then we pass a None. When using fixed stepping (to be precise, without rejections), then we could pass an integer. This would allow controls to handle this case and do something more efficient.

@lockwo WDYT of this? It's not quite 'stateful controls' – our previous proposal for solving issues of this sort – but would this be the right abstraction to solve some of our historical issues around efficiency of SDE solves?

I've been following a bit those questions, and I think that something like "stateful controls", or some small extensions when dealing with a known number of steps would be useful! (to me at least)