[Question] Understanding terminology reference, embedding

[DanielRoeder1]

I am getting a bit mixed up with the terminology (reference, embedding, query) and want to ensure that I pass the correct parameters.

I got a Bi Encoder architecture trained using CrossBatchMemory & NTXent loss. Each training sample contains one anchor document and three subsamples (positives). The intended inference task will be that the model retrieves an anchor document given a subsample (positive) like input.

The documentation states that the embedding -> anchor, ref_emb-> positive regarding the inputs for loss functions.

pytorch-metric-learning/src/pytorch_metric_learning/losses/cross_batch_memory.py

Lines 83 to 92 in 691a635

	indices_tuple = self.create_indices_tuple(
	embeddings,
	labels,
	E_mem,
	L_mem,
	indices_tuple,
	do_remove_self_comparisons,
	)
	loss = self.loss(embeddings, labels, indices_tuple, E_mem, L_mem)
	return loss

Now when using cross-batch memory the samples in the buffer (E_mem) are passed as ref_emb thus as positives/negatives to the loss function.

This indicates that I should ensure the enqueue mask feeds only the positive samples (i.e. not buffer my anchors) into the buffer.

And finally, given my inference task, I would pass the subsamples (i.e. positives during training) as queries to the AccuracyCalculator and the documents as the reference for evaluating the accuracy?

Let me know if my rational is correct & thanks for any feddback

[KevinMusgrave]

Now when using cross-batch memory the samples in the buffer (E_mem) are passed as ref_emb thus as positives/negatives to the loss function.

This indicates that I should ensure the enqueue mask feeds only the positive samples (i.e. not buffer my anchors) into the buffer.

In your approach, the positive pairs will be (anchor, positive1), (anchor, positive2)..., and everything in the buffer (i.e. positive1 and positive2) will be detached from gradient computation.

Another approach is to not use enqueue_mask at all. In this case, the anchors and subsamples will all be added to memory, and also used as the embeddings. Then the loss will be computed from all positive pairs, like

(anchor, positive1), (anchor, positive2), ..., (positive1, anchor), (positive1, positive2), ..., (positive2, anchor), (positive2, positive1), ...

The 1st element of each tuple comes from the current batch, and is used for gradient computation. The 2nd element of each tuple comes from the buffer and is therefore detached and not used in gradient computation.

One thing to be careful about in both approaches is to make sure your labels are unique across batches. For example, in some self-supervised scenarios, labels are not part of the dataset and are created per batch instead:

for d in data:
   embeddings = model(d)
   labels = torch.arange(batch_size)

But the above snippet would result in each batch's embeddings[0] to form positive pairs with all the previous batch's embeddings[0] that were added to the buffer, since they all have the same label, labels[0].

And finally, given my inference task, I would pass the subsamples (i.e. positives during training) as queries to the AccuracyCalculator and the documents as the reference for evaluating the accuracy?

Yes

[DanielRoeder1]

Now when using cross-batch memory the samples in the buffer (E_mem) are passed as ref_emb thus as positives/negatives to the loss function.
This indicates that I should ensure the enqueue mask feeds only the positive samples (i.e. not buffer my anchors) into the buffer.

In your approach, the positive pairs will be (anchor, positive1), (anchor, positive2)..., and everything in the buffer (i.e. positive1 and positive2) will be detached from gradient computation.

Another approach is to not use enqueue_mask at all. In this case, the anchors and subsamples will all be added to memory, and also used as the embeddings. Then the loss will be computed from all positive pairs, like

(anchor, positive1), (anchor, positive2), ..., (positive1, anchor), (positive1, positive2), ..., (positive2, anchor), (positive2, positive1), ...

The 1st element of each tuple comes from the current batch, and is used for gradient computation. The 2nd element of each tuple comes from the buffer and is therefore detached and not used in gradient computation.

One thing to be careful about in both approaches is to make sure your labels are unique across batches. For example, in some self-supervised scenarios, labels are not part of the dataset and are created per batch instead:

for d in data:
   embeddings = model(d)
   labels = torch.arange(batch_size)

But the above snippet would result in each batch's embeddings[0] to form positive pairs with all the previous batch's embeddings[0] that were added to the buffer, since they all have the same label, labels[0].

And finally, given my inference task, I would pass the subsamples (i.e. positives during training) as queries to the AccuracyCalculator and the documents as the reference for evaluating the accuracy?

Yes

I still have the following two questions:
As the second element of each tuple is from the buffer and thus not part of the gradient calculation, would this mean that when storing only the positives in the buffer which are encoded using a separate model from the anchors, this model would not get updated?

Why does switching the embeddings and ref_emb in a pairwise loss affect the loss value? (NtXent loss temp set to 1 to medicate effect)

loss_func = NTXentLoss(temperature=1)
embeds = torch.rand(2,100)
ref = torch.rand(4,100)
labels = torch.arange(2)
ref_label = labels.clone().repeat_interleave(2)
loss_func(embeddings = ref,labels = ref_label, ref_labels = labels, ref_emb = embeds)
# Indice tuples: (tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([1, 1, 0, 0]))
# Loss value: tensor(0.6727)
loss_func(embeddings = embeds,labels = labels, ref_labels = ref_label, ref_emb = ref)
# Indice tuples: (tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([2, 3, 0, 1]))
# Loss value: tensor(1.0712)

The indice tuples suggest that the same pairs were formed. Where is this difference created?

[KevinMusgrave]

Sorry for the late reply.

I still have the following two questions: As the second element of each tuple is from the buffer and thus not part of the gradient calculation, would this mean that when storing only the positives in the buffer which are encoded using a separate model from the anchors, this model would not get updated?

Yes I believe so. If you want the buffer elements to be included in the gradient calculation, then you can calculate this loss separately, in addition to using CrossBatchMemory:

loss_fn1 = NTXentLoss()
loss_fn2 = CrossBatchMemory()

# all embeddings in current batch will have gradients
loss1 = loss_fn1(emb, labels, ref_labels=ref_labels, ref_emb=ref_emb)
# embeddings in queue will not have gradients
loss2 = loss_fn2(emb, labels)
total_loss = loss1 + loss2

Why does switching the embeddings and ref_emb in a pairwise loss affect the loss value? (NtXent loss temp set to 1 to medicate effect)

loss_func = NTXentLoss(temperature=1)
embeds = torch.rand(2,100)
ref = torch.rand(4,100)
labels = torch.arange(2)
ref_label = labels.clone().repeat_interleave(2)
loss_func(embeddings = ref,labels = ref_label, ref_labels = labels, ref_emb = embeds)
# Indice tuples: (tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([1, 1, 0, 0]))
# Loss value: tensor(0.6727)
loss_func(embeddings = embeds,labels = labels, ref_labels = ref_label, ref_emb = ref)
# Indice tuples: (tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([2, 3, 0, 1]))
# Loss value: tensor(1.0712)

The indice tuples suggest that the same pairs were formed. Where is this difference created?

The reason is that in the loss calculation, the negatives for each positive pair are collected based on anchor index. This is probably best illustrated in this comment: #6 (comment) and at this line in the test for the loss function:

pytorch-metric-learning/tests/losses/test_ntxent_loss.py

Line 115 in 691a635

if a2 == a1:

For your example:

Case 1:

• indices_tuple: (tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([1, 1, 0, 0]))
• For positive pair (0,0), the negative pairs are [(0,1)]
• For positive pair (1,0), the negative pairs are [(1,1)]
• For positive pair (2,1), the negative pairs are [(2,0)]
• For positive pair (3,1), the negative pairs are [(3,0)]

Case 2:

• indices_tuple: (tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([0, 0, 1, 1]), tensor([2, 3, 0, 1]))
• For positive pair (0,0), the negative pairs are [(0,2), (0,3)]
• For positive pair (0,1), the negative pairs are [(0,2), (0,3)]
• For positive pair (1,2), the negative pairs are [(1,0), (1,1)]
• For positive pair (1,3), the negative pairs are [(1,0), (1,1)]

[DanielRoeder1]

Thank you for this answer this finally made me understand where the difficulty in my loss calculations originates!

I tried to achieve the following by setting custom indices_tuple but given this explanation, I think it isn't possible maybe you can suggest a solution:

I am trying to train a PolyEncoder-like architecture i.e. there is an extra late-stage interaction in which the query embedding are conditioned on the document embeddings before being scored against said document embeddings.

Pseudocode example for understanding:

query_embeds =  [num_que, dim]
doc_embeds = [num_doc, dim]
query_cond_embeds = PolyEncoder(query_embeds, doc_embeds) -> [num_que, num_doc, dim]
query_cond_embeds = query_cond_embeds.reshape(-1,dim) -> [num_que*num_doc, dim]
loss_func(query_cond_embeds, doc_embeds, custom_indice_tuples)

Now here are the kind of pos neg pairs I would like to form (note that q0d0 indicates query 0 conditioned on doc 0)
In this example, we have 3 documents, and 6 queries i.e. 2 queries each belong to one document

• For positive pair (q0d0,d0), the negative pairs are [(q0d1,d1), (q0d2,d2)]
• For positive pair (q1d0,d0), the negative pairs are [(q1d1,d1), (q1d2,d2)]
• For positive pair (q2d1,d1), the negative pairs are [(q2d0,d0), (q2d2,d2)]
• For positive pair (q3d1,d1), the negative pairs are [(q3d0,d0), (q3d2,d2)]
• for positive pair (q4d2,d2), the negative pairs are [(q4d0,d0),(q4d1,d1)]
• for positive pair (q5d2,d2), the negative pairs are [(q5d0,d0),(q5d1,d1)]

I wrote code that can create the custom indices_tuples but as you mentioned negative pairs are formed based on matching anc_idx. As in this example there is no chance to create a matching anc_idx as each conditioned query embedding as well as each document embedding has a unique index, is there any way to implement this without rewriting the complete loss functions?

Thanks for the help so far.

[KevinMusgrave]

As in this example there is no chance to create a matching anc_idx as each conditioned query embedding as well as each document embedding has a unique index, is there any way to implement this without rewriting the complete loss functions?

Unfortunately I don't think there's a way to do that, so you'll have to write a custom loss function.

[KevinMusgrave]

I'm guessing for pair (a1, p) you want to find all the negatives where labels[a1] == labels[a2]. Perhaps changing the condition of this line could work:

pytorch-metric-learning/src/pytorch_metric_learning/losses/ntxent_loss.py

Line 26 in 691a635

n_per_p = c_f.to_dtype(a2.unsqueeze(0) == a1.unsqueeze(1), dtype=dtype)

Something like

n_per_p = c_f.to_dtype(labels[a2].unsqueeze(0) == labels[a1].unsqueeze(1), dtype=dtype)

But I realize that _compute_loss doesn't take in labels, so you'd have to pass that in from GenericPairLoss if you take that approach.

[DanielRoeder1]

I looked into adapting the superclasses of NTXent loss to allow for this behavior. Unfortunately, this would have meant changing the comput_loss functions of all other losses that inherit from them to expect the input parameter.

I settled on coding a custom n_per_p in the NTXentLoss:

num_pos_pairs = pos_pairs.shape[0]
num_neg_pairs = neg_pairs.shape[0]
neg_per_pos = num_neg_pairs // num_pos_pairs
n_per_p = torch.zeros((num_pos_pairs,num_neg_pairs))
i0 = torch.arange(num_pos_pairs).repeat_interleave(neg_per_pos)
i1 = torch.arange(num_neg_pairs)
n_per_p[i0,i1] = 1
n_per_p = n_per_p.to(torch.device("cuda:0"))

Here I am exploiting the fact that I know that the first x pairs of the negative pairs belong to positive pair y and so on.

Thank you for your help, amazing how quickly you respond given the number of issues opened in this repo!

[KevinMusgrave]

I'm glad to hear you figured out a workaround.