How Do You Handle Imbalanced Datasets in Deep Learning?

How do you handle imbalanced datasets in deep learning

Every time I’ve worked on a fraud detection, medical diagnosis, or rare-event prediction project, I’ve run into the same wall almost immediately: the data is wildly imbalanced. Ninety-nine point nine percent of transactions are legitimate; only a tiny sliver are fraudulent. Ninety-eight percent of scans are healthy; only two percent show disease. If you train a naive deep learning model on data like this, it will happily achieve 99% accuracy — by simply predicting “no fraud” or “healthy” every single time, and being completely useless in practice.

In this article, I want to unpack exactly why imbalanced datasets cause problems for deep learning models, and walk through the practical techniques — from simple data-level tricks to sophisticated algorithm-level fixes — that actually work in production.

1. What Is an Imbalanced Dataset?

A dataset is considered imbalanced when the classes you’re trying to predict are not represented equally. In binary classification, this usually means one class (the majority class) vastly outnumbers the other (the minority class).

We can quantify this with the imbalance ratio (IR):

$$ IR = \frac{N_{\text{majority}}}{N_{\text{minority}}} $$

An IR of 1 means perfectly balanced classes. An IR of 100 means the majority class has 100 times more samples than the minority class — a fairly common scenario in fraud detection, network intrusion detection, and rare disease diagnosis.

2. Why Imbalance Breaks Naive Training

Deep learning models are typically trained by minimizing a loss function like cross-entropy, averaged over the entire training set:

$$ L = -\frac{1}{N} \sum_{i=1}^{N} \left[ y_i \log(\hat{y}_i) + (1 – y_i) \log(1 – \hat{y}_i) \right] $$

If 99% of samples belong to class 0, the gradient signal from class 1 samples gets drowned out by the sheer volume of class 0 samples. The optimizer finds it easy to minimize the loss simply by biasing predictions toward the majority class — even though that’s not what we actually want.

This leads to models with:

3. Choosing the Right Evaluation Metric First

Before even touching the data or model, it’s critical to abandon accuracy as your primary metric when dealing with imbalance. Instead, use metrics that reflect performance on the minority class specifically.

MetricFormulaWhat It Tells You
Precision$\frac{TP}{TP + FP}$Of all predicted positives, how many were correct
Recall (Sensitivity)$\frac{TP}{TP + FN}$Of all actual positives, how many were caught
F1-Score$2 \cdot \frac{\text{Precision} \cdot \text{Recall}}{\text{Precision} + \text{Recall}}$Harmonic mean of precision and recall
AUC-ROCArea under the ROC curveOverall ranking quality across thresholds
AUC-PRArea under Precision-Recall curveMore informative than AUC-ROC under heavy imbalance
Balanced Accuracy$\frac{1}{2}\left(\frac{TP}{TP+FN} + \frac{TN}{TN+FP}\right)$Average recall across both classes

For heavily imbalanced problems, AUC-PR and F1-score are generally more informative than plain accuracy or even AUC-ROC, since ROC curves can look deceptively good even when precision on the minority class is poor.

4. Data-Level Techniques

a) Random Oversampling

Duplicate examples from the minority class until the classes are more balanced. Simple, but risks overfitting since the model may just memorize duplicated examples.

b) Random Undersampling

Remove examples from the majority class to balance the dataset. Simple and fast, but risks discarding potentially useful information, especially when the majority class is diverse.

c) SMOTE (Synthetic Minority Oversampling Technique)

Rather than duplicating minority samples exactly, SMOTE generates synthetic samples by interpolating between existing minority class points and their nearest neighbors:

$$ x_{\text{new}} = x_i + \lambda \cdot (x_{\text{nn}} – x_i), \quad \lambda \sim U(0, 1) $$

where $x_i$ is a minority class sample, $x_{\text{nn}}$ is one of its $k$-nearest minority-class neighbors, and $\lambda$ is a random interpolation factor. This produces new, plausible synthetic examples rather than exact duplicates, reducing overfitting compared to naive oversampling.

d) ADASYN (Adaptive Synthetic Sampling)

A refinement of SMOTE that generates more synthetic samples for minority class instances that are harder to learn (i.e., those near the decision boundary with many majority-class neighbors), focusing the model’s attention where it struggles most.

e) Combining Over- and Under-Sampling

Techniques like SMOTEENN or SMOTETomek combine oversampling the minority class with cleaning up noisy or overlapping majority-class samples, often yielding better results than either technique alone.

from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)

smote = SMOTE(random_state=42)
X_resampled, y_resampled = smote.fit_resample(X_train, y_train)

print("Original class distribution:", y_train.value_counts())
print("Resampled class distribution:", y_resampled.value_counts())

5. Algorithm-Level Techniques

a) Class Weighting

Instead of resampling the data, we can modify the loss function to penalize misclassifications of the minority class more heavily. For binary cross-entropy:

$$ L = -\frac{1}{N} \sum_{i=1}^{N} \left[ w_1 \cdot y_i \log(\hat{y}_i) + w_0 \cdot (1 – y_i) \log(1 – \hat{y}_i) \right] $$

A common choice is to set weights inversely proportional to class frequency:

$$ w_c = \frac{N}{K \cdot N_c} $$

where $N$ is the total number of samples, $K$ is the number of classes, and $N_c$ is the number of samples in class $c$.

In PyTorch, this is straightforward:

import torch
import torch.nn as nn

# Suppose class 0 has 9900 samples and class 1 has 100 samples
class_weights = torch.tensor([100/9900, 9900/9900])  # inverse frequency (simplified)
criterion = nn.CrossEntropyLoss(weight=class_weights)

b) Focal Loss

Originally introduced for object detection (where background examples vastly outnumber object examples), focal loss down-weights easy, well-classified examples and focuses training on hard, misclassified ones:

$$ FL(p_t) = -\alpha_t (1 – p_t)^\gamma \log(p_t) $$

where:

When $\gamma = 0$, focal loss reduces to standard weighted cross-entropy. As $\gamma$ increases, the loss contribution from well-classified examples (where $p_t$ is close to 1) shrinks rapidly, forcing the model to focus on harder, often minority-class examples.

import torch
import torch.nn as nn
import torch.nn.functional as F

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, logits, targets):
        bce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
        p_t = torch.exp(-bce_loss)
        focal_loss = self.alpha * (1 - p_t) ** self.gamma * bce_loss
        return focal_loss.mean()

c) Threshold Adjustment

Many classifiers output a probability, and a default threshold of 0.5 is used to convert it into a class label. Under imbalance, this default threshold is often suboptimal. Adjusting the decision threshold based on the precision-recall trade-off you care about (using validation data) can substantially improve minority-class performance without retraining the model at all.

d) Ensemble Methods

Techniques like Balanced Random Forest, EasyEnsemble, and RUSBoost combine undersampling with ensemble learning — training multiple models on different balanced subsets of the majority class, then combining their predictions. This retains more information from the majority class than a single undersampled dataset would.

6. Visualizing the Overall Workflow

flowchart TD
    A[Raw imbalanced dataset] --> B{Choose strategy}
    B -->|Data-level| C[Oversample minority - SMOTE/ADASYN]
    B -->|Data-level| D[Undersample majority]
    B -->|Algorithm-level| E[Class-weighted loss / Focal loss]
    B -->|Algorithm-level| F[Ensemble methods]
    C --> G[Train deep learning model]
    D --> G
    E --> G
    F --> G
    G --> H[Evaluate using F1 / AUC-PR / Recall]
    H --> I{Satisfactory minority-class performance?}
    I -->|No| B
    I -->|Yes| J[Deploy model]

7. Advantages and Disadvantages of Each Approach

TechniqueAdvantagesDisadvantages
Random OversamplingSimple, preserves all majority dataRisk of overfitting to duplicated samples
Random UndersamplingFast, reduces training timeDiscards potentially useful majority-class information
SMOTE / ADASYNGenerates diverse synthetic examples, reduces overfitting vs. duplicationCan generate unrealistic samples in high-dimensional or noisy data
Class WeightingNo change to data, easy to implement, keeps all dataMay not fully compensate for extreme imbalance ratios
Focal LossFocuses learning on hard examples, effective for extreme imbalanceRequires tuning of $\gamma$ and $\alpha$ hyperparameters
Ensemble MethodsRetains majority-class diversity, often strong performanceMore complex to implement and train, higher computational cost

8. A Complete Example: Handling Imbalance in a Neural Network Classifier

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, WeightedRandomSampler, TensorDataset

# Assume X_train, y_train are already prepared tensors
class_counts = torch.bincount(y_train)
class_weights = 1.0 / class_counts.float()
sample_weights = class_weights[y_train]

sampler = WeightedRandomSampler(sample_weights, len(sample_weights), replacement=True)

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=64, sampler=sampler)

class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 2)
        )

    def forward(self, x):
        return self.net(x)

model = SimpleClassifier(input_dim=X_train.shape[1])
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(20):
    for batch_X, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()

This example combines two techniques at once: a WeightedRandomSampler (a data-level fix that oversamples minority-class batches during training) and a class-weighted loss function (an algorithm-level fix). Combining complementary techniques like this often works better than relying on just one.

9. Real-World Use Cases

10. Best Practices

  1. Never rely on accuracy alone. Use F1-score, AUC-PR, recall, and precision as your primary evaluation metrics.
  2. Always split your data before resampling. Apply SMOTE or oversampling only to the training set, never to validation or test sets, to avoid data leakage and inflated performance metrics.
  3. Combine data-level and algorithm-level techniques when the imbalance ratio is extreme (e.g., SMOTE plus class-weighted loss).
  4. Tune the decision threshold post-training based on your specific cost trade-off between false positives and false negatives.
  5. Use stratified sampling when splitting data into train/validation/test sets, to ensure class proportions are preserved across splits.
  6. Consider anomaly detection framing for extremely rare events (IR > 1000), since these problems sometimes behave more like outlier/anomaly detection than standard classification.
  7. Monitor per-class metrics during training, not just aggregate loss, to catch cases where the model is ignoring the minority class.
  8. Be cautious with synthetic oversampling in high-dimensional spaces (e.g., raw text or images), where interpolation between samples (as in SMOTE) may not produce meaningful or realistic examples — consider domain-specific augmentation instead.

11. Summary

Imbalanced datasets are one of the most common — and most consequential — challenges in applied deep learning. When left unaddressed, they lead to models that look impressive on paper (high accuracy) but fail catastrophically at the task that actually matters (detecting the rare, important cases).

We covered:

The key takeaway: there’s no single “correct” fix for imbalanced data. The right approach depends on your imbalance ratio, your data’s dimensionality, and — most importantly — the real-world cost of false positives versus false negatives in your specific application. Start with the right metrics, experiment with a combination of data-level and algorithm-level techniques, and always validate on a held-out set that reflects the true, imbalanced distribution you’ll see in production.

References

Exit mobile version