Encoder, decoder, or both?

Predict each target token using both the earlier target prefix and the available source.

The encoder reads x once and creates source memory. The decoder builds target-side states from the prefix. Each decoder block can then retrieve source information through cross-attention.

Our decoder block has three sublayers: causal self-attention, cross-attention, and the FFN. Each has a pre-norm input and a residual update. For the cross-attention operation, target-side states supply Q, while encoder memory supplies K and V.

Source and target can differ in length, vocabulary, or even modality. An image encoder can supply visual features to a text decoder. An audio encoder can supply acoustic representations. The source representation must fit the cross-attention projections; it need not consist of the same kind of tokens as the target.

4. Implement the decoder block

This extends the encoder-style block from lesson 6 with a second attention sublayer. Read the arguments to each call before reading the variable names. They tell you which stream is being queried.

class DecoderBlock:
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
        self.self_attn = MultiHead(width, heads, seed)
        self.cross_attn = MultiHead(width, heads, seed + 1)
        self.ffn = ffn_parameters(width, inner, seed + 2)

    def __call__(self, y: FloatArray, memory: FloatArray, target_visible: BoolArray, source_visible: BoolArray) -> FloatArray:
        normalized = layer_norm(y)
        y = y + self.self_attn(normalized, normalized, target_visible)[0]
        y = y + self.cross_attn(layer_norm(y), memory, source_visible)[0]
        return y + feed_forward(layer_norm(y), *self.ffn)

class DecoderBlock(nn.Module):
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
        super().__init__()
        self.self_attn = MultiHead(width, heads, seed)
        self.cross_attn = MultiHead(width, heads, seed + 1)
        self.norm1, self.norm2, self.norm3 = [nn.LayerNorm(width) for _ in range(3)]
        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))

    def forward(self, y: Tensor, memory: Tensor, target_visible: Tensor, source_visible: Tensor) -> Tensor:
        normalized = self.norm1(y)
        y = y + self.self_attn(normalized, normalized, target_visible)[0]
        y = y + self.cross_attn(self.norm2(y), memory, source_visible)[0]
        return y + self.ffn(self.norm3(y))

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

The memory tensor remains the encoder output throughout the decoder block. The target state y is updated after each sublayer. In our model, source memory has already passed through a final encoder normalization; the target stream uses a separate normalization before forming cross-attention queries.

Repeated decoder blocks would have their own parameters. “Stacking identical blocks” means identical structure, not necessarily shared weights. Our small model uses one block per side so its complete computation remains easy to inspect.

5. The same block structure can serve different tasks

The following forward demonstrations show an encoder-only classifier and a decoder-only next-token predictor. The latter reuses our two-sublayer block with a causal mask. Its class is called EncoderBlock because that is where we introduced it, but the mask changes the information it can access. The class name alone does not make it bidirectional.

def architecture_demo() -> tuple[FloatArray, FloatArray]:
    source = np.array([[4, 7, 9], [6, 8, 0]])
    prefix = np.array([[1, 4, 7], [1, 6, 8]])
    width, vocabulary = 32, 12
    rng = np.random.default_rng(8)
    embedding = rng.normal(size=(vocabulary, width))
    source_mask, target_mask = attention_masks(source, prefix)
    x = embedding[source] + positions(source.shape[1], width)
    y = embedding[prefix] + positions(prefix.shape[1], width)

    # Encoder-only: pool only real source positions for classification.
    encoded = layer_norm(EncoderBlock()(x, source_mask))
    real = (source != 0)[..., None]
    pooled = (encoded * real).sum(1) / real.sum(1)
    class_logits = pooled @ rng.normal(size=(width, 3))

    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
    next_token_logits = causal_hidden @ rng.normal(size=(width, vocabulary))
    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)

def architecture_demo() -> tuple[Tensor, Tensor]:
    source = torch.tensor([[4, 7, 9], [6, 8, 0]])
    prefix = torch.tensor([[1, 4, 7], [1, 6, 8]])
    width, vocabulary = 32, 12
    embedding = nn.Embedding(vocabulary, width)
    source_mask, target_mask = attention_masks(source, prefix)
    x = embedding(source) + positions(source.shape[1], width)
    y = embedding(prefix) + positions(prefix.shape[1], width)

    # Encoder-only: pool only real source positions for classification.
    encoded = layer_norm(EncoderBlock()(x, source_mask))
    real = (source != 0)[..., None]
    pooled = (encoded * real).sum(1) / real.sum(1)
    class_logits = nn.Linear(width, 3)(pooled)

    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
    next_token_logits = nn.Linear(width, vocabulary)(causal_hidden)
    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

Run architecture_demo() from either full source file. The outputs have shapes (2, 3) for three sentence classes and (2, 3, 12) for twelve vocabulary scores at each prefix position. They are random, untrained logits; the example demonstrates architecture and shape, not task performance.

The PyTorch helper creates its modules inside the function for a one-off forward demonstration. A reusable trainable model stores its modules on an nn.Module instance, as our final TinyTransformer does. Recreating layers on every training call would reset their weights and prevent learning across calls.

6. Choosing a starting point

These are starting points, not exclusive capability boundaries. For example, a decoder-only model can classify through a prompted label or a learned head, and an encoder can participate in iterative or non-autoregressive generation systems. What matters is the actual computation, objective, and inference procedure.

A useful engineering distinction is reuse. An encoder–decoder can encode the source once, then repeatedly query that memory during generation. A decoder-only model can reuse prefix keys and values through caching. Neither design removes all costs associated with long context.

7. “Decoder” does not always mean causal text generation

Outside autoregressive language modeling, the same word is used for other transformations from internal representations to task outputs. DETR ↗ , for example, uses object queries to decode a set of detections. Its object-query decoder is not simply a left-to-right text generator with a vocabulary head.

When reading another paper, inspect the query source, the masks, and the output head. Do not infer causality from the word “decoder” alone. This is especially useful when moving between language and computer vision.

Our choice for the final experiment

We will use an encoder–decoder for copying integer sequences. It gives us an explicit source/target separation and exercises all three kinds of attention from lesson 4. This is a pedagogical choice: copying could be formulated in other ways, but this version makes every component we have learned visible in one model.