[Python] Correct way of providing PRNGKeyArray to Dropout layer using JAX and Equinox

I'm using JAX + Equinox for training a machine learning model. The model contains a Dropout layer (eqx.nn.Dropout) which expects a random key in its __call__ method. My question is about what the correct way is to provide this key during training time.

My model basically looks like this

import equinox as eqx
from jaxtyping import PRNGKeyArray, Array

class Model(eqx.Module):
...
dropout_layer: eqx.nn.Dropout

def __init__(self, ...):
...
self.dropout_layer = eqx.nn.Dropout()
...

def __call__(self, x:Array, *, keyRNGKeyArray) -> Array:
...
x = self.dropout(x, key=key). # <-- MUST pass a key here
...
return x



And here is the basic structure of my training loop


key = jax.random.key(0)
key, subkey = jax.random.split(key)
model = Model()
optimizer = ...
optimizer_state = ...

for x, y in dataset:
model, opt_state, loss = _step_fn(model, optimizer, opt_state, x, y, key=subkey)



where

from functools import partial
import equinox as eqx
import optax

@eqx.filter_jit
def _step_fn(
model: Callable,
optim: optax.GradientTransformation,
optim_state: optax.OptState,
x: Float[Array, ""],
y: Float[Array, ""],
key: PRNGKeyArray,
):
def _loss_with_keys(model: Callable) -> Float[Array, ""]:
return loss_fn(model, x, y, key)

value_and_grad_fn = eqx.filter_value_and_grad(_loss_with_keys)

loss, updates = value_and_grad_fn(model)
updates, optim_state = optim.update(
updates, optim_state, params=eqx.filter(model, eqx.is_inexact_array)
)
model = eqx.apply_updates(model, updates)

return model, optim_state, loss

def _loss_fn(
model: Callable,
x: Float[Array, ""],
y: Float[Array, ""],
key: PRNGKeyArray,
) -> Float[Array, ""]:
pred = eqx.filter_vmap(partial(model, key=key))(x)
loss = jnp.mean(optax.squared_error(pred, y))

return loss


I tried several variations (creating the key in the training loop, the step function, the loss function), basically all of them work somehow, but none of them seems to work well as jax is incredibly slow (about 3-15 times slower than PyTorch, which in my experience is a clear indication that I'm doing something wrong here). It is also not clear to me, why some solutions are slower than others (maybe there is some recompilation happening).

I'm still quite new to Jax, thus may be missing some obvious points. To sum up:

1. What is the correct way to pass key to the dropout layer via the __call__ method?
2. Are there some obvious reasons why my implementation is so slow?

Continue reading...