PYTORCH · BUILD IT YOURSELF

Encoder-decoder transformers and teacher forcing

Introduction

For source 3, 4, 5, the target reversal task returns 5, 4, 3. Teacher forcing feeds beginning, 5, 4, 3 to the decoder while training against 5, 4, 3, end.

Learning goal: Build an encoder-decoder sequence model with correct target shifting, padding masks, cross-attention, and autoregressive evaluation.
Before you start: Transformer blocks, attention masks, special token IDs, cross-entropy, and sequence batches.

Lesson plan

Encode the source and trace how decoder cross-attention reads that representation.
Shift decoder inputs and targets correctly while handling unequal padded lengths.
Generate without answer tokens and compare teacher-forced accuracy with exact sequence accuracy.

Full build path and setup · Download transformer parts · Download training runner · Download correctness tests

A concrete sequence task

Send the source tokens [3,4,5]. Ask for the reversed sequence [5,4,3], followed by an end token. This small task makes ordering visible. It does not pretend to measure translation or language understanding.

The encoder reads the complete source. The decoder produces output one token at a time. During training we know the desired output, so we can supply its earlier tokens. During generation we must supply the model's own earlier predictions.

Build the source representation

Integer token IDs start as (B,S). An embedding lookup turns them into (B,S,C). IDs are row addresses, not numerical measurements. Token 6 is not twice as meaningful as token 3. The learned embedding row supplies its features.

The encoder passes these features through non-causal self-attention and feed-forward blocks. Every real source token may use all real source tokens. There is no future-answer leak here: the full source is already available when prediction starts.

Store the final encoder representation as memory. It has one vector per source position, not one vector for the whole sentence. The decoder can select different source positions for different output steps. Padding positions remain excluded as keys.

Build the target-side calculation

The decoder first attends to earlier target positions, including its current input position. Then it attends to encoder memory. Finally its feed-forward branch transforms each position. Each branch has a residual path. The final projection produces one score per vocabulary entry.

Self-attention gets queries, keys, and values from the target stream. Cross-attention gets queries from the target stream, but keys and values from source memory. Consequently its score matrix is (B,H,T,S). Source length S and target length T do not need to match.

Branch	Query source	Key/value source	Blocked positions
Encoder self-attention	Source	Source	Source padding
Decoder self-attention	Target	Target	Future target tokens and target padding
Decoder cross-attention	Target	Encoder memory	Source padding

Shift the target correctly

Write one row: for target tokens [BOS, A, B, EOS], decoder input is [BOS, A, B] and expected output is [A, B, EOS]. The shift is exactly one position.

Let 1 mean beginning-of-sequence and 2 mean end-of-sequence. For target [5,4,3,2], decoder input is [1,5,4,3]. The prediction at each input position is compared with the next desired token. This is teacher forcing: training uses known earlier target tokens.

import torch
target = torch.tensor([[5, 4, 3, 2]])
beginning = torch.full((1, 1), 1, dtype=torch.long)
decoder_input = torch.cat((beginning, target[:, :-1]), dim=1)
print(decoder_input.tolist())  # [[1, 5, 4, 3]]
assert decoder_input.shape == target.shape

Supplying the unshifted target allows the current position to see its own answer. A causal mask does not prevent that because the diagonal is allowed. Very low training loss can therefore be evidence of a broken task, not an excellent model.

Handle unequal lengths explicitly

The target causal mask hides later target tokens. The padding mask hides empty positions. These masks solve different problems. A causal mask alone still lets attention read padding; a padding mask alone still leaks future target answers.

A batch may contain sources of lengths 3 and 5. Pad both to length 5 with ID 0. The key mask excludes ID 0 from attention. Set the classification loss to ignore target ID 0. Ignoring padding in only one place is insufficient.

A padded query may still produce a vector. It is not a real prediction, so exclude it from metrics and loss. Ensure every real query has at least one allowed key. Reject empty sources instead of allowing an all-masked softmax row.

Generate without the answer

Encode the source once. Start the decoder input with the beginning token. Run the decoder. Read only the last position's scores. Choose a token, append it, and repeat. Stop at the end token or a configured maximum length.

Greedy decoding chooses the largest score. Sampling chooses using a probability distribution. This course uses greedy decoding for the sequence task so tests are easy to reproduce. When batching generation, remember which sequences have ended. Do not count later filler predictions as useful output.

Separate two kinds of accuracy

Teacher-forced token accuracy asks whether predictions are correct with the true previous tokens. Free-running sequence accuracy asks whether the whole generated answer is correct using predicted history. The second is harder. An early mistake can change every later input.

Test source-padding invariance by changing masked memory values while keeping the mask fixed. Test causality by changing a future target input and comparing earlier outputs. Use evaluation mode for both. These checks are more informative than inspecting a diagram.

Exercise: can the decoder attend to the last source token at its first step?

Yes, through cross-attention. The complete source is known. It cannot attend to future target tokens through causal self-attention. The two restrictions belong to different streams.

Read the executable source

These source listings are copied from the runnable files by scripts/sync-build-code.mjs. The training runner and module belong in the same folder. Read a section, run the tests, then make one small change.

TransformerEncoder — executable Python source

class TransformerEncoder(nn.Module):
    """Stack of non-causal encoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [EncoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, valid_tokens)
        return self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)

TransformerDecoder — executable Python source

class TransformerDecoder(nn.Module):
    """Stack of causal decoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [DecoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        for layer in self.layers:
            x = layer(
                x,
                memory,
                target_valid_tokens=target_valid_tokens,
                memory_valid_tokens=memory_valid_tokens,
            )
        return self.final_norm(x) * target_valid_tokens.unsqueeze(-1).to(x.dtype)

Seq2SeqTransformer — executable Python source

class Seq2SeqTransformer(nn.Module):
    """Complete encoder-decoder transformer for token IDs."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.source_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.target_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.encoder = TransformerEncoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.decoder = TransformerDecoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.output_projection = nn.Linear(d_model, vocab_size, bias=False)

    def encode(
        self, source_ids: torch.Tensor, source_valid: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if source_valid is None:
            source_valid = source_ids.ne(self.padding_id)
        _check_valid_mask(source_valid, tuple(source_ids.shape), "source_valid")
        if not bool(source_valid.any(dim=1).all()):
            raise ValueError("each source sequence needs at least one valid token")
        embedded = self.source_embedding(source_ids) * math.sqrt(self.d_model)
        return self.encoder(embedded, source_valid), source_valid

    def decode(
        self,
        target_ids: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid: torch.Tensor | None = None,
        memory_valid: torch.Tensor,
    ) -> torch.Tensor:
        if target_valid is None:
            target_valid = target_ids.ne(self.padding_id)
        embedded = self.target_embedding(target_ids) * math.sqrt(self.d_model)
        hidden = self.decoder(
            embedded,
            memory,
            target_valid_tokens=target_valid,
            memory_valid_tokens=memory_valid,
        )
        return self.output_projection(hidden)

    def forward(
        self,
        source_ids: torch.Tensor,
        target_input_ids: torch.Tensor,
        *,
        source_valid: torch.Tensor | None = None,
        target_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        memory, source_valid = self.encode(source_ids, source_valid)
        return self.decode(
            target_input_ids,
            memory,
            target_valid=target_valid,
            memory_valid=source_valid,
        )

    @torch.no_grad()
    def generate(
        self,
        source_ids: torch.Tensor,
        *,
        bos_id: int,
        eos_id: int,
        max_new_tokens: int,
        source_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Greedy decode and return generated IDs, including BOS and possible EOS."""
        self.eval()
        memory, source_valid = self.encode(source_ids, source_valid)
        generated = torch.full(
            (source_ids.shape[0], 1), bos_id, dtype=torch.long, device=source_ids.device
        )
        finished = torch.zeros(source_ids.shape[0], dtype=torch.bool, device=source_ids.device)
        for _ in range(max_new_tokens):
            logits = self.decode(generated, memory, memory_valid=source_valid)
            next_ids = logits[:, -1].argmax(dim=-1)
            next_ids = torch.where(finished, torch.full_like(next_ids, eos_id), next_ids)
            generated = torch.cat((generated, next_ids[:, None]), dim=1)
            finished |= next_ids.eq(eos_id)
            if bool(finished.all()):
                break
        return generated