Vector Quantization Layer Prototype

[Simsso]

This PR contains a vector quantization (VQ) prototype that was developed and tested in four Jupyter notebooks.

The code is not production-ready and only there for the sake of experimenting (it is located in /experiments). The well-tested and functionally extended version can now be developed based on this prototype (#51). Merging of this PR resolves #25.

The work can be found here (please review the four files + README.md):

• Layer function definition
• Projection testing; whether the mapping of inputs to the closest vector in the embedding space works
• Embedding space training testing; whether the gradient-based training of the embedding space vectors works
• Embedding space training (2) testing; same as above except the loss is applied to all vectors, not just the closest one

[Simsso]

Great notebook, @FlorianPfisterer.
Here is my review on the vq_layer function definition.

• Why do you need a transpose operations in this line of code?

  scores = tf.reduce_sum(tf.square(tf.expand_dims(x, 2) - tf.transpose(embed_space)), axis=1)

  Update: okay, got it. You can save the transpose when expanding on axis=1.
• You are adding a stop_gradient over here; I forgot that.

  vq_loss = tf.reduce_mean(tf.norm(tf.stop_gradient(x) - chosen_embeddings, ord='euclidean'))

• This loss term is similar to the VQ-VAE paper:

  commit_loss = BETA * tf.reduce_mean(tf.norm(x - tf.stop_gradient(chosen_embeddings), ord='euclidean'))

  It's a very good idea to add it, even though I'm not sure whether we're going to train our conv layers. But the production VQ layer version should have it!

[FlorianPfisterer]

This loss term is similar to the VQ-VAE paper

Well, that's because my experiments where based on the paper 😆. We're not going to train the conv layers, but it still adds a gradient flow that might be relevant for another VQ-layer further upstream.

[FlorianPfisterer]

Minor things, mostly questions. Really great experiments, I especially love the plots where one can see how the embeddings move!

[FlorianPfisterer]

As you noted yourself, add the tf.stop_gradient() here for y.

[Simsso]

Rather for x, I'd say. Can you triple-check, please?

[FlorianPfisterer]

If by alpha-loss you're referring to the second term in the following loss function (from the paper), then yes: