[NripeshN]

Proposed changes

Have added focal_loss, dice_loss, iou_loss, contrastive_loss, tversky_loss

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

FYI we already decided this one was a bit niche to put in nn in #224 .

[awni]

Could you say a bit more about why these specific losses should be included in nn? Some of them seem a bit niche / I don't find them in other frameworks. But maybe we can have a discussion about which if any are worth including...

[NripeshN]

Could you say a bit more about why these specific losses should be included in nn? Some of them seem a bit niche / I don't find them in other frameworks. But maybe we can have a discussion about which if any are worth including...

Hey there,

Great question! I totally get why some of these loss functions might seem a bit niche at first glance. Here's my rationale for including them:

1. Focal Loss: This one's a game-changer for dealing with class imbalance, which is a super common issue in classification tasks. It's especially useful in cases where you've got a lot of background noise dominating the few examples of the class you're actually interested in. By focusing more on the harder, misclassified examples, it helps models perform better in these tricky scenarios.

2. Dice and IoU Losses: These are bread-and-butter for segmentation tasks. They shine in scenarios where we need to evaluate the overlap between predicted and actual segments, making them super relevant for medical imaging, satellite image analysis, and other similar fields. Since these areas are growing rapidly, having robust loss functions like Dice and IoU can really empower models to perform better.

3. Contrastive Loss: This is a bit of a specialized one, I agree. But it's super useful in learning embeddings and similarity. Think face recognition, recommendation systems, or any scenario where understanding the relative similarity between data points is key. It's a bit of a unique tool, but in the right context, it's incredibly powerful.

4. Tversky Loss: It's like Dice Loss but with an extra tuning knob. This makes it super adaptable for datasets where you've got an imbalance between different types of errors (false positives vs false negatives). It's a bit more of a fine-tuning tool, but for datasets where precision is key, it can make a significant difference.

So, in summary, while some of these losses might not be everyday go-to's, they each have their own superpowers for specific, yet increasingly common, scenarios in the ML world. Having them in the toolkit could really bolster the versatility and effectiveness of the mlx framework, especially in tackling more complex or specialized tasks.

Looking forward to hearing your thoughts on this!

[JayYun98]

review from my pr #336

[JayYun98]

suggestion

        inputs (mx.array): Predicted probabilities for the positive class.
        targets (mx.array): The target values (binary).

[JayYun98]

suggestions. It's for binary target value so it's binary cross entropy, and you could use pre-defined bce loss for concise understanding.

    BCE_loss = binary_cross_entropy(inputs, targets)
    pt = mx.exp(-BCE_loss)
    loss = -alpha * ((1 - pt) ** gamma) * BCE_loss
    return _reduce(loss, reduction)

[JayYun98]

It is better to use original test examples used above.

Also, it is likely to use pre-defined expected output rather than copy of your loss fuction at losses.py

  ce_loss = mx.logaddexp(0.0, inputs) - targets * inputs
  pt = mx.exp(-ce_loss)
  expected_loss = -alpha * ((1 - pt) ** gamma) * ce_loss
  expected_loss = mx.mean(expected_loss)
  loss = nn.losses.focal_loss(
      inputs, targets, alpha=alpha, gamma=gamma, reduction="mean"
  )

These are my codes that uses like [[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]]. You can extend test functions to tversky_loss and iou. However, I am not sure why less common loss are added on this pr.

Nice work!

cc @awni

    def test_dice_loss(self):
        inputs = mx.array([[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]])
        targets = mx.array([[1, 0, 1, 0], [0, 1, 0, 1]])

        # Test with reduction 'none'
        losses_none = nn.losses.dice_loss(inputs, targets, reduction="none")
        expected_none = mx.array([0.658536, 0.529412])  # Example expected values
        self.assertTrue(mx.allclose(losses_none, expected_none))

        # Test with reduction 'mean'
        losses_mean = nn.losses.dice_loss(inputs, targets, reduction="mean")
        expected_mean = mx.mean(expected_none)
        self.assertTrue(mx.allclose(losses_mean, expected_mean))

        # Test with reduction 'sum'
        losses_sum = nn.losses.dice_loss(inputs, targets, reduction="sum")
        expected_sum = mx.sum(expected_none)
        self.assertTrue(mx.allclose(losses_sum, expected_sum))

    def test_focal_loss(self):
        inputs = mx.array([[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]])
        targets = mx.array([[1, 0, 1, 0], [0, 1, 0, 1]])
        alpha = 0.25
        gamma = 2.0

        # Test with reduction 'none'
        losses_none = nn.losses.focal_loss(
            inputs, targets, alpha, gamma, reduction="none"
        )
        expected_none = mx.array(
            [
                [0.0433217, 0.0433217, 0.25751, 0.466273],
                [0.000263401, 0.147487, 0.0433217, 0.0433217],
            ]
        )
        self.assertTrue(mx.allclose(losses_none, expected_none))

        # Test with reduction 'mean'
        losses_mean = nn.losses.focal_loss(
            inputs, targets, alpha, gamma, reduction="mean"
        )
        expected_mean = mx.mean(expected_none)
        self.assertTrue(mx.allclose(losses_mean, expected_mean))

        # Test with reduction 'sum'
        losses_sum = nn.losses.focal_loss(
            inputs, targets, alpha, gamma, reduction="sum"
        )
        expected_sum = mx.sum(expected_none)
        self.assertTrue(mx.allclose(losses_sum, expected_sum))

    def test_contrastive_loss(self):
        embeddings1 = mx.array([[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]])
        embeddings2 = mx.array([[0.6, 0.4, 0.3, 0.8], [0.2, 0.5, 0.6, 0.4]])
        targets = mx.array([1, 0])
        margin = 1.0

        # Test with reduction 'none'
        losses_none = nn.losses.contrastive_loss(
            embeddings1, embeddings2, targets, margin, reduction="none"
        )
        expected_none = mx.array([0.2, 0.735425])
        self.assertTrue(mx.allclose(losses_none, expected_none))

        # Test with reduction 'mean'
        losses_mean = nn.losses.contrastive_loss(
            embeddings1, embeddings2, targets, margin, reduction="mean"
        )
        expected_mean = mx.mean(expected_none)
        self.assertTrue(mx.allclose(losses_mean, expected_mean))

        # Test with reduction 'sum'
        losses_sum = nn.losses.contrastive_loss(
            embeddings1, embeddings2, targets, margin, reduction="sum"
        )
        expected_sum = mx.sum(expected_none)
        self.assertTrue(mx.allclose(losses_sum, expected_sum))

[NripeshN]

Hey @awni,
Could you please let me know which losses we could add and which losses would not make sense. I will make changes after your inputs. cc @Jyun1998
Thanks both!

[awni]

@NripeshN let's table these losses for now. Since they are slightly more niche I think it is better to wait to understand the API requirements before adding them.

Also the contrastive loss looks to me to be the same as the triplet loss which we already have (?).

[NripeshN]

@NripeshN let's table these losses for now. Since they are slightly more niche I think it is better to wait to understand the API requirements before adding them.

Also the contrastive loss looks to me to be the same as the triplet loss which we already have (?).

Hey there!

Yeah, you're right, the contrastive loss formula I've implemented looks pretty similar to the triplet loss already in place. The core idea in both is to compare the distances between anchor-positive and anchor-negative pairs. But, there's a slight twist in their focus.

Triplet loss simultaneously pulls the anchor closer to the positive and pushes it away from the negative. Contrastive loss, on the other hand, is more about just ensuring the anchor is closer to the positive than a defined margin, without explicitly considering the negative.

In practice, they can end up doing similar things, but their approach and emphasis differ a bit.