| --- |
| license: mit |
| language: en |
| library_name: flax |
| tags: [jax, flax-nnx, tinystories, language-model, from-scratch, rope, swiglu, rmsnorm] |
| datasets: [roneneldan/TinyStories] |
| --- |
| |
| # TinyStories-10M-JAX |
|
|
| A ~14.5M-parameter (≈6.3M non-embedding) decoder-only transformer trained **from scratch in JAX / Flax NNX** on |
| [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories), reproducing the setup of Eldan & Li (2023, |
| [arXiv:2305.07759](https://arxiv.org/abs/2305.07759)). |
|
|
| ## Results — held-out TinyStories validation (4.64M tokens) |
| | metric | value | |
| |---|---| |
| | val loss (nats/token) | **1.680** | |
| | perplexity | **5.364** | |
| | bits/token | 2.423 | |
|
|
| ## Architecture |
| Modern Llama/Mistral primitives, scaled down: |
| - d_model 256 · 6 layers · 8 heads · d_ff 1024 · context 512 |
| - **RoPE** · **RMSNorm** · **SwiGLU** FFN · **tied** input/output embeddings |
| - vocab 32,000 (byte-level BPE trained on TinyStories) |
|
|
| ## Training |
| - AdamW (β 0.9/0.95, wd 0.1), grad-clip 1.0 |
| - 1k-step warmup → cosine decay, peak LR 6e-4 |
| - 20,000 steps · batch 32 · context 512 · single Colab T4 |
|
|
| ## Usage |
| Weights are in `model.safetensors`. Reconstruct with the model code from the [GitHub |
| repo](https://github.com/Zayed024/tinystories-10m-jax) and `load_safetensors()` (see `sample.py`). Tokenizer: |
| `tokenizer.json`. |
|
|
| ## Limitations |
| Trained only on synthetic children's stories — coherent short English narratives, weak long-range consistency, no |
| factual/world knowledge. Not for general use. |
|
|
