Train a Tiny GPT to Copy

Train a small decoder-only transformer that learns to copy a sequence.

Each training row has 17 tokens: 8 random symbols, a separator, then the same 8 symbols again.

[7, 3, 3, 9, 2, 11, 5, 4,  1,  7, 3, 3, 9, 2, 11, 5, 4]
 <------- first half ----> sep <------- the copy ------>

Symbols use ids 2 to 11 and the separator is 1, so vocab_size is 12. The first half is random and cannot be predicted. The copy can. To learn it, the token at position t must attend back to position t - 8.

Write train_copy_model(train_tokens, vocab_size). It returns a trained nn.Module that maps a (B, T) tensor of ids to (B, T, vocab_size) next-token logits for any T up to 16. Train it on next-token prediction: the input is tokens[:, :-1] and the target is tokens[:, 1:]. Count the loss only on the copied half. Set the other targets to -100, which F.cross_entropy ignores by default.

The model must be causal: the logits at position t may depend only on tokens 0..t.

Function signature

def train_copy_model(train_tokens, vocab_size):

Example

Input

model = train_copy_model(train_tokens, 12)
prompt = torch.tensor([[7, 3, 3, 9, 2, 11, 5, 4, 1]])
greedy decoding: 8 times, append model(seq)[:, -1].argmax(-1)

Output

the 8 new tokens: [7, 3, 3, 9, 2, 11, 5, 4]

Reasoning

After the separator, each new token is the token 8 positions earlier. The model never saw this prompt in training. It has to learn the rule "look back 8 positions", and it can only do that with attention and position information.

Constraints and notes

• The tests decode 200 held-out prompts greedily and need at least 95% of the copied tokens right. They also check causality directly: changing a later token must not change an earlier logit.
• Without a causal mask, training looks perfect, because each position can see its own target. Generation then fails, because at inference time the future is not there.
• Use learned position embeddings (nn.Embedding(max_len, dim)). A model with no position information cannot tell which token is 8 steps back.
• A one or two layer model with dim = 64 learns this in a few epochs. Writing the block yourself with F.scaled_dot_product_attention(..., is_causal=True) is fast on a CPU. Return the model in eval mode.