Added quantisation layer, commitment loss, and associated unit tests

[RdoubleA]

Summary:

Vector Quantized Variational Autoencoder (VQVAE) is widely used to generate multimodal data (audio, image, or video samples). A well-known model implementing this is OpenAI's DALL-E model, which can generate images from text captions. The core component of the VQVAE is the quantisation layer, which is a discrete embedding of vectors that codifies the outputs of an encoder. Here, this layer is implemented as its own module, along with commitment loss which is used to fine-tune the encoder.

Test plan:

Unit tests were created and added in this PR. Testing was done with PyTest.

pytest -vv test/modules/losses/test_commitment.py

pytest -vv test/modules/layers/test_quantisation.py

[ankitade]

overall good start. some questions / suggestions

[ebsmothers]

Looks good! I like the way you structured the code, and adding the reshaping utils will come in handy as well.

[langong347]

Please include in the Test Plan:

• the exact command you used to run unit tests regarding your changes
• testing results

[langong347]

Sorry I probably didn't make the purpose of Test Plan clear: You can just run tests related to your changes, e.g., pytest -vv test.modules.layers.test_quantisation.py something like this instead of executing all tests from the library.

[langong347]

Thanks for the nice work! Going through the threads in this PR I can see a few points to follow up:

1. Learning of the codebook needs to be implemented by adding VQ loss or as moving average of the encoder output.
2. Re-design of the commitment loss (or vqvae loss as a whole) and clean up API of the quantization layer.

Also I have a question for you:

• The commitment loss is computed with the quantized vectors detached before returned by forward as in OpenAI and the Google example. In you current implementation, the quantized vectors are detached after (inside CommitmentLoss). Do you think the two steps are commutative?

Let me know your thoughts about those observation and if anything that might be missing.

Next step: click "import to fbsource" button below and wait for the internal CI to finish, fix any issues as needed.

[RdoubleA]

Thanks for the nice work! Going through the threads in this PR I can see a few points to follow up:

1. Learning of the codebook needs to be implemented by adding VQ loss or as moving average of the encoder output.
2. Re-design of the commitment loss (or vqvae loss as a whole) and clean up API of the quantization layer.

Also I have a question for you:

• The commitment loss is computed with the quantized vectors detached before returned by forward as in OpenAI and the Google example. In you current implementation, the quantized vectors are detached after (inside CommitmentLoss). Do you think the two steps are commutative?

Let me know your thoughts about those observation and if anything that might be missing.

Next step: click "import to fbsource" button below and wait for the internal CI to finish, fix any issues as needed.

I think as long as quantized is detached before computing loss then the gradients should be calculated correctly. Whether the loss is computed inside forward as in OpenAI's implementation or computed outside in a separate loss module as in CommitmentLoss the same tensors are being used, so it should have the same result.

[facebook-github-bot]

@RdoubleA has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@RdoubleA has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.