What you’ll learn

• why a word’s embedding isn’t enough — the same word can mean different things depending on its neighbors
• self-attention: how a model re-weights every token using every other token
• the query, key, value framing that all attention mechanisms share
• scaled dot-product attention — the actual formula, in NumPy
• multi-head attention: why one attention pass isn’t enough
• positional encoding — how an order-blind mechanism learns word order anyway
• how these pieces snap together into a Transformer encoder block
• why Transformers overtook RNNs for most sequence-modeling tasks

The problem: one vector per word isn’t enough

A keras.layers.Embeddingkeras.layers.Embedding gives every word a fixed vector — “bank” always gets the exact same vector, whether you mean a river bank or a savings bank. But that’s not how language works. Consider:

“The train left the station on time.”

Is “station” a train station? Could it be a radio station? The word alone doesn’t say — you need to look at its neighbors (“train”) to know. A smart model would give “station” a different vector depending on which sentence it appears in.

That’s exactly what self-attention does: it modulates the representation of a token using the representations of the other tokens around it, producing a context-aware vector for every word.

Self-attention, step by step

For the sentence above, self-attention asks: “how related is ‘station’ to every other word in the sentence?” It computes one relevancy score per pair of tokens (a plain dot product between their vectors), turns those scores into weights with a softmax, and then builds a new “station” vector as a weighted sum of every word’s vector — words that matter more (like “train”) contribute more.

  flowchart LR
  A["the"] --> S["Attention scores
vs station"]
  B["train"] --> S
  C["left"] --> S
  D["the"] --> S
  E["station"] --> S
  F["on"] --> S
  G["time"] --> S
  S --> W["Softmax weights"]
  W --> N["New context-aware
station vector"]

Here’s that exact idea in plain, runnable NumPy — no TensorFlow needed:

import numpy as np
 
def softmax(x):
    e = np.exp(x - np.max(x))
    return e / e.sum()
 
def self_attention(input_sequence):
    """input_sequence: shape (seq_len, embed_dim) -- one row per token."""
    output = np.zeros_like(input_sequence)
    for i, pivot_vector in enumerate(input_sequence):
        # 1. score this token against every token (including itself)
        scores = np.array([np.dot(pivot_vector, vector) for vector in input_sequence])
        # 2. scale, then softmax turns scores into weights that sum to 1
        scores = scores / np.sqrt(input_sequence.shape[1])
        weights = softmax(scores)
        # 3. weighted sum of every token's vector = new, context-aware vector
        output[i] = sum(w * v for w, v in zip(weights, input_sequence))
    return output
 
# 4 toy word vectors (3 dims each) standing in for "the cat sat down"
tokens = np.array([
    [1.0, 0.0, 0.2],
    [0.9, 0.1, 0.1],
    [0.1, 0.9, 0.0],
    [0.0, 0.8, 0.3],
])
 
context_aware = self_attention(tokens)
print("Original 'sat' vector:  ", tokens[2])
print("Context-aware 'sat':    ", np.round(context_aware[2], 3))

import numpy as np
 
def softmax(x):
    e = np.exp(x - np.max(x))
    return e / e.sum()
 
def self_attention(input_sequence):
    """input_sequence: shape (seq_len, embed_dim) -- one row per token."""
    output = np.zeros_like(input_sequence)
    for i, pivot_vector in enumerate(input_sequence):
        # 1. score this token against every token (including itself)
        scores = np.array([np.dot(pivot_vector, vector) for vector in input_sequence])
        # 2. scale, then softmax turns scores into weights that sum to 1
        scores = scores / np.sqrt(input_sequence.shape[1])
        weights = softmax(scores)
        # 3. weighted sum of every token's vector = new, context-aware vector
        output[i] = sum(w * v for w, v in zip(weights, input_sequence))
    return output
 
# 4 toy word vectors (3 dims each) standing in for "the cat sat down"
tokens = np.array([
    [1.0, 0.0, 0.2],
    [0.9, 0.1, 0.1],
    [0.1, 0.9, 0.0],
    [0.0, 0.8, 0.3],
])
 
context_aware = self_attention(tokens)
print("Original 'sat' vector:  ", tokens[2])
print("Context-aware 'sat':    ", np.round(context_aware[2], 3))

Original 'sat' vector:   [0.1 0.9 0. ]
Context-aware 'sat':     [0.419 0.523 0.147]

Original 'sat' vector:   [0.1 0.9 0. ]
Context-aware 'sat':     [0.419 0.523 0.147]

Notice the new “sat” vector has pulled a little away from its original position — it now carries a trace of every other word in the sentence, weighted by how related each one is.

Generalizing: query, key, value

So far both sides of every comparison came from the same sentence. But the same mechanism works when they come from different sequences — that’s the general query / key / value model:

• query — what you’re looking for
• keys — labels describing what’s available
• values — the actual content you’ll retrieve

Think of a search engine: you type a query (“dogs on the beach”). Every photo in the database has keys (its tags: “dog,” “beach,” “party”). The engine matches your query against every key, ranks the matches, and returns the corresponding values (the photos) — weighted by how well they matched.

outputs = sum(values * pairwise_scores(query, keys))

outputs = sum(values * pairwise_scores(query, keys))

For plain self-attention, query, keys, and values are all the same sequence — you’re comparing a sentence to itself. Keras exposes this directly through MultiHeadAttentionMultiHeadAttention, which is why you always pass the input three times:

import tensorflow as tf
 
num_heads = 4
embed_dim = 256
 
mha_layer = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
 
# query, keys, and values are the SAME sequence for plain self-attention
seq = tf.random.normal([2, 10, embed_dim])   # (batch, seq_len, embed_dim)
outputs = mha_layer(seq, seq, seq)
 
print("Output shape:", outputs.shape)

import tensorflow as tf
 
num_heads = 4
embed_dim = 256
 
mha_layer = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
 
# query, keys, and values are the SAME sequence for plain self-attention
seq = tf.random.normal([2, 10, embed_dim])   # (batch, seq_len, embed_dim)
outputs = mha_layer(seq, seq, seq)
 
print("Output shape:", outputs.shape)

Output shape: (2, 10, 256)

Output shape: (2, 10, 256)

Multi-head attention: why one pass isn’t enough

“Multi-head” means the layer doesn’t compute attention once — it splits the query/key/value vectors into several smaller, independent chunks (heads), runs attention separately on each one, then concatenates the results back together. Each head is free to learn a different kind of relationship between words (one head might track subject-verb agreement, another might track “which pronoun refers to which noun”), the same way different filters in a CNN learn different visual features.

Positional encoding: putting order back in

Here’s the catch: self-attention on its own is a set operation. It’s completely blind to whether “station” was the first word or the last — scramble every word in the sentence and you’d get the exact same pairwise scores. That means a plain self-attention layer, by itself, isn’t actually a sequence model at all.

The fix is positional encoding: add a vector that encodes where a token sits in the sequence directly onto its word embedding, so “position 0” and “position 5” look different to the model even for the same word.

import numpy as np
 
def positional_encoding(pos, d_model):
    i = np.arange(d_model)
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
    angles = pos * angle_rates
    pe = np.zeros(d_model)
    pe[0::2] = np.sin(angles[0::2])
    pe[1::2] = np.cos(angles[1::2])
    return pe
 
print("Position 0:", np.round(positional_encoding(0, 4), 3))
print("Position 1:", np.round(positional_encoding(1, 4), 3))

import numpy as np
 
def positional_encoding(pos, d_model):
    i = np.arange(d_model)
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
    angles = pos * angle_rates
    pe = np.zeros(d_model)
    pe[0::2] = np.sin(angles[0::2])
    pe[1::2] = np.cos(angles[1::2])
    return pe
 
print("Position 0:", np.round(positional_encoding(0, 4), 3))
print("Position 1:", np.round(positional_encoding(1, 4), 3))

Position 0: [0. 1. 0. 1.]
Position 1: [0.841 0.54  0.01  1.   ]

Position 0: [0. 1. 0. 1.]
Position 1: [0.841 0.54  0.01  1.   ]

The original “Attention Is All You Need” paper used this fixed sine/cosine trick. A simpler and equally effective alternative — and the one modern Keras code tends to use — is a learned PositionalEmbeddingPositionalEmbedding: an ordinary EmbeddingEmbedding layer whose input is just 0, 1, 2, ...0, 1, 2, ... (the position index), added on top of the word embedding, so position vectors get trained the same way word vectors do.

The Transformer encoder block

Put self-attention, a small feed-forward network, residual connections, and normalization together, and you get one Transformer encoder block — a reusable unit that ingests a sequence and returns a same-shaped sequence of richer, context-aware vectors:

  flowchart TD
  I["Input tokens
(word embedding + position)"] --> MHA["Multi-Head Self-Attention"]
  I --> ADD1["Add (residual)"]
  MHA --> ADD1
  ADD1 --> LN1["Layer Normalization"]
  LN1 --> FF["Feed-forward:
Dense -> ReLU -> Dense"]
  LN1 --> ADD2["Add (residual)"]
  FF --> ADD2
  ADD2 --> LN2["Layer Normalization"]
  LN2 --> O["Context-aware token vectors"]

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
 
class TransformerEncoder(layers.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
        super().__init__(**kwargs)
        self.attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.dense_proj = keras.Sequential([
            layers.Dense(dense_dim, activation="relu"),
            layers.Dense(embed_dim),
        ])
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
 
    def call(self, inputs, mask=None):
        if mask is not None:
            mask = mask[:, tf.newaxis, :]
        attention_output = self.attention(inputs, inputs, attention_mask=mask)
        proj_input = self.layernorm_1(inputs + attention_output)   # residual + norm
        proj_output = self.dense_proj(proj_input)
        return self.layernorm_2(proj_input + proj_output)           # residual + norm

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
 
class TransformerEncoder(layers.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
        super().__init__(**kwargs)
        self.attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.dense_proj = keras.Sequential([
            layers.Dense(dense_dim, activation="relu"),
            layers.Dense(embed_dim),
        ])
        self.layernorm_1 = layers.LayerNormalization()
        self.layernorm_2 = layers.LayerNormalization()
 
    def call(self, inputs, mask=None):
        if mask is not None:
            mask = mask[:, tf.newaxis, :]
        attention_output = self.attention(inputs, inputs, attention_mask=mask)
        proj_input = self.layernorm_1(inputs + attention_output)   # residual + norm
        proj_output = self.dense_proj(proj_input)
        return self.layernorm_2(proj_input + proj_output)           # residual + norm

Notice LayerNormalizationLayerNormalization, not BatchNormalizationBatchNormalization: it normalizes each sequence independently (pooling over the embedding axis), rather than pooling statistics across an entire batch — which matters because sequences in a batch aren’t statistically related to each other the way images are.

Why Transformers overtook RNNs

An RNN has to squeeze an entire sentence through a single hidden state, updated one word at a time — by word 100, it’s often forgotten word 1. A Transformer’s attention layer connects any two tokens directly, no matter how far apart, in a single step. That also makes it dramatically more parallelizable on a GPU: an RNN must finish word t before it can start word t + 1, while every token in a Transformer layer is processed simultaneously.

Visualize it

Below is the book’s own worked example: the sentence “The train left the station on time”, with attention scores computed between every pair of words. Click a row label to make that word the query and watch which keys it attends to most strongly:

Mini-checkpoint

Why does a Transformer need positional encoding, but an RNN doesn’t?

• An RNN processes tokens strictly one at a time, so word order is baked into how it computes — token 5 is only ever seen after tokens 1-4. Self-attention looks at all tokens simultaneously and treats them as an unordered set, so without positional encoding it has literally no way to tell “dog bites man” from “man bites dog.”

🧪 Try It Yourself

Exercise 1 – Scaled Dot-Product Attention Weights

Exercise 2 – Splitting a Vector Across Attention Heads

Exercise 3 – Sinusoidal Positional Encoding

Next

Continue to Sequence-to-Sequence Learning (Machine Translation) — see how an encoder and a decoder, built from either RNNs or Transformer blocks, work together to turn one sequence into another.