class imbalance
loss scaling
stochastic gradient descent
machine learning
data preprocessing

Tackling Class Imbalance scaling contribution to loss and sgd

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Class imbalance makes optimization misleading because the loss is dominated by whatever class appears most often. A model trained with plain stochastic gradient descent can look good on overall accuracy while learning almost nothing about the rare class you actually care about.

Why Imbalance Hurts SGD

In SGD, each mini-batch contributes a gradient estimate. If 95 percent of your samples belong to one class, then most batches mostly teach the model how to reduce error on that class. Even if minority examples are present, their contribution to the loss can be too small to noticeably steer the update.

That creates two separate problems:

  • The model sees too few minority examples
  • The loss function under-penalizes mistakes on those examples

Resampling addresses the first problem. Loss weighting addresses the second. In practice, strong training pipelines often use both.

Weighted Loss Changes Gradient Importance

The simplest fix is to multiply the loss contribution of each class by a weight. Rare classes receive larger weights so the optimizer treats their mistakes as more important.

For multi-class classification, libraries usually implement this directly. In PyTorch, CrossEntropyLoss accepts a class-weight vector:

python
1import torch
2import torch.nn as nn
3
4logits = torch.tensor([
5    [2.4, 0.3],
6    [1.8, 0.2],
7    [0.4, 1.7],
8], dtype=torch.float32)
9
10targets = torch.tensor([0, 0, 1], dtype=torch.long)
11class_weights = torch.tensor([1.0, 4.0], dtype=torch.float32)
12
13criterion = nn.CrossEntropyLoss(weight=class_weights)
14loss = criterion(logits, targets)
15print("Weighted loss:", float(loss))

Here, errors on class 1 count four times as much as errors on class 0. That does not magically fix the data, but it changes the gradient signal so minority mistakes are harder for the model to ignore.

Sampling Strategy Changes What SGD Sees

Loss weighting alone still leaves you with batches dominated by the majority class. A second improvement is to sample examples so minority classes appear more often during training.

PyTorch provides WeightedRandomSampler for this:

python
1import torch
2from torch.utils.data import TensorDataset, DataLoader, WeightedRandomSampler
3
4features = torch.randn(10, 4)
5labels = torch.tensor([0, 0, 0, 0, 0, 0, 0, 1, 1, 1])
6
7class_sample_count = torch.tensor([(labels == 0).sum(), (labels == 1).sum()])
8class_weights = 1.0 / class_sample_count.float()
9sample_weights = class_weights[labels]
10
11sampler = WeightedRandomSampler(
12    weights=sample_weights,
13    num_samples=len(sample_weights),
14    replacement=True,
15)
16
17dataset = TensorDataset(features, labels)
18loader = DataLoader(dataset, batch_size=4, sampler=sampler)
19
20for batch_features, batch_labels in loader:
21    print(batch_labels.tolist())
22    break

This does not change the underlying dataset; it changes which examples SGD sees and how often. In many cases that stabilizes early training because the minority signal stops disappearing from the first few hundred updates.

How to Choose Class Weights

A common starting point is inverse frequency:

weight_c = total_samples / number_of_classes / count_c

That gives larger weights to rarer classes while keeping the average scale reasonable. You can also normalize the weights so the loss magnitude stays closer to the unweighted baseline.

Be careful with extreme ratios. If one class is extremely rare, huge weights can make optimization unstable or cause the model to overfit noisy minority labels. In that situation, clipping the weights, using focal loss, or collecting more data can work better than blindly applying the raw inverse-frequency formula.

Weighting Versus Resampling

These techniques are related but not identical.

Loss weighting keeps the original data distribution but changes the objective. Resampling changes the training distribution presented to SGD. If you oversample heavily, the model may memorize a small set of rare examples. If you only weight the loss, batches may remain too unbalanced to learn useful minority features early enough.

A pragmatic recipe is:

  • Start with stratified train and validation splits
  • Add class weights to the loss
  • If batches are still badly skewed, add a weighted sampler
  • Evaluate with precision, recall, F1, and per-class confusion, not accuracy alone

A Minimal Training Loop

The example below combines both ideas in a toy binary classifier:

python
1import torch
2import torch.nn as nn
3from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler
4
5X = torch.randn(200, 6)
6y = torch.tensor([0] * 180 + [1] * 20)
7
8dataset = TensorDataset(X, y)
9
10class_counts = torch.tensor([(y == 0).sum(), (y == 1).sum()])
11class_weights = 1.0 / class_counts.float()
12sample_weights = class_weights[y]
13sampler = WeightedRandomSampler(sample_weights, num_samples=len(y), replacement=True)
14
15loader = DataLoader(dataset, batch_size=16, sampler=sampler)
16
17model = nn.Sequential(nn.Linear(6, 12), nn.ReLU(), nn.Linear(12, 2))
18criterion = nn.CrossEntropyLoss(weight=class_weights)
19optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
20
21for epoch in range(3):
22    for batch_X, batch_y in loader:
23        optimizer.zero_grad()
24        logits = model(batch_X)
25        loss = criterion(logits, batch_y)
26        loss.backward()
27        optimizer.step()
28
29    print("epoch", epoch, "loss", round(float(loss), 4))

This is small, but it demonstrates the central idea: rebalance both the gradient importance and the batch composition.

Common Pitfalls

The first mistake is optimizing only for accuracy. On imbalanced data, accuracy often flatters weak models. Always inspect recall and precision for the minority class.

Another mistake is applying weights on the training set but using a non-stratified validation split. That can make model comparison noisy because the minority class barely appears in validation.

A third issue is double-counting imbalance too aggressively. Heavy oversampling plus very large loss weights can overcorrect and produce many false positives. Tune both knobs together instead of maximizing each one independently.

Finally, class imbalance is sometimes a symptom of a data-collection problem. If the rare class is also mislabeled or poorly represented, weighting cannot invent signal that is not present in the data.

Summary

  • Imbalanced data weakens SGD because most updates are dominated by majority-class examples.
  • Loss weighting increases the gradient contribution of rare classes.
  • Weighted sampling changes batch composition so minority examples appear more often.
  • Use both techniques carefully and evaluate with class-aware metrics.
  • If ratios are extreme, check for label quality issues and consider focal loss or better data collection.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.