Batch Norm: Further Documentation and Simplified Definition

[shelhamer]

The current state of batch norm is somewhat unclear and requires a tedious specification of learning rates for correctness. This PR documents the batch norm layer in more detail, to clarify its blobs and how to handle the bias and shift, and ensures that the batch norm statistics are not mistakenly mangled by the solver. As the mean, variance, and bias correction are not learnable parameters to be optimized, they should be not be updated by the solver, and this PR enforces this exclusion.

A further PR could optionally include the scale and shift in the batch norm layer (which for now are handled by a separate layer), and this would align with the cuDNN interface, but this PR is helpful in itself to avoid accidents.

[bwilbertz]

@shelhamer Thanks a lot for these clarifications. It took me quite some time, when I started using Caffe's BatchNorm a few weeks ago, to get all these things clear by myself.

Since you are already looking into the BatchNorm implementation, please, please also revert the commit 0ad1d8a, since it makes it impossible to build up an accurate global estimator for the variance.

This commit was switching the calculation of Var(X) from EX^2 - (EX)^2 to E(X-EX)^2. Unfortunately, in each forward pass, the estimator m_b for EX is only based on the values of the current mini-batch (a rather bad estimator compared to the moving mean in blob[0]), which then adds the empirical mean of the non-linear transformation (X - m_b)^2 to the global stats.

In contrast, @cdoersch 's original implementation was storing EX^2 in blob[1], so that the final estimator for Var(X) was computed as EX^2 - m^2, where m is the global estimator for EX, which was computed using all batches (modulo the moving average factor).

This becomes especially an issue, if you want to compute some high accuracy estimators of mean and variance for further finetuning or inference and set moving_average_fraction = 1 (similar to what Kaiming He is doing for his ResNets (c.f. 5. in Disclaimer and known issues).

[bwilbertz]

Isn't blob[2] rather the cumulative moving average factor than a correction for the batch size?

It is updated as b(n+1) = alpha*b(n) + 1 with b(0) = 0.

[shelhamer]

Right, it is. Comment updated accordingly.

[shelhamer]

@bwilbertz

Since you are already looking into the BatchNorm implementation, please, please also revert the commit 0ad1d8a, since it makes it impossible to build up an accurate global estimator for the variance.

This commit was switching the calculation of Var(X) from EX^2 - (EX)^2 to E(X-EX)^2. Unfortunately, in each forward pass, the estimator m_b for EX is only based on the values of the current mini-batch (a rather bad estimator compared to the moving mean in blob[0]), which then adds the empirical mean of the non-linear transformation (X - m_b)^2 to the global stats.

Re: #3299 I'll try to review this once I have time for another round of batch norm reform. Note that it was addressing a numerical instability issue, so the solution isn't as simple as just reverting it.

[bwilbertz]

@shelhamer This kind of numerical issue only occurs if N is extreme large and in the same time the variance is very small.
But for that case, the parameter epsilon was introduced into the batch norm, so that one can easily avoid any numeric catastrophe by increasing epsilon (in case someone is really piping more or less constant data into the net).

Btw: epsilon correction is missing in MVNLayer (where this kind of numerical problem had its origin, see #3162)

[bwilbertz]

@shelhamer I think there is still an issue with the batch norm upgrade:

The operation new Net<Dtype>(net_param)->ToProto(...) is not idempotent any more.

That means if we pass an upgraded net_param into the Net constructor, we receive back a net, which now needs batchNormUpgrade again (because the first round LayerSetUp will add 3 params and the second one will refuse to work unless all params are removed).

Unless this is intended behaviour, I would suggest to maybe change those lines into something like:

  for (int i = 0; i < this->blobs_.size(); ++i) {
    if (this->layer_param_.param_size() == i) {
      ParamSpec* fixed_param_spec = this->layer_param_.add_param();
      fixed_param_spec->set_lr_mult(0.f);
    } else {
      CHECK_EQ(this->layer_param_.param(i).lr_mult(), 0.f)
          << "Cannot configure batch normalization statistics as layer parameters.";
    }
  }

[shelhamer]

@bwilbertz thanks for raising this, I will look into it soon.