Batch Normalization will disrupt multi-threading? Question:


In a Feed Forward network, assume we have a Mini-batch of 64 examples.
Our layer $l$ contains 20 neurons
Because every neuron in layer will require Mean and Variance from the whole minibatch, we must  to run all 64 'examples' in parallel (is this true?). We must wait until each such 'example' will have the appropriate sum (for that particular neuron!) available across all these minibatches. 


Wouldn't this break parallelism? Previously, each of our example-threads could have computed everything, up to the activation values for the entire layer $l$. However, now we have to: 


Wait for all example-threads to finish computing the sums for our layer $l$ (hence, our speed slows down to match the slowest example-thread)
Pause multithreading (!) and compute mean &amp; variance for the layer $l$
Apply it to the sums, resume example-threads, so each thread completes the job, by activating its corresponding adjusted sums


Do I miss a part of the puzzle, or it really has to be done like this?

$$$$
my understanding of Batch Norm:

Batch norm (link to original paper) makes the search space more even &amp; simpler to navigate during forward &amp; backward passes. 

It is usually applied to the sum $Z$ going into activation function (although there is a lot of debate (1) (2)). This is done for every neuron in a layer.

in usual settings, the neuron's output on layer $l$ is $$H_{i} = \phi(W\cdot X +b)$$  where W is a vector of weights leading into our $i$'th unit, $X$ is a vector of outputs from previous layer $l_{-1}$ and $b$ is a bias

With Batch norm, we still have our usual non-activated sum for neuron $i$:

$$Z_i = W\cdot X$$

but also have our mean for that neuron across all minibatches (Just for that neuron!)

$$\mu_i = \frac{1}{m} \sum_{p}^m Z_{ip} $$

Similar, for our Variance as:

$$\sigma^2_i = \frac{1}{m}\sum_{p}^m (Z_{ip} - \mu_i)^2$$

Our non-activated normalized sum for neuron $i$ (here, $\epsilon$ is a tiny number, in case if variance is too small, to avoid division by zero):

$$Z_{norm} = \frac{Z_i - \mu_i}{\sqrt{\sigma^2_i + \epsilon}}$$

We give a network a chance to learn the actual distribution it needs, so it's different from default distrib. We introduce 2 learnable parameters, $\gamma$ and $\beta$, and network can use it to 'squash and shift' the distributions as it needs (even canceling-out the batch-norm if required).

$$Z_{adjusted} = {\gamma}Z_{norm} + \beta $$

Finally neuron's output is computed as:

$$H_i =  \phi (Z_{adjusted}) $$ 

Bias isn't included, because we already have $\beta$ for the whole layer $l$.
My experience with batch and layer normalisation is that it is around 3 to 5 times slower per step than without it. Although the model may converge in less number of steps (which is what the authors report in the article) it takes much longer to calculate each step, and hence the whole training takes longer.

This has been reported in a few places like: 

https://www.reddit.com/r/MachineLearning/comments/4ufmxy/layer_normalization_implemented_in_tensorflow/

https://github.com/cooijmanstim/recurrent-batch-normalization/issues/2

https://github.com/keras-team/keras/issues/1309