anath2's picture
Update README.md
99dc748 verified
|
Raw History Blame Contribute Delete
1.2 kB
metadata
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:

{
    "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
}