Attention Mechanisms in NLP

The Story That Explains Attention

Before attention, sequence models like RNNs were forced to compress an entire input sentence into a single fixed-length vector — a bottleneck that made them struggle with long sentences. Attention broke that constraint by allowing the model to look at all input tokens simultaneously and dynamically focus on the relevant ones at each decoding step.

The Foundation — Why RNNs Needed Help

To understand why attention was revolutionary, you first need to feel the pain it solved. The classic encoder-decoder RNN architecture compressed an entire source sentence into a single vector — the final hidden state — before decoding.

Bahdanau Attention — The Original Breakthrough

The Three-Step Mechanism

Python Implementation — Bahdanau Attention

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

class BahdanauAttention(nn.Module):
    """
    Bahdanau (Additive) Attention — Bahdanau et al., 2014
    'Neural Machine Translation by Jointly Learning to Align and Translate'
    """

    def __init__(self, hidden_dim):
        super().__init__()
        # W_a: projects decoder state
        self.W_a = nn.Linear(hidden_dim, hidden_dim, bias=False)
        # U_a: projects encoder states
        self.U_a = nn.Linear(hidden_dim, hidden_dim, bias=False)
        # v_a: collapses to a scalar score
        self.v_a = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, decoder_state, encoder_outputs):
        # decoder_state:   (batch, hidden_dim)
        # encoder_outputs: (batch, seq_len, hidden_dim)

        # Step 1: Score — combine decoder state with each encoder state
        dec_proj = self.W_a(decoder_state).unsqueeze(1)  # (batch,1,hidden)
        enc_proj = self.U_a(encoder_outputs)               # (batch,seq,hidden)
        energy   = self.v_a(torch.tanh(dec_proj + enc_proj)).squeeze(-1)
        # energy: (batch, seq_len)  — raw alignment scores

        # Step 2: Weights — softmax over source positions
        attn_weights = F.softmax(energy, dim=-1)  # (batch, seq_len)

        # Step 3: Context — weighted sum of encoder states
        context = torch.bmm(
            attn_weights.unsqueeze(1),   # (batch, 1, seq_len)
            encoder_outputs               # (batch, seq_len, hidden)
        ).squeeze(1)                      # → (batch, hidden)

        return context, attn_weights

# ── Quick test ──────────────────────────────────────────────
batch, seq_len, hidden = 4, 12, 256
attn = BahdanauAttention(hidden)

decoder_state    = torch.randn(batch, hidden)
encoder_outputs  = torch.randn(batch, seq_len, hidden)

context, weights = attn(decoder_state, encoder_outputs)
print(f"Context shape : {context.shape}")   # (4, 256)
print(f"Weights shape : {weights.shape}")   # (4, 12)
print(f"Weights sum   : {weights.sum(-1)}") # tensor([1., 1., 1., 1.])

Luong Attention — The Streamlined Successor

Luong vs Bahdanau — Key Differences

Three Luong Scoring Functions

Python Implementation — Luong Attention

class LuongAttention(nn.Module):
    """
    Luong (Multiplicative) Attention — Luong et al., 2015
    'Effective Approaches to Attention-based Neural Machine Translation'
    Supports 'dot', 'general', and 'concat' scoring.
    """

    def __init__(self, hidden_dim, method="general"):
        super().__init__()
        self.method = method

        if method == "general":
            self.W_a = nn.Linear(hidden_dim, hidden_dim, bias=False)

        elif method == "concat":
            self.W_a = nn.Linear(hidden_dim * 2, hidden_dim, bias=False)
            self.v_a = nn.Linear(hidden_dim, 1, bias=False)

    def score(self, decoder_h, encoder_outputs):
        # decoder_h:       (batch, hidden)
        # encoder_outputs: (batch, seq_len, hidden)

        if self.method == "dot":
            # Simple dot product
            return torch.bmm(
                encoder_outputs,                  # (batch, seq, hidden)
                decoder_h.unsqueeze(-1)           # (batch, hidden, 1)
            ).squeeze(-1)                         # → (batch, seq)

        elif self.method == "general":
            # Learned transformation of encoder states
            transformed = self.W_a(encoder_outputs)  # (batch, seq, hidden)
            return torch.bmm(
                transformed,
                decoder_h.unsqueeze(-1)
            ).squeeze(-1)

        elif self.method == "concat":
            # Concatenate and project through MLP
            dec_exp = decoder_h.unsqueeze(1).expand_as(encoder_outputs)
            combined = torch.cat([dec_exp, encoder_outputs], dim=-1)
            return self.v_a(torch.tanh(self.W_a(combined))).squeeze(-1)

    def forward(self, decoder_h, encoder_outputs):
        # Step 1 + 2: Score → Softmax → Weights
        energy      = self.score(decoder_h, encoder_outputs) # (batch, seq)
        attn_weights = F.softmax(energy, dim=-1)             # (batch, seq)

        # Step 3: Context vector
        context = torch.bmm(
            attn_weights.unsqueeze(1),    # (batch, 1, seq)
            encoder_outputs                # (batch, seq, hidden)
        ).squeeze(1)                       # → (batch, hidden)

        return context, attn_weights

# ── Compare all three methods ───────────────────────────────
for method in ["dot", "general", "concat"]:
    attn   = LuongAttention(256, method=method)
    ctx, w = attn(decoder_state, encoder_outputs)
    print(f"[{method:8s}] context: {ctx.shape}, weights sum≈1: {w.sum(-1).allclose(torch.ones(4))}")

Self-Attention — When a Sentence Reads Itself

Self-attention (also called intra-attention) was the key architectural innovation in the Transformer (Vaswani et al., 2017). Unlike Bahdanau/Luong attention which connected encoder states to decoder states, self-attention lets every position in a sequence build a representation informed by every other position.

The Query–Key–Value Framework

Self-attention is built around three learned linear projections of every token embedding: Query (Q), Key (K), and Value (V).

The Scaled Dot-Product Attention Formula

Python Implementation — Scaled Dot-Product Self-Attention

import math

class SelfAttention(nn.Module):
    """
    Single-head Scaled Dot-Product Self-Attention.
    Vaswani et al. 2017 — 'Attention Is All You Need'
    """

    def __init__(self, d_model, d_k=None):
        super().__init__()
        self.d_k = d_k if d_k else d_model
        # Three learned projection matrices
        self.W_Q = nn.Linear(d_model, self.d_k, bias=False)
        self.W_K = nn.Linear(d_model, self.d_k, bias=False)
        self.W_V = nn.Linear(d_model, self.d_k, bias=False)

    def forward(self, x, mask=None):
        # x: (batch, seq_len, d_model)

        # Step 1: Project to Q, K, V spaces
        Q = self.W_Q(x)  # (batch, seq, d_k)
        K = self.W_K(x)  # (batch, seq, d_k)
        V = self.W_V(x)  # (batch, seq, d_k)

        # Step 2: Compute scaled dot-product scores
        # scores[i,j] = how much token i attends to token j
        scale  = math.sqrt(self.d_k)
        scores = torch.bmm(Q, K.transpose(1, 2)) / scale  # (batch, seq, seq)

        # Step 3: Optional mask (causal / padding)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # Step 4: Softmax → attention weights
        attn_weights = F.softmax(scores, dim=-1)  # (batch, seq, seq)

        # Step 5: Weighted sum of values
        output = torch.bmm(attn_weights, V)  # (batch, seq, d_k)

        return output, attn_weights

# ── Demo ─────────────────────────────────────────────────────
batch, seq, d_model = 2, 8, 64
sa = SelfAttention(d_model)

x_in = torch.randn(batch, seq, d_model)
out, weights = sa(x_in)

print(f"Input  shape: {x_in.shape}")      # (2, 8, 64)
print(f"Output shape: {out.shape}")       # (2, 8, 64)
print(f"Attn   shape: {weights.shape}")   # (2, 8, 8)
print(f"Row sums ≈ 1: {weights[0].sum(-1)}")

Attention Map — Visualising "it" Resolution

Multi-Head Attention — Parallel Perspectives

Why Multiple Heads?

The Full Multi-Head Attention Formula

Parameter Count — Where Does the Memory Go?

Python Implementation — Multi-Head Attention

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Scaled Dot-Product Attention.
    Vaswani et al. 2017 — 'Attention Is All You Need'
    """

    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model    = d_model
        self.num_heads  = num_heads
        self.d_k        = d_model // num_heads  # per-head dimension

        # Unified projection matrices — project all heads at once then split
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)  # output projection

    def split_heads(self, x):
        # x: (batch, seq, d_model) → (batch, num_heads, seq, d_k)
        batch, seq, _ = x.size()
        x = x.view(batch, seq, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # → (batch, heads, seq, d_k)

    def scaled_dot_product(self, Q, K, V, mask=None):
        scale  = math.sqrt(self.d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / scale
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        weights = F.softmax(scores, dim=-1)
        return torch.matmul(weights, V), weights

    def forward(self, query, key, value, mask=None):
        # For self-attention: query = key = value = x
        # For cross-attention: query from decoder, key/value from encoder

        batch = query.size(0)

        # Step 1: Project and split into h heads
        Q = self.split_heads(self.W_Q(query))  # (batch, heads, seq_q, d_k)
        K = self.split_heads(self.W_K(key))    # (batch, heads, seq_k, d_k)
        V = self.split_heads(self.W_V(value))  # (batch, heads, seq_v, d_k)

        # Step 2: Per-head attention (runs in parallel on all heads)
        x, attn_weights = self.scaled_dot_product(Q, K, V, mask)
        # x: (batch, heads, seq, d_k)

        # Step 3: Concatenate heads
        x = x.transpose(1, 2).contiguous()  # (batch, seq, heads, d_k)
        x = x.view(batch, -1, self.d_model)   # (batch, seq, d_model)

        # Step 4: Final output projection
        output = self.W_O(x)  # (batch, seq, d_model)

        return output, attn_weights

# ── Demo: BERT-base style (12 heads, d_model=768) ───────────
d_model, num_heads = 768, 12
mha = MultiHeadAttention(d_model, num_heads)

x_bert = torch.randn(2, 16, d_model)  # batch=2, seq=16
out, w  = mha(x_bert, x_bert, x_bert)  # self-attention

total_params = sum(p.numel() for p in mha.parameters())
print(f"Output shape   : {out.shape}")         # (2, 16, 768)
print(f"Attn map shape : {w.shape}")           # (2, 12, 16, 16)
print(f"Total params   : {total_params:,}")    # 2,359,296
print(f"d_k per head   : {d_model // num_heads}")  # 64

Three Types of Attention in a Transformer

The full Transformer architecture uses multi-head attention in three distinct roles, each with a different configuration of Q, K, V sources.

Attention Complexity — The Elephant in the Room

Multi-head attention is powerful but expensive. Understanding its complexity is essential for production systems and for understanding why models like Longformer and FlashAttention exist.

Positional Encoding — Giving Attention a Sense of Order

import numpy as np

def positional_encoding(max_seq, d_model):
    """
    Compute sinusoidal positional encodings.
    Returns tensor of shape (max_seq, d_model).
    """
    PE = np.zeros((max_seq, d_model))
    pos = np.arange(max_seq)[:, np.newaxis]           # (max_seq, 1)
    i   = np.arange(0, d_model, 2)[np.newaxis, :]      # (1, d_model/2)

    div_term = np.power(10000, (2 * i) / d_model)

    PE[:, 0::2] = np.sin(pos / div_term)   # even dims → sine
    PE[:, 1::2] = np.cos(pos / div_term)   # odd dims  → cosine

    return PE

PE = positional_encoding(100, 512)
print(f"PE shape: {PE.shape}")   # (100, 512)
print(f"Values range: [{PE.min():.3f}, {PE.max():.3f}]")   # [-1.000, 1.000]

Full Transformer Block — Putting It All Together

class TransformerEncoderBlock(nn.Module):
    """
    A single Transformer encoder block:
    Multi-Head Self-Attention → Add&Norm → FFN → Add&Norm
    """

    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        # ── Multi-head self-attention ───────────────────────────
        self.mha    = MultiHeadAttention(d_model, num_heads)
        self.norm1  = nn.LayerNorm(d_model)
        self.drop1  = nn.Dropout(dropout)

        # ── Feed-forward network ────────────────────────────────
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 1. Multi-head self-attention + residual
        attn_out, _ = self.mha(x, x, x, mask)
        x = self.norm1(x + self.drop1(attn_out))   # Add & Norm

        # 2. Position-wise feed-forward + residual
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.drop2(ffn_out))    # Add & Norm

        return x

# ── Stack 6 blocks (original Transformer encoder) ───────────
encoder = nn.Sequential(
    *[TransformerEncoderBlock(d_model=512, num_heads=8, d_ff=2048)
      for _ in range(6)]
)

src = torch.randn(2, 32, 512)  # batch=2, seq=32, d_model=512
out = encoder(src)
total = sum(p.numel() for p in encoder.parameters())
print(f"Encoder output : {out.shape}")    # (2, 32, 512)
print(f"Encoder params : {total:,}")      # ~25M

Comparison — Bahdanau vs Luong vs Self-Attention

Golden Rules — Attention in Practice

Modern Variants and Where Attention Is Going