The concept of module arguments

[manifest]

I'm trying to wrap my head around the concept of module arguments in Flax. I've read the design notes and got a few questions about Dropout layer and its deterministic argument.

1. What benefits does deterministic parameter have as a method argument comparing to it being a dataclass attribute, if it can only be used as a static argument of a jitted function. Is it our intention to trigger recompilation for the method arguments?

   from flax import linen as nn
   from flax.core import freeze
   from jax import numpy as jnp
   import jax
   
   model = nn.Dropout(rate=0.1)
   jax.jit(model.apply)(
     freeze({}),
     inputs=jnp.zeros((1,1)),
     deterministic=True,
   )

   in __call__
       if deterministic:
   jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
   The problem arose with the `bool` function. 
   While tracing the function apply at ../utils/python/env/lib/python3.8/site-packages/flax/linen/module.py:875 for jit, this concrete value was not available in Python because it depends on the value of the argument 'kwargs'.
   

2. Shouldn't we allow dynamic execution and avoid recompilation by always using jax.lax.cond instead of normal Python conditionals in layers accepting conditional method arguments (such as deterministic parameter in the Dropout layer)?

[andsteing]

You would usually use deterministic=False for training and deterministic=True for evaluation. And since you compile the entire training step and evaluation step, you'll get a recompile anyways.

Also note that if we used jax.lax.cond for implementation of the actual dropout, we would still have the codepaths for generating rngs etc (which would in turn require passing a dropout rng for evaluation etc) in Dropout.

[manifest]

Thank you for your reply. My initial intuition was that using the same jitted function for train and eval tasks, during model training, may speed up the training by avoiding recompilation. I now see that the resulting jaxpr is perhaps too small to cause any problem, and a better approach would be to use two separately jitted functions instead.

As for my 1st question, I still think that I don't fully comprehend when it's preferable to use method arguments instead of dataclass attributes .

You would usually use deterministic=False for training and deterministic=True for evaluation. And since you compile the entire training step and evaluation step, you'll get a recompile anyways.

During training, we would need to call the function with deterministic=True for training steps and deterministic=False for evaluation steps. We would trigger recompilation for each function call with the code like this:

train_fn = eval_fn = jax.jit(nn.Dropout(rate=0.1).apply, static_argnames='deterministic')

## ...in a training loop.
inputs=jnp.zeros((1,1))
rngkey_p, rngkey_d = jax.random.split(jax.random.PRNGKey(seed=0))
train_fn({}, inputs=inputs, deterministic=False, rngs=dict(params=rngkey_p,dropout=rngkey_d))
eval_fn({}, inputs=inputs, deterministic=True)

If I understand it correctly, we can simply avoid recompilation by changing the code a little.

train_fn = jax.jit(partial(nn.Dropout(rate=0.1).apply, deterministic=False))
eval_fn = jax.jit(partial(nn.Dropout(rate=0.1).apply, deterministic=True))

## ...in a training loop
inputs=jnp.zeros((1,1))
rngkey_p, rngkey_d = jax.random.split(jax.random.PRNGKey(seed=0))
train_fn({}, inputs=inputs, rngs=dict(params=rngkey_p,dropout=rngkey_d))
eval_fn({}, inputs=inputs)

The above code can also be written using dataclass attributes:

train_fn = jax.jit(nn.Dropout(rate=0.1, deterministic=False).apply)
eval_fn = jax.jit(nn.Dropout(rate=0.1, deterministic=True).apply)

It seems that the last example – the most concise. There also won't be any issues with parameter initialization during init function call, because deterministic can never change the initialization (the same logic seems to valid for any argument that may be seen as an argument for the apply method).

What would be a use case for deterministic being a method argument then?

[andsteing]

Note that if you have a jax.jit-decorated function inside another jax.jit-compiled function (e.g. apply_fn() inside a train_step()) then the function will be compiled only at the outermost level and the inner decorator does not have an effect (BTW the actual JAX tracing is quite fast, it's the XLA compilation that takes a lot of time).

As for whether deterministic should be a dataclass attribute or a function argument, we have some thoughts on this:

https://flax.readthedocs.io/en/latest/design_notes/arguments.html

It's obvious that we do not need to have deterministic as a dataclass attribute because the model initialization is the same. Now we have some advantage with the possibility of deterministic being a call argument, for example using the same apply_fn in train_state.TrainState that can then be passed to both evaluate() and train() functions.

[manifest]

It's obvious that we do not need to have deterministic as a dataclass attribute because the model initialization is the same. Now we have some advantage with the possibility of deterministic being a call argument, for example using the same apply_fn in train_state.TrainState that can then be passed to both evaluate() and train() functions.

I may be wrong, but I just see a single purpose for model dataclass to create init or apply function. With this high-order view, there isn't much difference between dataclass attributes and method arguments, meaning that T(deterministic=True).apply and partial(T().apply, deterministic=True) behave in the same way.

For example, the training in the SST example (TrainState.{create,apply})

state = TrainState.create(apply_fn=T().apply, tx=optimizer, ...)

## ..in the training loop.
state.apply_fn(..., deterministic=False, ...) ## train step
state.apply_fn(..., deterministic=False, ...) ## eval step

May also be written as

train_state = TrainState.create(apply_fn=T(deterministic=False).apply, tx=optimizer, ...)
eval_state = TrainState.create(apply_fn=T(deterministic=True).apply, tx=optimizer, ...)

## ..in the training loop.
train_state.apply_fn(...)
eval_state.apply_fn(...)

Or with some other implementation of the TrainState

state = TrainState.create(
  train_fn=T(deterministic=False).apply,
  eval_fn=T(deterministic=True).apply,
  tx=optimizer, ...)

## ..in the training loop.
state.train_fn(...)
state.eval_fn(...)

I now see only benefits of using data attributes because it provides just one obvious way to create init or apply function. If so, there would be no need in nn.merge_param, etc.

Can you please elaborate a bit more on the advantages of method arguments? Is it just a way how models supposed to be trained with TrainState or anything else I don't realize and need to be careful about?

[andsteing]

You're right even if deterministic was required for parameter initialization you could pass it in as an argument. When I said "model initialization" in my previous reply it's not very clear what it means actually - if it meant parameter initialization then we could use arguments if we wanted to, if it's only about setting up the dataclass fields then it really doesn't do so much...

While I think using a single state and specifying the deterministic argument differently in eval_step() and train_step() is the more elegant solution, it would also be possible to use two separate train states as you propose in your last reply. If you prefer that latter solution feel free to fork the class and modify accordingly.

But allowing deterministic to be an argument as well allows us to use the pattern sharing the state that we currently use.

[manifest]

Thank you for the clarification.