---
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.