I am trying to implement a toy VAE project. My goal is to use a VAE to model the moon dataset from scikit-learn, with an extra constant (but noisy) z-dimension.
To this end I use an approximate posterior with the form of a beta distribution and a uniform prior in a 1D latent space, because essentially the data is 1D. The decoder is a NN-parameterized gaussian.
I cannot get it to work using the simple ELBO.
I tried so far :
- Increasing the number of monte carlo samples in the SGVB
- Various deterministic pretrainings which tend to raise nans
- Increasing the width or depth of the networks
- Gradient clipping
- learning rate annealing
- Remove the noise in the data and perform Batch gradient descent instead of mini-batch
I use layers of residual blocks with Tanh nonlinearities, whose outputs are $\log \alpha$ and $\log \beta$ for the encoder, $\mu$ and $\log \sigma$ for the decoder.
I am starting to wonder whether the distribution is actually hard to model, because I ran out of bugs to fix and strategies to improve training.
Are some low dimensional distributions known to be hard to model this way ?
Additionally, what obvious or non obvious mistakes could I have made ?
Code to generate the data:
# Adapted from sklearn.dataset.make_moons def make_moons(n_samples=100, noise=None): generator = default_rng() n_samples_out = n_samples // 2 n_samples_in = n_samples - n_samples_out outer_circ_x = np.cos(np.linspace(0, np.pi, n_samples_out)) outer_circ_y = np.sin(np.linspace(0, np.pi, n_samples_out)) inner_circ_x = 1 - np.cos(np.linspace(0, np.pi, n_samples_in)) inner_circ_y = 1 - np.sin(np.linspace(0, np.pi, n_samples_in)) - .5 X = np.vstack([np.append(outer_circ_x, inner_circ_x), np.append(outer_circ_y, inner_circ_y), np.zeros(n_samples)]).T y = np.hstack([np.zeros(n_samples_out, dtype=np.intp), np.ones(n_samples_in, dtype=np.intp)]) if noise is not None: X += generator.multivariate_normal(np.zeros(3), np.diag([noise, noise, noise])**2, size=n_samples) return X, y # create dataset moon_coordinates, moon_labels = make_moons(n_samples=500, noise=.01) moon_coordinates = moon_coordinates.astype(np.float32) moon_labels = moon_labels.astype(np.float32) # normalize dataset moon_coordinates = (moon_coordinates-moon_coordinates.mean(axis=0))/np.std(moon_coordinates, axis=0)
I have found a mistake that can explain poor performance.
In my post I said that the data is basically 1D, yet when I create the dataset I normalize the standard deviation in every dimension. This increases the magnitude of the z noise, and all of a sudden the third dimension accounts for a lot of variance and my model tries to fit to this noise.
Removing the normalization dramatically increases the performance.