---
license: mit
language: en
tags:
- transformer
- language-model
- from-scratch
- tinystories
---

# TinyStories Transformer LM (From Scratch)

A small Transformer language model, implemented entirely from scratch (no `torch.nn` layers beyond `Parameter` and containers) as part of Stanford's CS336 (Language Modeling from Scratch) assignment.

## Model Details

- **Architecture**: Decoder-only Transformer with RoPE positional embeddings, causal multi-head self-attention, SwiGLU feed-forward layers, and RMSNorm (pre-norm configuration)
- **Parameters**: ~22.7M
- **Vocabulary**: 10,000 tokens (custom-trained byte-level BPE tokenizer)
- **Context length**: 256 tokens
- **Training data**: [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) (full training set, ~540M tokens)
- **Training steps**: 5,000
- **Final validation loss**: ~1.80 (cross-entropy)

## What this model can (and can't) do

This model was trained **only** on TinyStories — a dataset of simple children's stories. It is a **base language model**, not an instruction-tuned assistant. It can generate coherent, story-like text continuations (especially when prompted with story-like openings such as "Once upon a time..."), but it has **no** capability for question-answering, factual knowledge, instruction-following, or conversation. Prompting it with anything outside the TinyStories domain will typically produce fluent but nonsensical output, since it has only ever learned to continue text in that narrow style.

## Usage

This model uses a custom architecture (not directly compatible with `transformers`' `AutoModel`). To use it, you'll need the original implementation from [this repo](https://github.com/flbjcy/LM-from-scratch).

```python
import torch, pickle, json
from cs336_basics.model import TransformerLM, decode
from cs336_basics.tokenizer import Tokenizer

with open("config.json") as f:
    config = json.load(f)
with open("vocab.pkl", "rb") as f:
    vocab = pickle.load(f)
with open("merges.pkl", "rb") as f:
    merges = pickle.load(f)

tokenizer = Tokenizer(vocab, merges, special_tokens=["<|endoftext|>"])
model = TransformerLM(**config)
model.load_state_dict(torch.load("model_weights.pt"))

output = decode(model, tokenizer, "Once upon a time", max_tokens=100, temperature=0.8, top_p=0.9)
print(output)
```

## Training

Trained from scratch on a MacBook Air (Apple Silicon, MPS), using a custom implementation of AdamW, cosine learning rate scheduling with warmup, and gradient clipping. See the [source repository](https://github.com/flbjcy/LM-from-scratch) for full implementation details.