Add similarity loss functions for Keras 3?

[ma7555]

Hello,
Is there any interest in allowing training on embeddings using similarity loss functions? Currently the situation is tensorflow similarity is not working with Keras 3, with the latest working tf working version 2.15, TF-Similarity repo does not look very active to hope for an update anytime soon.

I was thinking if there would be some interest, I can rewrite even a single similarity loss function that would open the door for such tasks to take place on keras 3. I checked KerasHub too but it does not seem it has any loss functions.

I can go with the newer CircleLoss.

[sampathweb]

@ma7555 - Thanks for the suggestion. Please go ahead with a PR for CircleLoss. Thanks for your contribution!

[ma7555]

@sampathweb - I’ll be working on this.

We have three possible approaches to proceed:

1. Reuse LossFunctionWrapper:

   • This approach is straightforward and works well for typical (within-batch) training scenarios.
   • Here, the loss function will, as usual, accept only y_true and y_pred, limiting calculations to the same batch. This approach does not support cross-batch training (such as XBM), so distances and loss computations would be restricted to the current batch.
   • While this is the standard in most training configurations, it lacks flexibility for advanced training setups where cross-batch mining of positive and negative pairs might be beneficial.

2. Introduce a New Wrapper for Metric Losses:

   • To increase flexibility, we could create a new wrapper, such as MetricLossFunctionWrapper.
   • This setup would mirror the structure in MetricLoss and allow the loss function to support different query/key pairs by force-assigning self.fn(y_true, y_pred, y_true, y_pred, **self._fn_kwargs).
   • By supporting four parameters, this loss function could handle more advanced similarity training techniques while behaving like option 1 in standard configurations.

3. Modify the Loss Function to Accept Four Parameters:

   • This approach introduces the option of passing four parameters (query_embeddings, query_labels, key_embeddings, key_labels) to the loss function, with the latter two (key_embeddings, key_labels) defaulting to None.
   • Internally, if these additional parameters are not provided, the function would use query_embeddings and query_labels for both query and key values.
   • This makes the function compatible with both standard batch setups and advanced configurations without needing a separate wrapper, allowing flexibility for cross-batch or similarity-based training setups.

While cross-batch training theoretically offers a richer pool for positive/negative pair mining, the basic, more straightforward option (Option 1) may be sufficient for Keras Core. However, if there is interest in expanding to similarity-based training, either Option 2 or 3 might be better for future scalability with option 2 being the cleaner one.

I will prepare a draft PR for Option 1 and we can take it from there, if you prefer one of the other options, let me know and i will work on it.
You can also review this under tensorflow/similarity#222

[ma7555]

resolved by #20452