nn.TransformerEncoderLayer mismatch on batch size dimension?

username: edng
cooked: <p>In the forward function of nn.TransformerEncoderLayer, the input goes through MultiheadAttention, followed by Dropout, then LayerNorm. According to the documentation, the input-output shape of MultiheadAttention is (S, N, E) → (T, N, E) where S is the source sequence length, L is the target sequence length, N is the batch size, E is the embedding dimension. The input-output shape of LayerNorm is (N, *) → (N, *). Wouldn’t this cause a problem because the batch size of the MultiheadAttention output is the second dimension, while LayerNorm expects the batch size to be in the first dimension? Or am I missing something? Thanks!</p>

username: soulitzer
cooked: <p>From <a href="https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html" class="inline-onebox" rel="noopener nofollow ugc">LayerNorm — PyTorch 2.1 documentation</a>,</p>
<blockquote>
<p>If a single integer is used, it is treated as a singleton list, and this module will normalize over the last dimension which is expected to be of that specific size.</p>
</blockquote>
<p>So, TLDR it doesn’t care if seq len and batch size are permuted as long as the last dim is correct.</p>

username: pascal_notsawo
cooked: <p>In the case of <a href="https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html" rel="noopener nofollow ugc"><code>nn.TransformerEncoderLayer</code></a> (below I do not distinguish between <code>embed_dim</code> and <code>d_model</code> in the case of internal operations):</p>
<ul>
<li>if <code>batch_first=True</code>, then the input of the model will be of shape <code>(batch_size, seq_len, dim)</code> : here <code>N</code> of <code>nn.LayerNorm</code> corresponds to <code>batch_size</code>, <code>*</code> to <code>(seq_len, dim)</code> and <code>normalized_shape</code> to <code>dim</code>
</li>
<li>otherwise, the input of the model will be of shape <code>(seq_len, batch_size, dim)</code> : here <code>N</code> of <code>nn.LayerNorm</code> corresponds to <code>seq_len</code>, <code>*</code> to <code>(batch_size, dim)</code> and <code>normalized_shape</code> to <code>dim</code>
</li>
</ul>
<p>But in all cases, the normalization is done along the last dimension (<code>dim</code>), because according to the documentation of <a href="https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html" rel="noopener nofollow ugc"><code>nn.LayerNorm</code></a>, when a single integer is used for <code>normalized_shape</code>, it is treated as a singleton list, and this module will normalize over the last dimension which is expected to be of that specific size.</p>
<p>There is no need to talk about target length in the case of the encoder, we have only one input of length <code>seq_len</code> / <code>source_seq_len</code>. It is in the case of the <a href="https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoderLayer.html" rel="noopener nofollow ugc">nn.TransformerDecoderLayer</a> that we can even talk about <code>T</code> and <code>S</code> (target length and source length). And in this case an <code>nn.LayerNorm</code> will be applied to <code>(target_seq_len, batch_size, dim)</code> (or <code>(batch_size, target_seq_len, dim)</code> for <code>batch_first=True</code>) and others to <code>(source_seq_len, batch_size, dim)</code> (or <code>(batch_size, source_seq_len, dim)</code> for <code>batch_first=True</code>), always following the last dimension, no matter what value <code>batch_first</code> has.</p>
<p>See :</p>
<ul>
<li>TransformerEncoderLayer : <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/transformer.py#L255" class="inline-onebox" rel="noopener nofollow ugc">pytorch/transformer.py at master · pytorch/pytorch · GitHub</a>
</li>
<li>TransformerDecoderLayer : <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/transformer.py#L330" class="inline-onebox" rel="noopener nofollow ugc">pytorch/transformer.py at master · pytorch/pytorch · GitHub</a>
</li>
</ul>

username: edng
cooked: <p>Thank you both, that clarifies my confusion. For nn.LayerNorm, it doesn’t matter whether the batch size is in the first or second index for a 3D tensor, because it acts on the last dimension.</p>
<p>A related follow-up question to clarify my understanding of nn.MultiheadAttention: for nn.MultiheadAttention, unlike nn.LayerNorm, it does matter whether the batch size is in the first or second index, right? It doesn’t just act on the last dimension?</p>

username: pascal_notsawo
cooked: <p>In fact all the computations at the level of <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" rel="noopener nofollow ugc">nn.MultiheadAttention</a> are done with this shape <code>(seq_len, batch_size, dim)</code>, that is <code> batch_first=False</code>.<br>
Below, <code>seq_len</code> corresponds to <code>source_seq_len</code> in the case of the query, and to <code>target_seq_len</code> in the case of the key and value.</p>
<p>As you can see <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/activation.py#L1016-L1042" rel="noopener nofollow ugc">here</a>, when <code> batch_first=True</code>, i.e. the entries (<code>query</code>, <code>key</code>, <code>value</code>) are of the shape <code>(batch_size, seq_len, dim)</code>, they are first transposed into <code>(seq_len, batch_size, dim)</code> before the computations, and then the result (<code>attn_output</code>), of shape <code>(seq_len, batch_size, dim)</code>, is re-transposed to <code>(batch_size, seq_len, dim)</code> before being returned.</p>

username: edng
cooked: <p>That makes sense, thank you!</p>