Attention and Transformer Architecture

Attention and Transformer Architecture explains how modern sequence models compare tokens directly instead of processing them only left to right. In this lesson, the practical outcome is concrete: you will be able to read a transformer block, identify the query, key, value, mask, residual, normalization, and feed-forward parts, and verify that a PyTorch implementation is attending over the intended positions.

This topic connects directly to the sequence portion of the Deep Learning with PyTorch course. Recurrent layers carry a hidden state through time; attention replaces that bottleneck with pairwise token interactions. That change makes long-range relationships easier to model and makes training more parallel, but it also introduces quadratic memory cost, masking rules, and tensor layout mistakes that are easy to miss.

Purpose and Outcome

Attention answers a simple question for each token: which other tokens should influence this token’s representation, and by how much? A transformer uses attention repeatedly, then combines it with position information and small fully connected networks. For text, one token may attend to the noun that determines a verb form. For time series, one day may attend to earlier abnormal readings. For code, a variable use may attend to its assignment.

The output of an attention layer is still a tensor, usually shaped like batch, sequence, embedding when batch_first=True. What changes is how each output vector is formed. Instead of using only the current input vector, attention builds a weighted mixture of value vectors from the whole allowed context.

Mechanism: Queries, Keys, Values, and Scores

Scaled dot-product attention starts with three learned projections of the input. A query represents what the current token is looking for. A key represents what each token offers for matching. A value is the information that will be mixed into the output if the match is strong. In matrix form, attention computes softmax(QK^T / sqrt(d_k))V.

The division by sqrt(d_k) matters because dot products grow as the key dimension grows. Without scaling, the softmax can become overly sharp early in training, producing tiny gradients for most positions. The softmax is applied over source positions, so each target token receives attention weights that sum to one across the positions it is allowed to see.

A transformer block usually contains multi-head self-attention, a residual connection, layer normalization, a position-wise feed-forward network, another residual connection, and another normalization. Multi-head attention splits the embedding into several smaller attention spaces. One head might learn local syntax, another might learn delimiter matching, and another might track long-distance dependencies. The heads are concatenated and projected back to the model dimension.

API Anatomy in PyTorch

PyTorch provides nn.MultiheadAttention, nn.TransformerEncoderLayer, nn.TransformerEncoder, nn.TransformerDecoderLayer, and nn.Transformer. For most custom models, start with encoder or decoder layers instead of the monolithic nn.Transformer, because explicit layers make shapes and masks easier to test.

• embed_dim: the width of each token vector. It must be divisible by num_heads.
• num_heads: the number of independent attention heads.
• batch_first=True: makes tensors use batch, sequence, embedding, which is usually easier in application code.
• attn_mask: blocks particular target-source pairs, commonly future tokens in decoder-style language modeling.
• key_padding_mask: blocks padded source positions so padding tokens do not influence real tokens.
• dim_feedforward: the hidden width of the block’s token-wise MLP.
• dropout: regularizes attention probabilities, projections, and feed-forward activations during training.

Example 1: Scaled Dot-Product Attention by Hand

This first example exposes the internal calculation. Two query vectors compare against two key vectors. Because each query best matches the key at the same position, the output is mostly the corresponding value, with a smaller contribution from the other value.

import math
import torch

q = torch.tensor([[[1.0, 0.0], [0.0, 1.0]]])
k = torch.tensor([[[1.0, 0.0], [0.0, 1.0]]])
v = torch.tensor([[[10.0, 0.0], [0.0, 20.0]]])

scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
weights = torch.softmax(scores, dim=-1)
context = weights @ v

print(torch.round(weights, decimals=4))
print(torch.round(context, decimals=4))

The expected attention weights are approximately [[0.6698, 0.3302], [0.3302, 0.6698]]. The first output vector is therefore mostly [10, 0] but includes some of [0, 20]. This is the core transformer operation before learned projections, multiple heads, masking, and residual layers are added.

Example 2: Multi-Head Self-Attention Shapes

Real transformer layers learn the query, key, and value projections. This example uses PyTorch’s module and asks it to return attention weights. The important verification is shape: the output keeps the same embedding width, while the attention matrix records target-to-source relationships.

import torch
from torch import nn

torch.manual_seed(7)
attention = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
x = torch.randn(3, 4, 8)

out, weights = attention(x, x, x, need_weights=True)

print(tuple(out.shape))
print(tuple(weights.shape))

The deterministic output shapes are (3, 4, 8) and (3, 4, 4). The batch has three sequences, each sequence has four tokens, and each token keeps eight features. The attention weights have one row per target token and one column per source token. By default, PyTorch averages weights across heads when need_weights=True; request per-head weights when you need head-level diagnostics.

Example 3: Causal and Padding Masks

Masks define what information a token is allowed to use. A causal mask prevents a position from attending to future positions. A padding mask prevents real tokens from attending to placeholder tokens introduced for batching. Mixing these up is a common cause of suspiciously good validation loss or unstable generation.

import torch
from torch import nn

layer = nn.TransformerEncoderLayer(
    d_model=6,
    nhead=2,
    dim_feedforward=12,
    batch_first=True,
)
x = torch.randn(2, 5, 6)
causal_mask = nn.Transformer.generate_square_subsequent_mask(5)
padding_mask = torch.tensor([
    [False, False, False, False, False],
    [False, False, False, True, True],
])

y = layer(x, src_mask=causal_mask, src_key_padding_mask=padding_mask)
print(tuple(y.shape))

The expected output shape is (2, 5, 6). The first sequence has no padding. In the second sequence, positions three and four are padding and are blocked as source positions. The causal mask additionally prevents every target token from using later source positions. If this were a bidirectional encoder for classification, the causal mask would usually be removed; if this were next-token prediction, the causal mask is essential.

Example 4: A Small Transformer Classifier

A transformer does not inherently know token order, so sequence models need positional information. The following classifier embeds token ids, adds learned positional embeddings, runs an encoder, pools the first token, and predicts a class. This pattern is intentionally small but mirrors larger encoder classifiers.

import torch
from torch import nn

class TinyTransformerClassifier(nn.Module):
    def __init__(self, vocab_size: int, max_len: int, classes: int) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, 16)
        self.position_embedding = nn.Embedding(max_len, 16)
        layer = nn.TransformerEncoderLayer(
            d_model=16,
            nhead=4,
            dim_feedforward=32,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.classifier = nn.Linear(16, classes)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        batch, seq_len = token_ids.shape
        positions = torch.arange(seq_len, device=token_ids.device).expand(batch, seq_len)
        x = self.token_embedding(token_ids) + self.position_embedding(positions)
        encoded = self.encoder(x)
        return self.classifier(encoded[:, 0])

model = TinyTransformerClassifier(vocab_size=100, max_len=12, classes=3)
tokens = torch.randint(0, 100, (4, 12))
logits = model(tokens)
print(tuple(logits.shape))

The expected output shape is (4, 3): four examples and three class logits per example. The model pools the first position, so training data should reserve that position for a summary token or use a pooling strategy such as mean pooling over non-padding tokens.

Design Choices and Trade-Offs

Encoder versus decoder: encoders can attend bidirectionally and are natural for classification, tagging, retrieval embeddings, and representation learning. Decoders use causal attention and are natural for next-token generation. Encoder-decoder transformers combine both for tasks where an input sequence conditions an output sequence.

Number of heads: more heads can represent more relationship types, but each head gets a smaller slice of the embedding when the model dimension is fixed. Too many heads can waste capacity or slow inference without improving validation metrics.

Context length: self-attention over a sequence of length n builds an n by n attention pattern. Longer context improves access to distant information, but memory and compute grow quickly. Batching long sequences often requires truncation, chunking, sparse attention, gradient checkpointing, or smaller batch sizes.

Position encoding: learned embeddings are simple and effective up to the maximum length seen by the table. Sinusoidal or relative methods can behave better when positions beyond the training range matter. The correct choice depends on whether absolute location or relative distance carries the task signal.

Pre-norm versus post-norm: normalization placement changes optimization behavior. Many modern stacks prefer normalization before the attention and feed-forward sublayers because it often trains deeper transformers more stably.

Failure Modes and Troubleshooting

Symptom: validation is unrealistically strong for language modeling. Cause: future tokens may be visible because the causal mask is missing or has the wrong orientation. Diagnostic steps: inspect a small attention mask, verify that blocked future positions contain negative infinity or boolean blocks as expected, and run a toy example where the answer is only available in the future. Correction: pass a correctly shaped causal mask and keep a regression test around it.

Symptom: outputs change when only padding tokens change. Cause: the model is attending to padding positions. Diagnostic steps: create two batches with identical real tokens and different pad token ids, then compare logits. Correction: pass src_key_padding_mask or pool only non-padding positions.

Symptom: runtime error says the embedding dimension is not divisible by the number of heads. Cause: each head needs an equal feature slice. Diagnostic steps: print embed_dim and num_heads at model construction. Correction: choose compatible values such as d_model=128, nhead=8 or change the projection size.

Symptom: out-of-memory occurs on longer batches. Cause: attention matrices grow with sequence length squared. Diagnostic steps: log batch size, sequence length, dtype, and layer count for the failing batch. Correction: reduce batch size, cap sequence length, use mixed precision where appropriate, or redesign the input pipeline to avoid unnecessary padding.

Security, Performance, and Reliability Implications

Transformer models can memorize sensitive training data, so datasets should be reviewed, access-controlled, and deduplicated when privacy matters. At inference time, generated text systems need output filtering and task-specific evaluation because attention improves modeling capacity but does not provide truthfulness guarantees.

Performance work should start with sequence length histograms, not model code. Padding all examples to a large global maximum wastes memory. Bucketed batching, packed datasets, and careful truncation often improve throughput more than changing the number of layers. Reliability depends on saving the tokenizer or vocabulary with the model, because token ids are only meaningful under the same preprocessing rules used during training.

Hands-On Lab: Inspect a Transformer Encoder

Prerequisites: Python, PyTorch, and a local environment where you can run a short script. A GPU is optional; the lab is small enough for CPU execution.

1. Create a new script named inspect_attention.py.
2. Paste the multi-head attention example from this lesson and run it.
3. Change need_weights=True to request non-averaged attention weights if your installed PyTorch API supports that argument.
4. Add a padding mask for one sequence and confirm that padded source positions receive no attention.
5. Add a causal mask and verify that position zero cannot attend to later positions.
6. Record the printed tensor shapes and one attention row for a fixed random seed.

Verification: the output tensor should preserve batch, sequence, embedding. Attention weights should have a source-position dimension equal to the input sequence length, and masked positions should receive zero probability after softmax. Cleanup: remove the scratch script and any temporary environment if it was created only for the lab.

Assessment Exercises

1. Given tensors shaped (8, 32, 128) with batch_first=True, explain the shape of the averaged attention weights returned by self-attention and what each axis means.
2. A decoder model trains well but generates text that appears to copy later training labels. Design a minimal test that proves whether the causal mask is working.
3. Choose d_model and num_heads for a small classifier, then justify the choice in terms of divisibility, capacity, and speed.
4. Modify the classifier example so it mean-pools only non-padding tokens. Explain why pooling over padding can change predictions.
5. For a dataset with many short sequences and a few very long ones, propose a batching strategy and describe the trade-off between context preservation and memory use.

Summary

Attention builds each token representation from weighted value vectors selected by query-key similarity. Transformers stack that operation with multiple heads, position information, residual paths, normalization, and feed-forward networks. In PyTorch, the main engineering work is keeping tensor shapes, masks, positional encodings, and pooling rules explicit. Once those are testable, transformer models become a practical extension of the sequence modeling tools used throughout this course.