Proposal: Generic Triplet-Margin Loss

[ethch18]

🚀 Feature

We'd like to add a distance-agnostic version of the TripletMarginLoss, as the current version only supports l_p norms. This was something we implemented as part of a recent investigation for embedding learning and thought could be helpful to the PyTorch userbase at large.

Motivation

1. Extensibility: decoupling the distance function and the loss computation allow for more flexibility for users. PyTorch currently has a CosineEmbeddingLoss, but that serves a somewhat different purpose and doesn't really work for users wanting a triplet-margin loss with cosine distance.

2. Existing use cases: several papers have proposed triplet loss functions with cosine distance (1, 2) or have generally used cosine-based metrics (1, 2). PyTorch-BigGraph also does something similar with its ranking loss.

3. Products like Tensorflow decouple the distance functions and even allow for custom distance metrics

Pitch

We have a working implementation of this (see screenshots) and propose to use it directly, with adjustments to fit the existing PyTorch code structure. We'll tentatively call these torch.nn.GenericTripletMarginLoss and F.generic_triplet_margin_loss. In the proposed implementation, GenericTripletMarginLoss will be a stateful wrapper for F.generic_triplet_margin_loss, and the loss computation will occur in the latter.

The module signatures will resemble that of the production module, with an analogous signature for the functional module. The removed parameters (p and eps) are specific to l_p norm computation; size_average and reduce are removed because of deprecation. There are two new parameters: distance_function and is_similarity_function. distance_function is a function/module $$f: R^{n \cross d} \rightarrow R^n$$ that gives an index-wise distance metric between two tensors of embeddings. This can be, for example, torch.nn.PairwiseDistance or torch.nn.CosineSimilarity. is_similarity_function is a boolean that denotes whether distance_function is a distance or similarity function: this will indicate whether to flip the signs in distance computation.

Alternatives

Beyond the straightforward alternative of leaving this to the user to implement every time, we haven't come up with any other alternatives.

One open question is whether the Python implementation is performant enough, or if this would need to be added in aTen as well.

Additional context

cc @albanD @mruberry

[mruberry]

This is an interesting proposal and letting users specify their own distance functions makes sense.

I see two challenges with supporting this in PyTorch Core:

• PyTorch typically supports both a Python and C++ API, and I think this would only work in Python?
• Would this work with the JIT? I don't think scripting or tracing support passing callables.

[ethch18]

@mruberry thanks for your response!

Re: C++, would it be possible to pass in the distance function as a lambda? I'm not too familar with aTen but the distance function should just be another function. Or were there other concerns about C++?

Re: JIT, are there ways to work around the callable issue?

[mruberry]

Re: C++, would it be possible to pass in the distance function as a lambda? I'm not too familar with aTen but the distance function should just be another function. Or were there other concerns about C++?

Yes, you could have two implementations: one for Python and one for C++.

Re: JIT, are there ways to work around the callable issue?

Not that I know of.

Edit/Update: Passing a Python function to a module's init method and then calling that function in forward works fine with scripting inference.

[ethch18]

@mruberry I see that in this Q/A, it's possible to get around the JIT callable issue by wrapping it as a module. Since the implementation I proposed takes in an nn.Module and keeps it as an attribute, I think this should work fine as-is?

[mruberry]

@mruberry I see that in this Q/A, it's possible to get around the JIT callable issue by wrapping it as a module. Since the implementation I proposed takes in an nn.Module and keeps it as an attribute, I think this should work fine as-is?

Interesting! Good point. My mistake I should have thought of wrapping in a module / noticing that's how your callable was typed.

[mruberry]

Update from offline conversation with @ethch18:

This seems like a cool improvement on an existing module, but the work to implement it in core might be significant. If it's pursued, it's recommend to start by prototyping the C++ changes required in a BC-preserving way.

[mruberry]

We now have a "generic" triplet margin loss! Closing this issue.