Why does PyTorch's Transformer model implementation `torch.nn.Transformer` have an additional LayerNorm layer on Encoder/Decoder 's output?

username: Cyber_punk
cooked: <p>In the code of <code>torch.nn.Transformer</code> (<a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#TransformerEncoder" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.modules.transformer — PyTorch 2.1 documentation</a>),  the encoder part is implemented as follows:</p>
<pre><code class="lang-auto">encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,
                                        activation, layer_norm_eps, batch_first, norm_first,
                                        **factory_kwargs)
encoder_norm = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
</code></pre>
<p>From <code>torch.nn.TransformerEncoder</code> (<a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#TransformerEncoder" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.modules.transformer — PyTorch 2.1 documentation</a>) 's code we can see the output is passed through the above LayerNorm layer:</p>
<pre><code class="lang-auto">if self.norm is not None:
    output = self.norm(output)

return output
</code></pre>
<p>My questions is: <strong>why do we need to add another LayerNorm, considering we have LayerNorm layers applied to self-attention and FFNN respectively already?</strong> What’s the benefit? Is it necessary?</p>
<p><code>torch.nn.TransformerEncoderLayer</code> (<a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#TransformerEncoderLayer" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.modules.transformer — PyTorch 2.1 documentation</a>) 's output:</p>
<pre><code class="lang-auto">else:
    x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask, is_causal=is_causal))
    x = self.norm2(x + self._ff_block(x))

return x
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/9/89dc0abc4982e2cb92dd17b487edf238b02e2018.webp" data-download-href="https://discuss.pytorch.org/uploads/default/89dc0abc4982e2cb92dd17b487edf238b02e2018" title="attention_research_1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/89dc0abc4982e2cb92dd17b487edf238b02e2018_2_354x499.webp" alt="attention_research_1" data-base62-sha1="jFyQLIIu4F1giFhWlXmXGagxb0Q" width="354" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/89dc0abc4982e2cb92dd17b487edf238b02e2018_2_354x499.webp, https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/89dc0abc4982e2cb92dd17b487edf238b02e2018_2_531x748.webp 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/9/89dc0abc4982e2cb92dd17b487edf238b02e2018_2_708x998.webp 2x" data-dominant-color="E4E4E1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">attention_research_1</span><span class="informations">1320×1860 47.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: vdw
cooked: <p>Sorry for not having an answer!</p>
<p>However, asking for the purpose of a layer is generally a tricky question. For example, there a whole papers trying to figure out the purpose of individual components of the Transformer (e.g. just for the FF layer: <a href="https://aclanthology.org/2021.emnlp-main.446.pdf" rel="noopener nofollow ugc">Transformer Feed-Forward Layers Are Key-Value Memories</a>).</p>
<p>If you check the original Transformer paper (<a href="https://arxiv.org/pdf/1706.03762.pdf" rel="noopener nofollow ugc">Attention is all you need</a>), it’s mainly a description of the architecture and components, without any real theoretical justifications.</p>