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:

The "Aha!" Moment: Zooming in on a Single Question

Let's look at the exact moment the very first pair appears. The sequence starts with x₁ = 3 and the target answer key is y₁ = 3.

Grading the Entire Sequence at Once (The Forward Pass)

Transformers don't work step-by-step in a loop. They predict all 9 outputs simultaneously. Then, the loss function grabs all 9 predictions, throws out (masks) the ones where the model was asked to predict the next random x, and averages the penalty scores for the 5 actual questions (y targets).

The Mathematical Glass Ceiling

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.

The Mechanistic Paradox

If Head 2 is responsible for "recall", it should focus its attention completely on the correct $y$ target in the past context. But the attention maps proved it isn't. Across varying distances, the attention weight on the correct target often drops to near random (~8%). Instead of cleanly retrieving the right token via Attention, the model constructs an ingenious, hardware-level rescue operation using both layers simultaneously.

Information Flow Blueprint (Case: $x_5 = x_2 = 4$, Target $y = 3$)

By dissecting the neural activity exactly between Phase 1 and Phase 2 on raw test data, we catch this rescue operation in the act. Before the FFN fires, the system's probabilities are messy, entirely incorrect vectors. A single target neuron manually steps in and saves the sequence.

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.

Experiment: V = 5

Experiment: V = 10

Experiment: V = 15