🧠 LLM Distillation for Named Entity Recognition

Behavioral Knowledge Distillation for Fine-Grained NER: Compressing a 12B-parameter teacher model (Gemma 3 12B) into a 270M-parameter student (Gemma 3 270M) with no performance cliff, achieving strict-match F1 ~0.70 on 100+ entity types.

---

🎯 Key Results

Metric	Teacher (12B)	Distilled Student (270M)	Base Student (270M)
F1 (Two-Shot, Strict)	0.696	0.698	0.207
F1 (One-Shot, Strict)	0.677	0.667	0.224
Model Size	12B params	270M params	270M params
Inference Speed	~500ms/sample	~50ms/sample	~50ms/sample
VRAM (Inference)	~22GB	~1.5GB	~1.5GB

The Distilled Student Recovers ~97% of Teacher F1 While Being 44× Smaller and 10× Faster

---

📊 What's Inside

This repository implements a complete end-to-end behavioral fractional distillation pipeline:

1. Synthetic Data Generation: 3,000 training examples via Gemma 3 12B (zero/one/two-shot)
2. Teacher Annotation: Claude + ChatGPT provide gold-standard NER labels
3. Knowledge Distillation: KL Divergence (70%) + Cross-Entropy (30%) loss on logits
4. Comprehensive Evaluation: Strict-match + text-only metrics across 3 prompt modes
5. 100+ Fine-Grained Entities: People, orgs, locations, medicine, tech, events, and more

---

🚀 Quick Start

Prerequisites

pip install torch transformers accelerate bitsandbytes tqdm openai pandas

Generate Training Data

export OPENAI_API_KEY=sk-...
python training_data_generation.py --total 3000 --resume

Train Distilled Model

python NER_Distillation_v2.py
# Checkpoints saved to ./ner_distill_checkpoints/
# Final model saved to ./ner_distilled_model/

Evaluate

python evaluate_100_samples.py      # Normalised (with text processing)
python evaluate_100_exact.py        # Exact match (strict)

---

🔧 Architecture

┌─────────────────────────────────────────────────────────────────┐
│                    DISTILLATION PIPELINE                        │
│                                                                 │
│  Gemma 3 12B          Synthetic Sentence Generation             │
│  (Generator)   ──────▶ + Few-Shot Examples                      │
│                       (3,000 training records)                  │
│                                │                                │
│                                ▼                                │
│  Claude + GPT ──────▶ Gold-Standard NER Annotation              │
│  (Annotators)        (100+ entity types)                        │
│                                │                                │
│                                ▼                                │
│  Teacher 12B  ──────▶ Distillation Training                     │
│  (Logits KL)         KL + CE Loss (0.7 : 0.3)                   │
│                      T = 2.0 (Hinton softening)                 │
│                                │                                │
│                                ▼                                │
│  Gemma 3 270M ◀───── Fine-Tuned Student Model                  │
│  (Distilled)         (44× smaller, 10× faster)                  │
│                                │                                │
│                                ▼                                │
│  ┌──────────────────────────────────────────┐                   │
│  │ Evaluation (Strict + Text-Only Modes)     │                   │
│  │ Zero-Shot | One-Shot | Two-Shot          │                   │
│  │ P / R / F1 across 100-sample test set    │                   │
│  └──────────────────────────────────────────┘                   │
└─────────────────────────────────────────────────────────────────┘

---

📈 Performance Breakdown

Exact-Match Evaluation (Strict: Text + Type)

Two-Shot Mode (Best Performance)
┌─────────────────┬───────────┬───────────┬─────────┐
│ Model           │ Precision │ Recall    │ F1      │
├─────────────────┼───────────┼───────────┼─────────┤
│ Teacher (12B)   │  0.7131   │  0.6797   │ 0.6960  │
│ Distilled (270M)│  0.6923   │  0.7031   │ 0.6977* │
│ Base (270M)     │  0.1616   │  0.2891   │ 0.2073  │
└─────────────────┴───────────┴───────────┴─────────┘
                                          * Within 0.2% of teacher

Performance Across Modes

Mode	Teacher	Distilled	Base	Recovery
Zero-Shot	0.165	0.163	0.023	98.8%
One-Shot	0.677	0.667	0.224	98.5%
Two-Shot	0.696	0.698	0.207	100.3%

---

📁 Project Structure

.
├── 📓 Core Training
│   ├── NER_Distillation_v2.py          # Main distillation training (standalone)
│   ├── LLM_Distillation.ipynb           # Colab version
│   └── Test_Data_Generation.ipynb       # Multi-model inference
│
├── 🔧 Data Pipeline
│   ├── training_data_generation.py      # End-to-end: Gemma → GPT annotation
│   ├── generate_dataset.py              # From pre-generated sentences
│   ├── generate_fewshot_pool.py         # Few-shot example builder
│   └── assemble_dataset.py              # Merges dataset components
│
├── 🛠️ Dataset Curation
│   ├── fix_ner_dataset.py               # Label corrections + diversity
│   ├── analyze_fixed_dataset.py         # Statistics & QA
│   └── analyze_ner_quality.py           # Raw annotation analysis
│
├── 📊 Evaluation & Testing
│   ├── evaluate_100_samples.py          # Normalised matching (tolerant)
│   ├── evaluate_100_exact.py            # Exact matching (strict)
│   ├── evaluate_test_prompt.py          # Advanced fuzzy + strict
│   ├── generate_test_outputs.py         # Multi-model inference
│   ├── parse_to_csv.py                  # JSONL → prompt CSV
│   └── clean_outputs.py                 # JSON extraction from outputs
│
├── 📁 Training Data/
│   └── mixed_ner_dataset_final.jsonl    # 3,000 records (1.6 MB)
│
├── 📁 Test Data/
│   ├── Test_Data_100.csv                # 100-sample eval set (all outputs)
│   └── Test_Data_15.csv                 # 15-sample pilot set
│
└── 📄 Prompt_Structure.txt              # Example prompts

---

🎓 Pipeline Overview

Stage 1: Synthetic Data Generation

• Teacher Model: Gemma 3 12B generates diverse sentences across 12 writing styles
• Entity Hints: Dynamic selection of 1–4 entity types per sentence
• Annotation: Claude + ChatGPT perform two-turn (generate → self-review) validation
• Output: 3,000 JSONL records split equally into zero/one/two-shot

python training_data_generation.py --total 3000 --seed 42 --resume

Stage 2: Dataset Cleaning

Three-phase post-processing:

1. Hard label corrections (explicit mappings for known errors)
2. Schema normalization (canonical type enforcement)
3. Diversity maximization (domain-overlap scoring for few-shot replacement)

python fix_ner_dataset.py --input Training_Data/raw.jsonl \
                           --output Training_Data/cleaned.jsonl

Stage 3: Distillation Training

Loss Function: L = 0.7 × KL(student ∥ teacher) + 0.3 × CrossEntropy

• KL Divergence: Hinton temperature scaling (T=2.0) over top-50 logits
• Vocab Alignment: Sliced to min(teacher_vocab, student_vocab)
• Token Masking: Loss computed only over generated positions

python NER_Distillation_v2.py
# GPU Requirements: ≥24GB VRAM (e.g., A100 40GB, RTX 4090)

Stage 4: Test Data Generation

Runs inference for all three models sequentially:

python parse_to_csv.py              # JSONL → CSV prompts
python generate_test_outputs.py     # Inference (all 3 models)
python clean_outputs.py             # Extract JSON arrays

Stage 5: Evaluation

Two complementary metrics:

Normalised (text preprocessing, abbreviation expansion):

python evaluate_100_samples.py

Exact (strict character-level matching):

python evaluate_100_exact.py

---

📋 Entity Type Schema (100+)

Domain	Examples
People	PERSON, FICTIONAL_CHARACTER, DEITY, HISTORICAL_FIGURE
Organisations	COMPANY, STARTUP, UNIVERSITY, NGO, SPORT_TEAM, POLITICAL_PARTY
Geography	COUNTRY, CITY, STATE, RIVER, MOUNTAIN, OCEAN, GALAXY
Infrastructure	HOSPITAL, AIRPORT, MUSEUM, STADIUM, BRIDGE, MONUMENT
Medicine	DISEASE, DRUG, VACCINE, BODY_PART, CHEMICAL, PROTEIN, GENE
Technology	SOFTWARE, OS, AI_MODEL, PROGRAMMING_LANGUAGE, SPACECRAFT
Culture & Media	BOOK, MOVIE, SONG, ALBUM, LANGUAGE, RELIGION
Events	HISTORICAL_EVENT, WAR, SPORT_EVENT, FESTIVAL, DISASTER
Finance	CRYPTOCURRENCY, STOCK_TICKER
Temporal	DATE, TIME, MONEY, PERCENT, TEMPERATURE

---

🏋️ Training Configuration

Parameter	Value
Epochs	10
Batch Size	1
Learning Rate	5e-5 (cosine with 10% warmup)
Gradient Clipping	1.0
λ (KL weight)	0.7
KL Temperature	2.0
Top-K Logits	50
Max Sequence Length	768 tokens
Max Generation Length	256 tokens
Checkpoint Interval	Every 200 steps
Random Seed	42

---

📝 Prompt Format

All three modes follow the same structure:

Extract named entities from the sentence. Return a JSON array only.
Each item: {"entity": "exact text from sentence", "type": "ENTITY_TYPE"}
If no entities exist, return [].

Rules:
- Copy entity text exactly. Do not change it.
- Strip leading "a", "an", "the" from entity text only.
- Skip generic nouns, job titles alone, and abstract concepts.

[Examples: (one-shot / two-shot only)]
Example 1:
Sentence: """<annotated example>"""
Output: [...]

---

Sentence: """<target sentence>"""
Output:

Zero-shot records omit the Examples: block. See Prompt_Structure.txt for detailed examples.

---

💾 Requirements

torch>=2.0
transformers>=4.40
accelerate>=0.28
bitsandbytes>=0.43
tqdm
openai>=1.0
pandas

Hardware: GPU with ≥24GB VRAM for training (A100 40GB, RTX 4090, or Google Colab A100 instances recommended).

---

🔬 Experimental Details

Data Generation

• Teacher Sentences: Gemma 3 12B (temperature=0.7)
• Few-Shot Annotations: GPT-4/Claude (2-turn: generate @0.7, review @0.0)
• Ground Truth: Claude + ChatGPT (human-quality gold labels)
• Total Records: 3,000 (1,000 zero + 1,000 one + 1,000 two-shot)

Distillation Strategy

• No LoRA/PEFT: Full parameter updates for maximum knowledge transfer
• Prompt Alignment: Mode-aware masking (zero/one/two-shot)
• Logits Extraction: Teacher outputs logits + text in single forward pass
• Vocab Slicing: Both models use aligned vocabulary (50K tokens)
• Temperature Scaling: Hinton's softening (T=2.0) stabilizes KL divergence

Reproducibility Checklist

• Fixed random seed (42) across all stages
• Deterministic evaluation (no random sampling in scoring)
• Full hyperparameter documentation
• Dataset versioning (mixed_ner_dataset_final.jsonl)
• Checkpoint restoration logic (automatic best-model selection)
• Step-by-step pipeline documentation

---

🎯 When to Use This Approach

Ideal For:

• Deploying NER at edge/mobile with latency constraints (<100ms)
• Cost-sensitive inference (44× parameter reduction)
• Knowledge preservation from large proprietary models
• Multi-modal or resource-constrained environments

Trade-offs:

• Requires large teacher model during training (24GB+ VRAM)
• Synthetic data quality depends on teacher capabilities
• Fine-grained entity coverage limited by training set diversity
• Not ideal for zero-shot transfer to entirely new domains

Limitations:

1. Domain Specificity: Model trained on general English; may underperform on domain-specific text (legal, medical, code)
2. Out-of-Distribution Robustness: Distilled model may not generalize as well as teacher on unseen entity types
3. Long Sequences: Performance degrades on sequences >768 tokens
4. Rare Entity Types: Classes with <10 training examples may underperform

---

📊 Comparison with Baselines

Approach	Model Size	F1 (Two-Shot)	Inference Speed	Training Cost
Teacher Only (12B)	12B	0.696	~500ms	None
Base Student (untrained 270M)	270M	0.207	~50ms	None
Distilled Student ⭐	270M	0.698	~50ms	1× A100 (8h)
LoRA Fine-tuning (estimated)	270M	~0.620	~50ms	1× A100 (4h)

---

🚀 Usage Examples

Batch NER Inference

from transformers import pipeline

ner_pipeline = pipeline(
    "text2text-generation",
    model="./ner_distilled_model",
    device=0
)

text = "Apple's CEO Tim Cook visited Stanford University on Monday."
result = ner_pipeline(text, max_length=256)
# Returns: [{"entity": "Apple", "type": "COMPANY"}, ...]

Integrate into Production

# Load distilled model (low VRAM footprint)
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("./ner_distilled_model")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")

# Inference at scale
sentences = ["...", "...", "..."]
for sent in sentences:
    prompt = f"Extract NER: {sent}\nOutput:"
    tokens = tokenizer.encode(prompt)
    # ~50ms per sentence on CPU

---

📚 Citation

If you use this work, please cite:

@repo{intagliated2024ner_distillation,
  author       = {M.},
  title        = {{LLM Distillation for Named Entity Recognition}},
  url          = {https://github.com/intagliated/LLM-Distillation-For-NER},
  year         = {2024},
  note         = {Knowledge distillation of Gemma 3 12B into 270M for fine-grained NER}
}

---

🔗 Related Work

• Knowledge Distillation: Hinton et al. (2015), "Distilling the Knowledge in a Neural Network"
• LLM Compression: Student-teacher paradigm for language models (Sanh et al., Distil-BERT)
• NER Benchmarks: CoNLL03, OntoNotes, HateXplain, Few-NERD

---

🤝 Contributing

Contributions welcome! Areas for improvement:

• Multi-language support (extend to Tamil, Hindi, etc.)
• Domain-specific fine-tuning (medical, legal texts)
• Quantization to 4-bit for further compression
• Streaming/batched inference benchmarks
• Integration with Hugging Face Model Hub

---

📖 Troubleshooting

Q: Out of memory during training?
A: Reduce batch size further or use gradient accumulation. Adjust MAX_SEQUENCE_LENGTH in NER_Distillation_v2.py.

Q: Low F1 on my domain?
A: Regenerate training data with domain-specific prompts. Modify entity hints in training_data_generation.py.

Q: Model checkpoint not loading?
A: Ensure tokenizer and model architecture match. Re-download from ./ner_distilled_model/ or retrain.

---

📄 License

MIT License — See LICENSE file for details.

---

🙏 Acknowledgments

• Gemma Models: Google Research
• Gold Annotations: Claude (Anthropic), ChatGPT (OpenAI)
• Infrastructure: Google Colab, Hugging Face Transformers

---

📧 Contact

For questions, citations, or collaboration:

• GitHub: @intagliated
• Email: mariapault2000@gmail.com

---

Last Updated: August 2026
Status: Research artifact (actively maintained)