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:

The Memorization Probe Setup

To determine if the models have actually learned associative recall (matching $x \to y$ in memory) or are simply performing elimination, we construct a probe sequence on the repeated-prompt dataset. Here, the fifth $x$ token repeats $x_1$ ($x_5 = x_1$):

We trained **both** models on this repeated dataset. Yet, the architectural capacity of the model dictates whether it can solve it.

Proving Dual-Circuit Competence in the Larger Model (Model 2)

When evaluated on the test set with the repeated structure, Model 2 successfully deploys both elimination and recall at Step 5. The metrics show precise execution of both processes:

Understanding Attention Weight Routing

Each token produces Query (Q), Key (K), and Value (V) vectors. The attention weight between token $i$ and token $j$ represents the dot product of $Q_i \cdot K_j$, scaled and softmaxed. It governs the percentage of information transferred. By looking at these maps, we can mathematically unpack how the models allocate their heads.

Step 1: What Attention Does

The recall head (Head 2) creates tight local bonds between consecutive ($x, y$) pairs. Specifically:

This creates a chain: each $x$ token's post-attention representation is heavily colored by the $y$ token that preceded it. The ($x, y$) pair is encoded as a unit.

Step 2: What the FFN Does

After attention, each token passes through the Feed-Forward Network — two linear layers with a ReLU in between. The FFN is applied independently to each position (it does not mix information between positions).

Research on transformers (Geva et al., 2021: "Transformer Feed-Forward Layers Are Key-Value Memories") has shown that FFN layers act as associative memories: the first layer's weights are "keys" and the second layer's weights are "values." The FFN learns to map input patterns to output content.

The Full Recall Mechanism

Why we can't be 100% certain it's the FFN

We can observe the attention weights directly — that's what the heatmaps show. But the FFN's internal operation is less transparent. What we can say confidently:

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