Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Transformer Block Assembly › Component Inventory

Michael BrenndoerferUpdated June 13, 20259 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 110Transformer Block Assembly › Component Inventory

Reading section 3 of 16 · Component Inventory · 9 minSection 3 of 16 · 9 min

Section 3 of 16 · 9 min read

Transformer Block Assembly › Component InventoryOne section · Sections in this chapter3 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Before diving into implementation, let's catalog the components we'll need. Each has been covered in previous chapters, but seeing them together clarifies their roles.

Normalization layers ensure stable activations by normalizing each token's representation to have consistent statistics. We need one normalization layer before each sublayer, plus optionally a final normalization after all blocks.

Multi-head self-attention enables tokens to gather information from other positions in the sequence. It projects inputs to QKV representations, computes attention weights, and produces context-aware representations.

Feed-forward network transforms each position independently through a two-layer MLP with nonlinearity. Modern architectures often use gated variants like SwiGLU for improved expressiveness.

Residual connections add each sublayer's input directly to its output, creating gradient highways that enable training of deep networks.

Let's implement each component, then assemble them into a complete block:

In[3]:

Code

import numpy as np


def softmax(x, axis=-1):
    """Numerically stable softmax."""
    x_max = np.max(x, axis=axis, keepdims=True)
    exp_x = np.exp(x - x_max)
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)


def gelu(x):
    """Gaussian Error Linear Unit activation."""
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))


def silu(x):
    """SiLU (Swish) activation function."""
    return x * (1 / (1 + np.exp(-x)))

RMSNorm ImplementationLink Copied

Modern transformers predominantly use RMSNorm over standard LayerNorm due to its computational efficiency. While LayerNorm computes both mean and variance, RMSNorm normalizes by the root mean square alone:

\text{RMSNorm}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})}

where:

$\mathbf{x} \in \mathbb{R}^d$ : the input vector (one token's representation)
$\text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon}$ : the root mean square, with $\epsilon$ for numerical stability
$\boldsymbol{\gamma} \in \mathbb{R}^d$ : learnable scale parameters (initialized to ones)
$\odot$ : element-wise multiplication

RMSNorm saves compute by skipping the mean calculation and subtraction that LayerNorm requires. Here's the implementation:

In[4]:

Code

class RMSNorm:
    """
    Root Mean Square Layer Normalization.

    Normalizes inputs by their RMS value without mean centering.
    Used in LLaMA, Mistral, and most modern LLMs.
    """

    def __init__(self, dim, eps=1e-6):
        """
        Args:
            dim: Feature dimension to normalize over
            eps: Small constant for numerical stability
        """
        self.eps = eps
        self.weight = np.ones(dim)  # Learnable scale parameter

    def __call__(self, x):
        """
        Apply RMSNorm to input.

        Args:
            x: Input tensor of shape (..., dim)

        Returns:
            Normalized tensor of the same shape
        """
        # Compute RMS along last dimension
        rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)
        # Normalize and scale
        return self.weight * (x / rms)

Multi-Head Self-AttentionLink Copied

Attention is the mechanism that enables tokens to exchange information. The core computation projects each token into queries ( $Q$ ), keys ( $K$ ), and values ( $V$ ), computes attention scores between all query-key pairs, and uses these scores to create weighted combinations of values. Multi-head attention runs this process in parallel across multiple "heads," each learning different attention patterns:

\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \mathbf{W}^O

where:

$\mathbf{X} \in \mathbb{R}^{n \times d}$ : input sequence with $n$ tokens
$h$ : number of attention heads
$\text{head}_i = \text{Attention}(\mathbf{X}\mathbf{W}_i^Q, \mathbf{X}\mathbf{W}_i^K, \mathbf{X}\mathbf{W}_i^V)$ : attention computed for head $i$
$\mathbf{W}^O \in \mathbb{R}^{d \times d}$ : output projection that combines all heads

We'll implement a clean version that handles the full attention computation:

In[5]:

Code

class MultiHeadSelfAttention:
    """
    Multi-head self-attention mechanism.

    Splits the representation into multiple heads, computes attention
    independently for each head, then concatenates and projects.
    """

    def __init__(self, d_model, n_heads):
        """
        Args:
            d_model: Model dimension
            n_heads: Number of attention heads
        """
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # Initialize projections with Xavier/Glorot scaling
        scale = np.sqrt(2.0 / (d_model + d_model))
        self.W_q = np.random.randn(d_model, d_model) * scale
        self.W_k = np.random.randn(d_model, d_model) * scale
        self.W_v = np.random.randn(d_model, d_model) * scale
        self.W_o = np.random.randn(d_model, d_model) * scale

    def __call__(self, x, mask=None):
        """
        Apply multi-head self-attention.

        Args:
            x: Input tensor of shape (seq_len, d_model)
            mask: Optional attention mask of shape (seq_len, seq_len)

        Returns:
            Output tensor of shape (seq_len, d_model)
        """
        seq_len = x.shape[0]

        # Project to Q, K, V
        Q = x @ self.W_q  # (seq_len, d_model)
        K = x @ self.W_k
        V = x @ self.W_v

        # Reshape for multi-head attention: (seq_len, n_heads, d_head)
        Q = Q.reshape(seq_len, self.n_heads, self.d_head)
        K = K.reshape(seq_len, self.n_heads, self.d_head)
        V = V.reshape(seq_len, self.n_heads, self.d_head)

        # Transpose to (n_heads, seq_len, d_head) for batched computation
        Q = Q.transpose(1, 0, 2)
        K = K.transpose(1, 0, 2)
        V = V.transpose(1, 0, 2)

        # Scaled dot-product attention
        scale = np.sqrt(self.d_head)
        scores = (
            np.matmul(Q, K.transpose(0, 2, 1)) / scale
        )  # (n_heads, seq_len, seq_len)

        # Apply mask if provided (for causal attention)
        if mask is not None:
            scores = scores + mask  # mask contains -inf for positions to ignore

        # Softmax over key dimension
        attn_weights = softmax(scores, axis=-1)

        # Apply attention to values
        context = np.matmul(attn_weights, V)  # (n_heads, seq_len, d_head)

        # Transpose back and reshape: (seq_len, d_model)
        context = context.transpose(1, 0, 2).reshape(seq_len, self.d_model)

        # Output projection
        output = context @ self.W_o

        return output

Feed-Forward Network with SwiGLULink Copied

Modern architectures use gated linear units for improved expressiveness. SwiGLU, used in LLaMA and PaLM, combines a gating mechanism with the SiLU activation. The computation applies two parallel projections, one passed through SiLU (the "gate") and one linear (the "up" projection), then multiplies them element-wise:

\text{SwiGLU}(\mathbf{x}) = (\text{SiLU}(\mathbf{x}\mathbf{W}_{\text{gate}}) \odot (\mathbf{x}\mathbf{W}_{\text{up}})) \mathbf{W}_{\text{down}}

where:

$\mathbf{x} \in \mathbb{R}^d$ : input token representation
$\mathbf{W}_{\text{gate}}, \mathbf{W}_{\text{up}} \in \mathbb{R}^{d \times d_{\text{ff}}}$ : projection matrices to the hidden dimension
$\mathbf{W}_{\text{down}} \in \mathbb{R}^{d_{\text{ff}} \times d}$ : projection back to model dimension
$\text{SiLU}(z) = z \cdot \sigma(z)$ : the Swish activation, where $\sigma$ is the sigmoid function
$\odot$ : element-wise multiplication (the "gating" operation)

The gating allows the network to selectively pass or block information. This provides more expressive power than a simple nonlinearity. To see gating in action, let's visualize how the gate signal modulates the information flow:

Out[6]:

Visualization

Heatmap showing gate values for 4 tokens across 8 hidden dimensions. — Gate values (SiLU-activated projections) range from near-zero to high values.

Heatmap showing gated output for 4 tokens across 8 hidden dimensions. — Gated output shows how gate values selectively amplify or suppress the up-projection.

The heatmaps illustrate the gating mechanism's selective nature. Notice how dimensions with near-zero gate values (blue) suppress the corresponding up-projection values, while high gate values (red) allow information to pass through. This selective filtering gives SwiGLU its expressive advantage over simple activations.

Here's the implementation:

In[7]:

Code

class SwiGLUFFN:
    """
    Feed-forward network with SwiGLU activation.

    Uses gated linear units: output = (SiLU(x @ W_gate) * (x @ W_up)) @ W_down
    This is the standard FFN in LLaMA, Mistral, and similar models.
    """

    def __init__(self, d_model, d_ff):
        """
        Args:
            d_model: Model dimension
            d_ff: Hidden dimension of the FFN
        """
        self.d_model = d_model
        self.d_ff = d_ff

        # Initialize weights
        scale_up = np.sqrt(2.0 / (d_model + d_ff))
        scale_down = np.sqrt(2.0 / (d_ff + d_model))

        self.W_gate = np.random.randn(d_model, d_ff) * scale_up
        self.W_up = np.random.randn(d_model, d_ff) * scale_up
        self.W_down = np.random.randn(d_ff, d_model) * scale_down

    def __call__(self, x):
        """
        Apply SwiGLU feed-forward transformation.

        Args:
            x: Input tensor of shape (..., d_model)

        Returns:
            Output tensor of shape (..., d_model)
        """
        # Gate and up projections
        gate = silu(x @ self.W_gate)
        up = x @ self.W_up

        # Element-wise gating
        hidden = gate * up

        # Down projection
        output = hidden @ self.W_down

        return output

We'll also implement a standard (non-gated) FFN for comparison. The standard FFN uses a simpler two-layer structure:

\text{FFN}(\mathbf{x}) = \text{activation}(\mathbf{x}\mathbf{W}_1)\mathbf{W}_2

where:

$\mathbf{W}_1 \in \mathbb{R}^{d \times d_{\text{ff}}}$ : projects from model dimension to hidden dimension
$\mathbf{W}_2 \in \mathbb{R}^{d_{\text{ff}} \times d}$ : projects back to model dimension
$\text{activation}$ : typically GELU in GPT-style models

In[8]:

Code

class StandardFFN:
    """
    Standard two-layer feed-forward network.

    Implements: FFN(x) = activation(x @ W1) @ W2
    """

    def __init__(self, d_model, d_ff, activation=gelu):
        """
        Args:
            d_model: Model dimension
            d_ff: Hidden dimension
            activation: Nonlinear activation function
        """
        self.d_model = d_model
        self.d_ff = d_ff
        self.activation = activation

        # Initialize weights
        scale1 = np.sqrt(2.0 / (d_model + d_ff))
        scale2 = np.sqrt(2.0 / (d_ff + d_model))

        self.W1 = np.random.randn(d_model, d_ff) * scale1
        self.W2 = np.random.randn(d_ff, d_model) * scale2

    def __call__(self, x):
        """Apply feed-forward transformation."""
        hidden = self.activation(x @ self.W1)
        output = hidden @ self.W2
        return output