Correctly generating multivariate Gaussian random variable

username: Laya1
cooked: <p>I am generating a multivariate random variable from a Gaussian distribution with known mean and variance in the following two ways:</p>
<pre><code class="lang-auto">for t in range(T):
    samples[t, ...] = torch.normal(mean=means, std=torch.sqrt(variances))  # method 1
    samples[t, ...] = means + torch.mul(torch.sqrt(variances), torch.randn_like(means))  # method 2
</code></pre>
<p>I then perform training of a neural network based on <code>samples</code> as predictions and known <code>targets</code>. The losses that I get with these two methods and making no other changes are very different. If I use method 1, then the loss loss does not increase or decrease, i.e., the network does not learn anything. If I use method 2, the loss becomes <code>nan</code>.</p>
<p>According to my understanding, the two methods in the code snippet above are equivalent. I am not sure why the behaviour of the training changes so drastically. Is it a side effect of the loss, or are the two methods not equivalent?</p>

username: KFrank
cooked: <p>Hi Laya!</p>
<aside class="quote no-group" data-username="Laya1" data-post="1" data-topic="166050" data-full="true">
<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/l/8c91f0/48.png" class="avatar"> Laya1:</div>
<blockquote>
<p>According to my understanding, the two methods in the code snippet above are equivalent.</p>
</blockquote>
</aside>
<p>The two methods are not equivalent in that they do not handle<br>
backpropagation equivalently.</p>
<p>Specifically, <code>torch.normal()</code> doesn’t really backpropagate (although<br>
it acts like it does).  You can argue – as I suppose pytorch does – that<br>
<code>torch.normal()</code> <em>shouldn’t</em> backpropagate because there is no single,<br>
clear-cut definition of what it should mean to backpropagate through a<br>
sample from a probability distribution.  (I will say that it’s arguably a bug<br>
that the output of <code>torch.normal()</code> carries a <code>grad_fn</code> and produces<br>
gradients (that are zero) upon backpropagation.  To me it would make<br>
more sense for it simply not to support backpropagation.)</p>
<p>In “method 2” you can backpropagate just fine because <code>means</code> and<br>
<code>variances</code> only enter into standard differentiable functions – the output of<br>
<code>torch.randn_like()</code> are just “constants” (with <code>requires_grad = False</code>)<br>
as far as the computation involving <code>means</code> and <code>variances</code> is concerned.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.13.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; means = torch.randn (5, requires_grad = True)
&gt;&gt;&gt; vars = torch.rand (5, requires_grad = True)
&gt;&gt;&gt;
&gt;&gt;&gt; sample1 = torch.normal (mean = means, std = torch.sqrt (vars))              # method 1
&gt;&gt;&gt; sample2 = means + torch.mul (torch.sqrt (vars), torch.randn_like (means))   # method 2
&gt;&gt;&gt;
&gt;&gt;&gt; sample1                    # has grad_fn
tensor([-0.6208,  0.1995, -2.8783,  0.6203, -0.9839],
       grad_fn=&lt;NormalBackward3&gt;)
&gt;&gt;&gt; sample1.sum().backward()   # normal() pretends to backpropagate ...
&gt;&gt;&gt; means.grad                 # but with zero grad
tensor([0., 0., 0., 0., 0.])
&gt;&gt;&gt; vars.grad
tensor([0., 0., 0., 0., 0.])
&gt;&gt;&gt;
&gt;&gt;&gt; sample2                    # also has grad_fn
tensor([ 0.8952,  1.2811, -0.9495,  3.2830, -0.4096], grad_fn=&lt;AddBackward0&gt;)
&gt;&gt;&gt; sample2.sum().backward()   # and backpropagates "properly"
&gt;&gt;&gt; means.grad                 # meaningful grad
tensor([1., 1., 1., 1., 1.])
&gt;&gt;&gt; vars.grad
tensor([ 0.4397,  0.6463, -0.7891,  2.1393, -0.1648])
</code></pre>
<p>Whether this difference in backpropagation explains your loss blowing<br>
up would depend on the details of the rest of your computation.  But with<br>
different gradients the training you get with the two methods could well<br>
wander off in different directions with one of them becoming unstable while<br>
the other doesn’t.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Laya1
cooked: <p><a class="mention" href="/u/kfrank">@KFrank</a> Thank you very much for the explanation. That makes a lot of sense. I made a small change in the post - I was actually getting no improvement in loss with method 1 (it makes sense now since the grad is 0). I set <code>torch.autograd.set_detect_anomaly(True)</code> in my training routine and used method 2. The loss actually becomes <code>nan</code> and does not blow up (as mentioned earlier in the post). I get the following error:<br>
<code>RuntimeError: Function 'SqrtBackward0' returned nan values in its 0th output.</code></p>
<p>I think the error comes from the generation of the samples with method 2 that takes square root of the variances. The variable <code>variances</code> is generated from a <code>Linear</code> layer followed by <code>ReLU</code> activation, so it is never negative. I am now not sure why the loss is becoming <code>nan</code>.</p>
<p>To avoid this situation, I removed the <code>torch.sqrt()</code> and just used:<br>
<code>samples[t, ...] = means + torch.mul(variances, torch.randn_like(means))</code></p>
<p>The network now learns to predict the standard deviations (although the variable is named <code>variances</code>). When I do this, I get the following error:<br>
<code>RuntimeError: Function 'LogBackward0' returned nan values in its 0th output.</code></p>
<p>I am using a stochastic classification loss (eq 11 of this paper: <a href="https://proceedings.neurips.cc/paper/2017/file/2650d6089a6d640c5e85b2b88265dc2b-Paper.pdf" rel="noopener nofollow ugc">link</a>) that performs <code>torch.log</code> operations. I have the following implementation of the loss function:</p>
<pre><code class="lang-auto">class HeteroscedasticLoss(torch.nn.Module):
    def __init__(self, T=100, batch_size=128):
        super().__init__()
        self.T = T
        self.batch_size = batch_size

    def forward(self, outputs, targets):
                outputs = torch.swapaxes(outputs, 1, 0)  # convert to [batch_size, T, ..., num_classes]
        probs = torch.zeros_like(outputs, device=targets.device)
        i = 0
        for x in outputs:
            probs[i, ...] = torch.softmax(x, -1)
            i = i + 1
        probs = torch.mean(probs, 1)
        loss = torch.nn.functional.nll_loss(torch.log(probs), targets)
        return loss
</code></pre>
<p>I also used a different implementation (eq 12 of the paper):</p>
<pre><code class="lang-auto">class HeteroscedasticLoss(torch.nn.Module):
    def __init__(self, T=100, batch_size=128):
        super().__init__()
        self.T = T
        self.batch_size = batch_size

    def forward(self, outputs, targets):
        outputs = torch.swapaxes(outputs, 1, 0)  # convert to [batch_size, T, ..., num_classes]
        loss = 0
        i = 0
        for x in outputs:
            lsm = torch.logsumexp(input=x, dim=-1)  # [T,]
            arg = torch.sum(torch.exp(x[..., targets[i].item()] - lsm)) / self.T  # scalar
            loss = loss + torch.log(arg)
            i = i + 1
        loss = loss / self.batch_size
        return loss
</code></pre>
<p>In this case, the loss goes to <code>-inf</code>. I have changed the optimizer and learning rate, but I get the same error - only a little later with small learning rates. I am not sure now why this is happening and how to address this. Is there any way to handle this issue?</p>

username: KFrank
cooked: <p>Hi Laya!</p>
<aside class="quote no-group" data-username="Laya1" data-post="3" data-topic="166050" data-full="true">
<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/l/8c91f0/48.png" class="avatar"> Laya1:</div>
<blockquote>
<p>The loss actually becomes <code>nan</code> and does not blow up</p>
</blockquote>
</aside>
<p>I can’t tell from what you’ve posted what the problem might be.  You could<br>
have a bug somewhere or your training hyperparameters might be too<br>
“aggressive,” so your training might be becoming unstable.</p>
<p>(Consider starting with plain-vanilla <code>SGD</code> with a small learning rate, and<br>
if you can get that to train stably, try adding momentum.)</p>
<p>However …</p>
<aside class="quote no-group">
<blockquote>
<p>The variable <code>variances</code> is generated from a <code>Linear</code> layer followed by <code>ReLU</code> activation,</p>
</blockquote>
</aside>
<p>This may be the core cause of your problem.  You are trying to train a<br>
parameter (the variance, or perhaps standard deviation) that is naturally<br>
positive.  It is true that <code>ReLU</code> will clamp the parameter to be non-negative,<br>
but something abrupt (and technically singular) is happening right at zero.</p>
<p>To me a very small variance is “atypical” and an even smaller variance<br>
is even more “atypical,” but the output of <code>Linear</code> doesn’t know that getting<br>
close to zero is “atypical” and <code>Linear</code> can easily output negative values.</p>
<p>Consider passing the output of <code>Linear</code> through something like <a href="https://pytorch.org/docs/1.13/generated/torch.nn.Softplus.html" rel="noopener nofollow ugc">Softplus</a> to<br>
get your variance.  This maps the entire <code>(-inf, inf)</code> range of <code>Linear</code> to<br>
<code>(0.0, inf)</code> and does so in a way that seems to fit the “meaning” of a<br>
variance in the following sense:</p>
<p>A very large variance is “atypical” and would be achieved by having a very<br>
large positive output from <code>Linear</code>.  A very small variance (that is, positive,<br>
but very close to zero) would be achieved by having a very large <em>negative</em><br>
output from <code>Linear</code>.  Nothing abrupt nor singular occurs when the output<br>
of <code>Linear</code> happens to be zero.  (It is true that you get singular behavior as<br>
the variance approaches zero, but this happens when the output of <code>Linear</code><br>
approaches <code>-inf</code>, which is naturally a singular point.)</p>
<p>As long as your loss function penalizes small variances (skinny Gaussians),<br>
even if only weakly, it is unlikely that optimization will push the output of<br>
<code>Linear</code> to be a large enough negative value for the “singular” behavior of<br>
a variance near zero to be a problem.</p>
<p>(If you are also training the mean of your Gaussian, using the output of a<br>
<code>Linear</code> – without any further transformation – should be perfectly fine.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Laya1
cooked: <p><a class="mention" href="/u/kfrank">@KFrank</a> I tried changing the optimizer, learning rate, activation for predictions variances. Nothing seems to work at the moment. I get the error <code>RuntimeError: Function 'ExpBackward0' returned nan values in its 0th output.</code> now. I tried out a few checks and am still not able to figure out the problem. Thank you for your explanation and suggestions.</p>