[BUG] Data sets have a dangerous stage argument

[stefanradev93]

Describe the bug
Our data sets currently accept a stage argument which is training by default. Having validation and a stateful standardize transform will results in an error due to:

if self.adapter is not None: batch = self.adapter(batch, stage=self.stage)

since the running means and standard deviations have never been computed.

Expected behavior
I have come to the realization that BatchNorm layers should be part of the approximators and applied to all inference_conditions, summary_variables, and inference_variables. This will have the advantage that adapters will remain stateless and users will not have to deal with standardizing things explicitly. Still, we should keep the standardize transform with static means and stds for special cases.

Let me know what you think and I will provide an implementation.

[paul-buerkner]

Yes, I agree. That would make the adapter much cleaner I think.

[stefanradev93]

The question is only how to handle the ldj of the transformation.

[vpratz]

I think "standard" batch norm is not well suited for this, as during training it uses the current mean and std of a batch instead of the moving average, which is not what we want and will probably hurt performance. We could think about moving the standardization with momentum (as it is currently implemented in the adapter) to the approximator and standardize summary_conditions, inference_variables and (possibly) the outputs of the summary network.

A few other thoughts:

• we would usually not want to use different adapters for training set and validation set, but use the same adapter. Using different adapters will lead to different standardization, so different inputs to the network. The issue you describe then only arises if we use it on the validation set before the training, which might or might not make sense, depending on the context.
• we could in addition modify the behavior of the Standardize transform, so that it uses mean and variance of the batch in cases where it has momentum and was not called with stage="training" yet (maybe with a warning)

I see the problem though, and resolving it would be great. Another avenue we might explore is moving the adapter out of the datasets, e.g. by wrapping the datasets in the approximator. If I recall it correctly, @LarsKue has reflected on the general problem and possible difficulties with it, but I don't remember exactly which difficulties led to putting the adapter into the datsets in the first place. @LarsKue, could you comment on this?

[LarsKue]

If the adapter is in the dataset, calling it is "free" in terms of computational cost, since it will be done by a worker process. That is the primary reason why it is there.

[vpratz]

Thanks for the comment! Do you know whether the computational cost of the adapter is practically relevant here? I would assume the adapter to be cheap compared to the summary/inference networks in most cases. Did you do any benchmarking regarding this?

[stefanradev93]

@vpratz Batch norm was wrong terminology. Take a look at my initial changes in the ContinuousApproximator: https://github.com/bayesflow-org/bayesflow/tree/standardize_in_approx.

The idea is basically to have non-trainable running means and stds tracked for any of the relevant quantities.

[vpratz]

Closed by #486. The argument is still there (we might think about deprecating it going forward, but I'm not sure yet), but less dangerous as dynamic standardization is deprecated, and other transforms do not rely on being called with stage="training" first.