← Back to home
Case study · Machine learning

Churn prediction on a heavily imbalanced dataset

Predicting churn on millions of rows where positives are a tiny fraction of the population. The interesting work is not the model, it is what happens before and after the model.

PythonLightGBMImbalanced learningCalibrationPR-AUCFeature engineeringMLOps

01The problem

A dataset of several million user records with a strong class imbalance: churners were a small minority. A naive classifier reaches great accuracy by predicting the majority class and is useless in production. The business does not want a score, it wants a shortlist of accounts worth intervening on, within a fixed retention budget.

02The approach

I treated the problem as ranking under a budget rather than classification. That reframing changed every downstream choice: the metric, the sampling, the threshold, the monitoring.

03Pipeline

DATA Raw events (M rows) Feature store TRAINING Stratified resampling Gradient boosting (LGBM) Probability calibration Cost-sensitive threshold PR-AUC / recall@budget SERVING Batch scoring Retention actions Monitoring & drift

04Key technical decisions

  • Feature engineering over model tuning: recency, frequency and behavioural deltas moved the needle more than any hyperparameter search.
  • Stratified resampling, carefully: SMOTE on high-dimensional tabular data introduced noise. Undersampling the majority with stratification on segment kept the signal cleaner.
  • Gradient boosting (LightGBM): fast on millions of rows, handles categorical features natively, gives usable feature importance.
  • Probability calibration: raw boosting scores are miscalibrated after resampling. Isotonic calibration on a held-out fold made the scores usable for ranking and business cutoffs.
  • Cost-sensitive threshold, not 0.5: threshold picked to maximise recall at the retention team's daily capacity.
  • PR-AUC and recall@k, not accuracy: the only metrics that reflect what the business actually cares about on imbalanced data.

05What I would do differently

Ship a small monitoring dashboard alongside the model on day one: prediction drift, feature drift, capture rate against the retention actions actually taken. It closes the loop and turns the model from a one-off delivery into a system.

Next case studyMulti-agent RAG in Microsoft Teams