How to generate new data with a VAE?

Question

I have built the following function which takes as input some data and runs a VAE on them:

def VAE(data, original_dim, latent_dim, test_size, epochs):
    
    x_train, x_test = train_test_split(data, test_size=test_size, random_state=42)
    
    # Define the VAE architecture
    #Encoder
    encoder_inputs = tf.keras.Input(shape=(original_dim,))
    x = layers.Dense(64, activation='relu')(encoder_inputs)
    x = layers.Dense(32, activation='relu')(x)
    x = layers.Dense(8, activation='relu')(x)

    #--- Custom Latent Space Layer
    z_mean = layers.Dense(units=latent_dim, name='Z-Mean', activation='linear')(x)
    z_log_sigma = layers.Dense(units=latent_dim, name='Z-Log-Sigma', activation='linear')(x)
    z = layers.Lambda(sampling, name='Z-Sampling-Layer')([z_mean, z_log_sigma, latent_dim]) # Z sampling layer

    # Instantiate the encoder
    encoder = tf.keras.Model(encoder_inputs, [z_mean, z_log_sigma, z], name='encoder')

    #Decoder
    latent_inputs = tf.keras.Input(shape=(latent_dim,))
    x = layers.Dense(8, activation='relu')(latent_inputs)
    x = layers.Dense(32, activation='relu')(x)
    x = layers.Dense(64, activation='relu')(x)
    decoder_outputs = layers.Dense(1, activation='relu')(x)

    # Instantiate the decoder
    decoder = tf.keras.Model(latent_inputs, decoder_outputs, name='decoder')

    # Define outputs from a VAE model by specifying how the encoder-decoder models are linked
    # Instantiate a VAE model
    vae = tf.keras.Model(inputs=encoder_inputs, outputs=decoder(encoder(encoder_inputs)[2]), name='vae')
    
    # Reconstruction loss compares inputs and outputs and tries to minimise the difference
    r_loss = original_dim * tf.keras.losses.mse(encoder_inputs, decoder(encoder(encoder_inputs)[2]))  # use MSE

    # KL divergence loss compares the encoded latent distribution Z with standard Normal distribution and penalizes if it's too different
    kl_loss = -0.5 * K.mean(1 + z_log_sigma - K.square(z_mean) - K.exp(z_log_sigma), axis=-1)

    #VAE total loss
    vae_loss = K.mean(r_loss + kl_loss)

    # Add loss to the model and compile it
    vae.add_loss(vae_loss)
    vae.compile(optimizer='adam')
    
    # train the model
    vae.fit(x_train, x_train, epochs=epochs, validation_data=(x_test, x_test))

where

def sampling(args):
    
    z_mean, z_log_sigma, latent_dim = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=1., seed=42)
    return z_mean + K.exp(z_log_sigma) * epsilon

My question is, if I want to generate new data, by using the above VAE how can I achieve that ?

If I want to sample 100 new data, should I use this

   latent_mean = tf.math.reduce_mean(encoder(x_train)[2], axis=0) 
   latent_std = tf.math.reduce_std(encoder(x_train)[2], axis=0)
   tf.random.normal(shape=(100, latent_dim), mean=latent_mean, stddev=latent_std)

or

   latent_mean = tf.math.reduce_mean(encoder(x_train)[0], axis=0) 
   latent_std = tf.math.exp(tf.math.reduce_mean(encoder(x_train)[1], axis=0))
   tf.random.normal(shape=(100, latent_dim), mean=latent_mean, stddev=latent_std)

?

Basically, should I use the z_mean and z_log_sigma directly ? Or should I infer them from the z ?

Answer 1

At a high level of abstraction, there are three steps to VAEs. (The most common use-case of a VAE assumes that the $z$ are distributed as a multivariate normal distribution with a diagonal covariance matrix, so I focus on that.)

1. Encode. Use a neural network $f$ to map the input $x$ to a vector means and a vector of standard deviations: $f(x) = (\mu, \sigma)$
2. Sample. Draw a vector from the multivariate normal distribution defined by $(\mu, \sigma)$. This is $z \sim \mathcal{N}(\mu, \sigma^2 I)$.
3. Decode. Use a neural network to map the vector $z$ to something that "looks like" the input $x$. We could write this as $\hat x = g(z)$.

During, the loss is computed as the sum of the KLD and the reconstruction error: $\| \hat x - x \|_2^2 + KLD$, for example.

If you're trying to generate new data that "looks like" the input data, then you want $\hat x$. So neither of your proposals is correct -- you want to decode $z$.