Introduction

If your classification dataset is imbalanced, plain cross-entropy often under-penalizes mistakes on rare classes. A weighted loss fixes that by multiplying each example's loss by a class-dependent weight before reducing across the batch.

With tf.nn.sparse_softmax_cross_entropy_with_logits, the pattern is straightforward: compute the per-example sparse cross-entropy, look up a weight from the integer label, and multiply.

Why Weighting Works

Sparse softmax cross-entropy already gives one loss value per example. For imbalanced data, you can rescale those values so errors on minority classes count more during optimization.

The basic idea is:

• class 0 might use weight 1.0
• class 1 might use weight 4.0
• class 2 might use weight 2.5

That does not change the predicted probabilities directly. It changes how strongly the optimizer reacts to errors for each class.

A Runnable TensorFlow Example

1import tensorflow as tf
2
3logits = tf.constant([
4    [2.0, 0.5, 0.1],
5    [0.2, 1.5, 0.3],
6    [0.1, 0.2, 2.0],
7], dtype=tf.float32)
8
9labels = tf.constant([0, 1, 2], dtype=tf.int32)
10class_weights = tf.constant([1.0, 4.0, 2.0], dtype=tf.float32)
11
12per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
13    labels=labels,
14    logits=logits,
15)
16
17weights = tf.gather(class_weights, labels)
18weighted_loss = per_example_loss * weights
19
20loss = tf.reduce_mean(weighted_loss)
21
22print("per-example:", per_example_loss.numpy())
23print("weights:", weights.numpy())
24print("final loss:", loss.numpy())

This is the core recipe. tf.gather maps each integer class label to its corresponding weight, and the result scales the per-example loss before reduction.

Turning It Into a Keras Loss

For model training, wrap the logic in a callable:

1import tensorflow as tf
2
3class_weights = tf.constant([1.0, 4.0, 2.0], dtype=tf.float32)
4
5def weighted_sparse_ce(y_true, y_pred):
6    y_true = tf.cast(tf.reshape(y_true, [-1]), tf.int32)
7
8    per_example_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
9        labels=y_true,
10        logits=y_pred,
11    )
12
13    weights = tf.gather(class_weights, y_true)
14    return tf.reduce_mean(per_example_loss * weights)

Then compile with logits, not softmax probabilities:

1model.compile(
2    optimizer="adam",
3    loss=weighted_sparse_ce,
4)

If the model's last layer already applies softmax, then sparse_softmax_cross_entropy_with_logits is the wrong function because it expects raw logits. In that case, either remove the final softmax or use a different loss path that matches probabilities.

Choosing the Weights

Common strategies include:

• inverse class frequency
• median frequency balancing
• manually tuned values based on validation results

A simple start is:

weight_c = total_samples / (num_classes * samples_in_class_c)

That is not magic. It is just a reasonable baseline. Extremely large weights can destabilize training or hurt overall calibration.

When To Use Sample Weights Instead

If the weighting depends on more than class identity, use per-sample weights in the training loop or pass sample weights to model.fit. Class weighting and sample weighting are related, but class-based weighting is simpler when the only concern is label imbalance.

Reduction Strategy Matters

Most examples use tf.reduce_mean, which is fine for many training loops. In distributed training or mixed weighting setups, you may want to control normalization more explicitly so the effective loss scale remains stable across batches. The important point is that weighting should happen before the final reduction, not after it.

Common Pitfalls

• Passing softmax probabilities into a loss function that expects logits.
• Forgetting to cast labels to integer type before tf.gather.
• Using a class-weight vector whose indices do not line up with the encoded labels.
• Making rare-class weights so large that training becomes unstable.
• Reducing the loss before applying weights, which removes the per-example information you need.

Summary

• Compute sparse softmax cross-entropy per example, then multiply by class-specific weights.
• Use tf.gather to map integer labels to weights.
• Apply weights before the final reduction.
• Make sure the loss receives logits, not already-softmaxed outputs.
• Start with moderate inverse-frequency weights and tune from validation results.