Standardization: count is shared and gets updated multiple times, leading to wrong results

[vpratz]

Standardization._update_moments is called for each flattened input. While each input has its own mean and m2, the count is shared, and therefore updated multiple times, leading to an increase by n_inputs*batch_size instead of only batch_size. This leads to wrong results.

Possible fixes are:

1. Update count only once in the end of the updates
2. Have one count object for each input

The first is a bit more hacky, but does not break serialization. The latter would be in line with how we handle the mean and m2, and would fit better in the existing structure in my opinion, but would break the serialization.

I will provide a fix for 1., but 2. would be easy to implement as well.

[paul-buerkner]

Good catch. Fix for 1 sounds good. we can still implement 2 if we need to break serialization in standardization for some other reasons anyway

[vpratz]

Thanks for taking a look. I discovered a second bug, which unfortunately requires implementing option 2. Please refer to #525 for details. As this is a core component, and we discourage using the standardize transform already, maybe this fix would warrant a bug-fix release. The bugs cause:

• wrong standard deviations if the standardization layer has one input
• wrong means and standard deviations when the standardization layer has multiple inputs

I encountered the bug because I saw weird performance issues in an applied setting...

@paul-buerkner @stefanradev93 @LarsKue What do you think?

[paul-buerkner]

Glad you spotted this. I agree this warrants a bug-fix release and we should get it out asap.

[LarsKue]

@vpratz As far as I can see, this is fixed in #525, right?