Implement Focal Loss for Class-Imbalanced Classification

Problem Statement

In high-throughput observability platforms like Datadog, anomaly detection and alerting models frequently deal with extreme class imbalance where anomalies represent less than 0.1% of all ingested metrics. Standard binary cross-entropy (BCE) loss tends to be overwhelmed by easy negative examples.

Focal Loss addresses this by adding a modulating factor (1 - p_t)^gamma to standard cross-entropy, down-weighting well-classified examples and focusing training on hard negatives and false positives.

Implement binary_focal_loss_with_logits using PyTorch or NumPy operations from scratch. You must not call high-level library focal loss functions (such as torchvision.ops.sigmoid_focal_loss). Your implementation must compute the loss directly from unnormalized log-odds (logits) to avoid numerical instability (e.g., NaN or -inf values caused by log(0) or extreme exponentiation).

Mathematical Formulation

For a binary target y in {0, 1} and predicted logit z:

• Let p = sigmoid(z) = 1 / (1 + exp(-z))
• p_t = p if y == 1 else 1 - p
• alpha_t = alpha if y == 1 else (1 - alpha)
• Loss = -alpha_t * (1 - p_t)^gamma * log(p_t)

Your implementation must be numerically stable even when z takes extreme values such as z = -100 or z = 100.

Function Signature

import torch

def binary_focal_loss_with_logits(
    logits: torch.Tensor,
    targets: torch.Tensor,
    alpha: float = 0.25,
    gamma: float = 2.0,
    reduction: str = "mean"
) -> torch.Tensor:
    pass

Input & Output Format

• logits: Tensor of arbitrary shape (...) containing unnormalized prediction scores (float32/float64).
• targets: Tensor of the same shape as logits with values in {0.0, 1.0}.
• alpha: Weighting factor for the positive class (0.0 = 0.0). When gamma = 0, focal loss is equivalent to standard binary cross-entropy.
• reduction: Specifies reduction to apply: 'none', 'mean', or 'sum'.
• Output: Scalar tensor if reduction is 'mean' or 'sum'; otherwise, a tensor with the same shape as logits.

Examples

Example 1

logits = torch.tensor([0.0, 2.0, -2.0])
targets = torch.tensor([1.0, 1.0, 0.0])
loss = binary_focal_loss_with_logits(logits, targets, alpha=0.25, gamma=2.0, reduction="mean")
# Output: approx 0.0401

Explanation: For logits[1] = 2.0 and target = 1.0, p ≈ 0.8808. The modulating factor is (1 - 0.8808)^2 ≈ 0.0142, heavily discounting this easy positive example.

Example 2 (Numerical Stability)

logits = torch.tensor([-1000.0, 1000.0])
targets = torch.tensor([0.0, 1.0])
loss = binary_focal_loss_with_logits(logits, targets, alpha=0.5, gamma=2.0, reduction="none")
# Output: tensor([0.0, 0.0]) without NaN or Inf values

Constraints & Edge Cases

• logits and targets match in shape and data type.
• gamma = 0 and alpha = -1 should strictly match torch.nn.functional.binary_cross_entropy_with_logits.
• Must gracefully handle large positive and negative logits without returning NaN gradients or outputs.