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

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!

@vadmbertr

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.

@jpbrodrick89

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?

@lockwo

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?

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)

@lockwo

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?

My understanding of this is adding a feature such that controls can consume an index at call/evaluate time (which would allow it to index something like a precomputed array of noise values). This seems like it has similar challenges to the initial stateful approach, in that it would require a lot of impacts on function arguments/impact a lot of call signatures. Since in this index has to get threaded all the way to the control.eval, it has to go through solvers (so solver would need a new input parameters) etc. In this case, it seems like this is a more specific instance of the stateful controls (replacing index: None | Int[Array, ""] with index: None | PyTree[Array] seems like the original stateful approach), but that only narrows the impacts in some places. Or am I misinterpreting it?

If so, would it be "safe" to backpropagate through the solve if using UnsafeBrownianPath with a fixed step-size?

Maybe not officially safe, but I have done it ;). You can also just make a brownian path that takes keys smuggled in through the solver state pretty easily (you can find an implementation on one of my old branches)

This seems like it has similar challenges to the initial stateful approach, in that it would require a lot of impacts on function arguments/impact a lot of call signatures.

So I think passing the index is simpler in a few ways, actually:

crucially, there is no stateful output. Keeping track of this (esp. in the presence of step rejections etc.) was the most finickity part IMO.
the API impact really is far more limited: it affects AbstractSolver.step, AbstractTerm.contr, and our various brownian motion implementations, and that's about it. (And as a practical matter the diff for each solver is maybe 10 characters per class, to thread i in-and-out.) To highlight, it doesn't need threading around further than this, e.g. through the steps and states of diffeqsolve.

hmmm I see, it's just allowing the index only. Yes that could be more doable. Like this is sort of what you meant (obv very draft sketch): #763?

(for reference, speedup is

VBT: 0.058461
Old UBP: 0.012105
New UBP + Precompute: 0.000028
Pure Jax: 0.000530

)

Navigation Menu

Repository navigation

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

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

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

Description

Activity

jpbrodrick89 commented on Jun 11, 2026

vadmbertr commented on Jun 13, 2026

patrick-kidger commented on Jun 13, 2026

vadmbertr commented on Jun 15, 2026

lockwo commented on Jun 17, 2026

lockwo commented on Jun 17, 2026

patrick-kidger commented on Jun 27, 2026

lockwo commented on Jun 30, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions

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