Learn a normal distribution with MultivariateNormal(loc, covariance_matrix)

phan_phan: <p>Hi everyone,</p>
<p>I am trying to learn a multivariate normal distribution. To simplify things, let’s say the input is fixed, and the labels are some 3d points (x, y, z).</p>
<p>My approach was to create a NN to output 9 values, corresponding to 3 means, 3 variances and 3 co-variances. Then at each step, do a <strong>rsample</strong> from a</p>
<blockquote>
<p>MultivariateNormal(loc=m, covariance_matrix=c)</p>
</blockquote>
<p>, with c being the symmetric covariance matrix created with the 6 variances and co-variances. Then I get the loss by comparing the “rsampled” 3d point, to a point from the labels.</p>
<p>However, I sometimes get an error at the creation of the distribution, because the covariance matrix is not always positive definite.</p>
<ol>
<li>Is this the right approach ? Or is there a proper way to do this ?</li>
<li>If it is, how can I assure that the NN will produce valid distribution parameters ?</li>
<li>Should I use the other arguments (<em>precision_matrix</em> / <em>scale_tril</em>) instead of <em>covariance_matrix</em> at the creation of the distribution ? I don’t understand their mathematical meanings …</li>
</ol>
<p>Thank you,<br>
PL</p>

phan_phan: <p>Ended up using the 6 parameters to generate the <code>scale_tril</code> instead of the <code>covariance_matrix</code>.<br>
It worked really well.</p>

hanklin3: <p><a class="mention" href="/u/phan_phan">@phan_phan</a> can you explain your solution more in details? Do you do a regression output of 6 parameters and just send it to scale_tril?<br>
how do you exactly generate scale_tril?</p>

phan_phan: <p>Hi,<br>
Yes I eventually did a regression over the distribution parameters <code>loc</code> and <code>scale_tril</code>.<br>
So my last NN layer outputs a tensor of size <code>[batch, 9]</code>.</p>
<p>However, a few things to note:</p>
<ul>
<li>The diagonal of <code>scale_tril</code> must be <strong>strictly positive</strong>, so I used 1 + <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.ELU" rel="nofollow noopener">ELU</a>:</li>
</ul>
<pre><code class="lang-auto">loc, tril, diag = outputs.split(3, dim=-1)
diag = 1 + nn.functional.elu(diag)
# Each of loc, tril, diag is a tensor of size [batch, 3]
</code></pre>
<ul>
<li>I generated <code>scale_tril</code> the dirty way, maybe there’s a better one:</li>
</ul>
<pre><code class="lang-auto">z = torch.zeros(size=[loc.size(0)], device=device)
scale_tril = torch.stack([
    diag[:, 0], z         , z,
    tril[:, 0], diag[:, 1], z,
    tril[:, 1], tril[:, 2], diag[:, 2]
], dim=-1).view(-1, 3, 3)
# scale_tril is a tensor of size [batch, 3, 3]
dist = MultivariateNormal(loc=loc, scale_tril=scale_tril)
</code></pre>
<ul>
<li>In order to train, instead of minimizing a criterion between a distribution sample and the label (like I thought in the last comment), it is better to <strong>minimize the log-likelihood of the label</strong> with respect to the distrubution, with the <a href="https://pytorch.org/docs/stable/distributions.html#torch.distributions.multivariate_normal.MultivariateNormal.log_prob" rel="nofollow noopener">.log_prob()</a> method.</li>
</ul>
<pre><code class="lang-auto">loss = - dist.log_prob(labels).mean()
loss.backward()
</code></pre>

hanklin3: <p>Thank you so much, I will try it!</p>