Attention

LSTMs improved recurrent networks by learning what to remember and what to forget. But they still process a sequence one step at a time.

Suppose Jessie writes: “The burger that Jessie ordered after walking around Manhattan all afternoon was delicious.” To understand delicious, the model needs to connect it to burger.

Why force information to travel through every intermediate word? Why cannot delicious look directly at burger?

That question is the heart of attention. Attention lets each token decide which other tokens are relevant and retrieve information from them directly.

Why Memory Is Not Enough

An LSTM carries information through a hidden state and cell state. This is much better than a vanilla RNN, but it still compresses past context into a moving memory.

Attention changes the strategy. Instead of compressing everything into one memory, keep the token representations available and let each token look up what it needs.

When you read a confusing sentence, you do not reread every previous word in order. Your eyes jump to the relevant phrase. Your brain selectively focuses on useful information.

Attention gives neural networks that ability. Every token asks: which other tokens matter to me right now?

A Concrete Example

Consider the sentence:

To understand the word it, the model should look strongly at burger. Other words may matter a little, but burger is the key reference.

The model has learned where to look. That is the conceptual leap.

Library Analogy

An LSTM is like carrying one notebook and rewriting it after every book you read. Attention is more like keeping every book on the shelf. When you need something, you retrieve the relevant book directly.

Queries, Keys, And Values

Every token produces three vectors:

From the input embedding matrix $X$, the model learns three projections:

The matrices $W_Q$, $W_K$, and $W_V$ are learned parameters.

Measuring Relevance

A query is compared with every key using a dot product:

If a query and key point in similar directions, the dot product is large. If they are unrelated, it is small. This gives a matrix of raw attention scores.

Softmax Turns Scores Into Weights

Raw scores are not probabilities. Softmax converts each row into attention weights that sum to one:

Each row answers: for this query token, how much attention should go to every token?

Weighted Sum Of Values

The final output is a weighted average of value vectors:

A token's new representation is no longer based only on itself. It now includes information retrieved from other tokens.

Scaled Dot-Product Attention

The full equation divides by $\sqrt{d_k}$:

When vectors have many dimensions, dot products naturally become large. Large scores can make softmax too extreme, which weakens gradients. Dividing by $\sqrt{d_k}$ keeps scores at a more stable scale.

Worked Example

Suppose the word ordered attends to a short sentence:

The representation of ordered now includes information about who ordered and what was ordered. This is contextual representation.

Self-Attention

In self-attention, every token attends to other tokens in the same sequence. The word ordered might attend to Jessie and burger. The word burger might attend to ordered. Each token builds its own contextual view.

Why Attention Beat Recurrence

RNNs process tokens sequentially. Token 300 can only receive information from token 1 after it passes through every intermediate step. Attention allows token 300 to directly attend to token 1.

Attention also processes all tokens in parallel, which makes it much friendlier to GPUs than recurrence.

Multi-Head Attention

One attention mechanism can learn one kind of relationship. Multiple heads allow the model to attend in several ways at once: grammar, reference, sentiment, object relationships, or other patterns discovered during training.

Computational Cost

Attention compares every token with every other token. For sequence length $n$, the attention matrix has $n\times n$ entries:

This quadratic cost becomes expensive for very long documents and motivates research into efficient attention mechanisms.

Scaled dot-product attention is short to implement. The important part is the order of operations: project to Q, K, V; compute scores; scale; softmax; combine values.

import math
import torch
from torch import nn

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1)
    scores = scores / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ V
    return output, weights

batch_size = 2
sequence_length = 5
embed_dim = 16

X = torch.randn(batch_size, sequence_length, embed_dim)

W_q = nn.Linear(embed_dim, embed_dim)
W_k = nn.Linear(embed_dim, embed_dim)
W_v = nn.Linear(embed_dim, embed_dim)

Q = W_q(X)
K = W_k(X)
V = W_v(X)

output, attention_weights = scaled_dot_product_attention(Q, K, V)

print(output.shape)             # (2, 5, 16)
print(attention_weights.shape)  # (2, 5, 5)

The attention weights form a heat map. Row $i$ shows which tokens token $i$ attends to.