Pytorch doc about LayerNormalization is confusing

username: gdp
cooked: <p>According to my understanding, layer normalization is to normalize across the features (elements) of one example, so all the elements in that example should</p>
<p>(1)  use the same mean and variance computed over the example’s elements themselves.<br>
(2)  scale and bias via the same parameter gamma and beta</p>
<p>i.e. different elements  in one example should use the same normalization parameters, not different parameters for each element. (not per-element)</p>
<p>However, in the officail doc of nn.LayerNorm , a note says</p>
<p>"Unlike Batch Normalization and Instance Normalization, which applies scalar scale and bias for each entire channel/plane with the  <code>affine</code>  option, Layer Normalization applies per-element scale and bias with  <code>elementwise_affine</code> "</p>
<p>I think this is confusing… I wonder if there is a mistake in the doc or the implementation is truly as the note says, or if my understanding of layer normalization is wrong. (paper is at <a href="https://arxiv.org/abs/1607.06450" rel="nofollow noopener">https://arxiv.org/abs/1607.06450</a>)</p>
<p>Can anyone give me some suggestions ?</p>

username: SimonW
cooked: <p>PyTorch LN is correct. In Sec. 5.1 of the original paper,</p>
<blockquote>
<p>. They also learn an adaptive bias b and gain g for each neuron after the normalization.</p>
</blockquote>

username: gdp
cooked: <p>Thanks for your correction. So is the following understanding correct?<br>
Assuming the input minibatch is of size (N, C, H, W). For each example, the elements in all channels (CHW elements) are used to compute the mean and variance to normalize that example, and  there are NCHW scale and bias parameters for each element in each channel of each example.</p>

username: SimonW
cooked: <p>Mostly correct! But there are only CHW scale and bias parameters because N is the batch dimension.</p>
<p>This is all based on the assumption that you normalize over the last 3 dimensions. Actually in channeled data, it is unclear whether you should normalize over the channel dimension or not.</p>

username: gdp
cooked: <p>So if I do not normalize over channel dimension,  I would:<br>
(1) normalize over HW elements for each channel of each example in a minibatch (i.e. compute NC mean and variance values)<br>
(2) use scale and bias parameters for each element in the same channel of every examples in the minibatch (i.e. CHW scale and bias parameters)<br>
Am I correct ?</p>

username: SimonW
cooked: <aside class="quote no-group" data-username="gdp" data-post="5" data-topic="26977">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/a5b964/48.png" class="avatar"> gdp:</div>
<blockquote>
<p>(2) use scale and bias parameters for each element in the same channel of every examples in the minibatch (i.e. CHW scale and bias parameters)</p>
</blockquote>
</aside>
<p>No, you would have HW scale and bias parameters. Basically the number of such parameters is equal to the number of elements you normalized over (in each computation of mean and variance).</p>

username: gdp
cooked: <p>If I use HW scale and bias parameters, the same location in different channels will share the same scale and bias parameters, is it right ? while in a BatchNorm, different channels use different scale and bias parameters.</p>

username: SimonW
cooked: <p>Your understanding is correct.</p>

username: gdp
cooked: <p>Oh thanks a lot ! But the LayerNorm paper did not specify its usage in CNN, is the above usage a convention?</p>

username: SimonW
cooked: <p>It’s unclear how it should be applied on a CNN, so we provide a lot of flexibility and you can normalize over either C or not. But in either case, I think the paper is pretty clear in that you have a scale and bias for each normalized element.</p>

username: gdp
cooked: <p>OK, got it. Thanks !</p>

username: ado_sar
cooked: <p>Still I can’t understand from the docs (emphasis mine):</p>
<blockquote>
<p>Unlike Batch Normalization and Instance Normalization, which applies <strong>scalar scale and bias for each entire channel</strong>/plane with the <code>affine</code> option, Layer Normalization applies <strong>per-element scale and bias</strong> with <code>elementwise_affine</code></p>
</blockquote>
<p>In Batch Norm an adaptive bias <code>b</code> and gain <code>g</code> is also learned for each neuron after the normalization. In other words,  Batch Norm takes an input vector <code>x</code>, multiplies it element-wise with <code>g</code> and then adds the bias vector <code>b</code>. Doesn’t layer norm do the same thing with regards to the affine part?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="ado_sar" data-post="12" data-topic="26977">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ado_sar/48/59241_2.png" class="avatar"> ado_sar:</div>
<blockquote>
<p>Batch Norm takes an input vector <code>x</code>, multiplies it element-wise with <code>g</code> and then adds the bias vector <code>b</code>.</p>
</blockquote>
</aside>
<p>After broadcasting the <code>.weight</code> and <code>.bias</code> as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(1, 3, 224, 224)
bn = nn.BatchNorm2d(3)

print(bn.weight)
# Parameter containing:
# tensor([1., 1., 1.], requires_grad=True)
print(bn.bias)
# Parameter containing:
# tensor([0., 0., 0.], requires_grad=True)
ref = bn(x)

# manual approach
out = (x - x.mean([0, 2, 3], keepdim=True)) / torch.sqrt(x.var([0, 2, 3], unbiased=False, keepdim=True) + bn.eps)
out = out * bn.weight[None, :, None, None] + bn.bias[None, :, None, None]

print((out - ref).abs().max())
# tensor(4.7684e-07, grad_fn=&lt;MaxBackward1&gt;)
</code></pre>
<p>since both trainable parameters have <code>num_channels</code> values.</p>

username: ado_sar
cooked: <p>Apologize if misunderstood.</p>
<blockquote>
<p>Batch Norm takes an input vector <code>x</code>, multiplies it element-wise with <code>g</code> and then adds the bias vector <code>b</code>.</p>
</blockquote>
<p>I am referring to the original papers of BatchNorm and LayerNorm where they consider simple feed-forward networks (MLPs). The only difference is how the mean and the variance are calculated. That is, assuming a batched input <code>X</code> of shape <code>(N, n_feats)</code>, then BatchNorm calculates the mean (same for variance) as <code>X.mean(dim=0)</code> while LayerNorm calculates it as <code>X.mean(dim=1)</code>.</p>
<p>Other than that, <strong>they both learn <code>n_feats</code> gains and biases.</strong> Isn’t this correct?</p>