Deep Learning for Natural Language: Transformers

A static embedding gives a token one vector. A Transformer combines token representations with information about their positions and surrounding tokens, producing a contextual state at each input position. This lecture builds that transformation from attention weights to a complete encoder, then implements a small Keras text classifier.

Learning goals

You should finish able to calculate an attention-weighted representation, track query/key/value dimensions, explain positional information and masks, distinguish an encoder from a causal decoder, and build an encoder with residual connections, normalization, and a feed-forward network.

1. From token lookups to contextual states

Consider “the bank approved the loan” and “we rested beside the river bank.” A static word table starts both occurrences of bank with the same word-type vector. A contextual encoder can produce different final states because the surrounding inputs differ.

Three useful requirements for this encoder are:

• Context awareness: token states can incorporate information from permitted context positions.
• Order awareness: the model receives information distinguishing positions and arrangements.
• Length preservation: an input with L token positions produces L contextual states.

These describe the encoder transformation, not every possible NLP output. A classifier can pool the L states into one prediction. A translation model may generate an output of a different length. Tokens may be subwords rather than whole words, so preserving token length need not mean one output vector per written word.

The initial lookup table can be learned during task training or pretraining. Transformers do not require GloVe as input, and contextual models are not limited to Transformers: recurrent encoders can also incorporate context.

2. Self-attention as content-dependent mixing

Let X∈ℝ^(L×d) contain one input vector per row. A simplified, parameter-free mixing operation is:

A=softmax_rows(XXᵀ), followed by H=AX.

Each row of XXᵀ contains a token’s dot products with all token vectors. Row-wise softmax converts scores into weights; AX takes a separate weighted sum for every output position. Self-attention means that the queries and the information being attended to come from the same sequence.

This toy formula is useful intuition, but is not the full learned attention used in a standard Transformer. Its behavior still depends on X, which may itself come from trainable embeddings. Calling the mixing rule parameter-free does not make the surrounding network untrainable.

Softmax and weighted averaging

For scores s₁,…,s_L, softmax returns aⱼ=exp(sⱼ)/Σₖexp(sₖ). For finite, unmasked scores, weights are positive and sum to one. Numerically stable implementations subtract the largest score before exponentiating. Masked positions have zero weight in the ideal masked calculation; finite-precision arithmetic can also underflow small weights to zero.

For one query, suppose the scaled scores are [0,log 2,log 3]. Their exponentials are [1,2,3], giving attention weights [1/6,2/6,3/6]. Let the corresponding value vectors be:

v1 = [2,0]
v2 = [0,2]
v3 = [2,2]

The output is:

(1/6)[2,0]+(2/6)[0,2]+(3/6)[2,2]=[4/3,5/3].

If position 3 is masked, the remaining scores are renormalized to [1/3,2/3,0], producing [2/3,4/3]. Multiplying its old weight by zero without renormalizing would give a different, incorrect attention distribution.

Each single-head result is a convex combination of the value vectors before dropout. The complete Transformer also applies projections, residual paths, normalization, and nonlinear transformations, so its final output is not merely an average of the original word vectors.

3. Learned queries, keys, and values

Standard scaled dot-product attention introduces learned projections:

Q=XW_Q, K=XW_K, V=XW_V,

then computes:

Attention(Q,K,V)=softmax_rows(QKᵀ/√d_k+M)V.

M is an additive mask: zero for permitted pairs and ideally negative infinity for forbidden pairs. Biases can also be included in the projections. The terminology provides an intuition:

• A query describes what a position seeks.
• A key describes information against which queries are compared.
• A value supplies the content that is mixed.

These are trainable numerical roles, not manually labeled linguistic questions and answers. Unlike XXᵀ, QKᵀ need not be symmetric because queries and keys use different projections.

Quantity	Single-head shape
X	L×d_model
W_Q, W_K	d_model×d_k
W_V	d_model×d_v
Q, K	L×d_k
V	L×d_v
Scores and attention weights	L×L
Attention output	L×d_v

Batching adds a leading B dimension. With cross-attention, the query and source sequences can have different lengths: scores then have shape T×S instead of L×L.

Why divide by √d_k?

If query and key components are approximately independent, mean zero, and unit variance, their dot product has variance about d_k. Dividing by √d_k keeps that scale roughly stable as dimension grows. This helps avoid very large logits that make softmax sharply saturated. The assumptions explain the heuristic; they are not an exact description of every trained network.

All projections and attention operations participate in backpropagation. The training objective determines whether the resulting contextual states are useful for classification, masked-token prediction, or another task.

4. Multi-head attention

Several attention heads learn separate projections and mix information in parallel:

head_r=Attention(XW_Qʳ,XW_Kʳ,XW_Vʳ),

MultiHead(X)=Concat(head₁,…,head_h)W_O.

The output projection typically restores d_model features so residual addition is possible. Heads can learn different relationships, but there is no guarantee that one becomes a subject–verb head and another a modifier head. Nor should attention weights alone be treated as a complete causal explanation of a prediction.

A common configuration uses d_k=d_v=d_model/h. For d_model=32 and h=4, each head has key dimension 8. In Keras, key_dim means dimension per head, not total model width. Setting key_dim=32 with four heads would allocate a larger internal projection than this example intends.

With biases and total projected width 32, query, key, value, and output projections each have 32×32+32=1,056 parameters. Their total is 4,224. This count assumes the conventional widths just stated; other projection sizes change it.

5. Give attention access to position

Unmasked self-attention without positional information is permutation equivariant: permuting the input rows permutes the output rows correspondingly. It does not know that the first bank preceded the word loan. A mean-pooled classifier built only from such order-agnostic transformations cannot distinguish sequences solely by rearranging the same token multiset.

A simple solution adds a position vector to each token embedding:

Xᵢ=TokenEmbedding(tokenᵢ)+PositionEmbedding(i).

Both terms have d_model coordinates, so addition preserves width. Learned absolute position tables provide one trainable vector for each allowed position. They impose a configured maximum and do not automatically generalize to unseen positions.

The original Transformer also described fixed sinusoidal encodings:

PE(pos,2i)=sin(pos/10000^(2i/d_model)),

PE(pos,2i+1)=cos(pos/10000^(2i/d_model)).

For d_model=4 and pos=0, the vector is [0,1,0,1]. At pos=1 it is approximately [0.84147,0.54030,0.01000,0.99995]. Other architectures use relative-position mechanisms or rotary position representations, so additive learned vectors are one choice rather than the definition of all Transformers.

6. The complete encoder block

One post-normalization block can be written:

U=LayerNorm(X+Dropout(MultiHead(X))),

H=LayerNorm(U+Dropout(FFN(U))).

The same position-wise feed-forward network is applied independently to each token:

FFN(u)=W₂φ(W₁u+b₁)+b₂,

using a column-vector convention for that per-token equation. It expands to d_ff hidden units, applies a nonlinearity such as ReLU or GELU, then returns to d_model. Attention mixes positions; this feed-forward stage transforms features at each position.

Token IDs → Token vectors + Position vectors       (B,L,d_model)
                         │
                         ├── Multi-head self-attention
                         └── Residual add → LayerNorm
                                      │
                                      ├── Position-wise feed-forward network
                                      └── Residual add → LayerNorm
                                                     │
                                         Contextual token states
                                                     │
                              Optional masked pooling → Classifier

Residual paths preserve access to prior representations and aid gradient flow. Layer normalization usually normalizes across the feature dimension separately for each token, with learned scale and offset; it does not compute BatchNorm-style statistics across documents. Many modern models use pre-normalization, placing normalization before each sublayer. State which variant you implement instead of silently combining their equations.

Embedding dimension, feed-forward width, and head count are central hyperparameters. Depth, dropout, activation, sequence length, positional scheme, masking, and the training objective also matter. A block preserves L positions; pooling in an optional task head changes that length deliberately.

7. Masks define the permitted context

Mask / architecture	What a query can use	Typical purpose
Padding mask	Real source tokens, excluding padding	Prevent artificial padding from influencing representations
Bidirectional encoder	All permitted real positions on both sides	Encode an observed passage
Causal decoder	Current and earlier positions only	Predict subsequent tokens without future-token leakage

Bidirectional does not mean “two recurrent passes.” An encoder’s attention graph permits information from both directions. A causal mask creates a lower-triangular visibility pattern, including the diagonal. Padding and causal masks can be combined.

Parallelism also needs qualification. During teacher-forced decoder training, positions can be processed in parallel under the causal mask. Ordinary autoregressive generation still produces tokens step by step, often caching earlier keys and values. Removing recurrence does not mean that every generation task finishes in one parallel operation.

Standard dense attention computes O(L²d_k) score interactions per head and represents L×L pair weights. Doubling length quadruples the pair count. At L=128 there are 16,384 pairs per head; at L=256 there are 65,536. Efficient implementations can reduce memory traffic or avoid storing all weights, while sparse or alternative attention mechanisms can change the computation. Parallel matrix operations are powerful but not free.

8. A small Keras encoder with explicit padding masks

Use Keras 3 with a compatible TensorFlow installation and NumPy in a separate environment. This example accepts already indexed sequences so the encoder is easy to inspect. Zero is padding, one is unknown, and valid IDs must be smaller than vocab_size. It does not load BERT weights.

import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import numpy as np
import keras
from keras import ops

keras.utils.set_random_seed(42)
vocab_size = 5000
max_length = 128
d_model, d_ff, num_heads = 32, 64, 4
assert d_model % num_heads == 0

ids = keras.Input(shape=(max_length,), dtype="int32", name="token_ids")
valid = ops.not_equal(ids, 0)
# Each row is [0, 1, ..., 127], computed inside the model graph.
position_ids = ops.cumsum(ops.ones_like(ids), axis=-1) - 1
token_vectors = keras.layers.Embedding(vocab_size, d_model)(ids)
position_vectors = keras.layers.Embedding(max_length, d_model)(position_ids)
x = token_vectors + position_vectors

# Shape (B,1,L): exclude padding keys for every query, broadcasting over queries.
key_mask = ops.expand_dims(valid, axis=1)
attention = keras.layers.MultiHeadAttention(
    num_heads=num_heads, key_dim=d_model // num_heads,
    dropout=0.1, name="self_attention",
)
attended = attention(query=x, value=x, key=x, attention_mask=key_mask)
u = keras.layers.LayerNormalization(epsilon=1e-6)(
    x + keras.layers.Dropout(0.1)(attended)
)
ff = keras.layers.Dense(d_ff, activation="relu")(u)
ff = keras.layers.Dense(d_model)(ff)
states = keras.layers.LayerNormalization(epsilon=1e-6)(
    u + keras.layers.Dropout(0.1)(ff)
)
encoder = keras.Model(ids, states, name="small_encoder")

# Padded query states may be nonzero; exclude them explicitly at pooling.
pooled = keras.layers.GlobalAveragePooling1D()(states, mask=valid)
outputs = keras.layers.Dense(3, activation="softmax")(pooled)
classifier = keras.Model(ids, outputs, name="encoder_classifier")
classifier.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss=keras.losses.SparseCategoricalCrossentropy(),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
)
assert classifier.count_params() == vocab_size * d_model + 12739

This code uses an explicit mask rather than relying on automatic mask propagation through position addition and residual branches. Padding keys are hidden from attention, and padding query outputs are excluded from pooling. Padded states need not be zero. Reusing the key mask in every subsequent encoder block prevents those states from becoming visible as source content.

An all-padding row has no meaningful attention context or pooled mean. Reject such rows at the boundary:

def check_batch(batch):
    batch = np.asarray(batch)
    if batch.ndim != 2 or batch.shape[1] != max_length:
        raise ValueError("Expected a batch of fixed-length token sequences")
    if not np.issubdtype(batch.dtype, np.integer):
        raise ValueError("Token IDs must be integers")
    if (batch < 0).any() or (batch >= vocab_size).any():
        raise ValueError("Token ID outside embedding vocabulary")
    if not (batch != 0).any(axis=1).all():
        raise ValueError("Each document needs at least one nonpadding token")
    return batch.astype("int32")

example_ids = np.zeros((2, max_length), dtype="int32")
example_ids[0, :4] = [2, 3, 4, 5]
example_ids[1, :3] = [6, 7, 8]
example_ids = check_batch(example_ids)
assert tuple(encoder(example_ids, training=False).shape) == (2, 128, 32)
probabilities = classifier(example_ids, training=False).numpy()
np.testing.assert_allclose(probabilities.sum(axis=1), 1, atol=1e-6)
print("Output shapes verified; predictions are from random initialization.")

The checks above are included for you to execute with TensorFlow installed; a random output is not evidence of language understanding. The boundary check must also be applied to training and validation arrays before fitting. The model itself does not automatically call this Python helper.

Parameter and shape audit

The feed-forward stage has 32×64+64+64×32+32=4,192 parameters. Two layer normalizations add 2×(32+32)=128. Together with the attention projections, the block contains 4,224+4,192+128=8,544 parameters.

The token table contributes 5,000×32=160,000; learned positions contribute 128×32=4,096; the three-class head contributes 32×3+3=99. Total: 172,739. Sequence length affects the position table and attention workload, but does not multiply the shared attention projection count.

Connect the encoder to a real text task

Use a training-only adapted TextVectorization(output_mode="int", output_sequence_length=128) layer, as in the previous embedding lesson. Its IDs must match this model’s vocabulary size and token table. Do not load arbitrary token IDs from an unrelated tokenizer. Choose vocabulary and length before final evaluation; truncation can remove important context.

With nonempty, checked arrays x_train, x_val, and labels y_train, y_val in {0,1,2}, training follows the familiar pattern:

# Prepare these arrays from your own corpus before running this block.
x_train = check_batch(x_train)
x_val = check_batch(x_val)
history = classifier.fit(
    x_train, y_train,
    validation_data=(x_val, y_val),
    batch_size=32, epochs=20,
    callbacks=[keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=3, restore_best_weights=True,
    )], verbose=2,
)

For lyrics, retain grouped song/artist splits from the earlier lessons. Compare against majority, bag-of-words, and mean-embedding baselines. A small Transformer trained on little data may lose to simpler models. Record the vocabulary, class order, masks, and sequence limit with the trained model; report held-out metrics only after selecting settings on validation data.

9. BERT: a pretrained bidirectional encoder

BERT means Bidirectional Encoder Representations from Transformers. It stacks encoder blocks and is pretrained to learn context-dependent representations. Its input is normally subword tokens, not the whitespace word vocabulary used in the introductory lessons.

Original BERT combines token, position, and segment/type embeddings. Its key pretraining task is masked language modeling: selected tokens are corrupted and the model predicts their original identities using permitted context on both sides. Original BERT also used a next-sentence prediction objective; not every later encoder retains that objective.

A [MASK] token used for pretraining is different from an attention mask. The token substitutes content at selected positions; an attention mask restricts which positions are visible. Ordinary real tokens can still attend to a masked-token position under BERT’s bidirectional visibility rules.

For classification, a task head can use the contextual [CLS] state and be fine-tuned with the encoder. Other tasks use token-level outputs, such as named-entity recognition, or start/end position scores for extractive question answering. Sentence pooling choices must match the checkpoint and training objective; arbitrary averages of BERT states are not automatically optimal sentence embeddings.

When using pretrained BERT, use its matching tokenizer, special tokens, attention mask, and checkpoint configuration. A randomly initialized Keras block with the same broad architecture does not contain BERT’s learned knowledge. BERT is also not inherently an ordinary left-to-right text generator: its bidirectional masked-token objective differs from a causal next-token model’s objective.

10. Exercises with answers

1. Attention weights: Scores are [0,log 3]. What is softmax? [1/4,3/4].
2. Weighted values: With those weights and values [4,0] and [0,4], what is the output? [1,3].
3. Scaling: If d_k=64, what divides the query/key dot products? √64=8.
4. Heads: For d_model=96 and six equal-width heads, what key_dim implements the conventional split? 16, not 96.
5. Shapes: For B=8, L=50, h=4, what shape holds ordinary per-head attention weights? (8,4,50,50).
6. Positions: Why is content-only attention insufficient for order-sensitive tasks? It is permutation equivariant without order information or an order-imposing mask, so it cannot by itself identify absolute sequence arrangement.
7. Length: Does a classifier’s single output contradict encoder length preservation? No. The encoder retains L states; a separate pooling/classification head intentionally reduces them.
8. Padding: Why is masking attention keys alone insufficient for mean pooling? Padded query outputs can still be nonzero; exclude those positions from the pooled average too.
9. Complexity: What happens to the dense attention pair count when sequence length triples? It grows by nine times.
10. BERT: Can a bidirectional BERT token attend to later real tokens? Yes, subject to its attention mask. A causal decoder prevents that future access.
11. Pretraining: Does adding a [MASK] token make a randomly initialized encoder pretrained? No. Pretraining requires learning parameters through the appropriate objective and data.

Further reading and sources

• Vaswani et al., Attention Is All You Need: scaled attention, multiple heads, positional encodings, and the original encoder–decoder architecture.
• Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding: bidirectional encoder pretraining and downstream adaptation.
• Keras MultiHeadAttention: per-head dimensions, visibility masks, and output shape.
• Keras LayerNormalization: feature-axis normalization and learned scale/offset.
• Keras Transformer text-classification example: another encoder implementation to compare after understanding its masking and dimension choices.

The worked calculations are teaching examples. The Keras snippets illustrate an encoder’s structure and include runtime checks for a compatible environment; they do not report an executed training run, a pretrained model, or measured test accuracy.

Feynman study lab

Explain it simply

Self-attention builds a token’s representation by mixing information from accessible tokens.

Concrete example. If attention weights on two values 10 and 20 are 0.25 and 0.75, their weighted sum is 17.5.

Teach it back without looking

Close the lesson and explain the idea aloud to a beginner in two minutes. Define each technical term you use, walk through the example, and state one assumption or limitation. Then answer:

Why must a next-token predictor mask future tokens during training?

Check your explanation after you try

Otherwise it could use information unavailable when generating. A causal mask restricts attention to allowed earlier positions and the current position.

Find the gap and try again

If you hesitated, relied on a formula you could not explain, or missed the self-check, write down the exact gap. Revisit the relevant section of this lesson and work through its example. Close the notes, explain it again in simpler language, then change one number or assumption and predict what changes. Recalculate or check the lesson to test your prediction. Try once more tomorrow from memory.