`masked_fill` to `-inf` causes loss to not drop

Paulo_Costa: <p>i have a model that produces user embeddings, i batch them so that each row is a user, and the columns are the dimensions of their embeddings. the batch can have multiple of the same user. i’m training the model so that the user embeddings are close to the positive embeddings in embedding space. i want to use two sources of negatives—in batch negatives and random negatives. i aim to create a tensor like this:</p>
<pre data-code-wrap="python"><code class="lang-python"># u -&gt; user
# p -&gt; positives
# both tensors are 2d, rows are instances, and columns are dimension of embeddings
# rows in both correspond, so p[0] is the positive example for u[0]
# so when multiplied the diagonal values are the true positives
# the user should be close to that product in embedding space
&gt;&gt;&gt; user @ positives.t()
[ 
   [u0p0, u0p1, u0p2],
   [u1p0, u1p1, u1p2],
   [u2p0, u2p1, u2p2],
]
# so the diagonals are the positives here, and all other values are the negatives
</code></pre>
<p>when using in batch negatives—where the positives from some users are treated as negatives for other users in batch, since we can have multiple instances of the same user per batch, we run the risk of marking some instances as negatives when they actually should be near the user in embedding space.</p>
<p>to combat this, i create a mask that marks these false positives, and replace these false positive values with <code>-torch.inf</code>. this causes problems. without the mask, and nothing else changed. the model learns fine! but as soon as the mask is applied, the model learns nothing, and the loss never drops. something about using <code>masked_fill</code> is messing with my learning.</p>
<p>minimal code example for loss function:</p>
<pre data-code-wrap="python"><code class="lang-python">def _create_false_negative_mask(
    user: Tensor,
    user_ids: Tensor,
) -&gt; Tensor:
    """Create a mask where True if the position is not a false negative."""
    # n x m mask where true if same retailer in element n appears in element m
    is_same_retailer = user_ids[None, :] == user_ids[:, None]
    diagonal_true = torch.eye(
        user.size(dim=0), dtype=torch.bool, device=user.device
    )
    return ~(is_same_retailer &amp; ~diagonal_true)


def _calculate_diagonal_positive_logits(
    user: Tensor,
    positives: Tensor,
    user_ids: Tensor,
) -&gt; Tensor:
    """Calculate positive logits, masking out false negatives.

    Creates tensor where diagonals are the true positives, and others
    may be treated as negatives.
    """
    false_positive_mask = _create_false_negative_mask(
        user=user,
        user_id=user_ids,
    )
    positive_logits = user @ positives.t()
    # problem area!!!!!!!!!!!!!! if use use masked_fill here, then
    # model does not learn at all, if we simply return positive_logits,
    # loss drops as expected
    return positive_logits.masked_fill(~false_positive_mask, value=-torch.inf)


class SampledSoftmax(nn.Module):

    def __init__(self) -&gt; None:
        super().__init__()

    def forward(  # noqa: PLR0913
        self,
        user: Tensor,  # (B, D) user embeddings
        positives: Tensor,  # (B, D) positive embeddings
        negatives: Tensor,  # (N, D) negative embeddings
        user_ids: Tensor,  # (B,)
        tau: Tensor,
    ) -&gt; Tensor:
        """Compute sampled softmax loss."""
        positive_logits = _calculate_diagonal_positive_logits(
            user=user,
            positives=positives,
            user_ids=user_ids,
        )  # (B, B) with -inf for false negatives
        negative_logits = user @ negatives.t()  # (B, N)
        all_logits = (
            torch.cat((positive_logits, negative_logits), dim=1) / tau
        )  # (B, B+N)
        # We create a diagonal matrix where each row's true positive is
        # at the diagonal
        target = torch.arange(
            user.size(dim=0), dtype=torch.long, device=user.device
        )
        return functional.cross_entropy(input=all_logits, target=target)

</code></pre>
<p>things i’ve checked:</p>
<ul>
<li>one risk is that you mask all the values in the row except for one, this is a non issue here, <code>negatives</code> is large, batch size (number of rows) is 250, and <code>negatives</code> has 500 rows. so at most we get 9 / (250 + 500) masked values per row</li>
<li>we have checks that diagonals are never masked</li>
<li>we check that there is always at least one non-inf value per row in the final <code>all_logits</code></li>
</ul>