Variational Autoencoder in Keras: How to achieve different output of a Keras Layer at the time of training and prediction?

Question

We're implementing a paper titled - "Variational Autoencoders for Collaborative Filtering" in TF 2.0. The sample implementation of the above paper in TF 1.0 is given here.

The paper proposes an implementation of a Variational Autoencoder for collaborative filtering. As the output of the encoder, it uses the reparametrization trick to sample the latent vector Z at the time of training the network.

The reparametrization trick samples ϵ ∼ N (0, IK) and reparametrize the latent vector Z as: Zu = µϕ(xu ) + ϵ ⊙ σϕ(xu) where µϕ and σϕ are calculated from the output of the encoder.

But, at the time of prediction, the paper proposes to use only µϕ for sampling Z.

In our implementation, we used a custom tf.keras.layers.Layer to sample the latent vector Z. The following is the code of the architecture:

class Reparameterize(tf.keras.layers.Layer):
  """
  Custom layer.

  Reparameterization trick, sample random latent vectors Z from 
  the latent Gaussian distribution.

  The sampled vector Z is given by
  sampled_z = mean + std * epsilon
  """
  def call(self, inputs):
    Z_mu, Z_logvar = inputs

    Z_sigma = tf.math.exp(0.5 * Z_logvar)
    epsilon = tf.random.normal(tf.shape(Z_sigma))
    return Z_mu + Z_sigma * epsilon


class VAE:
  def __init__(self, input_dim, latent_dim=200):

    # encoder
    encoder_input = Input(shape=input_dim)
    X = tf.math.l2_normalize(encoder_input, 1)
    X = Dropout(0.5)(X)
    X = Dense(600, activation='tanh')(X)

    Z_mu = Dense(latent_dim)(X)
    Z_logvar = Dense(latent_dim)(X)

    sampled_Z = Reparameterize()([Z_mu, Z_logvar])

    # decoder
    decoder_input = Input(shape=latent_dim)
    X = Dense(600, activation='tanh')(decoder_input)
    logits = Dense(input_dim)(X)

    # define losses
    """
    custom loss function
    def loss(X_true, X_pred)
    """

    # create models
    self.encoder = Model(encoder_input, [Z_logvar, Z_mu, sampled_Z], name='encoder')
    self.decoder = Model(decoder_input, logits, name='decoder')

    self.vae = Model(encoder_input, self.decoder(sampled_Z), name='vae')
    self.vae.add_loss(kl_divergence(Z_logvar, Z_mu))

    # compile the model
    self.vae.compile(optimizer='adam', loss=loss, metrics=[loss])

Now, I am looking for a way to change the implementation of the custom Reparameterize layer at the time of prediction to use only µϕ (Z_mu) for sampling Z so as to achieve what is proposed by the paper mentioned above.

Or if there's another way of doing so in Tf 2.0, kindly recommend.

Answer 1

You could do:

# create your VAE model
my_vae = VAE(input_dim = my_input_dim)
# Train it as you wish
# .....

When training is done, you could use it as follows:

inp = Input(shape = my_input_dim)
_, Z_mu,_ = my_vae.encoder(inp) # my_vae is your trained model, get its outputs
decoder_output = my_vae.decoder(Z_mu) # use the Z_mu as input to decoder
vae_predictor = Model(inp, decoder_output) # create your prediction time model

You could use the vae_predictor model now for predictions.