Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Encoder Architecture › Building an Encoder Layer

Michael BrenndoerferUpdated June 15, 202512 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 111Encoder Architecture › Building an Encoder Layer

Reading section 5 of 13 · Building an Encoder Layer · 12 minSection 5 of 13 · 12 min

Section 5 of 13 · 12 min read

Encoder Architecture › Building an Encoder LayerOne section · Sections in this chapter5 / 13 ↓

Sections in this chapter13

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

Having understood the attention mechanism, we can now assemble a complete encoder layer. But attention alone isn't enough. Deep networks need architectural support to train successfully, and encoder layers combine several components that work together to enable learning.

An encoder layer has two major sublayers, each solving a different problem:

Multi-head self-attention enables tokens to gather information from each other. This is where bidirectional context mixing happens.
Feed-forward network (FFN) processes each token independently through a nonlinear transformation. This adds capacity that pure attention lacks.

Both sublayers are wrapped with two additional mechanisms that make deep stacking possible:

Residual connections add the input directly to the output of each sublayer. This creates "gradient highways" that allow learning signals to flow through many layers without vanishing.
Layer normalization stabilizes the distribution of activations, preventing them from exploding or collapsing as they pass through many transformations.

The modern pre-norm configuration places normalization before each sublayer rather than after. This ordering provides cleaner gradient flow and has become the standard for deep transformers. The complete computation unfolds in two stages:

Stage 1: Attention sublayer

\mathbf{h} = \mathbf{x} + \text{MultiHeadAttention}(\text{LayerNorm}(\mathbf{x}))

First, we normalize the input $\mathbf{x}$ . Then attention computes context-aware representations. Finally, we add the original input back (the residual connection). The result $\mathbf{h}$ blends the original representation with information gathered from other positions.

Stage 2: Feed-forward sublayer

\mathbf{y} = \mathbf{h} + \text{FFN}(\text{LayerNorm}(\mathbf{h}))

The same pattern repeats: normalize, transform, add residual. The FFN applies the same transformation independently to each token position, adding nonlinearity and additional learnable capacity.

Together, these stages form the encoder layer's output:

\begin{aligned} \mathbf{h} &= \mathbf{x} + \text{MultiHeadAttention}(\text{LayerNorm}(\mathbf{x})) \\ \mathbf{y} &= \mathbf{h} + \text{FFN}(\text{LayerNorm}(\mathbf{h})) \end{aligned}

where:

$\mathbf{x} \in \mathbb{R}^{n \times d}$ : input to the encoder layer, containing $n$ token representations
$\text{LayerNorm}(\cdot)$ : normalizes activations to have zero mean and unit variance per token
$\text{MultiHeadAttention}(\cdot)$ : bidirectional multi-head self-attention (no causal mask)
$\text{FFN}(\cdot)$ : position-wise feed-forward network with nonlinear activation
$\mathbf{h}$ : intermediate representation after the attention sublayer
$\mathbf{y}$ : final output of the encoder layer, same shape as input

Notice that the output $\mathbf{y}$ has the same shape as the input $\mathbf{x}$ . This dimensional consistency is what allows us to stack encoder layers: the output of one layer becomes the input to the next, with each layer progressively refining the representations.

Let's implement each component, starting with the building blocks:

In[13]:

Code

class LayerNorm:
    """Layer Normalization."""

    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.gamma = np.ones(dim)
        self.beta = np.zeros(dim)

    def __call__(self, x):
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta


class MultiHeadAttention:
    """Multi-head self-attention for encoder (bidirectional)."""

    def __init__(self, d_model, n_heads):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # Initialize projections
        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 attention.

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

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

        # Project to Q, K, V
        Q = x @ self.W_q
        K = x @ self.W_k
        V = x @ self.W_v

        # Reshape for multi-head: (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 for attention: (n_heads, seq_len, d_head)
        Q = Q.transpose(1, 0, 2)
        K = K.transpose(1, 0, 2)
        V = V.transpose(1, 0, 2)

        # Compute attention scores: (n_heads, seq_len, seq_len)
        scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(self.d_head)

        # Apply mask if provided (for padding)
        if mask is not None:
            scores = scores + mask

        # Note: No causal mask! This is bidirectional attention
        weights = softmax(scores, axis=-1)

        # Apply attention: (n_heads, seq_len, d_head)
        attended = np.matmul(weights, V)

        # Reshape back: (seq_len, d_model)
        attended = attended.transpose(1, 0, 2).reshape(seq_len, self.d_model)

        # Output projection
        output = attended @ self.W_o

        return output, weights


class FeedForward:
    """Position-wise feed-forward network."""

    def __init__(self, d_model, d_ff):
        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.b1 = np.zeros(d_ff)
        self.W2 = np.random.randn(d_ff, d_model) * scale2
        self.b2 = np.zeros(d_model)

    def __call__(self, x):
        """Apply two-layer FFN with GELU activation."""
        hidden = x @ self.W1 + self.b1
        # GELU activation
        hidden = (
            0.5
            * hidden
            * (
                1
                + np.tanh(np.sqrt(2 / np.pi) * (hidden + 0.044715 * hidden**3))
            )
        )
        return hidden @ self.W2 + self.b2

The feed-forward network uses GELU (Gaussian Error Linear Unit) activation, which has become the standard for transformer models. Unlike ReLU which hard-thresholds at zero, GELU provides smooth gating that depends on the input's magnitude:

Out[14]:

Visualization

Line plot comparing GELU and ReLU activation functions. — GELU activation compared to ReLU. GELU smoothly gates values based on their magnitude. This gives non-zero gradients for negative inputs. This helps transformers learn more nuanced representations than the hard cutoff of ReLU.

Let's visualize how layer normalization transforms the activation distributions:

Out[15]:

Visualization

Histogram of activation values before layer normalization with the mean marked by a dashed line and a callout reporting mean and standard deviation. — Activation distribution before layer normalization. The simulated activations have a non-zero mean and a standard deviation well above one, illustrating the varying location and scale that can destabilize deep networks.

Histogram of activation values after layer normalization with the mean marked by a dashed line and a callout reporting mean and standard deviation. — Activation distribution after layer normalization. Normalizing each token's representation produces an aggregate distribution centered near zero with unit standard deviation, stabilizing training dynamics.

Now we assemble these components into a complete encoder layer:

In[16]:

Code

class EncoderLayer:
    """
    Single transformer encoder layer.

    Uses pre-norm configuration with bidirectional self-attention.
    """

    def __init__(self, d_model, n_heads, d_ff):
        """
        Args:
            d_model: Model dimension
            n_heads: Number of attention heads
            d_ff: Feed-forward hidden dimension
        """
        self.d_model = d_model

        # Sublayer 1: Multi-head self-attention
        self.norm1 = LayerNorm(d_model)
        self.attention = MultiHeadAttention(d_model, n_heads)

        # Sublayer 2: Feed-forward network
        self.norm2 = LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)

    def __call__(self, x, mask=None):
        """
        Forward pass through encoder layer.

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

        Returns:
            Output of shape (seq_len, d_model)
        """
        # Sublayer 1: Attention with residual
        normed = self.norm1(x)
        attended, attn_weights = self.attention(normed, mask)
        x = x + attended  # Residual connection

        # Sublayer 2: FFN with residual
        normed = self.norm2(x)
        ffn_out = self.ffn(normed)
        x = x + ffn_out  # Residual connection

        return x, attn_weights


# Test the encoder layer
d_model = 256
n_heads = 8
d_ff = 1024
seq_len = 10

layer = EncoderLayer(d_model, n_heads, d_ff)
x = np.random.randn(seq_len, d_model)
output, attn_weights = layer(x)

Out[17]:

Console

Encoder Layer Test
----------------------------------------
Configuration:
  Model dimension (d_model): 256
  Attention heads: 8
  Head dimension: 32
  FFN dimension: 1024

Input shape: (10, 256)
Output shape: (10, 256)
Attention weights shape: (8, 10, 10)

Input norm: 51.3758
Output norm: 62.4308

The encoder layer preserves the input shape: a sequence of $n$ tokens with $d$ -dimensional representations goes in, and an identically shaped sequence comes out. The attention weights have shape (n_heads, seq_len, seq_len), showing that each head computes its own full attention matrix over all position pairs.

Visualizing Multi-Head Attention PatternsLink Copied

Different attention heads learn to focus on different aspects of the input. Some might attend to adjacent tokens (capturing local syntax), while others span long distances (capturing semantic relationships). Let's visualize how the 8 attention heads in our encoder layer distribute their attention differently:

Out[18]:

Visualization

Heatmap for attention head 1 showing a strong diagonal local-attention pattern. — Head 1 illustrates strongly local attention, with most probability concentrated on each token's diagonal self-position.

Heatmap for attention head 2 showing a mostly diagonal attention pattern with some off-diagonal weight. — Head 2 remains mostly local but distributes slightly more probability across other token positions.

Out[19]:

Visualization

Heatmap for attention head 3 showing a moderate diagonal with broader contextual weight. — Head 3 retains a visible diagonal while allocating more attention to non-local context.

Heatmap for attention head 4 showing balanced diagonal and contextual attention. — Head 4 balances local self-attention with broader contextual attention.

Out[20]:

Visualization

Heatmap for attention head 5 showing broad attention with a modest diagonal preference. — Head 5 is broadly distributed, with only a modest preference for the diagonal.

Heatmap for attention head 6 showing diffuse attention across token positions. — Head 6 illustrates diffuse attention across most token positions.

Out[21]:

Visualization

Heatmap for attention head 7 showing nearly global attention with a faint diagonal. — Head 7 approaches a nearly global pattern with only a faint diagonal preference.

Heatmap for attention head 8 showing nearly uniform global attention. — Head 8 illustrates fully global attention, spreading probability almost uniformly across token positions. Together, the eight patterns show how multiple heads can cover local and global relationships simultaneously.

The variation across heads illustrates why multi-head attention is so powerful. No single head needs to capture all relationships. The ensemble of heads provides coverage across local patterns, position-specific attention, and long-range dependencies. During training, each head specializes for patterns that help the overall task.