Using bind within loss functions instead of apply?

[dibyaghosh]

In my code, I have a module that has nested submodules, something like the following:

class Actor(nn.Module):
	def setup(self):
		...
	def __call__(self, x):
		...

	def compute_loss(self, batch):
		...

class Model:
    actor: Actor
    critic: # let's suppose this looks the same as Actor
	

Right now, in my train step, to make sure I can compute the losses on both submodules, I have to jump through some hoops by defining an intermediate function, and calling nn.Module.apply using this function

model_def = Model(...)

@jax.jit
def update(params, batch):
	def loss(params):
		def f(model):
                       return model.actor.compute_loss(batch) + model.critic.compute_loss(batch)
		return model_def.apply(params, method=f)
	grads = jax.grad(loss)(params)
	new_params = ...

It seems that if nn.Module.bind works as advertised, then I can write the code a lot simpler

@jax.jit
def update(params, batch):
	def loss(params):
		model = model_def.bind(params)
                return model.actor.compute_loss(batch) + model.critic.compute_loss(batch)
	grads = jax.grad(loss)(params)
	new_params = ...

but it seems from the documentation that module.bind() is highly discouraged and only for interactive sessions?

Why is bind discouraged within more traditional training workflows? Does this type of design pattern with module.bind() not work in Flax? What are the gotchas / where it fails?

[chiamp]

Could you define the compute_loss within the Model class, like this:

class Actor(nn.Module):
  def setup(self):
    self.dense = nn.Dense(2)
  def __call__(self, x):
    return self.dense(x)
  def compute_loss(self, x):
    return jnp.mean(self(x))

class Critic(nn.Module):
  def setup(self):
    self.dense = nn.Dense(3)
  def __call__(self, x):
    return self.dense(x)
  def compute_loss(self, x):
    return jnp.mean(self(x)**2)

class Model(nn.Module):
  def setup(self):
    self.actor = Actor()
    self.critic = Critic()
  def __call__(self, x):
    return self.critic(self.actor(x))
  def compute_loss(self, x):
    return self.actor.compute_loss(x) + self.critic.compute_loss(x)

m = Model()
batch = jax.random.uniform(jax.random.key(0), (4, 2))
v = m.init(jax.random.key(0), batch)
m.apply(v, batch, method='compute_loss')

Flax encourages a functional programming paradigm to make it compatible with Jax. Using .bind will turn the module into a stateful object and may introduce side effects depending on what you do to it. For example, if you were to use BatchNorm, you would need to remember that the batch_stats are updated every forward pass during training and return them as part of the variable update:

class MLP(nn.Module):
  @nn.compact
  def __call__(self, x, train):
    x = nn.Dense(3)(x)
    x = nn.BatchNorm()(x, use_running_average=not train)
    x = nn.Dense(1)(x)
    return x

x = jax.random.uniform(jax.random.key(0), (4, 2))
y = jnp.ones((4, 1))
module = MLP()
v = module.init(jax.random.key(0), x, train=False)

@jax.jit
def update(v):
  params = v['params']
  def loss(params):
    bound_module = module.bind({'params': params, 'batch_stats': v['batch_stats']}, mutable=['batch_stats'])
    out = bound_module(x, train=True) # bound_module.variables['batch_stats'] is updated
    return jnp.mean((y-out)**2), bound_module.variables['batch_stats']
  grads, batch_stats = jax.grad(loss, has_aux=True)(params)
  new_params = jax.tree_map(lambda p, g: p - 1e-3 * g, params, grads)
  return {'params': new_params, 'batch_stats': batch_stats}

Whereas writing it functionally is a little more explicit and clear that the batch_stats has been updated:

@jax.jit
def update(v):
  params = v['params']
  def loss(params):
    out, new_batch_stats = module.apply({'params': params, 'batch_stats': v['batch_stats']}, x, train=True, mutable=['batch_stats'])
    return jnp.mean((y-out)**2), new_batch_stats
  grads, batch_stats = jax.grad(loss, has_aux=True)(params)
  new_params = jax.tree_map(lambda p, g: p - 1e-3 * g, params, grads)
  return {'params': new_params, 'batch_stats': batch_stats}

I think in the example you provided using .bind is okay, since your update function binds the Module within the scope of the function and only takes in parameters and outputs updated parameters; this is similar to how .apply works where we create a copy of a bounded module and destroy it when the function exits (see the Module Lifecycle for more detail on this). But you wouldn't be able to pass a bounded instance of a Module and output an updated Module into a jitted function.