Tabular Imbalance Benchmark

A complete end-to-end machine learning pipeline for handling highly imbalanced tabular data using Neural Networks and LightGBM, with advanced techniques such as threshold optimization, probability calibration, and model stacking.

---

Overview

This project demonstrates how to properly design, train, and evaluate machine learning models for extreme class imbalance problems (e.g., fraud detection).

It includes:

• Neural Network (TensorFlow)
• LightGBM (tree-based boosting)
• PR-AUC optimization (critical for imbalance)
• Threshold tuning (F1 / F2)
• Probability calibration (Isotonic Regression)
• Model stacking (NN + LightGBM)

---

Key Features

• End-to-end pipeline (data → model → evaluation → deployment artifacts)
• Proper handling of imbalanced data
• Comparison of deep learning vs gradient boosting
• Production-oriented evaluation (not just accuracy)
• Clean and reproducible experimental setup

---

Models Implemented

1. Neural Network (TensorFlow)

• Fully connected architecture
• Batch Normalization + Dropout
• AdamW optimizer with LR scheduling
• Class weighting for imbalance

2. LightGBM

• Raw tabular input (no scaling)
• scale_pos_weight for imbalance
• Early stopping based on PR-AUC
• Hyperparameter tuning (small search)

3. Advanced Techniques

• Threshold tuning using F-beta (F1 / F2)
• Probability calibration (Isotonic Regression)
• Stacking (meta-model with Logistic Regression)

---

Evaluation Metrics

• ROC-AUC (ranking performance)
• PR-AUC / Average Precision (primary metric for imbalance)
• Precision / Recall / F1-score
• Confusion Matrix (threshold-based evaluation)

---

Why PR-AUC?

In highly imbalanced datasets:

• Accuracy is misleading
• ROC-AUC can be overly optimistic

PR-AUC focuses on the minority class and is more reliable.

---

Pipeline Steps

1. Generate imbalanced dataset
2. Train Neural Network (scaled data)
3. Train LightGBM (raw data)
4. Evaluate baseline models
5. Optimize LightGBM (PR-AUC tuning)
6. Tune decision threshold (F-beta)
7. Apply probability calibration
8. Build stacking model (NN + LightGBM)
9. Compare all models

---

Results

• LightGBM achieves performance comparable to Neural Networks
• PR-AUC is significantly improved after tuning
• Threshold optimization drastically improves recall
• Stacking provides additional performance gains

---

Project Structure

tabular-imbalance-benchmark.ipynb
artifacts_advanced/
    ├── scaler.joblib
    ├── nn_model.keras
    ├── lgbm_base.joblib
    ├── lgbm_best.joblib
    ├── lgbm_isotonic.joblib
    ├── meta_logreg.joblib
    ├── y_test.npy
    ├── predictions...

---

Tech Stack

• Python
• TensorFlow / Keras
• LightGBM
• Scikit-learn
• NumPy / Pandas
• Matplotlib

---

Use Cases

• Fraud Detection
• Credit Risk Modeling
• Medical Diagnosis
• Anomaly Detection
• Rare Event Prediction

---

Key Takeaways

• Tree-based models (LightGBM) are extremely strong for tabular data
• Neural Networks can match performance with proper tuning
• Threshold tuning is critical in real-world systems
• Calibration improves decision-making reliability
• Stacking combines strengths of multiple models

---

Author

Mohammad Amiri
AI / ML Engineer

---

If you found this useful

Give it a star and feel free to fork!