VAE with Multivariate Gaussian problems

username: banikr
cooked: <p>I am using Variational Autoencoders(VAE) to reconstruct some spectral data, to find lower(latent) dimensional representation out of the encoder.<br>
With mu and std from two FC layer outputs, I am getting the normal distribution <code>q(z|x)</code>:</p>
<pre><code class="lang-auto">q = torch.distributions.Normal(mu, std)
z = q.rsample()
</code></pre>
<p>and optimizing KL divergence with <code>p(x_hat|z)</code>:</p>
<pre><code class="lang-auto">p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
</code></pre>
<p>KL divergence:</p>
<pre><code class="lang-auto">log_qzx = q.log_prob(z)
log_pz = p.log_prob(z)
kl = (log_qzx - log_pz)
kl = kl.sum(-1)
</code></pre>
<p>Also, I am using, reconstruction loss as a regularizer:</p>
<pre><code class="lang-auto">def gaussian_likelihood(x_hat, logscale, x):
    scale = torch.exp(logscale) # simply unit
    mean = x_hat
    dist = torch.distributions.Normal(mean, scale)
    # measure p(x|z)
    log_pxz = dist.log_prob(x)
    return log_pxz.sum(dim=(1))
</code></pre>
<p>The final loss is:<br>
<code>KL + recon</code><br>
which I am minimizing with Adam optimizer with standard hyperparameters.<br>
To be fair, I am getting good results and expected outcomes from the above.</p>
<p>I wanted to implement the same strategy with Multivariate Normal distribution. To do this I did the following:</p>
<pre><code class="lang-auto">def kl_divergence( z, mu, lt_d_Chol):
    b, l = mu.size()
    p = torch.distributions.multivariate_normal.MultivariateNormal(loc=torch.zeros_like(mu),
                            covariance_matrix=torch.eye(l).reshape((1, l, l)).repeat(b, 1, 1).to(mu.get_device()))
    q = torch.distributions.multivariate_normal.MultivariateNormal(loc=mu, scale_tril=lt_d_Chol)
    log_qzx = q.log_prob(z)
    log_pz = p.log_prob(z)
    kl = (log_qzx - log_pz)
    return kl
</code></pre>
<p><code>lt_d_chol</code> is a rearrangement of an FC layer into lower-diagonal Cholesky decomposition, taken as:</p>
<pre><code class="lang-auto">def get_low_cholesky(self, latent_dim, logSigma, lt_indices, d_indices):
    """
    output -&gt; torch from ReLU in GPU
    lt_indices -&gt; numpy -&gt; np.tril_indices(latent_dim, k=-1, m=latent_dim)
    d_indices -&gt; numpy -&gt; np.diag_indices(latent_dim)
    """
    # elements = int(latent_dim * (latent_dim - 1) / 2) + latent_dim
    Sigma = logSigma.cpu()  # torch.exp(0.5 * logSigma) # todo: needed?
    bSize = Sigma.size(0)
    low_tri_Chol = torch.zeros(bSize, latent_dim, latent_dim)
    for b in range(bSize):
        b_lower_mat = low_tri_Chol[b, ...]
        b_lower_mat[lt_indices] = Sigma[b, :-latent_dim]
        b_lower_mat[d_indices] = 1 + nn.functional.elu(Sigma[b, -latent_dim:]) # elu coz diagonal values must be positive... 
        low_tri_Chol[b, ...] = b_lower_mat
    low_tri_Chol = low_tri_Chol.to(logSigma.get_device())
    return low_tri_Chol
</code></pre>
<p>Similarly the recon loss:</p>
<pre><code class="lang-auto">def gaussian_likelihood(self, x_hat, x):
    mean = x_hat
    b, l = mean.size()
    dist = torch.distributions.multivariate_normal.MultivariateNormal(loc=mean,
                    covariance_matrix=torch.eye(l).reshape((1, l, l)).repeat(b, 1, 1).to(mean.get_device()))
    log_pxz = dist.log_prob(x)
    return log_pxz
</code></pre>
<p>But the same hyperparameter and optimizer fail to minimize loss similarly to the univariate normal distribution.</p>
<ul>
<li>Is there anything wrong conceptually to assume that 1D vector/spectral data can be better reconstructed from multivariate distributions than univariate?</li>
<li>The output <code>logSigma</code> or <code>Sigma</code> comes from ReLU-&gt; FC layer so is it necessary to <code>ELU</code> as Cholesky factorized matrix of covariance-variance diagonal must be positive?</li>
<li>The runtime of the multivariate codes is pretty slower compared to univariate… Is it because there is a bunch of switching of tensors from CPU to GPU and vice versa?</li>
</ul>
<p>Please see the loss curve:<br>
for univariate:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d8bfa21f47533bf0baa77443987ca517824a7701.png" alt="Image" data-base62-sha1="uVrHM2IGJQYEcglxFKrpsVRct9f" width="330" height="250"><br>
for multivariate(too slow so only 80 epochs ran in the same time for 1000 univariate)<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/f/2ffbf1a1898b2c3c106920cc12b40445b8744231.png" alt="Image" data-base62-sha1="6Que9PTgmPBNsqq1LLNPfCgLms9" width="330" height="250"></p>

username: banikr
cooked: <p>Reason for slow run could be the large data size of cov-var or lower triangular matrix, for example, these are of dimensions:<br>
<code>batchSize x latent_dim x latent_dim</code></p>