Moving to numerically stable log-sum-exp leads to extremely large loss values

Umang_Sharma: <p>I am working on a network that uses a LSTM along with MDNs to predict some distributions. The loss function I use for these MDNs involve trying to fit my target data to the predicted distributions. I am trying to compute the log-sum-exp for the log_probs of these target data to compute the loss. When I use standard log-sum-exp, I get reasonable initial loss values (around 50-70) even though later it encounters some NaNs and breaks. Based on what I have read online, a numerically stable version of log-sum-exp is required to avoid this problem. However as soon I use the stable version, my loss values shoot up to the order of 15-20k. They do come down upon training but eventually they also lead to NaNs.</p>
<p>NOTE : I did not use the logsumexp function in PyTorch, since I needed to have a weighted summation based on my mixture components.</p>
<pre><code class="lang-auto">def log_sum_exp(self,value, weights, dim=None):
        eps = 1e-20
        m, idx = torch.max(value, dim=dim, keepdim=True)
        return m.squeeze(dim) + torch.log(torch.sum(torch.exp(value-m)*(weights.unsqueeze(2)),
                                       dim=dim) + eps)

def mdn_loss(self, pi, sigma, mu, target):
        eps = 1e-20
        target = target.unsqueeze(1)
        m = torch.distributions.Normal(loc=mu, scale=sigma)
        probs = m.log_prob(target)
        # Size of probs is batch_size x num_mixtures x num_out_features
        # Size of pi is batch_size x num_mixtures 
        loss = -self.log_sum_exp(probs, pi, dim=1)
        return loss.mean()
</code></pre>
<p>Upon adding anomaly_detection, the NaNs seem to occur at :<br>
probs = m.log_prob(target)</p>
<p>Seeing these huge initial loss values just by moving to the numerically stable version have led me to believe I have some bug in my current implementation. Any help please.</p>

KFrank: <p>Hi Umang!</p>
<aside class="quote no-group" data-username="Umang_Sharma" data-post="1" data-topic="61938" data-full="true">
<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/umang_sharma/48/18061_2.png" class="avatar"> Umang_Sharma:</div>
<blockquote>
<p>I am trying to compute the log-sum-exp for the log_probs of these target data to compute the loss. When I use standard log-sum-exp, I get reasonable initial loss values (around 50-70) even though later it encounters some NaNs and breaks. Based on what I have read online, a numerically stable version of log-sum-exp is required to avoid this problem. However as soon I use the stable version, my loss values shoot up to the order of 15-20k. They do come down upon training but eventually they also lead to NaNs.</p>
<p>NOTE : I did not use the logsumexp function in PyTorch, since I needed to have a weighted summation based on my mixture components.</p>
<pre><code class="lang-auto">def log_sum_exp(self,value, weights, dim=None):
        eps = 1e-20
        m, idx = torch.max(value, dim=dim, keepdim=True)
        return m.squeeze(dim) + torch.log(torch.sum(torch.exp(value-m)*(weights.unsqueeze(2)),
                                       dim=dim) + eps)
</code></pre>
</blockquote>
</aside>
<p>I haven’t checked your weighted, stable log-sum-exp implementation<br>
carefully, but, at first glance, it looks correct to me.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        m = torch.distributions.Normal(loc=mu, scale=sigma)
        probs = m.log_prob(target)
</code></pre>
<p>Upon adding anomaly_detection, the NaNs seem to occur at :<br>
probs = m.log_prob(target)</p>
</blockquote>
</aside>
<p>If your issues really do show up first with <code>probs = m.log_prob(target)</code>,<br>
then your <code>log_sum_exp()</code> is likely not the problem.</p>
<p>Note that the log in <code>log_prob()</code> itself can overflow, depending on<br>
the value of <code>target</code> and on the distribution, in particular, the value<br>
of <code>sigma</code>.</p>
<p>Here is a brief pytorch (version 0.3.0) session that illustrates this:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.distributions.Normal (0.0, 1.0).log_prob (torch.FloatTensor ([0, 10, 1.e5, 1.e10, 1.e19, 1.e20]))

-9.1894e-01
-5.0919e+01
-5.0000e+09
-5.0000e+19
-5.0000e+37
       -inf
[torch.FloatTensor of size 6]

&gt;&gt;&gt; torch.distributions.Normal (0.0, 1.e-15).log_prob (torch.FloatTensor ([0, 10, 1.e5, 1.e10, 1.e19, 1.e20]))

 3.3620e+01
-5.0000e+31
       -inf
       -inf
       -inf
       -inf
[torch.FloatTensor of size 6]
</code></pre>
<p>The normal distribution falls off exponentially, so it is possible have<br>
“reasonable” values of <code>target</code> (and <code>sigma</code>) for which the normal<br>
distribution “underflows” to 0.0, at which point the log in <code>log_prob()</code><br>
will return <code>-inf</code>.</p>
<p>I assume that <code>target</code> is some well-defined, static data in your training<br>
dataset.  I would first check that your <code>target</code> data doesn’t contain any outlandish values.</p>
<p>Then I assume that your training is adjusting your value of <code>sigma</code><br>
(and <code>mu</code>).  If your training, for some reason, is pushing <code>sigma</code><br>
towards zero, then your <code>log_prob()</code> will be getting pushed into<br>
the far tails of the normal distribution (even for reasonable <code>target</code><br>
values), and you will get <code>-inf</code>, as illustrated above.</p>
<p>So I would check for small values of <code>sigma</code>, perhaps with something<br>
as simple as</p>
<pre data-code-wrap="python"><code class="lang-python">    if  sigma &lt; 1.e-10:
        print ("Yikes!  Tiny sigma!)
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>Seeing these huge initial loss values just by moving to the numerically stable version have led me to believe I have some bug in my current implementation.</p>
</blockquote>
</aside>
<p>If I understand your description correctly, it sounds like moving to your<br>
numerically-stable <code>log_sum_exp()</code> postponed your training issues,<br>
so I would say that that could be viewed as an improvement.  (I don’t<br>
see it as evidence for a bug in <code>log_sum_exp()</code>.)</p>
<p>My guess is that your training is driving your model in an unhealthy<br>
direction (e.g., <code>sigma --&gt; 0</code>).  I would check for that first.</p>
<p>Good luck.</p>
<p>K. Frank</p>

Umang_Sharma: <p>Hi Frank,<br>
Thanks a lot for your suggestions.<br>
“I would first check that your  <code>target</code>  data doesn’t contain any outlandish values.”<br>
Great advice, turns out I did have some outlandish target values. Removing them and normalizing the data has fixed the issue. However, I am still a bit confused as to why the loss gave values of such different magnitudes just by introducing the numerically stable log-sum-exp.</p>
<pre><code class="lang-auto">m, idx = torch.max(value, dim=dim, keepdim=True)
        return m.squeeze(dim) + torch.log(torch.sum(torch.exp(value-m)*(weights.unsqueeze(2)),
                                       dim=dim) + eps)
</code></pre>
<p>Does it have something to do with the max we are taking in this function?</p>
<p>Thanks,<br>
Umang</p>

KFrank: <p>Hello Umang!</p>
<aside class="quote no-group" data-username="Umang_Sharma" data-post="3" data-topic="61938" data-full="true">
<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/umang_sharma/48/18061_2.png" class="avatar"> Umang_Sharma:</div>
<blockquote>
<p>However, I am still a bit confused as to why the loss gave values of such different magnitudes just by introducing the numerically stable log-sum-exp.</p>
</blockquote>
</aside>
<p>This is just speculation on my part.  I don’t have a clear picture of<br>
the exact sequence of events, but it sounded to me like when you<br>
used the naive log-sum-exp, your training broke early with NaNs.<br>
And when you used the more stable log-sum-exp, your training<br>
gave very large loss values, and then later broke with NaNs.</p>
<p>I assume that the naive version broke before you got to the large<br>
loss values, so they were “hidden.”  The stable version was able<br>
to run further, so it was able to reach the large loss values.  (Again,<br>
just speculation).</p>
<p>Best.</p>
<p>K. Frank</p>