| --- |
| license: mit |
| library_name: pytorch |
| datasets: |
| - roneneldan/TinyStories |
| language: |
| - en |
| tags: |
| - tinystories |
| - positional-embeddings |
| - nope |
| --- |
| |
| # 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`: |
|
|
| ```python |
| { |
| "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 |
| } |
| ``` |
