Multi-Head Attention in Transformers

[oops-shreya]

why should be output size in multi head should be same as input size?

after we concat multiple self attention output then we need to apply linear transformations to the generated output why?
why cant be the size be kept same after concatenation

[Alanturner2]

Great question! Let’s break it down step by step to understand why this design choice is made in the Transformer architecture.

1. Why Keep the Output Size the Same as the Input?

The main reason is residual connections (skip connections). In Transformers, the input of a layer is added directly to its output before moving to the next layer. For this addition to work, the dimensions of the input and output need to match. Keeping the output size the same ensures everything works seamlessly.

2. Why Do We Need a Linear Transformation After Concatenation?

When using Multi-Head Attention, we divide the input into several “heads.” Each head learns different patterns or relationships in the input, and their outputs are concatenated to form a larger tensor. For example:

• Suppose each head outputs a vector of size (d_k).
• If there are (h) heads, the concatenated output size becomes (h \times d_k).

This size is larger than the input, so we use a linear transformation (a simple matrix multiplication with learned weights) to reduce it back to the original size. But there’s more to it than just resizing:

• Mixing Information: The linear transformation helps combine the outputs from all heads into a single, cohesive representation. Without it, the heads would remain independent, and the model wouldn’t fully leverage the information each head captures.
• Compatibility: Reducing the size back to the original ensures that the data can move smoothly through the rest of the Transformer without increasing computational costs unnecessarily.
• Efficiency: Keeping dimensions consistent avoids breaking the architecture and keeps the Transformer modular.

3. Why Not Skip the Linear Transformation?

If we skipped the linear transformation and left the size as (h \times d_k), a few issues would arise:

• Dimension Mismatch: The larger size would break residual connections and make it harder to stack layers.
• Inefficiency: The model would need to handle larger tensors, increasing computation and memory usage without much benefit.
• Less Flexibility: The linear transformation is learnable, so it gives the model extra power to optimize how the information from multiple heads is combined.

TL;DR:

The linear transformation after concatenation:

• Makes sure the output size matches the input size for residual connections.
• Helps mix information from all heads to create a unified representation.
• Keeps the model efficient and compatible with the rest of the architecture.

I hope this helps clear things up! Let me know if you have any follow-up questions.