🎓 A1Learner
🔍

PyTorch Encoder–Decoder Project

An encoder-decoder model is a transformer that reads one sequence and writes another. This project trains one on a copy task, where the right output is the input itself.

Teacher forcing shifts the target

The decoder receives the target one step early, starting with BOS. It is scored against the same target ending with EOS.

import torch
from torch import nn
PAD, BOS, EOS = 0, 1, 2
seq = [3, 5, 7]
tgt_in = [BOS] + seq
tgt_out = seq + [EOS]
print("Decoder input:", tgt_in)
print("Target output:", tgt_out)

The mask hides the future

The causal mask adds minus infinity above the diagonal. Each position can only read itself and the positions before it.

mask = nn.Transformer.generate_square_subsequent_mask(4)
print(mask)

One pass, eight scores per position

Embeddings plus positions feed the transformer. The head turns each output position into one score per vocabulary token.

torch.manual_seed(0)
embedding = nn.Embedding(8, 32, padding_idx=PAD)
position = nn.Embedding(8, 32)
transformer = nn.Transformer(d_model=32, nhead=4, num_encoder_layers=1, num_decoder_layers=1, dim_feedforward=64, dropout=0.0, batch_first=True)
head = nn.Linear(32, 8)
src = torch.tensor([[3, 5, 7]])
tgt = torch.tensor([[BOS, 3, 5, 7]])
src_h = embedding(src) + position.weight[:3]
tgt_h = embedding(tgt) + position.weight[:4]
out = transformer(src_h, tgt_h, tgt_mask=nn.Transformer.generate_square_subsequent_mask(4))
print("Logits shape:", tuple(head(out).shape))

Example

import torch
from torch import nn

PAD, BOS, EOS = 0, 1, 2
VOCAB = 8
LENGTH = 3

held = torch.tensor([[3, 5, 7], [6, 3, 4], [7, 7, 3]])
gen = torch.Generator().manual_seed(11)
pool = torch.randint(3, VOCAB, (200, LENGTH), generator=gen)
keep = [row for row in pool.tolist()
        if row not in held.tolist()]
train_seqs = torch.tensor(keep[:64])


class Copier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(VOCAB, 32, padding_idx=PAD)
        self.position = nn.Embedding(8, 32)
        self.transformer = nn.Transformer(
            d_model=32, nhead=4, num_encoder_layers=1,
            num_decoder_layers=1, dim_feedforward=64,
            dropout=0.0, batch_first=True)
        self.head = nn.Linear(32, VOCAB)

    def forward(self, src, tgt_in):
        src_h = self.embedding(src) + self.position.weight[:src.size(1)]
        tgt_h = self.embedding(tgt_in) + self.position.weight[:tgt_in.size(1)]
        mask = nn.Transformer.generate_square_subsequent_mask(tgt_in.size(1))
        out = self.transformer(src_h, tgt_h, tgt_mask=mask)
        return self.head(out)


torch.manual_seed(0)
model = Copier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
loss_fn = nn.CrossEntropyLoss()

tgt_in = torch.cat([torch.full((len(train_seqs), 1), BOS), train_seqs], dim=1)
tgt_out = torch.cat([train_seqs, torch.full((len(train_seqs), 1), EOS)], dim=1)

first_loss = None
for step in range(100):
    optimizer.zero_grad()
    logits = model(train_seqs, tgt_in)
    loss = loss_fn(logits.reshape(-1, VOCAB), tgt_out.reshape(-1))
    if first_loss is None:
        first_loss = loss.item()
    loss.backward()
    optimizer.step()

print("Training sequences:", len(train_seqs))
print("First loss:", round(first_loss, 4))
print("Final loss:", round(loss.item(), 4))


def copy_one(seq):
    out = [BOS]
    with torch.no_grad():
        for _ in range(LENGTH + 1):
            nxt = model(seq.unsqueeze(0),
                        torch.tensor([out]))[0, -1].argmax().item()
            out.append(nxt)
            if nxt == EOS:
                break
    return out[1:-1]


model.eval()
for seq in held:
    copied = copy_one(seq)
    print("Input:", seq.tolist(), "-> copied:", copied,
          "| perfect:", copied == seq.tolist())

Key points

  • BOS and EOS: the begin and end tokens that frame every target sequence
  • Teacher forcing: training a decoder on the true earlier tokens instead of its own guesses
  • The decoder input is the target shifted right, starting with the BOS token
  • Use generate_square_subsequent_mask to hide later target positions during training
  • Greedy decoding: always taking the highest scoring next token after training

In this example, the vocabulary holds three control tokens and the digits 3 to 7. The training pool is drawn with a seeded generator, and any sequence equal to a held-out one is filtered out. Sixty-four sequences of length 3 remain. The Copier class embeds the source and the shifted target, adding position vectors to both. Its nn.Transformer runs one encoder layer and one decoder layer. Teacher forcing feeds the true target, shifted right, while the causal mask hides later positions. Training runs 100 full passes, and the loss falls from 2.1207 to 0.0021. After training, each held-out sequence is copied by greedy decoding, one token at a time. All three copies match their inputs exactly.

Note

The held-out sequences never appear in the training pool, because the builder filters them out first

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does teacher forcing feed the decoder during training?