How can I vectorize this for-loop loss function?

jhuteau: <p>I want an efficient (vectorized) implementation instead of my current for-loop implementation.</p>
<p>I have a MxN matrix named Sim corresponding to the similarity scores of M anchors with N documents.<br>
I have multiple negative (not matching the anchors) documents as well as multiple positive (matching) documents for each anchor, with the number of negative/positive documents not being the same for each anchor.</p>
<p>A triplet margin loss is defined on a triplet (anchor, positive, negative) and is computed as such:</p>
<pre><code class="lang-auto">H(a,p,n) = max(0, margin + Sim(a,n) - Sim(a,p))
</code></pre>
<p>The anchor corresponds to a row index, while the positive and negative are column indices.</p>
<p>For each anchor, my loss function is defined as the mean of the triplet losses with the f% least similar positives and the f% most similar positives, with f being an hyper-parameter (the same for all queries).</p>
<p>Say for the anchor <code>i</code> we have 10 positives and 100 negatives, with f=0.5, we want to select the 5 minimum Sim(i,p), the 50 maximum S(i,n), and compute the mean of the triplet loss for all pairs of positive/negatives (the mean of a 5x50 matrix). Then for anchor <code>j</code> we have 20 positives and 90 negatives, resulting in a 10x45 matrix of triplet losses to be averaged.</p>
<p>I have an implementation which computes the loss of each anchor in a for loop, but as you can imagine this is slow (with a small sim scores network, 50% of the training time is spent computing the loss).<br>
If I constrain my problem to only have 1 positive document per anchor, I can make an efficient vectorized implementation (negligible time cost). I suppose doing an efficient implementation with a different (but still constant) number of positives would not be too hard.</p>
<p>How can I efficiently implement the loss with a varying number of positive documents for each anchor ?</p>
<p>Here is the code for the loop version:</p>
<pre><code class="lang-auto">def hardest_triplet_margin_loss(preds, targets, margin, hardest_fraction):
    """
        preds is a vector of N scores.
        targets is a vector of N floats being the 
	'positiveness' of the corresponding element in preds.
    """
    targets = targets &gt;= 0.5
    positives = preds[targets]
    negatives = preds[~targets]

    hardest_positives = top_fraction(positives, hardest_fraction, False)\
            .unsqueeze(1)
    hardest_negatives = top_fraction(negatives, hardest_fraction, True)\
            .unsqueeze(0).expand(hardest_positives.numel(), -1) 

    losses = (margin + hardest_negatives - hardest_positives).clamp(min=0)
    loss = losses.mean()
    return loss

def top_fraction(x, fraction, largest=True):
    k = max(1, math.ceil(fraction * x.numel()))
    out = x.topk(k, largest=largest)[0]
    return out

num_anchors = 100
num_documents = 200

margin = 0.2
hardest_fraction = 0.5

# Matrix of similarity scores of anchors with documents (predicted by a NN)
sim = torch.rand(num_anchors, num_documents)

# Matrix of 'positiveness' of documents for each query. 
# (&gt;=0.5 means it's a positive) (given as input in the data)
targets = torch.rand(num_anchors, num_documents)

anchor_losses = torch.zeros(num_anchors)
for i in range(num_anchors):
    anchor_losses[i] = hardest_triplet_margin_loss(
        sim[i], targets[i], margin, hardest_fraction) 

loss = anchor_losses.mean()
</code></pre>
<p>If it’s not possible to fully vectorize the computation, I suppose doing it for blocks of anchors having the same number of positives would be the next best thing, and I think I can handle that on my own.</p>
<p>Thanks for reading me!</p>

KFrank: <p>Hi Jhuteau!</p>
<aside class="quote no-group quote-modified" data-username="jhuteau" data-post="1" data-topic="122407" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/439d5e/48.png" class="avatar"> jhuteau:</div>
<blockquote>
<p>I want an efficient (vectorized) implementation instead of my current for-loop implementation.<br>
…<br>
How can I efficiently implement the loss with a varying number of positive documents for each anchor ?</p>
</blockquote>
</aside>
<p>Easy peasy:</p>
<pre data-code-wrap="python"><code class="lang-python">import math

import torch

print (torch.__version__)

torch.manual_seed (20212021)


# code with tensor operations

def hardest_triplet_margin_lossB (preds, targets, margin, hardest_fraction):
    # preds:             the similarity scores, of shape [num_anchors, num-documents]
    # targets:           the "positiveness" of the scores, of the same shape
    # margin:            the margin-loss margin
    # hardest_fraction:  fraction of weakest "positive" and "negative" scores to analyze
    
    n_doc = preds.shape[1]
    
    pos_mask = (targets &gt;= 0.5).int()
    neg_mask = 1 - pos_mask
    max_preds = preds.max (dim = 1)[0]
    min_preds = preds.min (dim = 1)[0]
    
    pos_sort = pos_mask * preds  +  neg_mask * (1 + max_preds).unsqueeze (1)
    pos_sort = pos_sort.sort (dim = 1, descending = True)[0]
    pos_k = (hardest_fraction * pos_mask.sum (dim = 1)).ceil().maximum (torch.ones (1))
    pos_hard_mask = ((n_doc - pos_k).unsqueeze (1) &lt;= torch.arange (n_doc)).int()
    
    neg_sort = neg_mask * preds  +  pos_mask * (min_preds - 1).unsqueeze (1)
    neg_sort = neg_sort.sort (dim = 1)[0]
    neg_k = (hardest_fraction * neg_mask.sum (dim = 1)).ceil().maximum (torch.ones (1))
    neg_hard_mask = ((n_doc - neg_k).unsqueeze (1) &lt;= torch.arange (n_doc)).int()
    
    margin_mask = neg_hard_mask.unsqueeze (1) * pos_hard_mask.unsqueeze (2)
    margins = ((neg_sort.unsqueeze (1) - pos_sort.unsqueeze (2)) + margin).clamp (min = 0)
    
    losses = (margin_mask * margins).sum (dim = (1, 2)) / margin_mask.sum (dim = (1, 2))
    loss = losses.mean()
    return loss


# code with loop

def hardest_triplet_margin_loss(preds, targets, margin, hardest_fraction):
    """
        preds is a vector of N scores.
        targets is a vector of N floats being the 
	'positiveness' of the corresponding element in preds.
    """
    targets = targets &gt;= 0.5
    positives = preds[targets]
    negatives = preds[~targets]

    hardest_positives = top_fraction(positives, hardest_fraction, False)\
            .unsqueeze(1)
    hardest_negatives = top_fraction(negatives, hardest_fraction, True)\
            .unsqueeze(0).expand(hardest_positives.numel(), -1) 

    losses = (margin + hardest_negatives - hardest_positives).clamp(min=0)
    loss = losses.mean()
    return loss

def top_fraction(x, fraction, largest=True):
    k = max(1, math.ceil(fraction * x.numel()))
    out = x.topk(k, largest=largest)[0]
    return out

num_anchors = 100
num_documents = 200

margin = 0.2
hardest_fraction = 0.5

# Matrix of similarity scores of anchors with documents (predicted by a NN)
sim = torch.rand(num_anchors, num_documents)

# Matrix of 'positiveness' of documents for each query. 
# (&gt;=0.5 means it's a positive) (given as input in the data)
targets = torch.rand(num_anchors, num_documents)

anchor_losses = torch.zeros(num_anchors)
for i in range(num_anchors):
    anchor_losses[i] = hardest_triplet_margin_loss(
        sim[i], targets[i], margin, hardest_fraction) 

loss = anchor_losses.mean()

###

print ('loss =', loss)

lossB = hardest_triplet_margin_lossB (sim, targets, margin, hardest_fraction)
print ('lossB =', lossB)

print ('torch.allclose (lossB, loss) =', torch.allclose (lossB, loss))
print ('(lossB == loss) =', (lossB == loss))
</code></pre>
<pre><code class="lang-plaintext">1.7.1
loss = tensor(0.6941)
lossB = tensor(0.6941)
torch.allclose (lossB, loss) = True
(lossB == loss) = tensor(True)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

jhuteau: <p>I benchmarked different implementations on CPU, namely KFrank’s solution, <code>chunked_hardest_fraction_triplet_margin_loss</code> and <code>hardest_triplet_margin_loss</code>.<br>
<code>chunked</code> splits the preds into groups having the same number of positives, while <code>hardest</code> only work if you want the absolute hardest triplet.</p>
<pre><code class="lang-auto"># Multiple positives, hardest fraction, chunk vectorized O(n^2) memory
def chunked_hardest_fraction_triplet_margin_loss(
        preds, targets, margin, hardest_fraction):
    targets = (targets &gt;= 0.5)

    num_positives = targets.sum(dim=1)

    unique_num_positives = num_positives.unique()

    loss = torch.tensor(0.0)
    for i, single_num_positives in enumerate(unique_num_positives):
        group_indices = (num_positives == single_num_positives)\
            .nonzero(as_tuple=True)

        group_loss = constant_hardest_fraction_triplet_margin_loss(
                preds[group_indices], targets[group_indices], 
                margin, hardest_fraction)

        loss += group_loss * group_indices[0].size(0)
    loss /= preds.size(0)

    return loss

# Multiple (same # of) positives, hardest fraction, vectorized O(n^2) memory
def constant_hardest_fraction_triplet_margin_loss(
        preds, targets, margin, hardest_fraction):

    num_docs = preds.size(0)
    targets = (targets &gt;= 0.5)

    positive_indices = targets.nonzero(as_tuple=True)
    negative_indices = (~targets).nonzero(as_tuple=True)

    positives = preds[positive_indices].view(num_docs, -1)
    negatives = preds[negative_indices].view(num_docs, -1)

    pos_k = max(1, math.ceil(hardest_fraction * positives.size(1)))
    neg_k = max(1, math.ceil(hardest_fraction * negatives.size(1)))

    hardest_positives = positives.topk(pos_k, dim=1, largest=False)[0]
    hardest_negatives = negatives.topk(neg_k, dim=1, largest=True)[0]

    hardest_positives = hardest_positives.unsqueeze(1)
    hardest_negatives = hardest_negatives.unsqueeze(2)

    triplet_losses = (hardest_negatives + margin - hardest_positives).clamp(min=0)
    anchor_losses = triplet_losses.mean(dim=1)

    loss = anchor_losses.mean() / margin
    return loss

# Multiple positives, hardest triplet only, vectorized with O(n^2) memory
def hardest_triplet_margin_loss(
        preds, targets, margin):
    max_pred = preds.max(dim=1)[0]
    min_pred = preds.min(dim=1)[0]

    pos_mask = (targets &gt;= 0.5).int()
    neg_mask = 1 - pos_mask

    positives = preds * pos_mask + neg_mask * (max_pred + 1).unsqueeze(1)
    hardest_pos = positives.min(dim=1)[0]

    negatives = preds * neg_mask + pos_mask * (min_pred - 1).unsqueeze(1)
    hardest_neg = negatives.max(dim=1)[0]

    anchor_losses = (hardest_neg + margin - hardest_pos).clamp(min=0)
    loss = anchor_losses.mean() / margin
    return loss

</code></pre>
<p><code>naive</code> is a python for-loop implementation, <code>anchor</code> the implementation in the original post of this thread (mine), and <code>vectorized</code> KFrank’s implementation. I normalized results by the fastest one. MP_HF is for multiple positives with a fraction of the hardest, and MP_H for multiple positives with only the hardest.<br>
The tests were run on a 4-core CPU.<br>
I expect KFrank’s implementation to fare better on a GPU compared to the chunked implementation (but it requires O(num_anchors^3) memory). It also depends on how much the preds can be grouped (in the worst case <code>chunked==anchor</code>, and <code>anchor</code> is slower than <code>vectorized</code>).</p>
<p>Starting MP_HF.<br>
naive: 430.700<br>
anchor: 9.194<br>
chunked: 1.000<br>
vectorized: 2.516</p>
<p>Starting MP_H.<br>
naive: 1198.327<br>
anchor: 48.262<br>
chunked: 3.912<br>
vectorized: 11.033<br>
hardest: 1.000</p>