Create README.md

9483064 verified 4 months ago

5.52 kB

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

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