Different gradients from `torch.distributions.MultivariateNormal`/`torch.diag_embed` vs `torch.distributions.Normal`

jakelevi1996: <ul>
<li>I’m looking at 2 different ways of backpropagating through the log-probability of samples from a Gaussian RV WRT to its parameters:
<ul>
<li>With <code>torch.distributions.MultivariateNormal</code> and <code>torch.diag_embed</code>
</li>
<li>With <code>torch.distributions.Normal</code>, broadcasting and summing over the last dimension of the result</li>
</ul>
</li>
<li>The calculated log-probability is the same in each case</li>
<li>The gradient of the log probability WRT the mean is the same in each case</li>
<li>The gradient of the log probability WRT the standard deviation gives different gradients (seemingly by a factor of 2), as shown below</li>
<li>This is a <a href="https://en.wikipedia.org/wiki/Minimal_reproducible_example" rel="noopener nofollow ugc">MRE</a> for the behaviour of a PPO implementation, which gives meaningfully different results according to how I define the standard deviation and calculate the log probability (<code>torch.distributions.MultivariateNormal</code> gives significantly better results, which is unfortunate because I believe it is also less efficient)</li>
<li>Is this expected behaviour or is this a bug? I would have expected both approaches to be effectively doing the same thing and therefore to give the same answer</li>
</ul>
<pre><code class="lang-python">import torch
torch.manual_seed(0)
n = 10
mean = torch.zeros([n])
std = torch.ones([n], requires_grad=True)
a = torch.normal(mean, std)
print(a)
d1 = torch.distributions.MultivariateNormal(mean, torch.diag_embed(std))
p1 = d1.log_prob(a)
print(p1)
d2 = torch.distributions.Normal(mean, std)
p2 = d2.log_prob(a).sum(dim=-1)
print(p2)
p1.backward()
print(std.grad)
std.grad.zero_()
p2.backward()
print(std.grad)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">tensor([ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845, -1.3986,  0.4033,  0.8380,
        -0.7193, -0.4033], grad_fn=&lt;NormalBackward3&gt;)
tensor(-15.2935, grad_fn=&lt;SubBackward0&gt;)
tensor(-15.2935, grad_fn=&lt;SumBackward1&gt;)
tensor([ 0.6873, -0.4569,  1.8736, -0.3384,  0.0881,  0.4780, -0.4187, -0.1489,
        -0.2413, -0.4187])
tensor([ 1.3747, -0.9139,  3.7471, -0.6769,  0.1762,  0.9561, -0.8373, -0.2977,
        -0.4827, -0.8373])
</code></pre>

KFrank: <p>Hi Jake!</p>
<aside class="quote no-group quote-modified" data-username="jakelevi1996" data-post="1" data-topic="182263" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jakelevi1996/48/75757_2.png" class="avatar"> jakelevi1996:</div>
<blockquote>
<ul>
<li>I’m looking at 2 different ways of backpropagating through the log-probability of samples from a Gaussian RV WRT to its parameters:<br>
…</li>
</ul>
<pre data-code-wrap="python"><code class="lang-python">d1 = torch.distributions.MultivariateNormal(mean, torch.diag_embed(std))
...
d2 = torch.distributions.Normal(mean, std)
</code></pre>
</blockquote>
</aside>
<p>When you instantiate <code>d1</code>, you pass <code>MultivariateNormal</code> the covariance<br>
matrix (the diagonal elements of which are variances) that defines your<br>
distribution.  <code>Normal</code>, in contrast, takes the standard deviations (which are<br>
the square roots of the variances).</p>
<p>Try:</p>
<pre data-code-wrap="python"><code class="lang-python">d1 = torch.distributions.MultivariateNormal (mean, torch.diag_embed (std**2))
</code></pre>
<p>and you should get the results you expect.</p>
<p>(The power of 2 that relates the variance to the standard deviation explains<br>
the factor of two in your two computations of the gradient.  The fact that you<br>
got the same <code>log_prob()</code> in both cases is an artifact of <code>std</code> being a vector<br>
of <code>1.0</code>s whose square is equal to itself.)</p>
<p>Best.</p>
<p>K. Frank</p>

jakelevi1996: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a> ,</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="182263">
<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"> KFrank:</div>
<blockquote>
<p>When you instantiate <code>d1</code>, you pass <code>MultivariateNormal</code> the covariance<br>
matrix (the diagonal elements of which are variances) that defines your<br>
distribution. <code>Normal</code>, in contrast, takes the standard deviations (which are<br>
the square roots of the variances).</p>
</blockquote>
</aside>
<p>Thanks for explaining, makes 100% perfect sense <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="182263">
<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"> KFrank:</div>
<blockquote>
<pre><code class="lang-auto">d1 = torch.distributions.MultivariateNormal (mean, torch.diag_embed (std**2))
</code></pre>
</blockquote>
</aside>
<p>This works. Equivalently, I can apply <code>torch.sqrt</code> before passing <code>std</code> (which should be renamed to EG to <code>action_var</code>) to <code>torch.distributions.Normal</code>, keep everything else the same, and get the same results, even when sampling from the 2 distributions (as long as I reset the random seed before each sample):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.manual_seed(0)
n = 10
mean = torch.zeros([n])
std = torch.ones([n], requires_grad=True)
a = torch.normal(mean, std)
print(a)
d1 = torch.distributions.MultivariateNormal(mean, torch.diag_embed(std))
p1 = d1.log_prob(a)
print(p1)
d2 = torch.distributions.Normal(mean, torch.sqrt(std))
p2 = d2.log_prob(a).sum(dim=-1)
print(p2)
p1.backward()
print(std.grad)
std.grad.zero_()
p2.backward()
print(std.grad)
torch.manual_seed(0)
print(d1.sample())
torch.manual_seed(0)
print(d2.sample())
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">tensor([ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845, -1.3986,  0.4033,  0.8380,
        -0.7193, -0.4033], grad_fn=&lt;NormalBackward3&gt;)
tensor(-15.2935, grad_fn=&lt;SubBackward0&gt;)
tensor(-15.2935, grad_fn=&lt;SumBackward1&gt;)
tensor([ 0.6873, -0.4569,  1.8736, -0.3384,  0.0881,  0.4780, -0.4187, -0.1489,
        -0.2413, -0.4187])
tensor([ 0.6873, -0.4569,  1.8736, -0.3384,  0.0881,  0.4780, -0.4187, -0.1489,
        -0.2413, -0.4187])
tensor([ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845, -1.3986,  0.4033,  0.8380,
        -0.7193, -0.4033])
tensor([ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845, -1.3986,  0.4033,  0.8380,
        -0.7193, -0.4033])
</code></pre>
<p>Ideally I’d like to use <code>torch.distributions.Normal</code>, <code>sqrt</code>, and <code>log_prob.sum(dim=-1)</code> instead of <code>torch.distributions.MultivariateNormal</code> and <code>diag_embed</code> in my PPO implementation, because the former should be more efficient for high-dimensional actions. The above demo suggests both implementations should give the same results, although weirdly I get different results from the 2 different implementations (given the same random seed in each experiment). I don’t understand why this is the case, although I guess it’s a bit off-topic from my original question.</p>

KFrank: <p>Hi Jake!</p>
<aside class="quote no-group" data-username="jakelevi1996" data-post="3" data-topic="182263" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jakelevi1996/48/75757_2.png" class="avatar"> jakelevi1996:</div>
<blockquote>
<p>I get different results from the 2 different implementations</p>
</blockquote>
</aside>
<p>This is most likely not a problem.</p>
<p>It would not be surprising for the two implementation to give results that<br>
differ by numerical round-off error (even if mathematically equivalent).</p>
<p>Round-off error can accumulate, so, depending on how complicated a single<br>
iteration is, you might see the results of the two methods differ by, say, ten<br>
times typical round-off error after a single iteration.  As you continue to iterate,<br>
the two computations may well start to diverge.  This may or may not be a<br>
problem; if it is a problem, the problem isn’t the round-off error, but rather<br>
a sign that your problem is unstable in some way.</p>
<p>If your results differ by a lot more than typical round-off error after a single<br>
iteration, you likely have a real problem that you should track down.</p>
<p>A good way to test whether your discrepancy is due to round-off error is to<br>
repeat your computation in double-precision and check that the discrepancy<br>
drops by several orders of magnitude.</p>
<p>Best.</p>
<p>K. Frank</p>