Attention Mechanism

The Birth of Attention

Challenges of the Seq2Seq Model

In Sequence-to-Sequence (Seq2Seq) tasks such as machine translation, the encoder compressed the input sequence into a fixed-length context vector. This caused a problem.

The Solution Brought by Attention

The Attention mechanism (Bahdanau et al., 2014) allowed the decoder to refer to all of the encoder's hidden states at each step.

Computing Attention

Types of Score Functions

Self-Attention

Self-Attention is a scheme in which the input sequence computes Attention with respect to itself. It can directly model the relationship between any pair of positions within the sequence.

Why Scaling Is Necessary

Multi-Head Attention

Because a single Attention has limited expressive power, multiple "heads" compute Attention from different perspectives, and the results are combined.

How Are the Heads Split?

1. Linearly transform the input (each token is \(d_{model}\)-dimensional) with the weight matrices \(W^Q, W^K, W^V\) to produce the full Q, K, V (\(d_{model}\)-dimensional).
2. Split that \(d_{model}\)-dimensional vector into \(h\) equal parts, with each head handling a small Q, K, V of dimension \(d_k = d_{model}/h\) (in implementation this is a reshape: \((\text{seq length},\, d_{model}) \to (\text{seq length},\, h,\, d_k)\)). Rather than preparing \(h\) sets of separate weights, this simply splits one large projection into \(h\) parts, and the \(W_i^Q\) in the formula corresponds to the columns of \(W^Q\) that each head is responsible for.
3. Each head independently computes Scaled Dot-Product Attention, and the resulting outputs (each \(d_k\)-dimensional) are concatenated back to \(d_{model}\) dimensions, then finally mixed together by \(W^O\).

How to choose the number of heads \(h\): \(h\) is a hyperparameter chosen by a human, selected from the divisors of \(d_{model}\) so that \(d_k = d_{model}/h\) is an integer. Increasing the number of heads keeps the compute roughly constant since each head becomes thinner, but increasing it too much reduces the dimension per head and lowers expressive power. Empirically, 8 to 16 is the standard choice.

Typical Settings

• Number of heads h: 8 to 16
• Dimension of each head: \(d_k = d_v = d_{model} / h\)
• e.g., \(d_{model} = 512, h = 8 \Rightarrow d_k = 64\)

Positional Encoding

Self-Attention does not consider order (it is permutation-invariant). Therefore, positional information must be added explicitly.

Properties of Positional Encoding

• For any fixed offset \(k\), \(PE_{pos+k}\) can be expressed as a linear function of \(PE_{pos}\)
• This makes it easy to learn relative position information
• It can handle sequences longer than those in the trained model (extrapolation is possible)

Implementation in PyTorch

Scaled Dot-Product Attention

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Scaled Dot-Product Attention

    Args:
        query: (batch, heads, seq_len, d_k)
        key: (batch, heads, seq_len, d_k)
        value: (batch, heads, seq_len, d_v)
        mask: (batch, 1, 1, seq_len) or (batch, 1, seq_len, seq_len)

    Returns:
        output: (batch, heads, seq_len, d_v)
        attention_weights: (batch, heads, seq_len, seq_len)
    """
    d_k = query.size(-1)

    # QK^T / sqrt(d_k)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # Apply mask (optional)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Softmax
    attention_weights = F.softmax(scores, dim=-1)

    # Attention x V
    output = torch.matmul(attention_weights, value)

    return output, attention_weights

Multi-Head Attention

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

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Linear transformation layers
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def split_heads(self, x, batch_size):
        """(batch, seq_len, d_model) -> (batch, heads, seq_len, d_k)"""
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # Linear transformation
        Q = self.W_q(query)
        K = self.W_k(key)
        V = self.W_v(value)

        # Split into heads
        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)

        # Scaled Dot-Product Attention
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)

        # Concatenate the heads
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)

        # Final linear transformation
        output = self.W_o(attn_output)

        return output, attn_weights


# Usage example
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)
x = torch.randn(batch_size, seq_len, d_model)

output, attention = mha(x, x, x)  # Self-Attention
print(f"Output shape: {output.shape}")  # (32, 100, 512)
print(f"Attention shape: {attention.shape}")  # (32, 8, 100, 100)

Positional Encoding

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Precompute the positional encoding
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2).float() *
                            (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)  # even dimensions
        pe[:, 1::2] = torch.cos(position * div_term)  # odd dimensions

        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: (batch, seq_len, d_model)
        """
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


# Usage example
pe = PositionalEncoding(d_model=512)
x = torch.randn(32, 100, 512)
output = pe(x)
print(f"Output shape: {output.shape}")  # (32, 100, 512)

Visualizing Attention

By visualizing the Attention weights, we can understand "where the model is focusing."

Summary

Frequently Asked Questions (FAQ)