Reproduce LLM for Storytelling: Training TinyStories (3.6M & 19.3M) from Scratch

This chapter explores the construction, training, and deployment of ultra-compact Small Language Models (SLMs) trained entirely from scratch. Utilizing the TinyStories synthetic dataset and the Hugging Face ecosystem, we demonstrate how decoder-only Transformer models with as few as $3.6\text{M}$ to $19.3\text{M}$ parameters can learn to generate fluent, syntactically correct, and logically consistent English narratives without requiring millions of dollars in compute.

1. The Small Language Model (SLM) Paradigm Shift

Historically, the prevailing consensus in deep learning held that Large Language Models (LLMs) required hundreds of millions to hundreds of billions of parameters (e.g., GPT-3, LLaMA) to achieve coherent text generation, basic reasoning, and grammatical fluency. This belief was largely driven by the nature of web-scale training data: internet scrapes (such as Common Crawl) are filled with complex syntactic structures, domain-specific jargon, non-standard formatting, and noise.

To ingest and model this chaotic distribution, networks required immense capacity ($100\text{M}+$ parameters). Consequently, training an LLM from scratch was inaccessible to individual developers and small research teams due to massive GPU infrastructure costs.

The TinyStories paradigm fundamentally shifts this assumption. If the training distribution is constrained to clean, hyper-focused, synthetic narratives using a restricted vocabulary—specifically words easily understood by a 3-to-4-year-old child—a Transformer architecture with sub-10-million parameters can achieve remarkable linguistic fluency and plot reasoning.

2. Core Mental Model: Synthetic Data & The TinyStories Paper

The foundational work behind this approach stems from the research paper "TinyStories: How Small Can Language Models Be and Still Speak Coherent English?" by Ronen Eldan and Yuanzhi Li (Microsoft Research).

The Synthetic Data Engine

Instead of filtering raw internet text, the authors used GPT-3.5 and GPT-4 to generate a massive synthetic dataset of short stories (typically 200–500 words). The generation prompts explicitly instructed the teacher models to restrict their vocabulary, grammar, and thematic complexity to concepts understandable by a typical preschooler.

The Reasoning Benchmark

Despite their small size, models trained on this clean distribution demonstrate latent reasoning capabilities. For example, consider the prompt:

"Lily likes cats and dogs. She asked her mom for a dog and her mom said no. So instead she asked..."

To complete this logically, a model must maintain a mental state of the entities and constraints:

1. Lily likes set $S = \{\text{cat}, \text{dog}\}$.
2. $\text{Request}(\text{dog}) \implies \text{Denied}$.
3. Therefore, the logical alternative is $\text{Request}(\text{cat})$ or asking another authority figure (e.g., her dad).

Experimental results across parameter scales demonstrate how reasoning improves with model size:

3. Model Architecture & Sizing Dynamics

We construct two distinct model variants from scratch using a decoder-only Transformer architecture (specifically modeled after GPTNeoForCausalLM / GPT2).

Mathematical Foundations of Model Capacity

A Transformer's parameter count is primarily dictated by its hidden dimension ($d_{\text{model}}$), the number of layers ($L$), and the vocabulary size ($V$).

The token embedding matrix contributes $V \times d_{\text{model}}$ parameters, while each Transformer layer contains:

1. Multi-Head Self-Attention (MHSA): Projection matrices $W_q, W_k, W_v, W_o \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$, contributing $4 \cdot d_{\text{model}}^2$ parameters.
2. Feed-Forward Network (FFN): Typically expands $d_{\text{model}} \to 4 \cdot d_{\text{model}} \to d_{\text{model}}$, contributing $8 \cdot d_{\text{model}}^2$ parameters.

Structural Comparison: Variant S vs. Variant M

Note: Notice that in the $3.6\text{M}$ model, a substantial portion of the parameter budget ($\approx 3.2\text{M}$ parameters out of $3.65\text{M}$) is consumed by the Token Embedding matrix ($50,257 \times 64 \approx 3.21\text{M}$). The actual attention and feed-forward transformer blocks account for under $500\text{K}$ parameters!

4. Dataset Pipeline & Multi-Process Tokenization

The dataset used is roneneldan/TinyStories, hosted on Hugging Face. It contains approximately $2.1\text{M}$ training rows and $22\text{K}$ validation rows across single-column text documents.

Raw Text Story
   |
   | Byte-Pair Encoding (BPE Tokenizer: GPT-Neo-125M)
   v
Token ID Sequence: [15496, 996, 257, 2121, 2576, 3108, ...]
   |
   | Chunked & Truncated into Context Windows
   v
Fixed Tensor Shapes: (Batch Size x 512)

Multi-Process Tokenization Implementation

To prevent CPU bottlenecks during tokenization, we execute parallelized tokenization using datasets.map() with multiple worker processes (num_proc=16).

import datasets
from transformers import AutoTokenizer

# 1. Load the dataset from Hugging Face
dataset = datasets.load_dataset("roneneldan/TinyStories")

# 2. Instantiate the pre-trained Byte-Pair Encoding (BPE) tokenizer
tokenizer_name = "EleutherAI/gpt-neo-125m"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
tokenizer.pad_token = tokenizer.eos_token

# 3. Define the tokenization mapping function
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,
        padding="max_length"
    )

# 4. Tokenize dataset in parallel across CPU cores
tokenized_datasets = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=16,
    remove_columns=["text"]
)

print(f"Train split size: {len(tokenized_datasets['train'])}")
print(f"Validation split size: {len(tokenized_datasets['validation'])}")

5. Building & Configuring the Architecture in PyTorch

To build a model entirely from scratch (uninitialized weights), we define a GPTNeoConfig object with custom structural dimensions, then initialize an AutoModelForCausalLM via from_config().

import torch
from transformers import GPTNeoConfig, AutoModelForCausalLM

# Define structural hyperparameters for the 3.6M Small Variant
config = GPTNeoConfig(
    vocab_size=len(tokenizer),            # 50,257
    max_position_embeddings=512,          # Context window (tokens)
    hidden_size=64,                       # Embedding dimension (d_model)
    num_layers=8,                         # Stacked Transformer blocks
    num_heads=16,                         # Parallel attention heads
    attention_types=[[["global"], 8]],    # Global causal attention across all 8 layers
    embed_dropout=0.1,
    attention_dropout=0.1,
    resid_dropout=0.1,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id
)

# Instantiate the uninitialized base decoder model
model = AutoModelForCausalLM.from_config(config)

# Verify parameter distribution
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total Trainable Parameters: {total_params / 1e6:.2f}M")

Untrained Model Behavior

At step $0$, the model weights are randomly initialized. If evaluated immediately, the output distribution over the vocabulary $V$ is uniform random.

The initial cross-entropy loss $\mathcal{L}_{\text{initial}}$ can be estimated mathematically:

When prompted before training, an uninitialized model outputs gibberish:

Prompt: "Once upon a time"
Output: "Once upon a time find on a his full day was diden as never on him..."

6. End-to-End Training Execution & Checkpoint Management

Training is managed via Hugging Face’s Trainer API. Because cloud platforms like Google Colab enforcement limits can interrupt free T4 GPU sessions, adopting a resilient checkpoint recovery strategy is critical.

Python Training Script

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

# Configure training hyperparameters
training_args = TrainingArguments(
    output_dir="./tiny_stories_3M_checkpoints",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    logging_steps=500,                  # Log loss metrics every 500 steps
    save_steps=1000,                    # Persist disk checkpoint every 1,000 steps
    save_total_limit=3,                 # Retain only 3 latest checkpoints
    learning_rate=5e-4,                 # Optimizer initial learning rate
    weight_decay=0.01,
    fp16=True,                          # Mixed precision execution for speed/memory
    evaluation_strategy="steps",
    eval_steps=1000,
    push_to_hub=False
)

# Standard causal language modeling data collator (dynamically handles shifting)
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False                           # Causal LM (Autoregressive), not Masked LM
)

# Instantiate Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator
)

# Kickstart training (Pass resume_from_checkpoint=True if recovering from interruption)
trainer.train()

7. Benchmark Trajectories: Loss Curves & Compute Trade-Offs

The progress of language model training is tracked via Cross-Entropy Loss ($\mathcal{L}$). A lower loss indicates that the model assigns higher probability to the correct ground-truth next token in the sequence.

Empirical Metrics & Convergence Benchmarks

1. 3.6M Parameter Variant:

   • Starting Loss: $\approx 10.30$
   • Intermediate Loss (733 steps): $\approx 9.50$ (Output is still unintelligible)
   • Final Loss ($55,000$ steps / $0.22$ epochs): $2.17$
   • Convergence Note: Matches theoretical baseline established in the paper ($2.08$ minimum for $h=64$).

2. 19.3M Parameter Variant:

   • Final Loss ($265,000$ steps / $1.0$ full epoch): $1.40$
   • Training Time: $4\text{ hours } 40\text{ minutes}$ on an Nvidia A100 GPU.
   • Convergence Note: Matches paper target of $1.38$ for $h=256$.

8. Real-World Output Quality & Qualitative

Comparisons

To evaluate performance, we perform autoregressive text generation (inference) by supplying a short prompt and sampling successive tokens.

# Inference Setup
prompt = "Once upon a time there was a girl"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Autoregressive generation
outputs = model.generate(
    **inputs,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
    top_k=50,
    top_p=0.95
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Prompt Generation Output Matrix

Prompt: "Once upon a time there was a girl"

• 3.6M Parameter Variant Output:

  "Once upon a time there was a girl called Daisy. Daisy and her mommy were walking in the park. They saw a big pond and lots of pears. Daisy asked her mommy what's the lime? Her mommy said that's very yummy..."

  • Linguistic Assessment: Grammatically plausible, fluent vocabulary, correct punctuation. However, narrative context drifts rapidly (ponds $\to$ pears $\to$ lime) due to low attention dimension capacity ($d_{\text{model}}=64$).

• 19.3M Parameter Variant Output (Run 1):

  "Once upon a time there was a girl named Jane. She was 3 years old and very curious. One day Jane asked her mom what she should do. Her mom said okay let's go outside and hang a line back. Jane and her mom grabbed their jackets and went outside. When they were outside they saw a big black tree..."

  • Linguistic Assessment: Maintains cohesive character identities, temporal order, logical cause-and-effect (asking to go outside $\implies$ grabbing jackets $\implies$ stepping outside).

• 19.3M Parameter Variant Output (Run 2 - Testing Non-Memorization):

  "Once upon a time there was a girl named Amy. She was 3 years old and loved to explore. One day she decided to go on an adventure. She put on her bright red boots and walked outside. On her way she saw a big dark bush. She was curious so she leaned towards it to have a better look. Suddenly she heard a loud bark. It was a bear..."

  • Linguistic Assessment: Demonstrates dynamic generation and narrative tension (twist/climax) rather than raw memorization of static training text.

9. Production Pitfalls, Edge Cases & Debugging

Building and deploying SLMs introduces specific operational challenges:

1. High Initial Vocabulary Parameter Overhead

• Issue: When using standard tokenizers like gpt2 ($V=50,257$) with ultra-small embedding dimensions ($d_{\text{model}}=64$), embedding parameters completely dominate the compute architecture ($>85\%$ of the total model parameter count).
• Fix: If building a domain-specific model completely independent of Hugging Face defaults, train a custom BPE tokenizer restricted to a smaller vocabulary (e.g., $V=8,000$ to $16,000$). This frees up millions of parameters for deeper self-attention layers.

2. Context Boundary Memory Leaks

• Issue: Attempting to feed inputs exceeding max_position_embeddings ($512$) results in silent positional embedding truncations or tensor out-of-bound errors.
• Fix: Ensure the tokenizer explicitly sets truncation=True and max_length=512.

3. Premature Loss Halting

• Issue: Developers often stop training early when loss drops from $10.8 \to 6.0$, assuming the model has learned language.
• Fix: As proven by inference runs, models with loss $> 4.0$ still produce non-sensical, semi-random token sequences. Strict qualitative fluency in the TinyStories domain requires pushing loss below $2.2$.

10. Complete Architectural Flow Walkthrough

The following vertical diagram illustrates the complete end-to-end data processing workflow from raw story prompt to local deployment:

11. Summary and Key Takeaways

• Data Quality Over Model Scale: Extremely small language models ($3.6\text{M}$–$19.3\text{M}$ parameters) can generate coherent, grammatically sound, and logically consistent narratives if trained on clean, targeted synthetic data like TinyStories.
• Zero-Cost Accessibility: The $3.6\text{M}$ model variant can be trained from scratch on free GPU infrastructure (such as an Nvidia T4 on Google Colab) in approximately $3\text{ hours}$.
• Embedding Matrix Dominance: In ultra-small language models, the token embedding matrix ($V \times d_{\text{model}}$) often accounts for the vast majority of total parameters.
• Loss Thresholds Matter: Cross-entropy loss must drop below $\sim 2.2$ for $3.6\text{M}$ models and $\sim 1.4$ for $19.3\text{M}$ models before coherent text generation occurs. At loss values $>4.0$, outputs remain ungrammatical gibberish.
• Non-Memorization via Sampling: By utilizing temperature and top-$p$/top-$k$ sampling during inference, small models dynamically compose novel story variations rather than repeating static text from the training set.
• Lightweight Local Hosting: Trained model artifacts range from just $35\text{ MB}$ ($3.6\text{M}$) to $250\text{ MB}$ ($19.3\text{M}$), making them capable of fast CPU-based inference and local deployment.