Kahibaro
Discord Login Register

11.3 Transformer Building Blocks

High Level Picture of a Transformer Block

A Transformer is built by stacking many similar blocks that keep the same hidden size throughout the stack. Each block takes a sequence of vectors and returns a sequence of vectors with the same shape. If the input is a tensor $X$ with shape $(B, T, d_{model})$, where $B$ is batch size, $T$ is sequence length, and $d_{model}$ is the model hidden size, then every major sublayer in a standard Transformer encoder block is designed to output the same $(B, T, d_{model})$ shape so that blocks can be composed repeatedly.

A modern “pre norm” encoder style block can be summarized as two residual sublayers: self attention, then a position wise feedforward network. Each is wrapped with layer normalization and a skip connection.

Token and Position Representations

Transformers operate on token representations, not raw token ids. A token id sequence of shape $(B, T)$ is converted into dense vectors using an embedding table, producing $(B, T, d_{model})$. Because self attention itself does not encode order, a positional signal is added to token embeddings. In the classic design this is done by adding a position embedding to each token vector so the model can distinguish the same token appearing in different locations.

At this level the important building block detail is that after embedding and adding positional information, the model works in a continuous vector space with a fixed hidden size $d_{model}$.

Multi Head Self Attention

Self attention is the core sublayer that lets each position mix information from other positions. It is implemented by projecting the input $X$ into three learned views called queries, keys, and values. For one head, the projections are

$$
Q = XW_Q,\quad K = XW_K,\quad V = XW_V
$$

where $W_Q, W_K, W_V$ are learned matrices. Shapes are commonly arranged so that $Q, K, V$ become $(B, h, T, d_{head})$ for $h$ heads, with $d_{head} = d_{model}/h$.

Attention weights are computed using scaled dot products between queries and keys:

$$
A = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}}\right)
$$

and the output for a head is

$$
H = AV
$$

The outputs from all heads are concatenated and then projected back to $d_{model}$ with another learned matrix $W_O$. The result has shape $(B, T, d_{model})$, matching the block input.

Important rule: in multi head attention you must have $d_{model}$ divisible by the number of heads $h$, because $d_{head} = d_{model}/h$ is used for each head.

Important formula: scaled dot product attention uses the factor $\sqrt{d_{head}}$ in the denominator,
$$
\text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}}\right),
$$
to keep logits in a reasonable range as dimensions grow.

Attention Masks

Attention is usually paired with a mask that prevents attending to certain positions. In an encoder, padding tokens are masked so they do not contribute. In an autoregressive decoder, a causal mask prevents a token from looking at future positions. Conceptually, masking is applied by adding a large negative value to disallowed attention logits before the softmax, so their probabilities become near zero.

Important rule: masks are applied before the softmax. If you apply masking after the softmax, probabilities will not renormalize correctly.

Residual Connections and Layer Normalization

Transformers rely heavily on residual connections. After a sublayer computes some transformation $f(\cdot)$, the block adds the original input back:

$$
Y = X + f(X)
$$

This helps gradients flow through deep stacks.

Layer normalization stabilizes training by normalizing each token vector across its feature dimension. In the widely used pre norm design, normalization is applied before each sublayer:

$$
Y = X + \text{Sublayer}(\text{LN}(X))
$$

There is also a post norm design where layer norm is applied after the residual addition, but beginners will most often encounter pre norm in modern implementations.

Important rule: residual additions require identical shapes. Every sublayer inside a residual path must output $(B, T, d_{model})$ if the input is $(B, T, d_{model})$.

Dropout in Transformer Blocks

Dropout is typically used in multiple places, including on attention weights or attention outputs, and in the feedforward network. During training, dropout randomly zeros some elements and scales the remainder, during evaluation it is disabled. The main building block detail is placement: dropout is usually applied inside the sublayer, then the result is added via the residual connection.

Position Wise Feedforward Network

After self attention, each token position is processed independently by a small fully connected network, often called the MLP or FFN. It expands the hidden size to a larger intermediate dimension $d_{ff}$, applies a nonlinearity, then projects back to $d_{model}$:

$$
\text{FFN}(x) = W_2\,\sigma(W_1 x + b_1) + b_2
$$

where $W_1$ maps $d_{model}\to d_{ff}$ and $W_2$ maps $d_{ff}\to d_{model}$. The same FFN weights are used for every time step, which is why it is called position wise.

Common activations include ReLU and GELU, with GELU frequently used in Transformer language models.

Putting the Encoder Block Together

A standard pre norm Transformer encoder block can be written as

$$
X_1 = X + \text{MHA}(\text{LN}(X))
$$
$$
X_2 = X_1 + \text{FFN}(\text{LN}(X_1))
$$

and $X_2$ is passed to the next block. Stacking $N$ such blocks produces a deep model while keeping the representation size fixed.

Decoder Specific Building Blocks (When Present)

If the model has a decoder, a decoder block typically includes an extra attention sublayer. It first performs masked self attention over the decoder sequence, then cross attention where decoder queries attend to encoder outputs as keys and values, then the FFN. Cross attention uses the same scaled dot product mechanism, but $Q$ comes from the decoder hidden states and $K, V$ come from the encoder outputs.

The key building block idea is that encoder decoder attention is just attention with different sources for $Q$ versus $K, V$.

Shape Checklist for Implementation

In practice, most bugs in Transformer blocks come from shape mismatches and incorrect transposes. The conceptual shape flow is: start with $(B, T, d_{model})$, project into heads to get $(B, h, T, d_{head})$, compute attention weights of shape $(B, h, T, T)$, multiply by values to return to $(B, h, T, d_{head})$, then merge heads back to $(B, T, d_{model})$.

Important rule: attention weights have a square sequence dimension $(T, T)$. If this is not true, your query and key sequence dimensions are not aligned, which often means you mixed up transposes or batch first ordering.

Views: 58

Comments

Please login to add a comment.

Don't have an account? Register now!