Consider an autoencoder that encodes an input $\mathbf{x}\in\mathbb{R}^{D}$ to a code $\mathbf{z}\in\mathbb{R}^{K}$. Let $\mathbf{z}=e_{\theta}(\mathbf{x}):\mathbb{R}^{D}\rightarrow\mathbb{R}^{K}$ be the encoder parameterized by $\theta$ and $\mathbf{x}^{\prime}=d_{\phi}(\mathbf{z}):\mathbb{R}^{K}\rightarrow\mathbb{R}^{D}$ be the decoder parameterized by $\phi$. We define $e_{\theta}(\cdot)_{\leq i}:\mathbf{x}\in\mathbb{R}^{D}\mapsto(z_{1},z_{2},\cdots,z_{i},0,\cdots,0)^{\mkern-1.5mu\mathsf{T}}\in\mathbb{R}^{K}$, which truncates the representation to the first $i$ dimensions of the encoding $\mathbf{z}=e_{\theta}(\mathbf{x})$, masking out the remainder of the dimensions with a zero value. We define the ordered autoencoder objective as

We note that Eq$\ldotp$ (3) is equivalent to Rippel et al. (2014)’s nested dropout formulation using a uniform sampling of possible truncations. Moreover, when the encoder/decoder pair is constrained to be a pair of orthogonal matrices up to a transpose, then the optimal solution in Eq$\ldotp$ (3) recovers PCA.

Once we have learned an ordered autoencoder, we then train an autoregressive model on the full length codes in the ordered representation space, also referred to as ex-post density estimation (Ghosh et al., 2020). For each input $\mathbf{x}_{i}$ in a dataset $\mathbf{x}_{1},\cdots,\mathbf{x}_{N}$, we feed it to the encoder to get $\mathbf{z}_{i}=e_{\theta}(\mathbf{x}_{i})$. The resulting codes $\mathbf{z}_{1},\mathbf{z}_{2},\cdots,\mathbf{z}_{N}$ are used as training data. After training both the ordered autoencoder and autoregressive model, we can perform anytime sampling on a large spectrum of computing budgets. Suppose for example we can afford to generate $T$ code dimensions from the autoregressive model, denoted as $\mathbf{z}_{\leq T}\in\mathbb{R}^{T}$. We can simply zero-pad it to get $(z_{1},z_{2},\cdots,z_{T},0,\cdots,0)^{\mkern-1.5mu\mathsf{T}}\in\mathbb{R}^{K}$ and decode it to get a complete sample. Unlike the autoregressive part, the decoder has access to all dimensions of the latent code at the same time and can decode in parallel. The framework is shown in Fig$\ldotp$ 1(b). For the implementation, we use the ordered VQ-VAE (Section 4) as the ordered autoencoder and the Transformer (Vaswani et al., 2017) as the autoregressive model. On modern GPUs, the code length has minimal effect on the run-time of decoding, as long as the decoder is not itself autoregressive (see empirical verifications in Section 5.2.2).

To construct a VQ-VAE with code length of $K$ discrete latent variables , the encoder must first map the raw input $\mathbf{x}$ to a continuous representation $\mathbf{z}^{e}=e_{\theta}(\mathbf{x})\in\mathbb{R}^{K\times D}$, before feeding it to a vector-valued quantization function $q:\mathbb{R}^{K\times D}\to\{1,2,\cdots,C\}^{K}$ defined as

where $q(\mathbf{z}^{e})_{j}\in\{1,2,\cdots,C\}$ denotes the $j$-th component of the vector-valued function $q(\mathbf{z}^{e})$, $\mathbf{z}^{e}_{j}\in\mathbb{R}^{D}$ denotes the $j$-th row of $\mathbf{z}^{e}$, and $\mathbf{e}_{i}$ denotes the $i$-th row of the embedding matrix $\mathbf{E}\in\mathbb{E}^{C\times D}$. Next, we view $q(\mathbf{z}^{e})$ as a sequence of indices and use them to look up embedding vectors from the codebook $\mathbf{E}$. This yields a latent representation $\mathbf{z}_{d}\in\mathbb{R}^{K\times D}$, given by $\mathbf{z}^{d}_{j}=\mathbf{e}_{q(\mathbf{z}^{e})_{j}}$, where $\mathbf{z}^{d}_{j}\in\mathbb{R}^{D}$ denotes the $j$-th row of $\mathbf{z}^{d}$. Finally, we can decode $\mathbf{z}^{d}$ to obtain the reconstruction $d_{\phi}(\mathbf{z}^{d})$. This procedure can be viewed as a regular autoencoder with a non-differentiable nonlinear function that maps each latent vector $\mathbf{z}^{e}_{j}$ to $1$-of-$K$ embedding vectors $\mathbf{e}_{i}$.

During training, we use the straight-through gradient estimator (Bengio et al., 2013) to propagate gradients through the quantization function, i.e$\ldotp$, gradients are directly copied from the decoder input $\mathbf{z}^{d}$ to the encoder output $\mathbf{z}^{e}$. The loss function for training on a single data point $\mathbf{x}$ is given by

where $\mathtt{sg}$ stands for the stop_gradient operator, which is defined as identity function at forward computation and has zero partial derivatives at backward propagation. $\beta$ is a hyper-parameter ranging from $0.1$ to $2.0$. The first term of Eq$\ldotp$ (5) is the standard reconstruction loss, the second term is for embedding learning while the third term is for training stability (van den Oord et al., 2017). Samples from a VQ-VAE can be produced by first training an autoregressive model on its latent space, followed by decoding samples from the autoregressive model into the raw data space.

Since the VQ-VAE outputs a sequence of discrete latent codes $q(\mathbf{z}^{e})$, we wish to impose an ordering that prioritizes the code dimensions based on importance to reconstruction. In analogy to Eq$\ldotp$ (3), we can modify the reconstruction error term in Eq$\ldotp$ (5) to learn ordered latent representations. The modified loss function is an order-inducing objective given by

where $e_{\theta}(\mathbf{x})_{\leq i}$ and $\mathbf{z}^{d}_{\leq i}$ denote the results of keeping the top $i$ rows of $e_{\theta}(\mathbf{x})$ and $\mathbf{z}^{d}$ and then masking out the remainder rows with zero vectors. We uniformly sample the masking index $i\sim\mathcal{U}\{1,K\}$ to approximate the average in Eq$\ldotp$ (6) when $K$ is large. An alternative sampling distribution is the geometric distribution (Rippel et al., 2014). In Appendix C.3, we show that OVQ-VAE is sensitive to the choice of the parameter in geometric distribution. Because of the difficulty, Rippel et al. (2014) applies additional tricks. This results in the learning of ordered discrete latent variables, which can then be paired with a latent autoregressive model for anytime sampling.

In Section 4.1, we assume the encoder output to be a $K\times D$ matrix (i.e$\ldotp$, $\mathbf{z}^{e}\in\mathbb{R}^{K\times D}$). In practice, the output can have various sizes depending on the encoder network, and we need to reshape it to a two-dimensional matrix. For example, when encoding images, the encoder is typically a 2D convolutional neural network (CNN) whose output is a 3D latent feature map of size $L\times H\times W$. Here $L$, $H$, and $W$ stand for the channel, height, and width of the feature maps. We discuss below how the reshaping procedure can significantly impact the performance of anytime sampling and propose a reshaping procedure that facilitates high-performance anytime sampling.

The most common way of reshaping this 3D feature map, as in van den Oord et al. (2017), is to let $H\times W$ be the code length, and let the number of channels $L$ be the size of embedding vectors, i.e$\ldotp$, $K=H\times W$ and $D=L$. We call this pattern spatial-wise quantization, as each spatial location in the feature map corresponds to one code dimension and will be quantized separately. Since the code dimensions correspond to spatial locations of the feature map, they encode local features due to a limited receptive field. This is detrimental to anytime sampling, because early dimensions cannot capture the global information needed for reconstructing the entire image. We demonstrate this in Fig$\ldotp$ 2, which shows that OVQ-VAE with spatial-wise quantization is only able to reconstruct the top rows of an image with $\nicefrac{{1}}{{4}}$ of the code length.

To address this issue, we propose channel-wise quantization, where each channel of the feature map is viewed as one code dimension and quantized separately (see Fig$\ldotp$ 1(a) for visual comparison of spatial-wise and channel-wise quantization). Specifically, the code length is $L$ (i.e$\ldotp$, $K=L$), and the size of the embedding vectors is $H\times W$ (i.e$\ldotp$, $D=H\times W$). In this case, one code dimension includes all spatial locations in the feature map and can capture global information better. As shown in the right panel of Fig$\ldotp$ 2, channel-wise quantization clearly outperforms spatial-wise quantization for anytime sampling. We use channel-wise quantization in all subsequent experiments. Note that in practice we can easily apply the channel-wise quantization on VQ-VAE by changing the code length from $H\times W$ to $L$, as shown in Fig$\ldotp$ 1(a).

the encoding learned by a standard autoencoder (which we shall refer to as unordered). In addition to the theoretical analysis in Section 3.1.1, we provide further empirical analysis to characterize the difference between ordered and unordered codes. In particular, we compare the importance of the $i$-th code—as measured by the reduction in reconstruction error $\Delta(i)$—for PCA, standard (unordered) VQ-VAE, and ordered VQ-VAE. For VQ-VAE and ordered VQ-VAE, we define the reduction in reconstruction error $\Delta(i)$ for a data point $\mathbf{x}$ as

Averaging $\Delta_{\mathbf{x}}(i)$ over the entire dataset thus yields $\Delta(i)$. Similarly we define $\Delta(i)$ as the reduction on reconstruction error of the entire dataset, when adding the $i$-th principal component for PCA.

Fig$\ldotp$ 3 shows the $\Delta(i)$’s of the three models on the CIFAR-10. Since PCA and ordered VQ-VAE both learn an ordered encoding, their $\Delta(i)$’s decay gradually as $i$ increases. In contrast, the standard VQ-VAE with an unordered encoding exhibits a highly irregular $\Delta(i)$, indicating no meaningful ordering of the dimensions.

Fig$\ldotp$ 3 further shows how the reconstruction error decreases as a function of the truncated code length for the three models. Although unordered VQ-VAE and ordered VQ-VAE achieve similar reconstruction errors for sufficiently large code lengths, it is evident that an ordered encoding achieves significantly better reconstructions when the code length is aggressively truncated. When sufficiently truncated, we observe even PCA outperforms unordered VQ-VAE despite the latter being a more expressive model. In contrast, ordered VQ-VAE achieves superior reconstructions compared to PCA and unordered VQ-VAE across all truncation lengths.

We repeat the experiment on standard VAE to disentangle the specific implementations in VQ-VAE. We observe that the order-inducing objective has consistent results on standard VAE (Appendix C.1).

We test the performance of anytime sampling using OVQ-VAEs (Anytime + ordered) on several image datasets. We compare our approach to two baselines. One is the original VQ-VAE model proposed by van den Oord et al. (2017) without anytime sampling. The other is using anytime sampling with unordered VQ-VAEs (Anytime + unordered), where the models have the same architectures as ours but are trained by minimizing Eq$\ldotp$ (5). We empirically verify that 1) we are able to generate high quality image samples; 2) image quality degrades gracefully as we reduce the sampled code length for anytime sampling; and 3) anytime sampling improves the inference speed compared to naïve sampling of original VQ-VAEs.

We evaluate the model performance on the MNIST, CIFAR-10 (Krizhevsky, 2009) and CelebA (Liu et al., 2014) datasets. For CelebA, the images are resized to $64\times 64$. All pixel values are scaled to the range $[0,1]$. We borrow the model architectures and optimizers from van den Oord et al. (2017). The full code length and the codebook size are 16 and 126 for MNIST, 70 and 1000 for CIFAR-10, and 100 and 500 for CelebA respectively. We train a Transformer (Vaswani et al., 2017) on our VQ-VAEs, as opposed to the PixelCNN model used in van den Oord et al. (2017). PixelCNNs use standard 2D convolutional layers to capture a bounded receptive field and model the conditional dependence. Transformers apply attention mechanism and feed forward network to model the conditional dependence of 1D sequence. Transformers are arguably more suitable for channel-wise quantization, since there are no 2D spatial relations among different code dimensions that can be leveraged by convolutional models (such as PixelCNNs). Our experiments on different autoregressive models in Appendix C.2 further support the arguments.