Add binary focal cross entropy loss function

docs/src/python/nn/losses.rst

@@ -10,6 +10,7 @@ Loss Functions
    :template: nn-module-template.rst
 
    binary_cross_entropy
+   binary_focal_cross_entropy
    cosine_similarity_loss
    cross_entropy
    gaussian_nll_loss

python/mlx/nn/losses.py

@@ -121,6 +121,94 @@ def binary_cross_entropy(
     return _reduce(loss, reduction)
 
 
+def binary_focal_cross_entropy(
+    inputs: mx.array,
+    targets: mx.array,
+    with_logits: bool = True,
+    class_balancing: bool = False,
+    alpha: float = 0.25,
+    gamma: float = 2.0,
+    reduction: Reduction = "mean",
+) -> float:
+    r"""
+    Calculate the binary focal cross entropy loss between true labels and predicted probabilities.
+
+    According to `Lin et al., 2018`_, it helps to apply a "focal factor" to down-weight easy
+    examples and focus more on hard examples.
+
+    The loss is given by:
+
+    .. math::
+        \text{FL}(p,y)=
+        \begin{cases}
+            -\alpha(1-p)^\gamma\log(p), & \text{ if } y=1 \\
+            -(1-\alpha)p^\gamma\log(1-p), & \text{otherwise.}
+        \end{cases}
+
+    where :math:`p` is the predicted probability, :math:`y` is the target value, :math:`\alpha` is
+    the weight balancing factor for class 1, and :math:`\gamma` is the focusing parameter.
+
+    When ``class_balancing`` is ``False`` and ``gamma`` is ``0``, this function is equivalent to the
+    binary cross-entropy loss.
+
+    Args:
+        inputs (array): The predicted values. If ``with_logits`` is ``True``, then
+            ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities.
+        targets (array): Binary target values in {0, 1}.
+        with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``.
+        class_balancing (bool, optional): Whether to apply class balancing. Default: ``False``.
+        alpha (float, optional): A weight balancing factor for class 1.
+            Default: ``0.25`` (`Lin et al., 2018`_). The weight for class 0 is ``1.0 - alpha``.
+        gamma (float, optional): A focusing parameter used to compute the focal factor.
+            Default: ``2.0`` (`Lin et al., 2018`_).
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    .. _Lin et al., 2018: https://arxiv.org/pdf/1708.02002.pdf
+
+    Returns:
+        array: The computed binary focal cross entropy loss.
+
+    Examples:
+        >>> import mlx.core as mx
+        >>> from mlx.nn.losses import binary_focal_cross_entropy
+
+        >>> # inputs are logits, without class balancing
+        >>> logits = mx.array([[0.105361, 0.223144], [1.20397, 0.916291]])
+        >>> targets = mx.array([[0, 1], [1, 0]])
+        >>> binary_focal_cross_entropy(logits, targets, gamma=2.0)
+        array(0.244057, dtype=float32)
+
+        >>> # inputs are probabilities, with class balancing
+        >>> probs = mx.array([[0.1, 0.1], [0.4, 0.4]])
+        >>> targets = mx.array([[0, 1], [1, 0]])
+        >>> binary_focal_cross_entropy(probs, targets, class_balancing=True, alpha=0.25, gamma=2.0)
+        array(0.114078, dtype=float32)
+    """
+    if inputs.shape != targets.shape:
+        raise ValueError(
+            f"Inputs shape {inputs.shape} does not match targets shape {targets.shape}."
+        )
+
+    if with_logits:
+        inputs = mx.sigmoid(inputs)
+
+    focusing_pos = (1 - inputs) ** gamma
+    focusing_neg = inputs**gamma
+
+    loss_pos = -(focusing_pos * mx.log(inputs) * targets)
+    if class_balancing:
+        loss_pos *= alpha
+
+    loss_neg = -(focusing_neg * mx.log(1 - inputs) * (1 - targets))
+    if class_balancing:
+        loss_neg *= 1 - alpha
+
+    loss = loss_pos + loss_neg
+
+    return _reduce(loss, reduction)
+
+
 def l1_loss(
     predictions: mx.array, targets: mx.array, reduction: Reduction = "mean"
 ) -> mx.array:

python/tests/test_losses.py

@@ -160,6 +160,66 @@ def _test_probs_as_inputs():
         _test_logits_as_inputs()
         _test_probs_as_inputs()
 
+    def test_binary_focal_cross_entropy(self):
+        # Test ground truth values are calculated by tf.keras.losses.BinaryFocalCrossentropy
+
+        alpha = 0.25  # class balancing factor
+        gamma = 2.0  # focal factor
+
+        def _test_logits_as_inputs():
+            logits = mx.array([-1.6, 0.5, 2.9, -1.8])
+            targets = mx.array([0, 1, 0, 1])
+
+            # Test with reduction 'none', without class balancing
+            losses = nn.losses.binary_focal_cross_entropy(
+                logits, targets, gamma=gamma, reduction="none"
+            )
+            expected = mx.array([0.00518928, 0.06757348, 2.653519, 1.438211])
+            self.assertTrue(mx.allclose(losses, expected))
+
+            # Test with reduction 'none', with class balancing
+            losses_balanced = nn.losses.binary_focal_cross_entropy(
+                logits,
+                targets,
+                class_balancing=True,
+                alpha=alpha,
+                gamma=gamma,
+                reduction="none",
+            )
+            expected_balanced = mx.array(
+                [0.00389196, 0.01689337, 1.9901392, 0.35955274]
+            )
+            self.assertTrue(mx.allclose(losses_balanced, expected_balanced))
+
+        def _test_probs_as_inputs():
+            probs = mx.array([0.5, 0.6, 0.7, 0.8])
+            targets = mx.array([0, 1, 0, 1])
+
+            # Test with reduction 'none', without class balancing
+            losses = nn.losses.binary_focal_cross_entropy(
+                probs, targets, with_logits=False, gamma=gamma, reduction="none"
+            )
+            expected = mx.array([0.17328674, 0.08173206, 0.5899465, 0.00892573])
+            self.assertTrue(mx.allclose(losses, expected))
+
+            # Test with reduction 'none', with class balancing
+            losses_balanced = nn.losses.binary_focal_cross_entropy(
+                probs,
+                targets,
+                with_logits=False,
+                class_balancing=True,
+                alpha=alpha,
+                gamma=gamma,
+                reduction="none",
+            )
+            expected_balanced = mx.array(
+                [0.12996505, 0.02043301, 0.44245988, 0.00223143]
+            )
+            self.assertTrue(mx.allclose(losses_balanced, expected_balanced))
+
+        _test_logits_as_inputs()
+        _test_probs_as_inputs()
+
     def test_l1_loss(self):
         predictions = mx.array([0.5, 0.2, 0.9, 0.0])
         targets = mx.array([0.5, 0.2, 0.9, 0.0])