Implement and check a weighted binary cross-entropy loss from logits.

Instruction: For labels y in {0,1}, finite logits z, and positive-class weight p=3, define the batch-mean loss L=mean((1-y) softplus(z)+p*y*softplus(-z)). Implement a scalar reference, compute L and gradients for z=[0,0], y=[0,1], and explain how to preserve autograd in a training framework.

Context: Derive a numerically stable weighted loss, its reduction and a gradient sanity check.

Updated

Worked answer

For one example, ∂loss/∂z=(1−y)σ(z)−p y σ(−z). Divide by batch size for the stated mean reduction. At z=[0,0], y=[0,1], p=3, loss=(log 2+3 log 2)/2=2 log 2≈1.386294, and the gradient is [0.25, −0.75]. Raising the positive example's logit decreases its loss.

import math

def softplus(x):
    return max(x, 0.0) + math.log1p(math.exp(-abs(x)))

def weighted_bce(logits, labels, p=3.0):
    if len(logits) == 0 or len(logits) != len(labels):
        raise ValueError("nonempty matching batches required")
    if not math.isfinite(p) or p <= 0:
        raise ValueError("p must be finite and positive")
    if any(y not in (0, 1) for y in labels) or any(not math.isfinite(z) for z in logits):
        raise ValueError("binary labels and finite logits required")
    result = sum((1-y)*softplus(z) + p*y*softplus(-z)
                 for z, y in zip(logits, labels)) / len(logits)
    if not math.isfinite(result):
        raise ValueError("loss exceeds floating-point range")
    return result

z, y = [0., 0.], [0, 1]
assert abs(weighted_bce(z, y) - 2*math.log(2)) < 1e-12
for i, expected in enumerate([0.25, -0.75]):
    plus, minus = z.copy(), z.copy()
    plus[i] += 1e-5
    minus[i] -= 1e-5
    numeric = (weighted_bce(plus, y)-weighted_bce(minus, y))/2e-5
    assert abs(numeric-expected) < 1e-7
assert math.isfinite(weighted_bce([1000., -1000.], [0, 1]))

This standard-library reference checks arithmetic; it is not differentiable through a framework. In PyTorch use binary_cross_entropy_with_logits with a matching dtype/device pos_weight tensor and mean reduction, or build the same expression using tensor softplus operations. Do not detach logits or threshold probabilities inside the loss. Check framework gradients against this reference before training.

Positive weighting changes the optimization tradeoff; it does not guarantee a particular recall, calibration, or fairness outcome. Choose the weight and operating threshold using an appropriate validation design and report the resulting tradeoffs.

Reference: PyTorch BCEWithLogitsLoss.

Related Questions