Neural Discrete Representation Learning

The problem the paper starts from

Learning useful representations without supervision remains a key challenge in machine learning, and that is where the paper begins ^[1]. A representation is a compressed description of raw data — an image, a waveform, a sentence — that makes later work easier: if a model can turn raw input into a short code that keeps what matters and drops the rest, then generation, classification or compression can be done on the code instead of on the raw input. The autoencoder is the standard tool for this. An encoder maps the input to a code, a decoder maps the code back to the input, and training pushes the reconstruction towards the original. A variational autoencoder (VAE) adds probability to the picture: the encoder outputs a distribution over a continuous latent space, the decoder samples from it, and training maximises a bound on the likelihood of the data. Continuous latents are convenient, because gradients flow through them smoothly. The paper asks what happens if the code is not continuous at all, but a set of discrete symbols.

What VQ-VAE is

The paper proposes a generative model that learns discrete representations, the Vector Quantised-Variational AutoEncoder (VQ-VAE) ^[1]. It differs from VAEs in two key ways: the encoder outputs discrete rather than continuous codes, and the prior is learned rather than static ^[1]. To learn a discrete latent representation the authors incorporate ideas from vector quantisation ^[1].

Vector quantisation is a classical idea from signal compression. A continuous vector is replaced by the closest entry in a table of vectors — a codebook — and only the index of that entry is kept. The data is then described by a short sequence of indices instead of a long sequence of real numbers. VQ-VAE uses the same operation to turn the encoder's continuous output into discrete codes, which the decoder turns back into data.

The differentiability problem

Picking the nearest codebook entry is a lookup, not a smooth function of the input, so the ordinary chain rule does not apply at that step: the selection itself does not define a gradient with respect to the encoder's output. Making quantisation work inside a network trained by gradient descent therefore requires an explicit way of carrying the learning signal from the decoder back to the encoder. Any attempt to combine vector quantisation with neural networks has to face this difficulty, and it is where much of the design work in such a model sits.

Posterior collapse and the learned prior

The second difference from a VAE is that the prior is learned rather than static ^[1]. In a standard VAE the prior over the latent code is a fixed distribution, usually something simple, and the training objective pulls the encoder's distribution towards it. When the decoder is powerful, the model can end up ignoring the latent code altogether: the decoder reconstructs the data well on its own, the encoder's output carries little information, and the objective settles there. This is posterior collapse, the situation where latents are ignored when they are paired with a powerful decoder ^[1].

The authors say the VQ approach lets the model circumvent this problem ^[1]. The intuition is that a discrete code is not pulled towards a smooth fixed prior in the same way a continuous one is, and because the prior is learned instead of fixed, the model is free to fit a distribution over codes that actually describes the data rather than one chosen in advance.

Authors and publication

The paper, "Neural Discrete Representation Learning", was first submitted to arXiv on 2 November 2017 as 1711.00937 ^[1]. Its authors are Aaron van den Oord, Oriol Vinyals and Koray Kavukcuoglu ^[1].

Why discrete codes matter

Much data is discrete by nature: text is a sequence of words or word pieces, and a compressed signal can be described by a sequence of table indices. A generative model whose latent variable is a sequence of discrete codes is therefore a closer match to such data than one whose latent variable is a cloud of points in a continuous space. Modelling then becomes a matter of modelling sequences of symbols, with a learned prior sitting on top of the codes, instead of modelling densities in a continuous latent space. The paper's contribution is to set out how such a discrete latent space can be learned together with the encoder and the decoder, without supervision, and without the latents being quietly discarded along the way.