3.2 Methods to Mitigate Data Imbalance in Training Data

3.2 Methods to Mitigate Data Imbalance in Training Data

In enterprise machine learning applications—such as financial fraud detection, customer churn prediction, ad click-through rate modeling, and cyber threat identification—the target class distribution is often heavily skewed. Positive cases of interest may represent less than 1% (or even 0.01%) of total observations. Training models on severely imbalanced datasets without mitigation leads to catastrophic failure: standard loss functions incentivize algorithms to classify all records as the majority class to maximize raw accuracy. Addressing imbalance requires a multi-layered strategy across data resampling, algorithmic weighting, metric selection, and decision threshold optimization.

---

The Accuracy Paradox & Class Imbalance Dynamics

Consider a credit card fraud detection dataset with 1,000,000 transactions, where only 1,000 transactions (0.1%) are fraudulent ($y=1$) and 999,000 are legitimate ($y=0$).

+-----------------------------------------------------------------------------+
|                        THE ACCURACY PARADOX DEMONSTRATED                    |
|                                                                             |
|   TRIVIAL MAJORITY PREDICTOR: Predicts y = 0 for 100% of transactions       |
|                                                                             |
|   Confusion Matrix:                                                         |
|   +--------------------+--------------------+                               |
|   | TN = 999,000       | FP = 0             |  Accuracy  = 99.90% (Superb?) |
|   +--------------------+--------------------+  Precision = Undefined (0/0)  |
|   | FN = 1,000         | TP = 0             |  Recall    = 0.00%  (Useless!)|
|   +--------------------+--------------------+  F1-Score  = 0.00%            |
|                                                                             |
|   Business Impact: $10,000,000 in fraudulent transactions missed completely!|
+-----------------------------------------------------------------------------+

The Accuracy Paradox occurs because standard cross-entropy and 0-1 loss functions treat every misclassification equally. A naive model that predicts the majority class for every instance achieves 99.9% accuracy while having zero operational utility.

---

Multi-Level Imbalance Mitigation Strategies

Imbalance mitigation can be applied at three distinct stages of the machine learning lifecycle:

+-----------------------------------------------------------------------------+
|                   THREE-TIER IMBALANCE MITIGATION HIERARCHY                 |
|                                                                             |
|   [1] DATA-LEVEL STRATEGIES (Resampling before Training)                    |
|   +---------------------------------------------------------------------+   |
|   | - Random Undersampling (majority class reduction)                   |   |
|   | - Random Oversampling (minority class replication)                  |   |
|   | - SMOTE / ADASYN (synthetic minority interpolation)                 |   |
|   +---------------------------------------------------------------------+   |
|                                     |
|                                     v
|   [2] ALGORITHM-LEVEL STRATEGIES (Cost-Sensitive Loss Optimization)         |
|   +---------------------------------------------------------------------+   |
|   | - Scikit-learn: class_weight='balanced'                             |   |
|   | - XGBoost / LightGBM: scale_pos_weight = N_negative / N_positive     |   |
|   | - PySpark ML: weightCol="weight" (sample weight per row)             |   |
|   +---------------------------------------------------------------------+   |
|                                     |
|                                     v
|   [3] DECISION-LEVEL STRATEGIES (Post-Prediction Threshold Tuning)          |
|   +---------------------------------------------------------------------+   |
|   | - Shift classification cutoff tau from 0.5 to tau*                  |   |
|   | - Cost Matrix Optimization: Minimize C_FP * FP(tau) + C_FN * FN(tau)|   |
|   | - Precision-Recall Curve Cutoff Selection                           |   |
|   +---------------------------------------------------------------------+   |
+-----------------------------------------------------------------------------+

Data-Level Resampling Techniques

1. Random Undersampling: Randomly discards instances from the majority class until the desired ratio is reached.
   • Advantage: Decreases dataset size, significantly accelerating training speed.
   • Disadvantage: Discards potentially critical majority class boundary information, increasing model variance.
2. Random Oversampling: Randomly duplicates instances from the minority class.
   • Advantage: Preserves all majority information.
   • Disadvantage: Duplicating identical minority samples leads to severe overfitting on specific minority feature regions.
3. SMOTE (Synthetic Minority Over-sampling Technique): Synthesizes new minority instances by selecting a minority sample $\mathbf{x}i$, finding its $k$-nearest minority neighbors $\mathbf{x}{zi}$, and generating synthetic points along the connecting line segment: $\mathbf{x}_{\text{new}} = \mathbf{x}_i + \lambda (\mathbf{x}_{zi} - \mathbf{x}_i), \quad \lambda \sim \text{Uniform}(0, 1)$
   • Spark Consideration: SMOTE requires distance calculations in feature space, making it expensive on distributed datasets. In Databricks, SMOTE is typically executed on sampled single-node data or partition-by-partition using Pandas UDFs.

Algorithm-Level Cost-Sensitive Learning

Cost-sensitive learning modifies the objective loss function to penalize minority class errors more heavily than majority class errors:

• Scikit-learn class_weight='balanced': Automatically assigns inverse class frequency weights: $w_j = \frac{N}{K \cdot N_j}$ where $N$ is total samples, $K$ is number of classes, and $N_j$ is samples in class $j$.
• XGBoost / LightGBM scale_pos_weight: Scales the gradient and hessian of positive instances: $\text{scale\_pos\_weight} = \frac{\text{Number of Negative Instances}}{\text{Number of Positive Instances}}$
• PySpark ML weightCol: Most PySpark ML classification estimators (LogisticRegression, RandomForestClassifier) accept a weightCol parameter pointing to a computed column in the DataFrame.

---

Decision Threshold Optimization & Cost Matrix Analysis

By default, binary classifiers apply a decision threshold of $\tau = 0.5$ to predicted probabilities: $\hat{y} = \begin{cases} 1 & \text{if } P(y=1|\mathbf{x}) \ge 0.5 \\ 0 & \text{if } P(y=1|\mathbf{x}) < 0.5 \end{cases}$

However, $\tau = 0.5$ is only optimal when classes are balanced and the cost of a False Positive ($C_{FP}$) equals the cost of a False Negative ($C_{FN}$).

+-----------------------------------------------------------------------------+
|                     COST MATRIX & THRESHOLD OPTIMIZATION                    |
|                                                                             |
|                          ACTUAL: Fraud (y=1)      ACTUAL: Legitimate (y=0)  |
|   PREDICT: Fraud (y_hat=1)      TP (Cost = $0)          FP (Cost = $C_FP)   |
|                                 Transaction Caught      False Alarm/Friction|
|                                                                             |
|   PREDICT: Legitimate (y_hat=0) FN (Cost = $C_FN)       TN (Cost = $0)      |
|                                 Fraud Loss Incurred     Normal Transaction  |
|                                                                             |
|   Total Business Cost Function:                                             |
|   Cost(tau) = C_FP * FP(tau) + C_FN * FN(tau)                               |
+-----------------------------------------------------------------------------+

Impact of Adjusting Threshold $\tau$

• Lowering Threshold (e.g., $\tau = 0.15$): The model predicts positive more aggressively. Recall increases (fewer False Negatives; more fraud detected), but Precision decreases (more False Positives; more false alarms). Ideal when $C_{FN} \gg C_{FP}$ (e.g. cancer detection, severe fraud).
• Raising Threshold (e.g., $\tau = 0.85$): The model predicts positive only when highly confident. Precision increases (fewer False Positives), but Recall decreases (more False Negatives). Ideal when False Alarms are extremely costly (e.g. automated account termination).

---

Metric Selection Under Imbalance: PR-AUC vs. ROC-AUC

Evaluating imbalanced models with ROC-AUC can be dangerously misleading:

---

Implementation Code Examples

PySpark ML: Class Weighting with weightCol

from pyspark.sql import functions as F
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.evaluation import BinaryClassificationEvaluator

df = spark.table("lakehouse_gold.fraud_transactions")

# Calculate class distribution
total_count = df.count()
pos_count = df.filter(F.col("is_fraud") == 1).count()
neg_count = total_count - pos_count

# Calculate balanced class weights
pos_weight = total_count / (2.0 * pos_count)
neg_weight = total_count / (2.0 * neg_count)

# Add weight column
weighted_df = df.withColumn(
    "class_weight",
    F.when(F.col("is_fraud") == 1, pos_weight).otherwise(neg_weight)
)

train_df, test_df = weighted_df.randomSplit([0.8, 0.2], seed=42)

# Train cost-sensitive Logistic Regression
lr = LogisticRegression(
    featuresCol="features",
    labelCol="is_fraud",
    weightCol="class_weight",
    maxIter=50
)

lr_model = lr.fit(train_df)
predictions = lr_model.transform(test_df)

# Evaluate with PR-AUC
evaluator = BinaryClassificationEvaluator(
    labelCol="is_fraud",
    rawPredictionCol="rawPrediction",
    metricName="areaUnderPR"
)
pr_auc = evaluator.evaluate(predictions)
print(f"Weighted Model PR-AUC: {pr_auc:.4f}")

Scikit-learn: Cost-Based Threshold Optimization

import numpy as np
from sklearn.metrics import precision_recall_curve

# Predicted probabilities and ground truth
y_true = np.array([...])
y_prob = np.array([...])

# Business cost specification
COST_FP = 15.0   # $15 customer service friction per false alarm
COST_FN = 500.0  # $500 average loss per missed fraud transaction

precisions, recalls, thresholds = precision_recall_curve(y_true, y_prob)

best_threshold = 0.5
min_total_cost = float("inf")

for threshold in np.linspace(0.01, 0.99, 100):
    y_pred = (y_prob >= threshold).astype(int)
    
    fp = np.sum((y_pred == 1) & (y_true == 0))
    fn = np.sum((y_pred == 0) & (y_true == 1))
    
    total_cost = (COST_FP * fp) + (COST_FN * fn)
    
    if total_cost < min_total_cost:
        min_total_cost = total_cost
        best_threshold = threshold

print(f"Optimal Decision Threshold: {best_threshold:.3f}")
print(f"Minimized Financial Cost: ${min_total_cost:,.2f}")

---

Where Each Technique Belongs in the Pipeline

A recurring exam trap is applying a resampling technique in the wrong place.

Class weighting is the lowest-risk option because it neither discards majority rows nor invents minority ones; it simply tells the loss function that a minority error costs more. It is the technique Databricks material reaches for first, and the one that most often appears as the correct answer when a question asks how to "directly mitigate the model's bias toward the majority class".