MDN VAE

ptmelt/layers.py

@@ -347,3 +347,29 @@ def forward(self, input: torch.Tensor):
             )
 
         return self.weight * x_hat + self.bias
+
+
+class Reparameterization(nn.Module):
+    """Reparameterization trick for Gaussian Mixture Models."""
+
+    def forward(self, mix_coeffs, means, log_vars):
+        # Ensure mix_coeffs is a valid probability distribution
+        mix_coeffs = F.softmax(mix_coeffs, dim=-1)
+
+        # Sample component indices from the categorical distribution
+        component_indices = torch.multinomial(mix_coeffs, num_samples=1).squeeze(-1)
+
+        # Select the means and log variances of the sampled components
+        batch_size = means.size(0)
+        latent_dim = means.size(-1)
+
+        # Gather the means and log_vars based on sampled indices
+        selected_means = means[torch.arange(batch_size), component_indices, :]
+        selected_log_vars = log_vars[torch.arange(batch_size), component_indices, :]
+
+        # Reparameterization trick
+        stds = torch.exp(0.5 * selected_log_vars)
+        eps = torch.randn_like(stds)
+        z = selected_means + eps * stds
+
+        return z

ptmelt/losses.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 
 
@@ -8,7 +9,7 @@ def safe_exp(x):
     return torch.exp(x)
 
 
-class MixtureDensityLoss(torch.nn.Module):
+class MixtureDensityLoss(nn.Module):
     """
     Custom loss function for Mixture Density Network (MDN).
 
@@ -50,6 +51,7 @@ def forward(self, y_pred, y_true):
         log_var_preds = torch.clamp(log_var_preds, min=-10.0, max=10.0)
 
         # Ensure mixture coefficients sum to 1
+        # temperature = 1.0  # lower = sharper, higher = softer
         m_coeffs = F.softmax(m_coeffs, dim=1)
         # Convert log variance to variance
         var_preds = safe_exp(log_var_preds)
@@ -80,6 +82,13 @@ def forward(self, y_pred, y_true):
         # loss = -torch.mean(log_sum_exp)
         loss = log_sum_exp
 
+        # # add in entropy regularization
+        # lambd_reg = 1e-3
+        # entropy = -torch.sum(
+        #     m_coeffs * torch.log(torch.clamp(m_coeffs, min=1e-8)), dim=1
+        # )
+        # loss += lambd_reg * entropy
+
         # add in the mse as well
         if self.mse_weight > 0.0:
             mix_mean = (m_coeffs.unsqueeze(-1) * mean_preds).sum(dim=1)
@@ -94,3 +103,29 @@ def forward(self, y_pred, y_true):
         # else no reduction, return the full loss tensor
 
         return loss
+
+
+class VAELoss(nn.Module):
+    def __init__(self, reconstruction_loss_fn=nn.MSELoss()):
+        super(VAELoss, self).__init__()
+        self.reconstruction_loss_fn = reconstruction_loss_fn
+
+    def compute_reconstruction_loss(self, x, x_reconstructed):
+        return self.reconstruction_loss_fn(x_reconstructed, x)
+
+    def compute_kl_divergence(self, mix_coeffs, means, log_vars):
+        kl_div = -0.5 * torch.sum(1 + log_vars - means.pow(2) - log_vars.exp(), dim=-1)
+        kl_div = torch.mean(torch.sum(mix_coeffs * kl_div, dim=-1))
+        return kl_div
+
+    def forward(self, x, x_reconstructed, mix_coeffs, means, log_vars):
+        # Reconstruction loss
+        # reconstruction_loss = self.reconstruction_loss_fn(x_reconstructed, x)
+        reconstruction_loss = self.compute_reconstruction_loss(x, x_reconstructed)
+
+        # KL Divergence for Gaussian Mixtures
+        # kl_div = -0.5 * torch.sum(1 + log_vars - means.pow(2) - log_vars.exp(), dim=-1)
+        # kl_div = torch.mean(torch.sum(mix_coeffs * kl_div, dim=-1))
+        kl_div = self.compute_kl_divergence(mix_coeffs, means, log_vars)
+
+        return reconstruction_loss + kl_div