How to determine whether a Linen Module is still initializing?

[marcvanzee]

We currently check whether a module is initialized using self.has_variable, for instance in our definition of MultiHeadDotProductAttention:

...
# detect if we're initializing by absence of existing cache data.
is_initialized = self.has_variable('cache', 'cached_key')
...
if is_initialized:
  ...

I think this looks a bit strange to the average user. In the old API we were using self.is_initialized, and I think we should consider bringing this back in some form or another.

Below some remarks made by different users offline:

From @levskaya : "The reason we did it this way is that it's strictly more general if you imagine having separate variable collections that might have a more complicated initialization than just a single-shot "init" function. That said, everyone seems to hate it, so I'm not sure the generality is worth not having an is_initializing module-global state variable..."

From @jheek : "You need to be careful with the definition of initializing. The simplest one is we called Module.init and not Module.apply. But there are other cases like if I call the same Module in init for the second time it's already fully initialized so should self.is_initializing return True or False? Something like this:

def __call__(self, x):
  fc = Dense(5)
  x = fc(x) # this will init kernel, bias...
  x = fc(x) # already initialized. 

So to summarize: it seems useful to be able to have a boolean in Module specifying whether a module is initialized or not, but there are some tricky cases we should think about (see example above).

[jheek]

I'm not sure if cache is the best example to frame this discussion around. Ideally cache doesn't exist up until decoding stage. Right now we use a slightly odd workaround because we cannot init carry variables inside a loop (but we could just fix that).

[cgarciae]

In #2052 I had to use something like:

is_initializing = self.scope.is_collection_empty("batch_stats")

the module itself didn't have variables, initialization state depended on its sub-modules. This solution was very tricky (had to look for option in the Scope API), wondering if a method that checks if the module contain any variables at all would be easier to use and applicable in most common cases:

def variables_empty(self) -> bool:
  return not self.scope._variables

Alternatively is_collection_empty could be added exposed as a method in Module.

[marksandler2]

Not having a standard is_initializing behavior also leads to weird side effects when providing mutable=True as much of the libray code uses is_mutable_collection('params') to check whether it is initializing.

I've just been bitten badly by an API, where if i pass mutable=True to a model containing BatchNorm, then it won't update its batch norm statistics nor use global batch norm statistics, because, it thinks it is initializing. This led to a hard to detect discrepancies in our evaluation data that we spent few days chasing.

Checking whether "params" is mutable for initialization behavior seems like a source of endless hard to detect bugs. Not to mention that we might have a good reason to make params mutable (for partial reinitialization for isntance) without expecting library code to change their behavior.

For instance

foo = ResNet50(axis_name='batch')
weights = foo.init(...)
l1 =  foo.apply(weights, mutable=True)
l2 = foo.apply(weights, mutable=['batch_stats'])

l1 and l2 will produce slightly different results in distributed training scenario.

[cgarciae]

@marksandler2 I was looking into this, BatchNorm defines initialization like this:

initializing = self.is_mutable_collection('params')

Which results in batch statistics not being updated if params are mutable because of:

if not initializing:
  ra_mean.value = self.momentum * ra_mean.value + (self.momentum) * mean
  ra_var.value = self.momentum * ra_var.value + (1 - self.momentum) * var

@jheek from an outside perspective it's not clear why making params mutable results in BatchNorm not training properly. If we want to avoid this issue maybe additional context information about whether we are inside init or not is necessary.

[marksandler2]

Yes, it is obvious why it is not doing things properly. It is just a very sharp edge in flax (and not necessarily batchnorm specific). Re batch norm specifically this could be fixed by checking if batch_stats 'mean' exists instead. But overall i think it should be fixed at the API level. Potentially i think the problem is that today `mutable` conflates two different operations: "init" and "update", it would be nice if those were separate then we can make variable updateable, but guarantee that we won't by accident run "init" on it, and vice-versa.

…

On Wed, May 11, 2022 at 12:39 PM Cristian Garcia ***@***.***> wrote: @marksandler2 <https://github.com/marksandler2> I was looking into this, BatchNorm defines initialization like this: initializing = self.is_mutable_collection('params') Which results in batch statistics not being updated because of: if not initializing: ra_mean.value = self.momentum * ra_mean.value + (self.momentum) * mean ra_var.value = self.momentum * ra_var.value + (1 - self.momentum) * var @jheek <https://github.com/jheek> from the outside it's not clear why making params mutable results for BatchNorm not training properly. If we want to avoid this issue maybe additional context information about whether we are inside init or not is necessary. — Reply to this email directly, view it on GitHub <#652 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABVIDWGR6NA773BIK3P4WT3VJQEGRANCNFSM4T3IU6MQ> . You are receiving this because you were mentioned.Message ID: ***@***.***>