Significant performance difference: diffeqsolve vs. lax.scan - Expected Behavior?

Hi,

I appreciate all the work that has gone into Diffrax! I'm trying to use diffrax to simulate the interaction between an NN policy and an ode system model for a given horizon length. Specifically, the policy takes the system state as input at each step and generate the control action that is to be applied to the system. Then the ode is solved for obtaining the evolution of the state. For that I have two implementations. The first one is by using lax.scan, and in each iteration the ode is solved for one step using solver.init() and solver.step(). In the second implementation, I use Diffrax.diffeqsolve() which takes as arguments the final simulation time and saveAt(ts = jnp.linespace(t0,t1, num_sim_steps) in addition to term, solver, ..etc. For both implementations I use 'Diffrax.Euler' as solver and I augment the ODE with the policy network inside the vector-field. However, I noticed that the simulation with Diffrax.diffeqsolve() is almost 3 times slower than with lax.scan and single steps, and this difference gets even bigger when comparing Diffrax.diffeqsolve() with 'Adaptive steps' and lax.scan with 'fixed single steps'. The idea behind using Diffrax.diffeqsolve() is that I wanted to investigate if reducing the number of steps, by adapting step sizes, would improve the simulation speed even in the presence of the overheads resulted from intermediate solver-related calculations and rejected steps when exceeding tolerances. But what I don't understand is that when using Euler I would expect both implementations to have similar simulation speeds, which is not the case. My second question would be is there a way to utilize adaptive solver to make the simulation faster for this application.

Here is a benchmark example for this comparison in which I replaced the actual system with an arbitrary first-order ODE. I use Python (3.12.7), Jax (0.4.35), equinox (0.11.10), and diffrax (0.5.1). I run the code on CPU 13th Gen Intel i7-1355U.

Imports

import jax
import jax.numpy as jnp
import numpy as np

import equinox as eqx
import optax
import diffrax

import timeit

Define policy network class

class MLP(eqx.Module):
    """ class for a policy for providing control actions.
    """
    layers: list[eqx.nn.Linear]

    def __init__(self, layer_sizes, key):
        self.layers = []
        for fan_in, fan_out in zip(layer_sizes[:-1], layer_sizes[1:]):
            key, subkey = jax.random.split(key)
            self.layers.append(eqx.nn.Linear(fan_in, fan_out, use_bias=True, key=subkey))

    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = jax.nn.leaky_relu(layer(x))
        return self.layers[-1](x)

case 1: lax.scan and single-steps

def ode_step(init_state, ref, policy):
     """Method for simulating the policy-environment interaction and solving the ode for one step
 
     Args:
         init_state (jax.Array): state at the current step (y0)
         ref (jax.Array): reference state at the current step
         policy (MLP): for predicting the control action 
 
     Returns:
         jax.Array: updated state after one step simulation
     """
     args = (ref)
 
     # ode, including policy (MLP) and system: tau*dy/dt + y = u, aumming tau=1
     d_y = lambda t,y,args: (policy(jnp.concatenate([y, args]))-y)                                        
 
     term = diffrax.ODETerm(d_y)
     solver = diffrax.Euler()
     t0 = 0
     t1 = 1e-4
     y0 = init_state
     env_state = solver.init(term, t0, t1, y0, args)
     y, _, _, env_state, _ = solver.step(term, t0, t1, y0, args, env_state, made_jump=False)
     return y

def rollout_traj_scan(policy, init_states, ref_states, horizon_length):
     """rollout policy-environment interaction for 'horizon_length' for single sample using lax.scan
 
     Args:
         policy (MLP): predicting control actions
         init_states (jax.Array): initial environment state before rollout
         ref_states (jax.Array): reference to be tracked
         horizon_length (int): length of future predictions
 
     Returns:
         jax.Array: MSE tracking loss
     """
     # extending ref_states to horizon length
     ref_o = jnp.repeat(ref_states[None, :], horizon_length, axis=0)                                      
     
     def body_fun(carry, ref):
         state = carry
         state = ode_step(carry, ref, policy)
         return (state), (state)
 
     _, (states) = jax.lax.scan(body_fun, (init_states), ref_o, horizon_length)
 
     # error between the simulation state and reference
     error= states-ref_states                                                                           
 
     loss=jnp.mean((error)**2)                                                                               
     return jnp.clip(loss, max=1e5)

Case 2: Diffrax.diffeqsolve

def ode_diffeqsolve(policy, init_state, ref_state,  horizon_length):
      """Method for rollout, simulating the policy-environment interaction, and solving the ode for 'horizon_length' usin diffrax.diffeqsolve.
  
      Args:
          policy (MLP): predicting control actions
          init_state (jax.Array): initial environment state before rollout
          ref (jax.Array): reference to be tracked
          horizon_length (int): length of future predictions
  
      Returns:
          jax.Array: state trajectory at predefined time steps
      """
      args = (ref_state)
  
      # ode, including policy (MLP) and system: tau*dy/dt + y = u, assuming tau=1
      d_y = lambda t,y,args: (policy(jnp.concatenate([y, args]))-y)                                        
                                                                                                           
      return diffrax.diffeqsolve(
          terms = diffrax.ODETerm(d_y), 
          solver = diffrax.Euler(), 
          t0 = 0, 
          t1 = horizon_length*1e-4, 
          dt0=1e-4, 
          y0=init_state, 
          args=args, 
          saveat=diffrax.SaveAt(ts=jnp.linspace(1e-4, horizon_length*1e-4, horizon_length)), 
          stepsize_controller=diffrax.ConstantStepSize()
      ).ys 

#Roll out using diffrax.diffeqsolve until t1
def rollout_traj_diffeqsolve(policy, init_states, ref_states, horizon_length):
      """Calls 'ode_diffeqsolve' for the rollout and calculate MSE tracking loss
  
      Args:
          policy (MLP): for predicting control actions
          init_states (jax.Array): initial environment state before rollout
          ref_states (jax.Array): reference to be tracked
          horizon_length (int): length of future predictions
  
      Returns:
          jax.Array: MSE tracking loss
      """
  
      # get state trajectory through horizon_length
      states = ode_diffeqsolve(policy, init_states, ref_states, horizon_length)                          
  
      # error between the simulation state and reference
      error= states-ref_states                                                                           
  
      loss=jnp.mean((error)**2)
      return jnp.clip(loss, max=1e5)
      

Setup for the benchmark

jax_key = jax.random.PRNGKey(np.random.randint(0, 2**31))
jax_key, policy_key = jax.random.split(jax_key)

policy=MLP([2,20,20,20,1],key=jax_key)  # initialize policy network

state_key, ref_key = jax.random.split(jax_key)
init_states= jax.random.uniform(state_key, minval=0.0, maxval=30.0, shape=(1,))  # generate random initial state
ref_states =  jax.random.uniform(ref_key, minval=0.0, maxval=30.0, shape=(1,))  # generate random reference state

horizon_length = 25  # rollout length
train_steps = 5000  # number of iterations to be measured

Function for testing the speed of fwd and bwd propagations of the system

def speedtest(fcn, name):
    
    fwd = eqx.filter_jit(fcn)
    bwd = eqx.filter_jit(eqx.filter_grad(fcn))

    #Measure fwd time for train_steps iterations
    fwd_times = timeit.repeat(
        lambda: jax.block_until_ready(fwd(
        policy,
        init_states,
        ref_states,
        horizon_length,
    )), number=train_steps, repeat=10
    )                                                                                                      
    print(f"{name} fwd: {min(fwd_times)}")

    #Measure fwd+Bwd time for train_steps iterations
    bwd_times = timeit.repeat(
        lambda: jax.block_until_ready(bwd(
        policy,
        init_states,
        ref_states,
        horizon_length,
    )), number=train_steps, repeat=10
    )                                                                                                      
    print(f"{name} fwd+bwd: {min(bwd_times)}")

Run tests

speedtest(rollout_traj_scan, "scan")
speedtest(rollout_traj_diffeqsolve, "diffeqsolve")

Start by reproducing the benchmark comparing lax.scan with solver.init()/solver.step() against diffrax.diffeqsolve using Euler and the supplied speedtest. Read the diffrax.diffeqsolve and solver-step entry points, then determine whether the observed overhead is expected and whether adaptive stepping can improve this policy-and-ODE workload; done means documenting the cause and guidance for this use case.