E2E: Transformer Encoder - End-to-End Implementations

Implement E2E: Transformer Encoder in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

End-to-End Transformer Encoder

What you'll build

Every piece of the original (post-norm) Transformer encoder from "Attention Is All You Need", written as small functions and then stacked: layer norm, sinusoidal positions, scaled dot-product attention, multi-head attention, the feed-forward network, an encoder block and a multi-layer encoder.

Architecture

Step	Operation	Shape
input	`x`	`(batch, seq_len, d_model)`
positions	`x = x + positional_encoding(seq_len, d_model)`	same
per block: attention	`a = multi_head_attention(x, x, x, num_heads)`	same
per block: add & norm	`x = layer_norm(x + a, ln1_gamma, ln1_beta)`	same
per block: FFN	`f = relu(x @ W1 + b1) @ W2 + b2`	hidden `(…, d_ff)`, out same
per block: add & norm	`x = layer_norm(x + f, ln2_gamma, ln2_beta)`	same

This is post-norm (LayerNorm after the residual add), as in the 2017 paper. GPT-style models use pre-norm; see Mini GPT.

Function contract

Function	Formula / behaviour	Returns
`layer_norm(x, gamma, beta, eps=1e-5)`	`gamma · (x − mean) / sqrt(var + eps) + beta` over the last axis, population variance	array shaped like `x`
`positional_encoding(seq_len, d_model)`	`PE[p, 2i] = sin(p / 10000^(2i/d_model))`, `PE[p, 2i+1] = cos(…)`	`(seq_len, d_model)`
`scaled_dot_product_attention(Q, K, V, mask=None)`	`A = softmax(Q Kᵀ / sqrt(d_k) + mask)`; `mask` is additive (0 or −inf)	tuple `(A @ V, A)` for 3-D inputs `(batch, seq, d_k)`
`multi_head_attention(Q, K, V, num_heads)`	split `d_model` into `num_heads` heads of `d_k = d_model / num_heads`, attend per head, concatenate	output only, `(batch, seq_len, d_model)`
`feed_forward(x, W1, b1, W2, b2)`	`relu(x @ W1 + b1) @ W2 + b2`	`(…, d_model)`
`TransformerEncoderBlock(d_model, num_heads, d_ff)`	attributes below; `forward(x)` runs one post-norm block	`(batch, seq_len, d_model)`
`transformer_encoder(x, num_layers=2, d_model=64, num_heads=8, d_ff=256)`	add positions, then `num_layers` fresh blocks	`(batch, seq_len, d_model)`

Simplification: multi_head_attention has no learned projections: the heads are slices of the given Q, K, V and there is no W_O. Real layers project first (Q = x W_Q etc.); the PyTorch nn.MultiheadAttention and the Mini GPT problem do.

Block parameters (the tests use these names)

After np.random.seed(42) (in the starter), create, in this order:

ln1_gamma, ln2_gamma: ones (d_model,); ln1_beta, ln2_beta: zeros (d_model,)
W1 = np.random.randn(d_model, d_ff) * 0.02, then W2 = np.random.randn(d_ff, d_model) * 0.02; b1, b2 zeros

Because every block re-seeds, all blocks in transformer_encoder start with identical weights; that keeps the exercise deterministic.

Edge cases and conventions

Subtract the row max inside softmax (the provided softmax does).
d_model must be divisible by num_heads.

Constraints

NumPy only, no loops over batch, heads or positions (einsum or @ with transposes).

Follow-up

Attention costs O(seq_len² · d_model) time and O(seq_len²) memory per head. What do FlashAttention and sliding windows change?
Why does pre-norm train more stably than post-norm in deep stacks?

Examples

Input: layer_norm(np.array([[1.0, 2.0, 3.0, 4.0]]), np.ones(4), np.zeros(4))
Output: array([[-1.3416, -0.4472, 0.4472, 1.3416]])

Mean 2.5 and population variance 1.25, so each value becomes (x − 2.5) / 1.118

Input: positional_encoding(3, 4)[2]
Output: array([0.9093, -0.4161, 0.02, 0.9998])

Position 2: sin(2), cos(2), then the slower pair sin(2/100), cos(2/100)

Input: Q = K = [[[1, 0], [0, 1]]], V = [[[1, 2], [3, 4]]]; scaled_dot_product_attention(Q, K, V)
Output: (array([[[1.6605, 2.6605], [2.3395, 3.3395]]]), array([[[0.6698, 0.3302], [0.3302, 0.6698]]]))

Scores are I / sqrt(2), so each query puts softmax([0.707, 0]) = [0.67, 0.33] on itself and the other key

Practice in Python

import numpy as np

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


def layer_norm(x, gamma, beta, eps=1e-5):
    """Normalize over the last axis (population variance), then scale and shift."""
    # Your code here
    pass


def positional_encoding(seq_len, d_model):
    """Sinusoidal positional encoding, shape (seq_len, d_model): sin on even, cos on odd columns."""
    # Your code here
    pass


def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: (batch, seq_len, d_k). mask: optional additive mask (0 / -inf).

    Returns:
        (output, attention): (batch, seq_len, d_k) and (batch, seq_len, seq_len)
    """
    # Your code here
    pass


def multi_head_attention(Q, K, V, num_heads):
    """
    Split d_model into num_heads heads (no learned projections), attend per head,
    concatenate. Returns the output only: (batch, seq_len, d_model).
    """
    # Your code here
    pass


def feed_forward(x, W1, b1, W2, b2):
    """Position-wise feed-forward network: Linear -> ReLU -> Linear"""
    # Your code here
    pass


class TransformerEncoderBlock:
    def __init__(self, d_model: int, num_heads: int, d_ff: int):
        """
        Post-norm encoder block.

        Create: ln1_gamma, ln1_beta, ln2_gamma, ln2_beta (ones / zeros, shape (d_model,)),
        W1 = np.random.randn(d_model, d_ff) * 0.02, b1 = zeros(d_ff),
        W2 = np.random.randn(d_ff, d_model) * 0.02, b2 = zeros(d_model)  (W1 drawn before W2).
        """
        np.random.seed(42)
        self.d_model = d_model
        self.num_heads = num_heads

        # Initialize weights
        # Your code here
        pass

    def forward(self, x: np.ndarray) -> np.ndarray:
        """
        x = LN1(x + MHA(x)); x = LN2(x + FFN(x)).

        Args:
            x: Input (batch, seq_len, d_model)

        Returns:
            output: (batch, seq_len, d_model)
        """
        # Your code here
        pass


def transformer_encoder(x, num_layers=2, d_model=64, num_heads=8, d_ff=256):
    """Add positional encoding, then apply num_layers new TransformerEncoderBlocks."""
    # Your code here
    pass

The interactive editor loads with Python and NumPy. Run your implementation against 10 tests, or use the progressive hints.

More exercises in this section →