Formulating Metric Learning Problem

[kheyer]

I need some help figuring out how to structure labels for a metric learning problem.

I'm training a model for embedding based retrieval. In this context, there is one target document for each retrieval, and MRR/recall@1 is the desired metic.

I trained a base model with cosine similarity loss. Now I want to fine-tune with a ranking focused loss. For each target embedding, I also have the k nearest neighbors of the target embedding. Ideally I'd like some sort of ranking loss that drives the predicted embedding to be more similar to the target embedding and less similar to the target embedding's nearest neighbors.

I'm not sure how to structure the labels of this problem since there are no classes associated with the embeddings. What would be a good way to formulate this?

[KevinMusgrave]

In each batch you have predicted embeddings (pred) of size (N, D) and reference embeddings (ref) of size (M, D) where M > N.

For a specific pred[i], there's a ref[j] that you want it to be near (the target), so it should have the same label. All other embeddings in ref would be considered negatives then.

Here's one possible approach. It assumes you know which indices of ref you want each pred to be near.

loss_fn = ContrastiveLoss()

ref_labels = torch.arange(len(ref))
pred_labels = # I'm assuming you have this data already. pred_labels[i] is the index of `ref` that you want `pred` to be near

labels = torch.cat([ref_labels, pred_labels], dim=0)
embeddings = torch.cat([ref, pred], dim=0)
loss = loss_fn(embeddings, labels)

If you have multiple pred embeddings that map to the same ref embedding, then you'll have multiple pred embeddings with the same label. The loss function will explicitly push those pred embeddings together. If you don't want that, you can make it so that the relationship is only between pred and ref:

loss_fn = ContrastiveLoss()

ref_labels = torch.arange(len(ref))
pred_labels = # ... 

loss = loss_fn(pred, pred_labels, ref_emb=ref, ref_labels=ref_labels)