BCEWithLogitsLoss requires using pos_weight (not needed in TensorFlow/optax)

sammlapp: <p>I’m a bit confused by an empirical observation when using BCEWithLogitsLoss.</p>
<p>The scenario: I have very sparse multi-target labels for 1000 classes. For instance, each sample has on average 1 label. However, samples can have 0, 1, 2, or more labels, so it is a multi-target problem. I’m training any CNN, for instance resnet18, on a training dataset with a balanced dataset. The loss is evaluated on “full size” matrices of labels vs logits output by the CNN:</p>
<p><code>loss = BCEWithLogitsLoss( logits_2d, target_2d)</code></p>
<p><em>Observation</em>: If I train with BCEWithLogitsLoss (without specifying pos_weight), the model always quickly learns to predict 0 for all classes and all samples. It then cannot escape this local minimum. This makes sense, because the labels are 99.9% 0, so the loss will be decreased by learning to output low logit values for all classes. When training, all logits quickly become &lt;0. To avoid this problem, I can specify<br>
<code>BCEWithLogitsLoss(pos_weight=1/np.mean(labels))</code>.<br>
This makes the 1/1000 positive labels equally as important as the 999/1000 negative labels, and the model trains properly, with logit outputs averaging around 0, and the model learns.</p>
<p>This “trick” of heavily weighting the loss doesn’t seem to be necessary when training in TF or jax+optax: default parameters of binary cross entropy loss train fine with sparse multi-target labels.</p>
<p><em>Question</em>: Is there some implementation difference between BCEWithLogitsLoss and corresponding loss functions in other libraries (eg <a href="https://optax.readthedocs.io/en/latest/api.html#optax.sigmoid_binary_cross_entropy" rel="noopener nofollow ugc">optax BCE loss</a> or <a href="https://www.tensorflow.org/api_docs/python/tf/keras/losses/CategoricalCrossentropy" rel="noopener nofollow ugc">TF CategoricalCrossentropy loss</a> ) such that the default behavior in other packages acts like <code>BCEWithLogitsLoss(pos_weight=1/np.mean(labels))</code>?</p>
<p>I don’t know what <code>reduction=losses_utils.ReductionV2.AUTO</code> means in the TF implementation so I’m not sure if that could be involved.</p>
<p>Other users have reported similar behavior for PyTorch BCE Loss:</p><aside class="quote quote-modified" data-post="2" data-topic="103422">
  <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/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/multi-label-model-outputs-only-negative-values-while-trained-with-bcewithlogitsloss/103422/2">Multi-label model outputs only negative values while trained with BCEWithLogitsLoss()</a> 
  </div>
  <blockquote>
    Hi Amir - 

The short answer is that your dataset is “unbalanced,” so you 
should try using the pos_weight argument when you construct 
your BCEWithLogitsLoss loss criterion. 
Looking at your target, and naively assuming that all positive 
class labels appear about equally frequently, it appear that any 
giving class will be labelled positive only once in about every 
nine images. 
So your classifier could do a pretty good job by just predicting 
negative for all of your classes all of the time. …
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="1" data-topic="153313">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zahra_momeni/48/50247_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/multi-label-classification-all-predictions-become-zero/153313">Multi-label classification: all predictions become zero</a> 
  </div>
  <blockquote>
    I’m working on a multi-label classification. The number of labels are 122, and for each sample,  a maximum of 10 of these labels are one and the rest are zero. (so we have sparse multi-hot encoding) 
I used BCEWithLogitsLoss as a loss function in two ways: 
1- without weight 
criterion = nn.BCEWithLogitsLoss() 
outputs = self.model(inputs) 
loss = criterion(outputs, targets) 
2- with weight 
criterion = nn.BCEWithLogitsLoss(reduction=‘none’) 
outputs = self.model(inputs) 
loss = criterion(output…
  </blockquote>
</aside>