I benchmarked different implementations on CPU, namely KFrank’s solution, chunked_hardest_fraction_triplet_margin_loss and hardest_triplet_margin_loss.
chunked splits the preds into groups having the same number of positives, while hardest only work if you want the absolute hardest triplet.
# Multiple positives, hardest fraction, chunk vectorized O(n^2) memory
def chunked_hardest_fraction_triplet_margin_loss(
        preds, targets, margin, hardest_fraction):
    targets = (targets >= 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 >= 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 >= 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


naive is a python for-loop implementation, anchor the implementation in the original post of this thread (mine), and vectorized 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.
The tests were run on a 4-core CPU.
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 chunked==anchor, and anchor is slower than vectorized).
Starting MP_HF.
naive: 430.700
anchor: 9.194
chunked: 1.000
vectorized: 2.516
Starting MP_H.
naive: 1198.327
anchor: 48.262
chunked: 3.912
vectorized: 11.033
hardest: 1.000