Nonlinearity in TransformerEncoderLayer

username: Yoo
cooked: <p>Is nn.MultiheadAttention in pytorch just a linear transformation layer, while nn.TransformerEncoderLayer the combination of nn.MultiheadAttention and a feed forward layer?</p>
<p>In addition, here is an example of nn.TransformerEncoderLayer:</p>
<pre><code class="lang-auto">TransformerEncoderLayer(
          (self_attn): MultiheadAttention(
            (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)
          )
          (linear1): Linear(in_features=128, out_features=512, bias=True)
          (dropout): Dropout(p=0.1, inplace=False)
          (linear2): Linear(in_features=512, out_features=128, bias=True)
          (norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
          (norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
          (dropout1): Dropout(p=0.1, inplace=False)
          (dropout2): Dropout(p=0.1, inplace=False)
</code></pre>
<p>My question is, does the nonlinearity only come from LayerNorm and dropout layers? I did not see any activation function in this block.</p>

username: AlphaBetaGamma96
cooked: <p>Isn’t there a softmax function within the MultiHeadAttention <code>nn.Module</code>? That’s a non-linear function.</p>

username: Yoo
cooked: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="2" data-topic="165116">
<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/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>MultiHeadAttention <code>nn.Module</code></p>
</blockquote>
</aside>
<p>Hi thanks for your reply but I could not find any activation functions defined within <a href="https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" class="inline-onebox" rel="noopener nofollow ugc">MultiheadAttention — PyTorch 2.1 documentation</a></p>

username: AlphaBetaGamma96
cooked: <p>After a brief look through the source code, <code>MultiHeadAttention</code> calls <code>multi_head_attention_forward</code> which contains a <code>softmax</code> call as seen in the source <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L5180" rel="noopener nofollow ugc">here</a>. So, it does seem to have a non-linear function of some kind. <a class="mention" href="/u/ptrblck">@ptrblck</a> (apologizes for the tag) will be able to confirm any more specific details! But the source does show a non-linear function being used.</p>

username: Yoo
cooked: <p>Thanks again. I checked the source too and found there is a linear layer in MultiheadAttention function that maps derived attention output in embed_dim to embed_dim at the end:</p>
<blockquote>
<p><a href="https://github.com/pytorch/pytorch/blob/0a274c4b6c916363ce3e3f75b315ac66156f8ce6/torch/nn/functional.py#L5187" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/0a274c4b6c916363ce3e3f75b315ac66156f8ce6/torch/nn/functional.py#L5187</a></p>
</blockquote>
<p>And in MultiheadAttention module which calls the function, self.out_proj is defined as the following:</p>
<blockquote>
<p>self.out_proj = NonDynamicallyQuantizableLinear(embed_dim, embed_dim, bias=bias, **factory_kwargs)</p>
</blockquote>
<p>So I guess there is an extra linear transformation before pytorch outputs the attention than what I originally expected.<br>
<a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> May you have a look, thanks!</p>