anath2/tinystories-pos-embed

TinyStories toy transformer: learned vs no positional embeddings

Two variants, identical except for the positional embedding:

• learned
• none

Model

d_model	32
n_layers	4
n_heads	16
d_head	2
d_mlp	128
n_ctx	128
d_vocab	50257 (GPT-2 tokenizer)

Training

AdamW, lr 1e-3, weight decay 1e-2, batch size 32, 500 steps per "epoch" (each epoch is ~2M tokens, i.e. fresh data — no repeats). Test loss is mean next-token cross-entropy on a held-out 1000-sequence split.

checkpoint	epochs	steps	tokens	test loss
learned.pt	40	20k	80M	2.779
none.pt	40	20k	80M	2.808

Checkpoint format

Each .pt is a plain dict, loadable with weights_only=True:

{
    "cfg": {...},         # model config incl. "pos_embed": "learned" | "none"
    "state_dict": {...},  # CPU tensors
    "args": {...},        # training args (epochs, lr, seed, ...)
    "test_loss": float,   # final held-out loss
}