distil_Med42_8B_Llama-3.2-1B/README.md

---
license: llama3.2
base_model: meta-llama/Llama-3.2-1B
tags:
- medical
- healthcare
- knowledge-distillation
- distillation
- llama
- llama-3.2
- lightweight
- clinical-nlp
datasets:
- MohamedAhmedAE/Med_LLaMa3_fine-tuning_dataset
language:
- en
pipeline_tag: text-generation
library_name: transformers
---

# distil_Med42_8B_Llama-3.2-1B

**A 1B-parameter medical language model distilled from [Llama3-Med42-8B](https://huggingface.co/m42-health/Llama3-Med42-8B) into a [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) student**, using the framework from the paper *"DistilLLM-Med: A Lightweight Medical Language Model through Knowledge Distillation"* (IEEE ICICIS 2025).

- 📄 **Paper**: [IEEE Xplore, document 11313220](https://ieeexplore.ieee.org/document/11313220/)
- 💻 **Training notebook**: [`Train_Distil_Med42_LLama_Transformer.ipynb`](https://github.com/Mohamed-Ahmed-Abo-El-Enen/MasterPapers/blob/main/DistilLLM-Med%20A%20Lightweight%20Medical%20Language%20Model%20through%20Knowledge%20Distillation/Training/Train_Distil_Med42_LLama_Transformer.ipynb)
- 🗂️ **Training data**: [MohamedAhmedAE/Med_LLaMa3_fine-tuning_dataset](https://huggingface.co/datasets/MohamedAhmedAE/Med_LLaMa3_fine-tuning_dataset)
- 👨‍🏫 **Teacher**: [m42-health/Llama3-Med42-8B](https://huggingface.co/m42-health/Llama3-Med42-8B) (8.03B params)
- 🎓 **Student base**: [meta-llama/Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) (1.24B params)

## About the paper

**DistilLLM-Med: A Lightweight Medical Language Model through Knowledge Distillation**
Mohamed Abo El-Enen, Sally Saad, Taymoor Nazmy — Faculty of Computer and Information Sciences, Ain Shams University, Cairo, Egypt.
Published in *2025 IEEE Twelfth International Conference on Intelligent Computing and Information Systems (ICICIS)*.
[IEEE Xplore](https://ieeexplore.ieee.org/document/11313220/)

The paper distills medical expertise from two specialized teacher models — **MedGemma-4B** (referred to as "MedGemini-4B" in the paper text, citing Sellergren et al.'s MedGemma technical report) and **Llama3-Med42-8B** — into a single lightweight **LLaMA 3.2-1B** student, using temperature-scaled KL-divergence distillation, specialty-weighted losses, and attention-map alignment. The resulting student retains **89.3%** of teacher token-level accuracy while cutting parameters by **75%**, reaching **47.7%** average accuracy on MMLU-Medical (**67.8%** of the Med42-8B teacher's 72.6%) and **59.5 tokens/sec** inference throughput.

## Model family

This checkpoint is one of four sibling distilled models trained in this project, all with a Llama-3.2-1B student:

| Model | Teacher | Student base | Notebook |
|---|---|---|---|
| [distil_med42_8B_Llama-3.2-1B-Instruct](https://huggingface.co/MohamedAhmedAE/distil_med42_8B_Llama-3.2-1B-Instruct) | Med42-8B | Llama-3.2-1B-**Instruct** | `Train_Distil_Med42_LLama_Transformer.ipynb` |
| **distil_Med42_8B_Llama-3.2-1B** (this model) | Med42-8B | Llama-3.2-1B | `Train_Distil_Med42_LLama_Transformer.ipynb` |
| [distil_llama_3_8B_Llama-3.2-1B](https://huggingface.co/MohamedAhmedAE/distil_llama_3_8B_Llama-3.2-1B) | Meta-Llama-3-8B | Llama-3.2-1B | `Distil_LLama_LLama.ipynb` |
| [distil_MedGemma_4B_Llama-3.2-1B](https://huggingface.co/MohamedAhmedAE/distil_MedGemma_4B_Llama-3.2-1B) | MedGemma-4B | Llama-3.2-1B | `Train_Distil_MedGemma_LLama.ipynb` |

## Method

Because Med42-8B and Llama-3.2-1B are both built on the Llama-3 tokenizer (128,256 shared vocabulary tokens), this run distills **directly on logits** with no vocabulary-projection layer needed (unlike the MedGemma run, which has a mismatched vocabulary).

The training objective combines four components, following the paper's Eq. (1)–(7):

1. **Temperature-scaled KL distillation** — `L_KD = α·L_CE(y, p_S) + (1-α)·L_KL(p_T^τ, p_S^τ)`, with `α = 0.7` (favoring the soft-label/KL term over the hard-label cross-entropy term).
2. **Progressive temperature scheduling** — softmax temperature `τ` decays exponentially from `τ₀ = 8` to `τ_final = 1` over the first 50,000 steps (`τ(t) = τ₀·exp(-γt)`), so the student first learns broad relationships between similar conditions before sharpening toward confident predictions.
3. **Specialty-weighted loss** — per-specialty weights `w_s` re-balance the KL loss across medical subdomains (e.g., cardiology vs. dermatology) so that no single specialty dominates training.
4. **Attention-map alignment** — an auxiliary Frobenius-norm loss `L_att` pulls the student's per-layer, per-head attention matrices toward the teacher's, so the student learns *where* the teacher looks (e.g., key symptoms in a vignette), not just *what* it outputs.

Total loss: `L_total = L_weighted + β·L_att + λ‖W_S‖²` (`β ≈ 0.1–0.3`, `λ` = weight decay).

### Training setup (from the paper)

| | |
|---|---|
| Hardware | 2× NVIDIA T4 (15GB each, ~30GB total) |
| Optimizer | AdamW, lr = 1e-6, weight decay = 0.01 |
| Batch size | 1 (+ gradient accumulation) |
| Precision | Teacher: 4-bit NF4 quantized (bitsandbytes, double quantization, fp16 compute) · Student: full fp16 |
| Temperature schedule | τ₀ = 8 → τ_final = 1, exponential decay over 50,000 steps |
| α (KD/CE balance) | 0.7 |
| Training extent | ~0.6 epoch over the unified corpus (per the paper) |
| Checkpoint on this Hub repo | Global step **80,118** (`save_every_steps = 5000`) |

Config as logged for this checkpoint on the Hub:
```json
{
  "alpha": 0.7, "temperature": 4, "accumulation_steps": 1, "hidden_dim": 1024,
  "learning_rate": 1e-06, "num_epochs": 10,
  "teacher_vocab_size": 128256, "student_vocab_size": 128256,
  "save_every_steps": 5000, "attention_distill": true, "feature_matching": true,
  "progressive_unfreezing": true, "adaptive_temperature": true, "layer_wise_distill": true,
  "attention_loss_weight": 0.05, "feature_loss_weight": 0.1,
  "min_temperature": 1.0, "max_temperature": 8.0, "unfreeze_schedule": "linear"
}
```

Full implementation (loss functions, projection layers, training loop) is in [`Train_Distil_Med42_LLama_Transformer.ipynb`](https://github.com/Mohamed-Ahmed-Abo-El-Enen/MasterPapers/blob/main/DistilLLM-Med%20A%20Lightweight%20Medical%20Language%20Model%20through%20Knowledge%20Distillation/Training/Train_Distil_Med42_LLama_Transformer.ipynb).

## Training data

The unified training corpus merges **18 established medical benchmarks** (~1.5–1.65M samples after cleaning) into a single instruction–response format, including:

- **MMLU** (medical subtasks) and **MedMCQA** — multiple-choice medical exam questions
- **PubMedQA** and **COVID-QA** — context-grounded biomedical QA
- **ChatDoctor** and **MIMIC-III** — real-world clinical dialogues and notes
- **Medical Meadow** and **MedQuAD** — structured instructional/expert-curated Q&A

Preprocessing deduplicated samples, truncated to a 1024-token limit, normalized medical abbreviations (via MeDAL mappings), and segmented long clinical notes (e.g., MIMIC-III) into coherent chunks. Benchmark **test sets were excluded from training** to keep evaluation fair. The processed/published version of this corpus is on the Hub as [Med_LLaMa3_fine-tuning_dataset](https://huggingface.co/datasets/MohamedAhmedAE/Med_LLaMa3_fine-tuning_dataset) (1.64M train / 14.4k validation rows), with an `instruction` / `input` / `context` / `output` / `choices` / `type` (`QA` / `MCQ` / `CASE`) schema.

## Results (from the paper)

The paper reports two distilled variants ("Experiment 1" and "Experiment 2") evaluated against both teachers and the untrained baseline. **The paper does not explicitly label which experiment number maps to which Hub checkpoint** — treat the numbers below as the project's overall benchmark results rather than a guaranteed one-to-one match with this specific file. Distilled "Experiment 1" is the paper's best/full-framework model.

**MMLU-Medical accuracy:**

| Task | LLaMA 3.2-1B (base) | MedGemma-4B | Llama3-Med42-8B | Distilled Exp. 1 | Distilled Exp. 2 |
|---|---|---|---|---|---|
| Anatomy | 49.6% | 23.0% | 74.1% | 51.1% | 49.6% |
| Clinical Knowledge | 35.9% | 22.6% | 75.1% | 49.1% | 44.5% |
| College Biology | 38.2% | 26.4% | 82.6% | 50.0% | 41.0% |
| College Medicine | 33.5% | 24.3% | 68.8% | 33.0% | 38.2% |
| Medical Genetics | 41.0% | 34.0% | 77.0% | 45.0% | 49.0% |
| Nutrition | 42.2% | 20.3% | 73.9% | 56.2% | 51.6% |
| Professional Medicine | 37.9% | 17.7% | 77.6% | 55.2% | 48.9% |
| Virology | 38.6% | 25.9% | 51.8% | 42.2% | 34.9% |
| **Average** | **39.6%** | **24.3%** | **72.6%** | **47.7%** | **44.7%** |

The best distilled model reaches **67.8%** of the Med42-8B teacher's accuracy and a **20.5%** relative improvement over the base LLaMA 3.2-1B.

**MedAlpaca (medical-meadow-medical-flashcards) generation quality:**

| Metric | Med42-8B | MedGemma-4B | LLaMA 3.2-1B (base) | Distilled Exp. 1 | Distilled Exp. 2 |
|---|---|---|---|---|---|
| ROUGE-1 | 0.310 | 0.178 | 0.171 | 0.212 | 0.209 |
| ROUGE-2 | 0.183 | 0.093 | 0.094 | 0.139 | 0.130 |
| ROUGE-L | 0.267 | 0.149 | 0.150 | 0.179 | 0.170 |
| BLEU | 0.071 | 0.026 | 0.026 | 0.037 | 0.034 |
| Token accuracy | 0.537 | 0.482 | 0.495 | 0.505 | 0.510 |

Average retention of the Med42-8B teacher's quality: **63.1%** (up to 76.1% on ROUGE-2).

**Compute efficiency:**

| Model | Params (B) | Memory (GB) | Speed (tok/s) | Latency (ms) |
|---|---|---|---|---|
| Llama3-Med42-8B (teacher) | 8.03 | 5.32 | 49.7 | 1006.0 |
| MedGemma-4B (teacher) | 4.97 | 9.75 | 33.6 | 1486.6 |
| LLaMA 3.2-1B (base) | 1.24 | 7.68 | 102.7 | 487.0 |
| **Distilled LLaMA 3.2-1B** | **1.24** | **7.68** | **59.5** | **841.0** |

**Ablation study** (contribution of each component, MMLU-Medical average):

| Variant | MMLU Acc. | Δ vs. base |
|---|---|---|
| LLaMA 3.2-1B (base) | 39.6% | – |
| A: Baseline KD | 42.3% | +2.7 |
| B: A + Progressive Temperature | 44.7% | +5.1 |
| C: B + Specialty Weighting | 46.5% | +6.9 |
| D: C + Attention Alignment (full model) | 47.7% | +8.1 |

Independent medical-expert review of 89 model answers found most responses (70/89) more detailed and context-rich than reference answers, but flagged 19/89 with critical factual errors (misdiagnoses, incorrect mechanisms) — underscoring that the model **requires expert oversight** and is not a standalone diagnostic tool.

## Usage

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "MohamedAhmedAE/distil_Med42_8B_Llama-3.2-1B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
)

prompt = "What are some possible causes of low PTH and high calcium levels?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

> This is a **base/completion-style** checkpoint (no chat template registered on the Hub). For chat-style interaction, see [distil_med42_8B_Llama-3.2-1B-Instruct](https://huggingface.co/MohamedAhmedAE/distil_med42_8B_Llama-3.2-1B-Instruct).

## Intended use & limitations

- Research use in efficient/lightweight medical NLP and knowledge-distillation studies; suitable for edge/low-resource deployment experiments.
- **Not a certified clinical tool.** The paper's own expert review found critical errors in ~21% of sampled answers (misdiagnoses, incorrect mechanisms). Always require qualified human oversight before acting on any output.
- Distillation training only reached ~0.6 epoch over the corpus — treat as a research checkpoint, not a fully converged production model.
- Inherits any biases present in the teacher model (Med42-8B) and in the 18-source training corpus.

## License

Released under the **Llama 3.2 Community License** (matches the student base, [meta-llama/Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B)). Since Med42-8B was used as a distillation teacher, also review the license terms on [m42-health/Llama3-Med42-8B](https://huggingface.co/m42-health/Llama3-Med42-8B) (Llama 3 Community License family) before downstream/commercial use.

## Citation

```bibtex
@inproceedings{aboelenein2025distilllmmed,
  title     = {DistilLLM-Med: A Lightweight Medical Language Model through Knowledge Distillation},
  author    = {Abo El-Enen, Mohamed and Saad, Sally and Nazmy, Taymoor},
  booktitle = {2025 IEEE Twelfth International Conference on Intelligent Computing and Information Systems (ICICIS)},
  year      = {2025},
  publisher = {IEEE},
  url       = {https://ieeexplore.ieee.org/document/11313220/}
}
```