Beta-VAE Loss

Problem

Combine the reconstruction term and the KL term with a beta weight: mean squared reconstruction error + beta * KL, using the closed-form KL with log_var.

The signature is beta_vae_loss(x, x_hat, mu, log_var, beta). Return a scalar.

Hint

Higher beta trades reconstruction quality for a more disentangled latent space.

Test cases

case 1

in: [[1,2],[0,0],[0],[0],1]

expected: 2.5

case 2

in: [[0],[0],[1],[0],2]

expected: 1

case 3

in: [[1],[0],[0],[1],0.5]

expected: 1.1795704571

case 4

in: [[1,1],[0.5,0.5],[0.5],[-0.5],0.25]

expected: 0.2945663325

Starter codePython

import math
def beta_vae_loss(x, x_hat, mu, log_var, beta):
    # Your code here
    pass

4 test cases · run your code against them in the editor