ModernBERT DGA Classifier (Edge-Optimized)
This repository contains a fine-tuned sequence classification model based on
answerdotai/ModernBERT-base, specifically trained to detect Domain Generation Algorithms (DGA) in network traffic.
The model is optimized for real-time binary classification of domain names (Benign vs. Malicious) and is structurally designed to be exported to ONNX for lightweight execution on ARM-based edge hardware, such as a Raspberry Pi acting as a secure DNS sinkhole.
Model Details
- Architecture: ModernBERT (Encoder-only)
- Task: Binary Sequence Classification
Label 0: Benign DomainLabel 1: Malicious (DGA) Domain
- Maximum Sequence Length: 64 tokens (Dynamic Padding applied per batch)
- Primary Language: English / Agnostic (Domain Strings)
Intended Use
This model is intended to be integrated into network security appliances, specifically localized DNS sinkholes. By analyzing incoming DNS queries in real-time, the model categorizes requests.
Critical considerations for this deployment:
In a DNS sinkhole environment, False Positives (FPR) are highly disruptive, as they block legitimate web traffic. The training configuration and evaluation metrics were heavily weighted toward optimizing Precision and F1 scores to ensure a near-zero False Positive Rate.
Training Data
The model was fine-tuned using high-quality OSINT threat intelligence feeds (such as the Harpomaxx DGA dataset format), compiling domain strings into two primary classes:
- Benign: Sourced from top-ranking lists (e.g., Tranco).
- Malicious: Sourced from known DGA feeds (e.g., Bambenek / UMUDGA).
Training Procedure & Hardware
The model was trained using native Hugging Face libraries (transformers, datasets, Trainer) with aggressive hardware optimization for consumer-grade GPUs.
Compute Environment
- Hardware: NVIDIA GeForce RTX 4080 Laptop GPU (12GB VRAM)
- Optimization Engine: PyTorch with Native AMP (Automatic Mixed Precision)
Hyperparameters
Based on hardware stress-testing, the following hyperparameter configuration was utilized to maximize VRAM efficiency while maintaining mathematical stability:
- Train Batch Size: 128 (dynamically scaled for 12GB VRAM limits)
- Eval Batch Size: 128
- Learning Rate:
3e-5
- Weight Decay:
0.01
- Warmup Ratio:
0.1
- Precision:
FP16 (Tensor Core accelerated)
- Epochs: 2 (to prevent overfitting on sequential letter patterns)
- Best Model Selection: Based on highest
F1 score.
Evaluation Metrics
During evaluation, a custom metric computation block was utilized to track metrics beyond standard accuracy. The model was evaluated on:
- F1-Score (Primary optimization target)
- Precision & Recall
- False Positive Rate (FPR)
(Note: Add your final printed metrics from the terminal here once the training run completes).
Edge Deployment & ONNX
While these weights are native to PyTorch (.safetensors), this model is designed to be exported to ONNX via the Hugging Face optimum library. This allows the model to drop PyTorch dependencies entirely and run purely on ONNX Runtime for low-latency inference on ARM architectures (Raspberry Pi 4/5).
