Aashw1n/tinystories-22M

TinyStories 22M

A 22.7M parameter decoder-only transformer trained from scratch on TinyStories, with a custom byte-level BPE tokenizer. No transformers, no pretrained components — every layer written from primitives (Stanford CS336 Assignment 1).

Results

Validation loss	1.375 (perplexity 3.96)
Parameters	22.7M total, ~17M non-embedding
Training tokens	327.7M
Hardware	1x A100 80GB, ~40 min

Architecture

Pre-norm decoder. RMSNorm, RoPE (theta=10000), SwiGLU FFN, causal multi-head attention. 4 layers, 16 heads, d_model 512, d_ff 1344, context 256, vocab 10000.

Trained with hand-written AdamW (decoupled weight decay, bias correction), cosine LR schedule with 300-step warmup, peak LR 1e-3, global-norm gradient clipping at 1.0.

Usage

Not a transformers architecture. Load with the included model.py:

import json, torch
from safetensors.torch import load_file
from model import TransformerLM
from tokenizer import Tokenizer

cfg = json.load(open("config.json"))
model = TransformerLM(
    cfg["vocab_size"], cfg["context_length"], cfg["num_layers"],
    cfg["d_model"], cfg["num_heads"], cfg["d_ff"], cfg["theta"],
)
model.load_state_dict(load_file("model.safetensors"))
model.eval()

tok = Tokenizer.from_file("ts_vocab.json", "ts_merges.txt", ["<|endoftext|>"])

Limitations

Fluent within the TinyStories distribution — simple vocabulary, short narratives, dialogue. Loses track of entity state across paragraphs: characters swap names, actions get attributed to the wrong actor. No world knowledge, no instruction-following. It is a base LM trained on synthetic children's stories.

Sample

Prompt: Once upon a time

Once upon a time, there was a little boy named Tim. He liked to play outside in the sun. One day, he saw a small bird. The bird was sad because it could not fly. Tim wanted to help the bird.