NexEdge AI // Production-Grade Intelligence ~~ Enterprise AI Integration

Modern AI systems rely heavily on attention. Whether you're building LLMs, document processors, or domain-specific copilots, understanding attention is no longer optional — it is foundational.

This tutorial explains attention from the ground up, then shows how to implement it in PyTorch with clean, annotated code. The goal is not just to show what the code does, but to explain why each part exists and how it fits into the Transformer architecture.

What You Will Learn

By the end of this guide, you will understand:

Why attention replaced recurrence in modern sequence models.
How queries, keys, and values work.
Why positional encodings are needed.
How scaled dot-product attention is computed.
How multi-head attention extends the basic idea.
How these pieces combine into a Transformer block.
How to build a tiny language model on top.

Why Attention Matters

Traditional RNNs and LSTMs process tokens sequentially. That makes them harder to parallelize and less effective at capturing long-range dependencies.

Attention solves this by letting every token directly examine every other token in the sequence. Instead of compressing the past into a single hidden state, the model learns what to focus on at each step.

That is why attention became the backbone of Transformers, LLMs, and many multimodal systems.

The Core Attention Equation

The Transformer attention mechanism is defined as:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Where:

Q = queries, what each token is looking for.
K = keys, what each token offers.
V = values, the information to be aggregated.
d_k = key dimension used for scaling.

The sequence is simple:

Compute query-key similarity.
Normalize scores with softmax.
Use the resulting weights to mix the value vectors.

Step 1: Scaled Dot-Product Attention

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


class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int, d_v: int):
        super().__init__()
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_v, bias=False)
        self.d_k = d_k

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        """
        x: (batch, seq_len, d_model)
        mask: optional (seq_len, seq_len) or (batch, seq_len, seq_len)

        returns:
            out: (batch, seq_len, d_v)
            attn_weights: (batch, seq_len, seq_len)
        """
        # Project input embeddings into query, key, and value spaces.
        q = self.W_q(x)  # (batch, seq_len, d_k)
        k = self.W_k(x)  # (batch, seq_len, d_k)
        v = self.W_v(x)  # (batch, seq_len, d_v)

        # Compute similarity scores between every query and every key.
        scores = torch.matmul(q, k.transpose(-2, -1))  # (batch, seq_len, seq_len)

        # Scale by sqrt(d_k) to keep gradients stable.
        scores = scores / math.sqrt(self.d_k)

        # Optionally block future tokens or padding tokens.
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # Convert scores into probabilities.
        attn_weights = F.softmax(scores, dim=-1)

        # Use attention weights to combine value vectors.
        out = torch.matmul(attn_weights, v)

        return out, attn_weights

Step 2: Why Positional Encodings Are Needed

Self-attention does not know token order by itself. If you give it the same set of tokens in a different order, it has no built-in sense of sequence position.

To fix that, we add positional information to the token embeddings before attention. This can be learned or fixed.

Step 3: Sinusoidal Positional Encoding

python
def sinusoidal_positional_encoding(seq_len: int, d_model: int, device: torch.device):
    """
    Returns positional encodings of shape (seq_len, d_model).
    """
    position = torch.arange(seq_len, device=device, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(
        torch.arange(0, d_model, 2, device=device, dtype=torch.float)
        * (-math.log(10000.0) / d_model)
    )

    pe = torch.zeros(seq_len, d_model, device=device)
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

Step 4: Learned Token + Position Embeddings

python
class InputEmbedding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, max_seq_len: int):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)

    def forward(self, tokens: torch.Tensor):
        """
        tokens: (batch, seq_len)
        returns: (batch, seq_len, d_model)
        """
        seq_len = tokens.size(1)
        positions = torch.arange(seq_len, device=tokens.device).unsqueeze(0)
        return self.token_emb(tokens) + self.pos_emb(positions)

Step 5: Multi-Head Attention

python
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.num_heads = num_heads
        self.d_model = d_model
        self.head_dim = d_model // num_heads

        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)

    def split_heads(self, x):
        batch_size, seq_len, _ = x.shape
        x = x.view(batch_size, seq_len, self.num_heads, self.head_dim)
        return x.transpose(1, 2)

    def merge_heads(self, x):
        batch_size, num_heads, seq_len, head_dim = x.shape
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, self.d_model)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        q = self.split_heads(self.W_q(x))
        k = self.split_heads(self.W_k(x))
        v = self.split_heads(self.W_v(x))

        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores / math.sqrt(self.head_dim)

        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(0).unsqueeze(0)
            elif mask.dim() == 3:
                mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, float("-inf"))

        attn_weights = F.softmax(scores, dim=-1)
        out = torch.matmul(attn_weights, v)

        out = self.merge_heads(out)
        out = self.W_o(out)

        return out, attn_weights

Step 6: Transformer Block

python
class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        attn_out, attn_weights = self.attn(x, mask=mask)
        x = self.norm1(x + self.dropout(attn_out))

        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)

        return x, attn_weights

Step 7: Tiny Transformer Language Model

python
class TinyTransformerLM(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        num_heads: int,
        num_layers: int,
        d_ff: int,
        max_seq_len: int,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.embed = InputEmbedding(vocab_size, d_model, max_seq_len)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)

    def forward(self, tokens: torch.Tensor):
        batch_size, seq_len = tokens.shape
        device = tokens.device

        x = self.embed(tokens)
        mask = torch.tril(torch.ones(seq_len, seq_len, device=device))

        all_attn = []
        for block in self.blocks:
            x, attn_weights = block(x, mask=mask)
            all_attn.append(attn_weights)

        x = self.norm(x)
        logits = self.head(x)
        return logits, all_attn

Example Usage

python
vocab_size = 100
d_model = 32
num_heads = 4
num_layers = 2
d_ff = 64
max_seq_len = 16

model = TinyTransformerLM(
    vocab_size=vocab_size,
    d_model=d_model,
    num_heads=num_heads,
    num_layers=num_layers,
    d_ff=d_ff,
    max_seq_len=max_seq_len,
)

tokens = torch.randint(0, vocab_size, (2, 8))
logits, attn_maps = model(tokens)

print(logits.shape)
print(attn_maps.shape)

Training Sketch

python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = TinyTransformerLM(
    vocab_size=1000,
    d_model=128,
    num_heads=8,
    num_layers=4,
    d_ff=256,
    max_seq_len=64,
).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss()

dummy_batch = torch.randint(0, 1000, (4, 20), device=device)
x = dummy_batch[:, :-1]
y = dummy_batch[:, 1:]

logits, _ = model(x)
loss = criterion(logits.reshape(-1, logits.size(-1)), y.reshape(-1))

optimizer.zero_grad()
loss.backward()
optimizer.step()

print("Loss:", loss.item())

Why This Matters for AI Engineering

This architecture is the foundation behind many production AI systems. If you understand attention at this level, you can:

debug model behavior more effectively,
adapt architectures for custom clients,
reason about long-context tradeoffs,
and communicate clearly with technical and non-technical stakeholders.

That is especially valuable when building AI products that care about reliability, interpretability, and business impact.

text

[Target audience] 


AI engineers, ML engineers, applied researchers, and technical decision-makers  buidling and evaluating 
custom AI solutions.

Next read

How AI Engineers Build Reliable Agentic Systems

Attention Mechanisms in Transformers