Hello Simon!



 smonsays:

While tinkering with the official code example for Variational Autoencoders, I experienced some unexpected behaviour with regard to the Binary Cross-Entropy loss. When I use F.binary_cross_entropy in combination with the sigmoid function, the model trains as expected on MNIST.  However, when changing to the F.binary_cross_entropy_with_logits function, the loss suddenly becomes arbitrarily small during training and the model no longer produces meaningful results.
# For this loss function the loss becomes arbitrarily small 
BCE = F.binary_cross_entropy_with_logits(recon_x, x.view(-1, 784), reduction='sum') / x.shape[0]

# For this loss function, the training works as expected
BCE = F.binary_cross_entropy(torch.sigmoid(recon_x), x.view(-1, 784), reduction='sum') / x.shape[0]

To my understanding, the only difference between the two approaches should be numerical stability. Am I missing something?


I agree with your understanding, and looking at the two lines of
code you posted, I don’t see anything suspicious (although I miss
things all the time …).
If I had to guess, I would guess that you have a typo somewhere
else in your code that causes the two runs to differ.
However, we can start by checking out some basics:
Here I run the example given in the documentation for
torch.nn.functional.binary_cross_entropy.   (Please note
that I am running this test, for whatever reason, with pytorch 0.3.0.)
Here is the script:
import torch
print (torch.__version__)
torch.manual_seed (2019)
input =  torch.autograd.Variable (torch.randn ((3, 2)))
print (input)
target = torch.autograd.Variable (torch.rand ((3, 2)))
print (target)
loss_plain = torch.nn.functional.binary_cross_entropy (torch.sigmoid (input), target)
print (loss_plain)
loss_logits = torch.nn.functional.binary_cross_entropy_with_logits (input, target)
print (loss_logits)
print (loss_plain - loss_logits)

And here is the output:
>>> import torch
>>> print (torch.__version__)
0.3.0b0+591e73e
>>> torch.manual_seed (2019)
<torch._C.Generator object at 0x00000207759A60F0>
>>> input =  torch.autograd.Variable (torch.randn ((3, 2)))
>>> print (input)
Variable containing:
-0.1187  0.2110
 0.7463 -0.6136
-0.1186  1.5565
[torch.FloatTensor of size 3x2]

>>> target = torch.autograd.Variable (torch.rand ((3, 2)))
>>> print (target)
Variable containing:
 0.7628  0.0721
 0.2208  0.3979
 0.6338  0.1922
[torch.FloatTensor of size 3x2]

>>> loss_plain = torch.nn.functional.binary_cross_entropy (torch.sigmoid (input), target)
>>> print (loss_plain)
Variable containing:
 0.8868
[torch.FloatTensor of size 1]

>>> loss_logits = torch.nn.functional.binary_cross_entropy_with_logits (input, target)
>>> print (loss_logits)
Variable containing:
 0.8868
[torch.FloatTensor of size 1]

>>> print (loss_plain - loss_logits)
Variable containing:
1.00000e-08 *
 -5.9605
[torch.FloatTensor of size 1]

And, indeed, the two expressions are the same (up to floating-point
precision).
Is there any way you can capture a specific instance of your
recon_x and x.view() so that you can pump what you know
are the same values into your two cross-entropy expressions?
Is there any way you track your loss (your BCE) on a step-by-step
basis so you can see when they first diverge?
The thing that is noteworthy to me is that you say that the less
numerically stable version (regular bce) works, while the more
stable version (bce_with_logits) quits working at some point.
This is backwards of what one might expect.
Perhaps there is something about your model that puts it on
the edge of being poorly behaved.  In such a case it could be
plausible that, by happenstance, the bce version stays in a
well-behaved region, but small differences due to using the
bce_with_logits version cause it to drift into a poorly-behaved
region.
If this were the case I would expect that other perturbations
such as starting with different initial weights, or using a different
optimization algorithm or learning rate could cause the training
to “randomly” end up being well behaved or poorly behaved.
So:
I would

proofread the code to make sure there isn’t some outright error
check that your model and data are reasonably well behaved
and stable with respect to perturbations of the details of your
training
track your bce and bce_logits step by step to find out where
they first diverge, and drill down with the values that immediately
precede the divergence.

Have fun!
K. Frank