--- license: apache-2.0 --- # ScratchLM-95M A 95-million-parameter transformer language model built **entirely from scratch** in PyTorch — no pretrained weights, no `from_pretrained`. Trained on WikiText-103 and fine-tuned for question answering using two-stage LoRA. > Beats GPT-3 Small (PPL 26.0) on WikiText-103 with 30% fewer parameters. --- ## Model Details | Property | Value | |---|---| | **Parameters** | ~95M (88.7M base + embedding table) | | **Architecture** | 12-layer causal transformer | | **Position Encoding** | RoPE (Rotary Position Embedding) | | **FFN Activation** | SwiGLU | | **Attention** | Causal + Flash Attention fallback | | **Context Length** | 512 tokens | | **Tokenizer** | GPT-2 BPE (vocab size 50,257) | | **Framework** | PyTorch (from scratch) | | **Base Pretraining** | WikiText-103 | | **Fine-tuning Method** | Two-stage LoRA (R128 → R64) | | **Trainable Params (LoRA)** | ~1.6M (1.8% of base) | --- ## Model Configuration ```python VOCAB_SIZE = 50257 # GPT-2 tokenizer MAX_SEQ_LEN = 512 # Context window EMBED_DIM = 512 # Model dimension NUM_HEADS = 8 # Attention heads NUM_LAYERS = 12 # Transformer blocks D_FF = 2730 # SwiGLU intermediate dim DROPOUT = 0.1 ``` --- ## Performance | Checkpoint | PPL | |---|---| | Base model (WikiText-103, Epoch 17) | 24.40 | | GPT-3 Small (reference) | 26.0 | | Stage 1 LoRA — R128, α=256, Epoch 7 | 21.48 | | **Stage 2 LoRA — R64, α=128, Epoch 2** | **20.83** | - **Total PPL improvement:** 3.57 points (24.40 → 20.83) - **Train/Val gap:** 0.1216 (healthy, no overfitting) - **Fine-tune hardware:** Single RTX 5050 (8.5GB VRAM), ~6 hours --- ## Training Details ### Base Pretraining (V3) | Hyperparameter | Value | |---|---| | Dataset | WikiText-103 | | Effective batch size | 32 (8 × 4 grad accum) | | Learning rate | 2e-4 (cosine decay) | | Warmup steps | 800 | | Optimizer | AdamW (β1=0.9, β2=0.95) | | Gradient clipping | 0.8 | | Best checkpoint | Epoch 17 | ### Fine-Tuning (Two-Stage LoRA) Full fine-tuning (all 95M params) caused catastrophic forgetting — PPL rose from 24.40 to 35+. LoRA was adopted to freeze 99% of weights. **LoRA config:** ```python LORA_R = 32 # Final rank LORA_ALPHA = 64 # Scaling factor (2× rank) LORA_DROPOUT = 0.05 # Target modules: qkv_proj, out_proj ``` **Two-stage strategy:** - **Stage 1** — LoRA R=128, α=256 → best PPL 21.48 at Epoch 7 - **Stage 2** — Merge Stage 1 weights into base, restart with LoRA R=64, α=128 → best PPL **20.83** at Epoch 2 Weight merge equation: ```python W_merged = W_base + (B @ A) * (alpha / rank) ``` ### Dataset Engineering 355k clean QA pairs from three sources, after rejecting 460k noisy alternatives: | Source | Pairs | Notes | |---|---|---| | FreebaseQA / TriviaQA | 205k | Factual pairs | | ELI5 (cleaned) | 90k | Reddit markdown stripped | | SciQ | 60k | Science explanations | Cleaning steps: regex removal of filler phrases (`"the correct answer is..."`), Reddit markdown stripping, deduplication, length filtering (20–80 word answers), proper punctuation enforcement. ~22,000 entries removed. --- ## Usage ```python from transformers import AutoTokenizer, AutoModelForCausalLM import torch tokenizer = AutoTokenizer.from_pretrained("Debarun12/ScratchLM-95M") model = AutoModelForCausalLM.from_pretrained("Debarun12/ScratchLM-95M") model.eval() prompt = "What is the official language of Brazil?" inputs = tokenizer(prompt, return_tensors="pt") with torch.no_grad(): output = model.generate( **inputs, max_new_tokens=100, repetition_penalty=3.0, do_sample=False ) print(tokenizer.decode(output[0], skip_special_tokens=True)) # → "The official language is Portuguese, used in national and international contexts." ``` > **Note:** A repetition penalty of 3.0 (window 128 tokens) is strongly recommended — it is required to suppress repetition loops at this model scale. --- ## Limitations - Hallucinates facts for questions outside training data coverage (e.g., precise dates, counts) - Struggles with exact numerical facts - 95M parameter size caps final PPL around 20–22 without architectural scaling - Context window limited to 512 tokens --- ## Key Engineering Decisions **Why RoPE?** Superior relative position awareness over learned absolute embeddings; no position table to overfit. Same choice as LLaMA, Mistral, and Gemma. **Why SwiGLU?** Gated linear unit consistently outperforms GELU in practice; requires 3 linear layers (w1, w2, w3). **Why LoRA over full fine-tune?** Full fine-tune destroyed base knowledge (PPL 24.40 → 35+) due to domain shift between long encyclopedia text and short factual QA. LoRA preserves base weights while adapting ~1.8% of parameters. **Why two-stage LoRA?** Stage 1 (R128) showed widening train/val gap at Epoch 7. Merging and restarting with lower rank (R64) gave a fresh second stage without continued overfitting. --- ## Citation If you use this model or find this work helpful, please cite: ```bibtex @misc{scratchlm95m, author = {Debarun Das}, title = {ScratchLM-95M: A 95M Parameter Language Model Built From Scratch}, year = {2025}, publisher = {HuggingFace}, url = {https://huggingface.co/Debarun12/ScratchLM-95M} } ``` --- ## Author **Debarun Das** - GitHub: [github.com/Debarun12](https://github.com/Debarun12) - HuggingFace: [huggingface.co/Debarun12](https://huggingface.co/Debarun12)