The Loss Function

def mdn_loss(alpha, sigma, mu, target, eps=1e-8):
    target = target.unsqueeze(1).expand_as(mu)
    m = torch.distributions.Normal(loc=mu, scale=sigma)
    log_prob = m.log_prob(target)
    log_prob = log_prob.sum(dim=2)
    log_alpha = torch.log(alpha + eps)  # Avoid log(0) disaster
    loss = -torch.logsumexp(log_alpha + log_prob, dim=1)
    return loss.mean()

Here’s the breakdown:

1. target = target.unsqueeze(1).expand_as(mu): Expand the target to match the shape of mu.
2. m = torch.distributions.Normal(loc=mu, scale=sigma): Create a normal distribution.
3. log_prob = m.log_prob(target): Calculate the log probability.
4. log_prob = log_prob.sum(dim=2): Sum log probabilities.
5. log_alpha = torch.log(alpha + eps): Calculate log of mixing coefficients.
6. loss = -torch.logsumexp(log_alpha + log_prob, dim=1): Combine and log-sum-exp the probabilities.
7. return loss.mean(): Return the average loss.

The Neural Network

Let’s create a neural network that’s all set to handle the wizardry:

class MDN(nn.Module):
    def __init__(self, input_dim, output_dim, num_hidden, num_mixtures):
        super(MDN, self).__init__()
        self.hidden = nn.Sequential(
            nn.Linear(input_dim, num_hidden),
            nn.Tanh(),
            nn.Linear(num_hidden, num_hidden),
            nn.Tanh(),
        )
        self.z_alpha = nn.Linear(num_hidden, num_mixtures)
        self.z_sigma = nn.Linear(num_hidden, num_mixtures * output_dim)
        self.z_mu = nn.Linear(num_hidden, num_mixtures * output_dim)
        self.num_mixtures = num_mixtures
        self.output_dim = output_dim
def forward(self, x):
        hidden = self.hidden(x)
        alpha = F.softmax(self.z_alpha(hidden), dim=-1)
        sigma = torch.exp(self.z_sigma(hidden)).view(-1, self.num_mixtures, self.output_dim)
        mu = self.z_mu(hidden).view(-1, self.num_mixtures, self.output_dim)
        return alpha, sigma, mu

Notice the softmax being applied to 𝛼ᵢ alpha = F.softmax(self.z_alpha(hidden), dim=-1), so they sum up to 1, and the exponential to 𝜎ᵢ sigma = torch.exp(self.z_sigma(hidden)).view(-1, self.num_mixtures, self.output_dim), to ensure they remain positive, as explained earlier.

The Prediction

Getting predictions from MDNs is a bit of a trick. Here’s how you sample from the mixture model:

def get_sample_preds(alpha, sigma, mu, samples=10):
    N, K, T = mu.shape
    sampled_preds = torch.zeros(N, samples, T)
    uniform_samples = torch.rand(N, samples)
    cum_alpha = alpha.cumsum(dim=1)
    for i, j in itertools.product(range(N), range(samples)):
        u = uniform_samples[i, j]
        k = torch.searchsorted(cum_alpha[i], u).item()
        sampled_preds[i, j] = torch.normal(mu[i, k], sigma[i, k])
    return sampled_preds

Here’s the breakdown:

1. N, K, T = mu.shape: Get the number of data points, mixture components, and output dimensions.
2. sampled_preds = torch.zeros(N, samples, T): Initialize the tensor to store sampled predictions.
3. uniform_samples = torch.rand(N, samples): Generate uniform random numbers for sampling.
4. cum_alpha = alpha.cumsum(dim=1): Compute the cumulative sum of mixture weights.
5. for i, j in itertools.product(range(N), range(samples)): Loop over each combination of data points and samples.
6. u = uniform_samples[i, j]: Get a random number for the current sample.
7. k = torch.searchsorted(cum_alpha[i], u).item(): Find the mixture component index.
8. sampled_preds[i, j] = torch.normal(mu[i, k], sigma[i, k]): Sample from the selected Gaussian component.
9. return sampled_preds: Return the tensor of sampled predictions.

Practical Example: Predicting the ‘Apparent’ 🌡️

Let’s apply MDNs to predict ‘Apparent Temperature’ using a simple Weather Dataset. I trained an MDN with a 50-hidden-layer network, and guess what? It rocks! 🎸

Find the full code here. Here are some results: