Added a NTXent loss

[GrantMcConachie]

An normalized temperature scaled cross entropy (NTXent) loss for a contrastive learning objective. I am fairly new to submitting pull requests to public repos, so I didn't add a ton of tests for this outside a batched test. Let me know if there is anything else I should add!

[GrantMcConachie]

Oh also, I know NTXent isn't really a classification loss, but I didn't know where else to put it.

[vroulet]

Thank you @GrantMcConachie!
This would be a great addition. I left you some comments. In addition of those comments:

• you may create a file _contrastive.py where you would put this loss (so don't put this in _classification.py),
• you will need to add the loss in the __init__.py file of the losses folder,
• you will also need to add it in the docs (in docs/api/losses.rst).

Ping me if you have any difficulties and thank you again!

[vroulet]

Hello @GrantMcConachie!

Yes, it's a bug due to a new release of jax, it's already been pointed out in #904. I've upstreamed the bug internally to the jax team, I hope it can get solved. I don't see any quick hot fix (the bug shall affect many of our modules). I can ping you when it gets solved :)

[vroulet]

The tests have been fixed in #908. You should be able to proceed. Again it would be really nice if you could avoid passing through an exponential without careful tricks like the losumexptrick (see https://en.wikipedia.org/wiki/LogSumExp). Alternatively, you may use directly functions like jax.nn.log_softmax that implemented such a trick.
Thanks again for the PR!

[GrantMcConachie]

I am still working on this! It is trickier to implement than I thought due to the partitioning of positive and negative pairs. For example If there is 4 embeddings with labels like this:

e1 -> 0
e2 -> 0
e3 -> 0
e4 -> 1

then your cosine similarity matrix will be a 4x4 matrix like this:

[ 0          sim01+     sim02+     sim03-]
[sim10+        0        sim12+     sim13-]
[sim20+      sim21+       0        sim23-]
[sim30       sim31      sim32        0   ]

Where the pluses indicate positive pairs.

So we want the loss to be log(exp(sim01+) / (exp(sim01+) + exp(sim03-)) + log(exp(sim02+) / (exp(sim02+) + exp(sim03-)) + log(exp(sim10+) / (exp(sim10+) + exp(sim13-)) + .... However a row wise log_softmax on this matrix will return values log(exp(sim01+) / (exp(sim01+) + exp(sim02+) + exp(sim03-)) and will incorporate exp(sim02+) in the denominator when it shouldn't.

I think there's probably a way around this using the logsumexp or log_softmax, but I am still trying to figure out how to do it.

[GrantMcConachie]

Hi @vroulet! I believe that I implemented the loss function using the same trick that logsumexp and log_softmax. I did not use these functions explicitly, but I was able to "normalize" the cosine similarity values by subtracting the row wise maximum cosine similarity values from each cosine similarity value before exponentiating, summing, then taking the logarithm. I believe this is sufficient to avoid overflow/underflow problems, but please let me know if you see an issue with this!

[vroulet]

Perfect! Thank you for pushing this through!
I'm leaving last small comments and we'll merge :)
Also you'll need to add the loss to the main __init__.py file (otherwise it won't be caught in the doc; for an example see definition of e.g. convex_kl_divergence in __init__.py and its appearance in the definition of __all__ in the __init__.py file).

[GrantMcConachie]

Thanks for all the edits! I am happy that this will get added. Please let me know if there's more I can do!

[vroulet]

Thank you again @GrantMcConachie !