README.md

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