What is the attention mechanism in transformers?

Theoretical Background:

• Self-Attention: In a transformer model, self-attention is used to compute a representation of a word in the context of a sentence. It does this by creating three vectors for each word: Query (Q), Key (K), and Value (V). The output of the self-attention is a weighted sum of the value vectors, where the weights are determined by the similarity of the query vector to key vectors. This allows the model to dynamically focus on different parts of the input sequence.

• Multi-Head Attention: This involves running multiple self-attention mechanisms (heads) in parallel. Each head has its own set of learned projections for Q, K, and V, allowing the model to attend to different information subspaces. The outputs of these heads are concatenated and linearly transformed to produce the final output.

Practical Applications: Transformers and their attention mechanisms are widely used in NLP tasks. For example, in machine translation, attention enables the model to align words between languages dynamically. In text summarization, it helps focus on the most relevant parts of the text.

Code Example: Here's a conceptual snippet to illustrate multi-head attention using PyTorch:

import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, d_model):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        self.depth = d_model // num_heads
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.dense = nn.Linear(d_model, d_model)

    def split_heads(self, x, batch_size):
        x = x.view(batch_size, -1, self.num_heads, self.depth)
        return x.permute(0, 2, 1, 3)

    def forward(self, v, k, q):
        batch_size = q.size(0)
        q = self.split_heads(self.wq(q), batch_size)
        k = self.split_heads(self.wk(k), batch_size)
        v = self.split_heads(self.wv(v), batch_size)

        # Scaled dot-product attention here...
        # Concatenate and apply a final linear layer
        return self.dense(...)

External References:

• Vaswani et al., "Attention is All You Need," 2017: https://arxiv.org/abs/1706.03762
• Illustrated Transformer by Jay Alammar: http://jalammar.github.io/illustrated-transformer/

Diagrams: Here is a simplified diagram of a single self-attention mechanism:

graph TD;
    A[Input Sequence] --> B[Query/Key/Value]
    B --> C[Self-Attention Calculation]
    C --> D[Weighted Output Sequence]

This diagram shows how an input sequence is transformed through self-attention, highlighting the process of generating queries, keys, and values, and computing the attention-weighted output.