DistilBERT Phishing Site Classifier (Student)

A 4-layer DistilBERT trained via knowledge distillation from a fine-tuned BERT teacher (Akash-Sakala/bert-phishing-classifier_teacher) for binary phishing site URL classification.

Model Details

Property	Value
Base model	distilbert/distilbert-base-uncased
Architecture	DistilBertForSequenceClassification
Layers	4 (distilled from 12-layer BERT teacher)
Attention heads	8
Task	Binary classification (phishing / benign)
Parameters	~52M

Training — Distillation Setup

Hyperparameter	Value
Temperature	3.0
Alpha (KL weight)	0.6
Hard label weight	0.4
Learning rate	2e-5
Batch size	64
Epochs	4
Warmup steps	10% of total steps
Weight decay	0.01
Optimizer	AdamW
Scheduler	Linear with warmup
Mixed precision	fp16 (torch.amp)

Loss Function

Combined KL divergence (soft targets) + Cross-Entropy (hard labels):

loss = alpha * KL(student_soft || teacher_soft) * T^2 + (1 - alpha) * CrossEntropy(student, labels)

Test Set Results

Model	Accuracy	Precision	Recall	F1
BERT Teacher	0.8971	0.9136	0.8763	0.8945
DistilBERT Student	0.9601	0.9710	0.9483	0.9595

The student outperforms the teacher across all metrics while being smaller and faster.

Dataset

Dataset: Akash-Sakala/phishing-site-classification
Train: 154,000 | Validation: 33,000 | Test: 33,000
Labels: 0 = benign, 1 = phishing

Usage

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained('Akash-Sakala/bert-phishing-classifier_student')
model = AutoModelForSequenceClassification.from_pretrained('Akash-Sakala/bert-phishing-classifier_student')

url = 'http://suspicious-login.verify-account.com/secure'
inputs = tokenizer(url, return_tensors='pt', truncation=True, padding='max_length')

with torch.no_grad():
    logits = model(**inputs).logits

pred = torch.argmax(logits, dim=1).item()
print('Phishing' if pred == 1 else 'Benign')

Teacher Model

Akash-Sakala/bert-phishing-classifier_teacher