NLP · Seq2seq

Translation, summarisation and question answering map a sequence to a sequence of a different length. The encoder–decoder design (Sutskever, Vinyals and Le, 2014):

1. The encoder RNN reads the input and ends with a hidden state, a fixed-size summary.
2. The decoder RNN starts from that state and generates outputs one token at a time, feeding each prediction back in.
3. With teacher forcing, during training the decoder receives the true previous token instead of its own guess, which makes training faster and more stable.

As a toy task, I trained a GRU encoder–decoder to reverse strings of 5–10 digits. Reversal is hard for a single summary vector, because the first output depends on the last input. Accuracy drops as strings get longer. That’s the bottleneck the next post fixes.

import torch, torch.nn as nn
torch.manual_seed(0)
V, SOS, H = 12, 10, 128            # digits 0-9, SOS=10, PAD=11

def batch(n, lo=5, hi=10):
    L = torch.randint(lo, hi + 1, (1,)).item()
    x = torch.randint(0, 10, (n, L))
    return x, x.flip(1)

class Seq2Seq(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, 32)
        self.enc = nn.GRU(32, H, batch_first=True)
        self.dec = nn.GRU(32, H, batch_first=True)
        self.out = nn.Linear(H, V)
    def forward(self, x, y=None, steps=None):
        _, h = self.enc(self.emb(x))
        tok = torch.full((x.size(0), 1), SOS)
        logits = []
        for t in range(steps or y.size(1)):
            o, h = self.dec(self.emb(tok), h)
            logits.append(self.out(o))
            tok = y[:, t:t + 1] if y is not None else logits[-1].argmax(-1)   # teacher forcing
        return torch.cat(logits, 1)

m = Seq2Seq(); opt = torch.optim.Adam(m.parameters(), lr=2e-3)
for step in range(2500):
    x, y = batch(64)
    loss = nn.functional.cross_entropy(m(x, y).reshape(-1, V), y.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step()
print("final train loss:", round(loss.item(), 4))
with torch.no_grad():
    for L in [5, 7, 10]:
        x, y = batch(500, L, L)
        pred = m(x, steps=L).argmax(-1)
        print(f"length {L:>2}: exact-match {(pred == y).all(1).float().mean():.3f}, token acc {(pred == y).float().mean():.3f}")
    x, _ = batch(1, 8, 8)
    print("input", x.tolist()[0], "-> output", m(x, steps=8).argmax(-1).tolist()[0])

Output:

final train loss: 0.0608
length  5: exact-match 1.000, token acc 1.000
length  7: exact-match 0.982, token acc 0.995
length 10: exact-match 0.666, token acc 0.912
input [6, 7, 9, 5, 9, 0, 9, 0] -> output [0, 9, 0, 9, 5, 9, 7, 6]