GroupNorm(num_groups=1) and LayerNorm are not equivalent?

username: Jiachen
cooked: <p>Can anyone help me understand how GroupNorm(num_groups=1) and LayerNorm can be equivalent?</p>
<p>I tried the following code modified from the original GN document <a href="https://pytorch.org/docs/stable/generated/torch.nn.GroupNorm.html" rel="noopener nofollow ugc">link</a> and I found the two functions are not equivent: (I check the initialization for both functions and they are the same)</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; import torch.nn as nn
&gt;&gt;&gt; 
&gt;&gt;&gt; input = torch.randn(20, 6, 10, 10)
&gt;&gt;&gt; g = nn.GroupNorm(1, 6)
&gt;&gt;&gt; l = nn.LayerNorm(6)
&gt;&gt;&gt; print((g(input) - l(input.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)).norm())
tensor(55.1597, grad_fn=&lt;CopyBackwards&gt;)
</code></pre>
<p>Then I tried another code:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; import torch.nn as nn
&gt;&gt;&gt; 
&gt;&gt;&gt; input = torch.randn(20, 6, 10, 10)
&gt;&gt;&gt; g = nn.GroupNorm(1, 6)
&gt;&gt;&gt; l = nn.LayerNorm((6, 10, 10))
&gt;&gt;&gt; print((g(input) - l(input)).norm())
tensor(0., grad_fn=&lt;CopyBackwards&gt;)
&gt;&gt;&gt; print(g.weight.shape)
torch.Size([6])
&gt;&gt;&gt; print(l.weight.shape)
torch.Size([6, 10, 10])
</code></pre>
<p>Although it gives the same result, you can see that the dimension of the parameters of the two functions are not the same at all.</p>
<p>So I am wondering, in what sense, the GroupNorm(num_groups=1) and LayerNorm are equivalent? Do I miss something here?</p>

username: Atcold
cooked: <p>Hmm… Let’s do some checks.</p>
<pre><code class="lang-auto">import torch
from torch import nn

x = torch.randn(20, 6, 10, 10)
g = nn.GroupNorm(1, 6, affine=False)
l = nn.LayerNorm((6, 10, 10), elementwise_affine=False)

y_g = g(x)
y_l = l(x)
(y_g - y_l).pow(2).sum().sqrt()
#  tensor(5.1145e-06)
</code></pre>
<p>Now, when you type <code>nn.LayerNorm(6)</code>, you’re instructing <code>torch</code> to compute the normalisation over a single dimension, <em>i.e.</em> the last one. So, when you feed the permutated input to your <code>LayerNorm</code> module, it will compute a normalisation <strong>only</strong> over the 6 channels, per every location of the map. No wonder you got a difference of <code>55.1597</code> in your first snippet.</p>
<p>Finally, <code>GroupNorm</code> uses a (global) channel-wise learnable scale and bias, while <code>LayerNorm</code> has a (local) scale and bias for each location as well. Unless you share them across all locations for <code>LayerNorm</code>, <code>LayerNorm</code> will be more flexible than <code>GroupNorm</code> using a single group. You can see how their CPP implementation differs below.</p>
<p><a href="https://github.com/pytorch/pytorch/blob/bceb1db885cafa87fe8d037d8f22ae9649a1bba0/aten/src/ATen/native/cpu/group_norm_kernel.cpp#L67-L69" rel="noopener nofollow ugc"><code>group_norm_kernel.cpp</code></a></p>
<pre><code class="lang-auto">// global scale and bias

for (const auto k : c10::irange(HxW)) {
  Y_ptr[k] = scale * X_ptr[k] + bias;
}
</code></pre>
<p><a href="https://github.com/pytorch/pytorch/blob/52ccbf449495640e043289eae8521e987a6449b1/aten/src/ATen/native/cpu/layer_norm_kernel.cpp#L60-L68" rel="noopener nofollow ugc"><code>layer_norm_kernel.cpp</code></a></p>
<pre><code class="lang-auto">// per location scale and bias
vec::map3&lt;T&gt;(
  [scale, bias](Vec x, Vec gamma, Vec beta) {
    return (x * Vec(scale) + Vec(bias)) * gamma + beta;
  },
  Y_ptr,
  X_ptr,
  gamma_data,
  beta_data,
  N
);
</code></pre>
<p>Where <code>map</code> tells you that (<a href="https://en.wikipedia.org/wiki/Map_(parallel_pattern)" rel="noopener nofollow ugc">citing Wikipedia</a>):</p>
<blockquote>
<p>a simple operation is applied to all elements of a sequence, potentially in parallel <a href="https://en.wikipedia.org/wiki/Map_(parallel_pattern)#cite_note-1" rel="noopener nofollow ugc">[1]</a>. It is used to solve <a href="https://en.wikipedia.org/wiki/Embarrassingly_parallel" rel="noopener nofollow ugc">embarrassingly parallel</a> problems: those problems that can be decomposed into independent subtasks, requiring no communication/synchronization between the subtasks</p>
</blockquote>
<p>Answering your question, <code>GroupNorm(num_groups=1)</code> and <code>LayerNorm</code> are <strong>not</strong> equivalent, unless followed by a fully-connected layer.</p>