One block: attention and a feed-forward network

Attention scores, weights, and weighted sums

Each word is now a row of numbers — a [3 × 4] grid for “the cat sat” (rows = tokens, columns = vector entries). A block transforms that grid using attention and a feed-forward network.

Attention combines information across positions

So far each row started from a token and its position. Attention lets each word take a weighted mix of the others — learned projections determine the weights; a large weight is not a guarantee of relevance. In this decoder block, attention is the operation that mixes information across token positions.

How does it calculate the weights? Each row is projected into a query, key, and value. Dot a query with each allowed key, divide by the square root of the key width, and apply softmax to turn the scores into weights. Use those weights to combine the value rows. And it’s masked so a position can use itself and earlier positions but cannot use future tokens.

One weighted mixture, with numbers

Suppose two allowed positions have scores $0$ and $\ln3$ . Softmax turns their exponentials, 1 and 3, into weights $1/4$ and $3/4$ . If their value vectors are $(2,0)$ and $(0,4)$ , the attention output is

\frac14(2,0)+\frac34(0,4)=(0.5,3).

Queries and keys determined how much weight each position received. Values supplied what was mixed. The mixture is a vector with the same number of entries as each value vector.

The matrix notation for the same operation

For n positions, Q and K each have shape $n\times d_k$ . Their product $QK^\mathsf T$ has shape $n\times n$ : one score for each query–key pair. Divide by $\sqrt{d_k}$ , mask future positions, and apply softmax to each row. Multiplying those weights by V produces the weighted value vectors.

\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\!\left(\frac{QK^\mathsf T}{\sqrt{d_k}}+M\right)V.

The causal mask M contains zero at allowed positions and $-\infty$ at blocked positions, giving blocked positions zero weight. This is one head; multiple heads are combined with an output projection.

Optional contrast: an encoder can use later context

A useful contrast is bidirectional attention, as used in an encoder. In that setting a word can use later context too. Consider: “The chicken didn’t cross the road because it was too tired.” What is it?

Many heads at once

A block runs several attention heads in parallel, each with learned query, key, and value projections. They can represent different relationships; they are not assigned fixed linguistic jobs in advance.

see this in PyTorch

Each token builds a query, key, and value, scores how much it should attend to every earlier token, and mixes their values into its new representation.

# Single-head self-attention: every token decides how much to "look at" the others.
import torch
import torch.nn as nn
import torch.nn.functional as F

seq, d_model = 4, 8        # 4 tokens, each a vector of size 8
x = torch.randn(seq, d_model)   # x: (seq, d_model) -- our input tokens

# Each token asks a Question, offers a Key (label), and carries a Value (content).
to_q = nn.Linear(d_model, d_model)   # learns to build queries
to_k = nn.Linear(d_model, d_model)   # learns to build keys
to_v = nn.Linear(d_model, d_model)   # learns to build values

q = to_q(x)   # q: (seq, d_model) -- what each token is looking for
k = to_k(x)   # k: (seq, d_model) -- what each token offers to be matched on
v = to_v(x)   # v: (seq, d_model) -- the info each token will share

d_k = q.size(-1)                          # size of each query/key vector
scores = q @ k.mT / d_k ** 0.5            # scores: (seq, seq) -- how much each token looks at each other; divide by the square root of key width

# Causal mask: a token may only look at itself and tokens before it, never the future.
mask = torch.triu(torch.ones(seq, seq, dtype=torch.bool), diagonal=1)  # True = future
scores = scores.masked_fill(mask, float('-inf'))   # block the future with -inf

weights = F.softmax(scores, dim=-1)   # weights: (seq, seq) -- each row sums to 1 (attention %)
out = weights @ v                     # out: (seq, d_model) -- each token = weighted blend of values

Transform each row

After attention, the same feed-forward network is applied separately to each row. It transforms the contextual information already in that row; it does not read another row directly.

The shape never changes

A block takes a [3 × 4] grid and hands back a [3 × 4] grid — same shape out as in. So you can feed it into another block, and another.

Here is the route through the pre-normalized block used in this guide:

Operation	What changes	Shape for three tokens
Normalize, then attention	Mix information across allowed token positions	3 × 4 → 3 × 4
Add the original input	Add matching entries	3 × 4
Normalize, then feed-forward	Transform features within each row; widen to 8 internally	3 × 4 → 3 × 8 → 3 × 4
Add the input to that sublayer	Add matching entries again	3 × 4

The complete token-flow explorer lets you select a token and inspect every calculation. First, the next lesson explains what happens inside the feed-forward network.

see this in PyTorch

A transformer block stacks two simple moves, each added back onto the input: first the tokens look at each other (attention), then each token row is transformed separately (the feed-forward net).

# A transformer block: each token gets refined by (1) looking at other tokens,
# then (2) applying the same feed-forward function to each row. We keep both as submodules and just wire them up.
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, attention, ffn):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)  # steadies values before attention
        self.norm2 = nn.LayerNorm(d_model)  # steadies values before the ffn
        self.attention = attention          # tokens look at each other
        self.ffn = ffn                      # the same function for each row

    def forward(self, x):                   # x: (batch, seq, d_model)
        # pre-norm: normalize FIRST, then add the result back onto x (residual).
        # residual addition retains a direct path for x; later updates can still cancel information.
        x = x + self.attention(self.norm1(x))  # step 1: mix info across tokens
        x = x + self.ffn(self.norm2(x))         # step 2: transform each row
        return x                                # x: (batch, seq, d_model), same shape in = out

Next: inside the FFN.

Sources · 3

Full bibliography →