Understanding Attention: The Idea Behind Modern AI

Let V be a non-empty set and let K be a field (typically ℝ or ℂ).

We say that V is a vector space over K if the following operations are defined:

- Vector addition

- Scalar multiplication

and the following axioms hold for all and all .

Axioms of vector addition

Closure
For all

Commutativity

Associativity

Additive identity
There exists an element such that:

Additive inverse
For each , there exists an element such that:

What does each subspace represent?

When we project a token vector into , , and , we are not creating random vectors.

Each projection learns a different “view” of the same token.

You can think of it like this:

• The original embedding : a rich, compressed meaning of the token.

• The projections: specialized lenses for the attention mechanism.

1) Query subspace ()

This subspace represents what this token wants to find in other tokens.

If the token is “elephant”, its query might encode something like:

• “I’m looking for attributes that describe me”

• “I’m looking for words related to size, color, action, etc.”

So is an information request vector.

2) Key subspace (​)

This subspace represents what this token can offer to others.

A token’s key acts like a label or index:

• “I contain useful information about being an adjective”

• “I contain information about location/time/action”

During attention, other tokens compare their query with this key to decide:

“Is this token relevant for me right now?”

So is a matchability vector.

3) Value subspace (​)

This subspace represents the actual content to pass forward if the token is selected by attention.

Important distinction:

• Key is for matching.

• Value is for carrying information.

If attention weight is high, more of that token’s value is mixed into the output.

So is the payload vector.

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

class SelfAttention(nn.Module):
    """
    Minimal single-head self-attention implementation.

    Input shape:
        x: [batch_size, seq_len, d_model]

    Output shape:
        out: [batch_size, seq_len, d_model]
        attn_weights: [batch_size, seq_len, seq_len]
    """

    def __init__(self, d_model: int, d_k: int = None):
        super().__init__()
        # If d_k is not provided, use d_model (common in simple examples)
        self.d_model = d_model
        self.d_k = d_model if d_k is None else d_k

        # Trainable projection matrices:
        # W_q, W_k, W_v are implemented as Linear layers.
        # They map each token embedding from d_model -> d_k (for q, k)
        # and d_model -> d_model (for v in this simple version).
        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, d_model, bias=False)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
        """
        x: [B, T, D]
        mask (optional): [B, T, T] with 1 for allowed positions, 0 for blocked positions
                         (useful for padding mask or causal mask)

        Returns:
            out: [B, T, D]
            attn_weights: [B, T, T]
        """
        # 1) Project input embeddings into Query, Key, Value subspaces
        # q: [B, T, d_k], k: [B, T, d_k], v: [B, T, D]
        q = self.W_q(x)
        k = self.W_k(x)
        v = self.W_v(x)

        # 2) Compute raw attention scores with dot product q @ k^T
        # scores shape: [B, T, T]
        # Each position i compares its query q_i with every key k_j.
        scores = torch.matmul(q, k.transpose(-2, -1))

        # 3) Scale scores to stabilize training
        # Without scaling, dot products can grow too large for big d_k.
        scores = scores / math.sqrt(self.d_k)

        # 4) (Optional) Apply mask
        # Positions with mask == 0 are set to a large negative value,
        # so after softmax they become ~0 attention probability.
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # 5) Convert scores into probabilities with softmax
        # attn_weights[b, i, j] = how much token i attends to token j
        attn_weights = F.softmax(scores, dim=-1)

        # 6) Weighted sum of values
        # out_i = sum_j attn_weights[i, j] * v_j
        # out shape: [B, T, D]
        out = torch.matmul(attn_weights, v)

        return out, attn_weights


if __name__ == "__main__":
    # Reproducibility
    torch.manual_seed(42)

    # Example dimensions
    B = 2          # batch size
    T = 5          # sequence length
    D = 16         # embedding size
    d_k = 8        # query/key subspace size

    # Dummy token embeddings
    x = torch.randn(B, T, D)

    # Create attention module
    attn = SelfAttention(d_model=D, d_k=d_k)

    # Forward pass
    out, weights = attn(x)

    print("Input shape:   ", x.shape)        # [2, 5, 16]
    print("Output shape:  ", out.shape)      # [2, 5, 16]
    print("Weights shape: ", weights.shape)  # [2, 5, 5]

    # Optional: inspect attention of first batch item
    print("\nAttention matrix for sample 0:")
    print(weights[0])

1) What are Q, K, and V?

For each token representation xix_ixi​, the model computes:

• (query): what this token is looking for.

• (key): what this token offers.

• (value): the content this token contributes if attended.

2) Why dot product?

The compatibility score between token iii and token jjj is:

Dot product can be written as:

So it measures alignment:

• large positive similar direction high compatibility,

• near zero weak relation,

• negative opposite direction.

That is why is interpreted as “how well they match.”

3) Why does appear?

If we stack all queries and keys:

then:

and each entry is:

So one matrix multiplication computes all pairwise query-key scores at once.

4) Short linguistic intuition

In the sentence:
“The cat ate fish because it was hungry.”

the query for “it” may align more with “cat” than with “fish,” so the score for “cat” is higher, and attention puts more weight there.