Are some low dimensional distributions known to be hard to model with VAEs?

Question

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 ?

ADDENDA

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)

UPDATE

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.

Answer 1

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.