Attention and Normalization

MultiHeadAttention: the Transformer core

MultiHeadAttention is the layer the entire Transformer family is built around. It lets every position in a sequence look at every other position and decide what's relevant — "attention" — and it runs that lookup in several independent heads at once so the model can attend to different relationships in parallel. Two args carry most of the meaning: num_heads is how many parallel attention patterns to learn, and key_dim is the width of each head's query/key space.

Call it with the same tensor for query, key, and value and you get self-attention (a sequence relating to itself); pass a different source for the key/value and you get cross-attention (the encoder-decoder link).

Normalization: what keeps deep nets trainable

Normalization layers rescale activations to a stable distribution so gradients don't explode or vanish as depth grows. The one you pick depends on what axis you normalize over:

Layer	Normalizes	Best For
`BatchNormalization`	Per batch (across samples)	CNNs, general training
`LayerNormalization`	Per sample (across features)	Transformers, RNNs
`GroupNormalization`	Per group of channels	Small batch sizes

The split matters because BatchNormalization depends on batch statistics — fragile when the batch is tiny or the sequence length varies — while LayerNormalization normalizes each sample on its own, which is why Transformers and RNNs reach for it.

# Multi-head attention (Transformer key component) layers.MultiHeadAttention( num_heads=8, # Number of attention heads key_dim=64, # Dimension of each head flash_attention=True, # Enable Flash Attention (3.7+) )

Exercise

Build a Transformer encoder block using MultiHeadAttention + LayerNormalization in pre-norm style: x → LayerNorm → MHA → +residual → LayerNorm → FFN (Dense→Dense) → +residual. Run a tiny sanity fit on random data and confirm the loss actually decreases. Then delete one LayerNorm and watch training degrade — feel why the principle callout is non-negotiable.

Hint

Pre-norm means LayerNorm goes BEFORE the sublayer, then you add the original input back as the residual. Keep key_dim × num_heads equal to your model width for the residual add to line up.

Attention and Normalization

MultiHeadAttention: the Transformer core

Normalization: what keeps deep nets trainable

Code

External links

Exercise

Progress

Comments 0