[question] Computational complexity of integrating / backpropogating through SDE

Hi, I was a bit surprised by the difference in speed between backpropogating through ODEs vs SDEs but couldn't find any discussion in the documentation about the time complexity of either. With a quick look, I couldn't find any issues addressing this. In particular, consider this piece of code, where we vary the different terms being integrated:

import optax
import jax
import jax.numpy as jnp
import jax.random as jrandom
import equinox as eqx
import diffrax as dx
from tqdm import tqdm

def sample(model, num_steps, gamma, dim, key):
    drift, diff = model

    def f(t, y, args):
        return drift(jnp.concatenate([t[None], y]))

    def g(t, y, args):
        return diff(jnp.concatenate([t[None], y]))**2

    control = dx.VirtualBrownianTree(
        t0=0,
        t1=1,
        tol=1/(2*num_steps),
        shape=(dim,),
        key=key,
    )

    drift_term = dx.ODETerm(f)
    diffusion_term = dx.WeaklyDiagonalControlTerm(g, control)
    terms = dx.MultiTerm(drift_term, diffusion_term)
    solver = dx.Euler()
    y0 = jnp.zeros(dim)

    ts = jnp.linspace(0, 1, num_steps + 1)
    saveat = dx.SaveAt(ts=ts)

    sol = dx.diffeqsolve(
        terms,
        # diffusion_term,
        # drift_term,
        solver,
        0,
        1,
        1/num_steps,
        y0,
        saveat=saveat,
        max_steps=num_steps + 1,
    )

    return sol.ys

def loss(drift, num_steps, gamma, dim, key):
    path = sample(drift, num_steps, gamma, dim, key)
    final = path[-1]
    loss = jnp.sum(final**2)
    return loss

@eqx.filter_value_and_grad
def loss_mean(drift, num_steps, gamma, dim, key, batch_size):
    loss_vmapped = jax.vmap(loss, (None, None, None, None, 0), 0)
    key = jrandom.split(key, batch_size)
    return jnp.mean(loss_vmapped(drift, num_steps, gamma, dim, key))


if __name__=="__main__":

    key = jrandom.PRNGKey(0)

    init_drift_key, init_diff_key, train_key = jrandom.split(key, 3)

    dim = 500


    drift = eqx.nn.MLP(dim + 1, dim, 300, 2, key=init_drift_key)
    diff = eqx.nn.MLP(dim + 1, 1, 300, 2, key=init_diff_key)

    model = (drift, diff)

    optimizer = optax.adamw(1e-4)
    opt_state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))
    
    @eqx.filter_jit
    def make_step(model, num_steps, gamma, dim, key, batch_size, opt_state):
        loss, grads = loss_mean(model, num_steps, gamma, dim, key, batch_size)
        updates, opt_state = optimizer.update(
            grads, opt_state, eqx.filter(model, eqx.is_inexact_array)
        )
        model = eqx.apply_updates(model, updates)
        return loss, model, opt_state


    for step in tqdm(range(100)):
        step_key = jrandom.fold_in(train_key, step)
        loss, model, opt_state = make_step(
            model, 20, 0.1, dim, step_key, 32, opt_state
        )

On my machine locally (macbook m1) integrating terms or diffusion_term takes around 41s and integrating drift_term takes around 5s. What is the reason for this difference? Am I doing something wrong here? Note, that the computation in the diffusion term is simply multiplying BM by a scalar. Is VirtualBrownianTree the slow part here? I suspect implementing the euler solver in plain jax would give a faster solution -- would that be wrong? Maybe it's worth adding some documentation about this.

Start with the example's dx.VirtualBrownianTree, dx.MultiTerm, dx.Euler, and dx.diffeqsolve calls, and reproduce the reported timings for the drift, diffusion, and combined terms. Compare the integration and differentiated paths to identify the source of the SDE overhead. Done means documenting the relevant complexity and whether a plain-JAX Euler implementation changes the result.