File size: 1,195 Bytes
4632b87 c467ee0 4632b87 c467ee0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | ---
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
}
``` |