Can a Tiny Transformer Learn to Remember?

Three sub-skills the model needs

Why is $V=5$ the Optimal Choice?

Choosing a vocabulary size of $V=5$ and configuring the dataloader's epoch sizes represent a deliberate balance between computational simplicity, architectural capacity constraints, and optimization stability.

The Code

def __getitem__(self, idx):
    mapping_perm = random.choice(self.mappings)          # pick one of 96 train perms
    mapping = {x: y for x, y in zip(self.elements, mapping_perm)}

    X_order = list(self.elements)
    random.shuffle(X_order)                              # shuffle [1,2,3,4,5] → e.g. [3,4,2,1,5]

    seq = []
    for x in X_order:
        seq.append(x)
        seq.append(mapping[x])                           # interleave: [3,3, 4,5, 2,1, 1,4, 5,2]

    inputs  = torch.tensor(seq[:-1], dtype=torch.long) # tokens 0–8  (length 9)
    targets = torch.tensor(seq[1:],  dtype=torch.long) # tokens 1–9  (length 9, shifted by 1)
    mask = torch.zeros(len(inputs), dtype=torch.float32)
    mask[::2] = 1.0                                    # positions 0,2,4,6,8 → the y-prediction positions
    return inputs, targets, mask

The Raw Sequence (length 10)

Say the mapping for this example is 3→3, 4→5, 2→1, 1→4, 5→2 and the x-order is [3,4,2,1,5]. The full sequence is:

Positions 0,2,4,6,8 are x tokens (inputs/queries). Positions 1,3,5,7,9 are y tokens (answers).

Aligning inputs to targets — position by position

Let's line them up and see what the model is predicting at each position:

What is Cross-Entropy Loss?

At each $y$-prediction position, the model outputs a probability distribution over the 5 possible values. Cross-entropy measures how wrong that distribution is:

The Masking — Why We Don't Train on Every Position

# mask[::2] = 1.0 means positions 0, 2, 4, 6, 8 get mask=1 (active)
# positions 1, 3, 5, 7 get mask=0 (ignored)

loss_elementwise = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
# loss_elementwise shape: [batch*9] — one loss number per position per sample

masked_loss = (loss_elementwise * masks.view(-1)).sum()
# multiply by mask → zero out x-positions, keep only y-positions

loss = masked_loss / masks.sum()
# divide by number of active tokens (5 per sample) → average loss per y-prediction

Loss at Each Position — What a Perfect Model Would See

What is $d_{\text{model}}$ and why does it matter?

Every token is represented as a vector of $d_{\text{model}}$ numbers. Think of it as the "working memory" size for each position. Everything the model knows about a token — its identity, its position, the context it has gathered — must fit in this vector.

Quantitative Validation: The Perfect Bayesian Proof

By comparing the model's output distribution against the exact mathematical Bayesian posterior over the entire test set (5,000 samples), we can quantitatively measure the precision of its probabilistic reasoning:

Probe Design — Making Elimination and Recall Contradict Each Other

We take a real test sequence, keep the first 4 pairs intact, then replace x₅ with a repeat of x₁. The correct answer is x₁'s original y — which is already in the "used" set.

# Built programmatically from a real test sequence — no hardcoding
inputs, targets, _ = test_dataset[0]
full_seq = [inputs[0].item()] + targets.tolist()
mapping = {full_seq[i]: full_seq[i+1] for i in range(0, 10, 2)}
x_order = [full_seq[i] for i in range(0, 10, 2)]

probe_seq = []
for x in x_order[:4]:           # first 4 pairs — normal
    probe_seq.append(x)
    probe_seq.append(mapping[x])
probe_seq.append(x_order[0])       # x5 = repeat of x1
probe_seq.append(mapping[x_order[0]]) # correct answer if model recalls

The Result — d=16, 1 head, Normal Dataset

Probe sequence: [2, 4, 5, 3, 4, 5, 3, 2, 2, 4] — x=2 repeats, correct recall answer = y=4

Why This Happened — The Training Signal Was Never There

The original dataset contains only valid permutations — every x appears exactly once per sequence. The model was never penalized for ignoring (x→y) associations, because elimination always gave the right answer. There was literally zero gradient signal pushing it toward recall.

The New Dataset — One Class Change

We introduce a modified dataset where 30% of sequences have a repeated x at position 5. Everything else — the mapping split, the interleaving logic, the masking — stays identical.

class RepeatedPermutationSequenceDataset(Dataset):
    def __getitem__(self, idx):
        mapping_perm = random.choice(self.mappings)
        mapping = {x: y for x, y in zip(self.elements, mapping_perm)}
        X_order = list(self.elements)
        random.shuffle(X_order)

        if random.random() < self.repeat_ratio:   # 30% of the time:
            X_order[4] = random.choice(X_order[:4]) # x5 becomes a copy of x1–x4

        seq = []
        for x in X_order:
            seq.append(x)
            seq.append(mapping[x])   # same mapping — x still maps to the same y
        # inputs, targets, mask — identical to before

Retraining — New Model, New Dataset

We also upgrade the model: d=32, nhead=2. Everything else in the training loop is identical.

# New dataset with repeats
train_dataset = RepeatedPermutationSequenceDataset(train_mappings, V=V, num_samples=100000)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# Upgraded model
model = MinimalCausalTransformer(vocab_size=V, max_seq_len=9, d_model=32, nhead=2, num_layers=1).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
epochs = 50

# Training loop — identical to before
for epoch in range(1, epochs + 1):
    for inputs, targets, masks in train_loader:
        logits = model(inputs)
        loss_elementwise = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
        masked_loss = (loss_elementwise * masks.view(-1)).sum()
        loss = masked_loss / masks.sum()
        loss.backward()
        optimizer.step()

Memorization Probe — Task 2 Model

Same probe design: x₁=2 repeats at step 5. Correct recall answer = y=3. Only unused value = y=2.

Bayesian Accuracy — Formal Evaluation on 5000 Test Sequences

We evaluate against the theoretically perfect Bayesian distribution at each step — for both subgroups at step 5 (pure elimination vs. memorization).

What are we actually looking at?

Self-attention lets each token in the sequence ask: "which past tokens should I pay attention to?" The heatmap shows the answer. Each row is a token asking. Each column is a token being looked at. The number in each cell is the attention weight — how much the row token attends to the column token. Each row sums to 1.0.

The upper-right triangle is always zero — the causal mask prevents any token from seeing future positions. Only the lower-left triangle matters.

Why the Heads Sometimes Swap Between Runs

There is nothing in the architecture that forces Head 1 to do elimination and Head 2 to do recall. The two heads are mathematically symmetric. Across different random seeds, gradient descent sometimes assigns the roles the other way around. The functional split is always the same — only the head number changes. This is actually strong evidence that the specialization is robust: the model always converges to this two-role solution, it just doesn't care which head gets which job.

Why We Can't Be 100% Certain It's the FFN

The attention weights are directly readable — that's what the heatmaps show. The FFN's internal operation is less transparent. What we can say:

Multi-head attention discovers complementary functional roles without being told to. One head tracks what has been used. One head binds what belongs together. This emergent modularity is not designed — it is found. And it only appears when the training data actually requires it.