Attention: Explanation is All You Need

We all know the story of the Transformers. They’re the robots in disguise. And what is Attention? It's something Optimus Prime never lacked!

Okay, bye! Blog is over!

Haha! Classic Schmosby humor. But let’s get serious now.

We all have used ChatGPT, let's be real. So ChatGPT uses this "Attention". But to understand what attention really means, we have to rewind to a time before ChatGPT, when the world was a different place. A time when the world was ruled by RNNs and CNNs.

Back then, Recurrent Neural Networks (RNNs) was the go-to model for many natural language processing tasks. We had vanilla RNNs, LSTMs, and GRUs, all of which were pretty good at processing sequential data like sentences. But they weren’t perfect. They struggled with long-range dependencies—basically, if you had to predict the last word in a very long sentence, the early parts of the sentence would often fade into a distant memory. On top of that, RNNs had trouble being parallelized. You had to process tokens step by step, which slowed training.

Meanwhile, Convolutional Neural Networks (CNNs) tried to fix some issues by capturing local patterns in a more parallel-friendly way. But CNNs had their own quirks, like needing deeper layers to capture broader context. Enter the big question: “How can we let a model see the entire sequence all at once while focusing on the most relevant parts?”

Attention: The Attention Seeker

The concept of “attention” was actually around before the Transformer paper. Early versions of attention appeared in sequence-to-sequence models for tasks like machine translation (especially around 2014-2015). The gist was: instead of encoding the entire input sequence into a single hidden state (like older RNN-based models did), you could let the model “attend” to certain parts of the input sequence more strongly when predicting each output token.

Imagine you’re translating a long sentence from English to French. At each step, you might want to look more closely at the English words that correspond to the French phrase you’re about to generate. That’s exactly what attention did: a set of weights telling you how much to focus on each input token at every decoding step.

English to French Attention Heatmap

This was a game-changer: it helped preserve long-range context and gave more interpretability to the model, because you could visualize the attention weights and see where the model was “looking.”

The Recurrent Problem

But there was a catch. Even with attention, the underlying RNNs or LSTMs were still there. They were like the old chairs with fancy new cushions that looked good but still creaked when you sat on them. The key issues were:

1. Sequential Processing: The model still had to process tokens one by one, serially. These models were hard to parallelize.
2. Long-Range Dependencies: The RNNs or LSTMs were still responsible for capturing long-range dependencies, which they struggled with.
3. Exploding and Vanishing Gradients: Training was still difficult because of the vanishing and exploding gradient problem.

Transformers: The Attention Revolution

Then came the Transformer. The bold claim was that we could ditch recurrent connections entirely. Instead of stepping through the sequence one token at a time, the Transformer uses a mechanism called “self-attention,” which compares every token in the sequence with every other token simultaneously. That means for each word, the model can figure out which other words it should pay more attention to.

Self-Attention: The Key Idea

The formula for self-attention is simple: you take the dot product of every token with every other token, scale it, apply a softmax to get weights, and then take a weighted sum of the token embeddings. This weighted sum is the output of self-attention.

Formula:

1Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V

Okay, difficult! Let’s break it down:

• Queries, Keys, and Values (Q, K, V): Each token in the input is projected into three vectors: a query vector (Q), a key vector (K), and a value vector (V).
• QK^T: This is the dot product of the Query and Key vectors. It tells you how much each token should attend to every other token.
• sqrt(d_k): This is a scaling factor. It’s the square root of the dimension of the Key vectors (Also, it is the embedding dimension). It helps stabilize the gradients during training.
• softmax: This function normalizes the weights so they sum to 1. It tells you how much each token should attend to every other token.
• QK^T / sqrt(d_k): This is the scaled dot product. It tells you how much each token should attend to every other token.
• V: This is the Value vector. It’s the output of self-attention. It’s a weighted sum of the Value vectors of all tokens.

Implementation in PyTorch

1import torch
2import torch.nn as nn
3from torch.nn import functional as F
4
5
6class Head(nn.Module):
7    """One head of self-attention."""
8
9    def __init__(self, n_embd, head_size, block_size, dropout):
10        super().__init__()
11        self.key = nn.Linear(n_embd, head_size, bias=False)
12        self.query = nn.Linear(n_embd, head_size, bias=False)
13        self.value = nn.Linear(n_embd, head_size, bias=False)
14        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
15        self.dropout = nn.Dropout(dropout)
16
17    def forward(self, x):
18        B, T, C = x.shape
19        k = self.key(x)
20        q = self.query(x)
21        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
22        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
23        wei = F.softmax(wei, dim=-1)
24        wei = self.dropout(wei)
25        v = self.value(x)
26        out = wei @ v
27        return out

Multi-Head Attention: Doing More

Instead of computing one single attention distribution, Transformers do this multiple times in parallel (each is called a “head”). Each head learns to focus on different aspects of the sequence. Then, we concatenate the heads back together and pass them through a linear layer.

This multi-head trick makes the model more expressive. It can learn different relationships in the data at the same time.

Formula:

1MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O

Implementation in PyTorch

1import torch
2import torch.nn as nn
3from torch.nn import functional as F
4
5
6class MultiHeadAttention(nn.Module):
7    """Multiple heads of self-attention in parallel."""
8
9    def __init__(self, n_embd, n_head, block_size, dropout):
10        super().__init__()
11        assert n_embd % n_head == 0, f"n_embd ({n_embd}) must be divisible by num_heads ({n_head})"
12
13        self.n_embd = n_embd
14        self.n_head = n_head
15        self.head_size = n_embd // n_head
16
17        self.heads = nn.ModuleList([Head(n_embd, self.head_size, block_size, dropout) for _ in range(n_head)])
18        self.proj = nn.Linear(n_embd, n_embd)
19        self.dropout = nn.Dropout(dropout)
20
21    def forward(self, x):
22        out = torch.cat([h(x) for h in self.heads], dim=-1)
23        out = self.dropout(self.proj(out))
24        return out

Now this is the original version of the Attention provided in the paper“Attention is All You Need.”

There are infact many other versions of the Attention Mechanism that were built on top of this original version. Some of them are:

• Self-Attention: The original version of the Attention Mechanism.
• Bidirectional Attention: Each token can look both left and right in the sentence.
• Multi-Query Attention: Instead of having multiple sets of queries, keys, and values, we share keys and values across heads while keeping separate queries.
• Infini-Attention: an efficient method to infinitely long inputs with bounded memory and computation
• Sliding Window Attention (Local Attention): This is a variant of the Attention Mechanism where each token can only attend to a fixed window of tokens around it.
• Flash Attention: This is the same attention mechanism as the vanilla attention, but it uses efficient memory layouts and GPU-friendly operations to compute attention with less overhead.

Alright, I am not going to bore you with the details of all these versions. But I have another blog coming up where I will be explaining these versions in detail. So, stay tuned!

Why Attention Matters?

So, why all the fuss about attention? Here’s the bottom line:

• Better Long-Range Dependencies: RNNs struggled with forgetting old information. Attention can look at the entire sequence, so it doesn’t lose track of distant tokens as easily.
• Parallelization: Transformers process all tokens at once (in the self-attention block), which means you can harness the power of GPUs to do it in parallel.
• Flexibility: You can adapt the attention mechanism to different tasks—machine translation, summarization, image captioning, code generation, and so on.
• Interpretability: You can visualize attention weights to see which words are influencing each other. Although nobody really knows what the model is thinking, attention gives you some clues.

Attention started as a small improvement to help models focus on relevant parts of a sentence. Then it exploded into a full-blown revolution with the Transformer architecture, drastically changing how we build language models, speech models, and even some vision models. Over time, new flavors of attention have popped up to handle large sequences, reduce computation, and speed up inference (like Flash Attention and multi-query attention). And to make generation efficient, we introduced the KV cache.

Right now it really feels like Attention is really all you need. But who knows, maybe in the future we will have a new architecture that will make us rethink the importance of Attention!

References:

• Attention is All You Need
• Raschka Forever 💯

Source Code

You can find the source code for my implementation of the Attention Mechanism in PyTorch here.

Checkout my Decoder Only Transformer implementation: https://github.com/NotShrirang/QuillGPT/