The update equation in examples/generative /real_nvp.py

[amida47]

in your implementation of a realNVP, here is a part of the model class

class RealNVP(keras.Model):
    def __init__(self, num_coupling_layers):
        super().__init__()

        self.num_coupling_layers = num_coupling_layers

        # Distribution of the latent space.
        self.distribution = tfp.distributions.MultivariateNormalDiag(
            loc=[0.0, 0.0], scale_diag=[1.0, 1.0]
        )
        self.masks = np.array(
            [[0, 1], [1, 0]] * (num_coupling_layers // 2), dtype="float32"
        )
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.layers_list = [Coupling(2) for i in range(num_coupling_layers)]

    @property
    def metrics(self):
        """List of the model's metrics.

        We make sure the loss tracker is listed as part of `model.metrics`
        so that `fit()` and `evaluate()` are able to `reset()` the loss tracker
        at the start of each epoch and at the start of an `evaluate()` call.
        """
        return [self.loss_tracker]

    def call(self, x, training=True):
        log_det_inv = 0
        direction = 1
        if training:
            direction = -1
        for i in range(self.num_coupling_layers)[::direction]:
            x_masked = x * self.masks[i]
            reversed_mask = 1 - self.masks[i]
            s, t = self.layers_list[i](x_masked)
            s *= reversed_mask
            t *= reversed_mask
            gate = (direction - 1) / 2
            x = (
                reversed_mask
                * (x * tf.exp(direction * s) + direction * t * tf.exp(gate * s))
                + x_masked
            )
            log_det_inv += gate * tf.reduce_sum(s, [1])

        return x, log_det_inv

    # Log likelihood of the normal distribution plus the log determinant of the jacobian.

my problem is with the values the variable direction takes, if training is true direction = -1, but with this value if we calculate the new value of x (in this case lets take x=[x1, x2] a two dimension data vector) going forward (from data to latent space), I found that given the current mask is [1, 0] and that s = [s1,s2] and t = [t1, t2] and :

x = (
                reversed_mask
                * (x * tf.exp(direction * s) + direction * t * tf.exp(gate * s))
                + x_masked
            )

will give me x = [x1 , x2 * exp(-s2) - t2 *exp(-s2)] which different than what the formula says x should be, equal to [x1, x2 * exp(s2) + t2]

if you can please clarify the choice for setting direction values and so that were on the same page what I understand to go forward ( from data to latent) we use this equation and for the inverse pass (from latent to data) we use this and that during training we take a data point, and go from data to latent (for me this is a forward pass)

[amida47]

since @mbrukman was last to modify this file, I hope you can clarify my confusion