Implementing The Diagonal Multivariate Gaussian Distribution

Merlin136: <p>Hello there,</p>
<p>I need to use - or implement - a means of calculating the probability density function of a diagonal, multivariate Gaussian distribution.</p>
<p>The context is that I am building an “Active Inference” agent based model to solve a version of the CartPole-v1 environment:  <a href="https://gymnasium.farama.org/environments/classic_control/cart_pole/" class="inline-onebox" rel="noopener nofollow ugc">Cart Pole - Gymnasium Documentation</a></p>
<p>The agent receives a noisy observation vector and must operate on the basis of this noisy/imperfect observation.</p>
<p>I will omit the details as to what my method actually is, I’ll only say that it necessitates the ability to calculate the probability density function (pdf) of a diagonal multivariate Gaussian distribution. I know that PyTorch offers this functionality out of the box with the MultivariateNormal class: <a href="https://pytorch.org/docs/stable/distributions.html#multivariatenormal" class="inline-onebox" rel="noopener nofollow ugc">Probability distributions - torch.distributions — PyTorch 2.0 documentation</a>. Indeed, I have attempted to use this class, in addition to my own “from scratch” implementation, but I get the same problem in any case.</p>
<p>After importing:</p>
<pre><code class="lang-python">import torch
import torch.distributions.multivariate_normal as mvn
import numpy as np
</code></pre>
<p>My implementation is as follows - note that the print statements are for debugging only:</p>
<pre><code class="lang-python">def diag_mvg_pdf_torch(samples, means, diag_vars):
    """
    Compute the PDF of a diagonal multivariate Gaussian distribution using PyTorch.

    Parameters:
    - samples (torch.Tensor): The samples tensor (batch_size x num_features).
    - means (torch.Tensor): The means tensor (batch_size x num_features).
    - diag_vars (torch.Tensor): The diagonal elements of the covariance matrix (batch_size x num_features).

    Returns:
    - pdf (torch.Tensor): The average PDF value over the batch.
    """
    batch_size = len(samples)

    final_pdf = 0

    for i in range(batch_size):
        mvn_dist_i = mvn.MultivariateNormal(
            loc=means[i],
            covariance_matrix=torch.diag(diag_vars[i])
        )

        pdf_i = torch.exp(mvn_dist_i.log_prob(samples[i]))
        final_pdf += pdf_i

        print(f"\npdf_i: {pdf_i}")
        print(f"samples[i]: {samples[i]}")
        print(f"means[i]: {means[i]}")
        print(f"diag_vars[i]: {diag_vars[i]}\n")

    # Compute the average PDF over the batch
    pdf = final_pdf / batch_size

    return pdf
</code></pre>
<p>This function is supposed to instantiate a multivariate Gaussian pdf, parameterised by the mean vector and variances in each row of “means” and “diag_vars” (respectively). The pdf value of each respective sample in “samples” is computed and the average value of these pdfs should be returned.</p>
<p>The batch size is 64, so this code should instantiate 64 diagonal multivariate Gaussians, compute the pdf of each sample with respect to exactly one, unique, instantiated Gaussian and then return the average of these pdf values.</p>
<p>My “samples”, “means” and “diag_vars” tensors look like this - I’ve concatenated their full extent to save space:</p>
<pre><code class="lang-python"># reparameterized samples
samples = torch.tensor([
                [ 0.0317,  0.0163, -0.0062,  0.0694],
                [ 0.1244, -0.1245,  0.0240,  0.0954],
               ...
                [-0.1707,  0.4822,  0.0498,  0.1098],
                [ 0.1741,  0.0472,  0.1415,  0.1108]
        ])

        # predicted mean vectors
        means = torch.tensor([
                [-0.0248,  0.1706,  0.0314,  0.1002],
                [-0.0248,  0.1706,  0.0314,  0.1002],
                ...
                [ 0.0237,  0.1656,  0.0014,  0.1111],
                [ 0.0237,  0.1656,  0.0014,  0.1111]
        ])

        # predicted variances 
        diag_vars = torch.tensor([
                [8.1372e-03, 2.6815e-02, 3.7143e-02, 1.4694e-03],
                [8.1372e-03, 2.6815e-02, 3.7143e-02, 1.4694e-03],
                ...
                [1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07],
                [1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07]
        ])
</code></pre>
<p>Now the output of this following code:</p>
<pre><code class="lang-python">pdf = diag_mvg_pdf_torch(samples, means, diag_vars)
print(f"\npdf: {pdf}\n")
</code></pre>
<p>is this:</p>
<pre><code class="lang-python">...
pdf_i: 1966.2965087890625
samples[i]: tensor([0.0185, 0.1974, 0.2239, 0.1114])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])


pdf_i: 39.062076568603516
samples[i]: tensor([ 0.1559,  0.2052, -0.1978,  0.1238])
means[i]: tensor([0.0253, 0.1559, 0.0055, 0.1171])
diag_vars[i]: tensor([2.0956e-02, 6.7220e-02, 2.8870e-02, 9.5371e-06])


pdf_i: 250.84556579589844
samples[i]: tensor([-0.1707,  0.4822,  0.0498,  0.1098])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])


pdf_i: 1444.880859375
samples[i]: tensor([0.1741, 0.0472, 0.1415, 0.1108])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])


pdf: 485.08538818359375
</code></pre>
<p>where this is only the last few lines of the output.</p>
<p>Clearly this is not working correctly, since the final “averaged” pdf value is about 485. Indeed, each of the 64 “pdf_i” values are usually vastly larger than 1.0.</p>
<p>I have tried several things, like implementing the pdf calculation “from scratch”:</p>
<pre><code class="lang-python">def diag_mvg_pdf_from_scratch(samples, means, diag_vars):
    """
    Compute the PDF of a diagonal multivariate Gaussian distribution using PyTorch.

    Parameters:
    - samples (torch.Tensor): The samples tensor (batch_size x num_features).
    - means (torch.Tensor): The means tensor (batch_size x num_features).
    - diag_vars (torch.Tensor): The diagonal elements of the covariance matrix (batch_size x num_features).

    Returns:
    - pdf (torch.Tensor): The average PDF value over the batch.
    """
    batch_size = len(samples)

    final_pdf = 0

    for i in range(batch_size):

        exponent = (-(samples[i] - means[i]) ** 2)/(2 * diag_vars[i])
        normalization_term = (2*np.pi*diag_vars[i]) ** (-1/2)
        pdf_i_dims = normalization_term * torch.exp(exponent)
        pdf_i = torch.prod(pdf_i_dims)

        print(f"\nexponent: {exponent}")
        print(f"normalization_term: {normalization_term}")
        print(f"pdf_i_dims: {pdf_i_dims}")
        print(f"pdf_i: {pdf_i}")
        print(f"samples[i]: {samples[i]}")
        print(f"means[i]: {means[i]}")
        print(f"diag_vars[i]: {diag_vars[i]}\n")

    # Compute the average PDF over the batch
    pdf = final_pdf / batch_size

    return pdf
</code></pre>
<p>but this has exactly the same result.</p>
<p>I suspect that the issue will have to do with the “samples”, “means” and “log_vars” themselves, since “means” and “log_vars” are actually predicted by a feed-forward neural network. “samples” are samples drawn from the this distribution, by means of the “reparameterization trick”.</p>
<p>I would appreciate any thoughts/assistance that anyone might have to offer. I’m very happy to elaborate on anything that you desire.</p>
<p>Many thanks for reading this far!</p>

KFrank: <p>Hi Ronnie!</p>
<aside class="quote no-group" data-username="Merlin136" data-post="1" data-topic="187609" 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/merlin136/48/51387_2.png" class="avatar"> Merlin136:</div>
<blockquote>
<p>I need to use - or implement - a means of calculating the probability density function of a diagonal, multivariate Gaussian distribution.</p>
</blockquote>
</aside>
<p>Your probability-density-function implementation is correct.</p>
<aside class="quote no-group">
<blockquote>
<p>I have attempted to use this class, in addition to my own “from scratch” implementation, but I get the same problem in any case.</p>
</blockquote>
</aside>
<p>Your “from scratch” implementation is also correct and is a successful<br>
cross-check on your first implementation.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">...
pdf_i: 1966.2965087890625
samples[i]: tensor([0.0185, 0.1974, 0.2239, 0.1114])
means[i]: tensor([0.0237, 0.1656, 0.0014, 0.1111])
diag_vars[i]: tensor([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])
</code></pre>
<p>Clearly this is not working correctly, since the final “averaged” pdf value is about 485. Indeed, each of the 64 “pdf_i” values are usually vastly larger than 1.0.</p>
</blockquote>
</aside>
<p>You haven’t said exact;ly why you think this is wrong, but is seems that you<br>
are surprised by the “large” value of the pdf.</p>
<p>In this case, I think that your intuition is misleading you.  As the variance<br>
(or, similarly, the standard deviation) of a gaussian becomes small, the peak<br>
value of that gaussian’s pdf becomes large.  This effect is enhanced in multiple<br>
dimensions because the peak value of the multidimensional pdf is (in a<br>
well-defined sense) the product of the one-dimensional peaks.  And your<br>
variances are small.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.1'
&gt;&gt;&gt; import math
&gt;&gt;&gt; diag_vars = torch.tensor ([1.7298e-02, 6.7676e-02, 3.3873e-02, 8.5967e-07])
&gt;&gt;&gt; peak_pdf = (1.0 / (2.0 * math.pi * diag_vars)**0.5).prod()
&gt;&gt;&gt; peak_pdf
tensor(4338.4268)
</code></pre>
<p>So for this particular example, the pdf you calculate – <code>1966.2965</code> – is less than<br>
half of the peak value (which would be obtained if your <code>samples</code> were equal to<br>
your <code>means</code>), which is to say that it’s not unreasonably large.</p>
<p>Best.</p>
<p>K. Frank</p>

Merlin136: <p>Thank you very much <a class="mention" href="/u/kfrank">@KFrank</a></p>
<p>Apologies for the late reply.</p>
<p>Yes I see now that I have just confused probability with probability density value. I have now fixed the issue. I didn’t actually need the pdf value at all in the end, to do what I wanted.</p>
<p>Thank you for your time!</p>