---
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
}
```