Original: A NYTIME S AMPLING FOR A UTOREGRESSIVE M ODELS VIA O RDERED A UTOENCODING
Revised: A NYTIME S AMPLING FOR A UTOREGRESSIVE M ODELS VIA O RDERED A UTOENCODING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Autoregressive models are widely used for tasks such as image and audio generation.
Revised: Autoregressive models are widely used for tasks such as image and audio generation.

Original: The sampling process of these models, however, does not allow interruptions and cannot adapt to real-time computational resources.
Revised: The sampling process of these models, however, does not allow interruptions and cannot adapt to real-time computational resources.

Original: This challenge impedes the deployment of powerful autoregressive models, which involve a slow sampling process that is sequential in nature and typically scales linearly with respect to the data dimension.
Revised: This challenge impedes the deployment of powerful autoregressive models, which involve a slow sampling process that is sequential in nature and typically scales linearly with respect to the data dimension.

Original: To address this difﬁculty, we propose a new family of autoregressive models that enables anytime sampling.
Revised: To address this difﬁculty, we propose a new family of autoregressive models that enables anytime sampling.

Original: Inspired by Principal Component
Revised: Inspired by Principal Component

Original: Analysis, we learn a structured representation space where dimensions are ordered based on their importance with respect to reconstruction.
Revised: Analysis, we learn a structured representation space where dimensions are ordered based on their importance with respect to reconstruction.

Original: Using an autoregressive model in this latent space, we trade off sample quality for computational efﬁciency by truncating the generation process before decoding into the original data space.
Revised: Using an autoregressive model in this latent space, we trade off sample quality for computational efﬁciency by truncating the generation process before decoding into the original data space.

Original: Experimentally, we demonstrate in several image and audio generation tasks that sample quality degrades gracefully as we reduce the computational budget for sampling.
Revised: Experimentally, we demonstrate in several image and audio generation tasks that sample quality degrades gracefully as we reduce the computational budget for sampling.

Original: The approach suffers almost no loss in sample quality (measured by
Revised: The approach suffers almost no loss in sample quality (measured by

Original: FID) using only 60% to 80% of all latent dimensions for image data.
Revised: FID) using only 60% to 80% of all latent dimensions for image data.

Revised: Code is available at https://github.com/Newbeeer/Anytime-Auto-Regressive-Model.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Autoregressive models are a prominent approach to data generation, and have been widely used to produce high quality samples of images (Oord et al., 2016b; Salimans et al., 2017; Menick & Kalchbrenner, 2018), audio (Oord et al., 2016a), video (Kalchbrenner et al., 2017) and text (Kalchbrenner et al., 2016; Radford et al., 2019).
Revised: Autoregressive models are a prominent approach to data generation, and have been widely used to produce high quality samples of images (Oord et al., 2016b; Salimans et al., 2017; Menick & Kalchbrenner, 2018), audio (Oord et al., 2016a), video (Kalchbrenner et al., 2017) and text (Kalchbrenner et al., 2016; Radford et al., 2019).

Original: These models represent a joint distribution as a product of (simpler) conditionals, and sampling requires iterating over all these conditional distributions in a certain order.
Revised: These models represent a joint distribution as a product of (simpler) conditionals, and sampling requires iterating over all these conditional distributions in a certain order.

Original: Due to the sequential nature of this process, the computational cost will grow at least linearly with respect to the number of conditional distributions, which is typically equal to the data dimension.
Revised: Due to the sequential nature of this process, the computational cost will grow at least linearly with respect to the number of conditional distributions, which is typically equal to the data dimension.

Original: As a result, the sampling process of autoregressive models can be slow and does not allow interruptions.
Revised: As a result, the sampling process of autoregressive models can be slow and does not allow interruptions.

Original: Although caching techniques have been developed to speed up generation (Ramachandran et al., 2017; Guo et al., 2017), the high cost of sampling limits their applicability in many scenarios.
Revised: Although caching techniques have been developed to speed up generation (Ramachandran et al., 2017; Guo et al., 2017), the high cost of sampling limits their applicability in many scenarios.

Original: For example, when running on multiple devices with different computational resources, we may wish to trade off sample quality for faster generation based on the computing power available on each device.
Revised: For example, when running on multiple devices with different computational resources, we may wish to trade off sample quality for faster generation based on the computing power available on each device.

Original: Currently, a separate model must be trained for each device ( i.e . , computational budget) in order to trade off sample quality for faster generation, and there is no way to control this trade-off on the ﬂy to accommodate instantaneous resource availability at time-of-deployment.
Revised: Currently, a separate model must be trained for each device ( i.e . , computational budget) in order to trade off sample quality for faster generation, and there is no way to control this trade-off on the ﬂy to accommodate instantaneous resource availability at time-of-deployment.

Original: To address this difﬁculty, we consider the novel task of adaptive autoregressive generation under computational constraints .
Revised: To address this difﬁculty, we consider the novel task of adaptive autoregressive generation under computational constraints .

Original: We seek to build a single model that can automatically trade-off sample quality versus computational cost via anytime sampling , i.e . , where the sampling process may be interrupted anytime ( e.g . , because of exhausted computational budget) to yield a complete sample whose sample quality decays with the earliness of termination.
Revised: We seek to build a single model that can automatically trade-off sample quality versus computational cost via anytime sampling , i.e . , where the sampling process may be interrupted anytime ( e.g . , because of exhausted computational budget) to yield a complete sample whose sample quality decays with the earliness of termination.

Original: In particular, we take advantage of a generalization of Principal Components Analysis (PCA) proposed by Rippel et al.
Revised: In particular, we take advantage of a generalization of Principal Components Analysis (PCA) proposed by Rippel et al.

Original: (2014), which learns an ordered representations induced by a structured application of dropout to the representations learned by an autoencoder.
Revised: (2014), which learns an ordered representations induced by a structured application of dropout to the representations learned by an autoencoder.

Original: Such a representation encodes raw data into a latent space where dimensions are sorted based on their importance for reconstruction.
Revised: Such a representation encodes raw data into a latent space where dimensions are sorted based on their importance for reconstruction.

Original: Autoregressive modeling is then applied in the ordered representation space instead.
Revised: Autoregressive modeling is then applied in the ordered representation space instead.

Original: This approach enables a natural trade-off between quality and computation by truncating the length of the representations: When running on devices with high computational capacity, we can afford to generate the full representation and decode it to obtain a high quality sample; when on a tighter computational budget, we can generate only the ﬁrst few dimensions of the representation and decode it to a sample whose quality degrades smoothly with truncation.
Revised: This approach enables a natural trade-off between quality and computation by truncating the length of the representations: When running on devices with high computational capacity, we can afford to generate the full representation and decode it to obtain a high quality sample; when on a tighter computational budget, we can generate only the ﬁrst few dimensions of the representation and decode it to a sample whose quality degrades smoothly with truncation.

Original: Because decoding is usually fast and the main computation bottleneck lies on the autoregressive part, the run-time grows proportionally relative to the number of sampled latent dimensions.
Revised: Because decoding is usually fast and the main computation bottleneck lies on the autoregressive part, the run-time grows proportionally relative to the number of sampled latent dimensions.

Original: Through experiments, we show that our autoregressive models are capable of trading off sample quality and inference speed.
Revised: Through experiments, we show that our autoregressive models are capable of trading off sample quality and inference speed.

Original: When training autoregressive models on the latent space given by our encoder, we witness little degradation of image sample quality using only around 60% to 80% of all latent codes, as measured by Fréchet Inception Distance (Heusel et al., 2017) on CIFAR-10 and
Revised: When training autoregressive models on the latent space given by our encoder, we witness little degradation of image sample quality using only around 60% to 80% of all latent codes, as measured by Fréchet Inception Distance (Heusel et al., 2017) on CIFAR-10 and

Original: CelebA. Compared to standard autoregressive models, our approach allows the sample quality to degrade gracefully as we reduce the computational budget for sampling.
Revised: CelebA. Compared to standard autoregressive models, our approach allows the sample quality to degrade gracefully as we reduce the computational budget for sampling.

Original: We also observe that on the VCTK audio dataset (Veaux et al., 2017), our autoregressive model is able to generate the low frequency features ﬁrst, then gradually reﬁne the waveforms with higher frequency components as we increase the number of sampled latent dimensions.
Revised: We also observe that on the VCTK audio dataset (Veaux et al., 2017), our autoregressive model is able to generate the low frequency features ﬁrst, then gradually reﬁne the waveforms with higher frequency components as we increase the number of sampled latent dimensions.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Autoregressive Models Autoregressive models deﬁne a probability distribution over data points x ∈ R D by factorizing the joint probability distribution as a product of univariate conditional distributions with the chain rule.
Revised: Autoregressive Models Autoregressive models deﬁne a probability distribution over data points x ∈ R D by factorizing the joint probability distribution as a product of univariate conditional distributions with the chain rule.

Original: Using p θ to denote the distribution of the model, we have:
Revised: Using p θ to denote the distribution of the model, we have:

Original: [Equation]
Revised: [Equation]

Original: The model is trained by maximizing the likelihood:
Revised: The model is trained by maximizing the likelihood:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where p d ( x ) represents the data distribution.
Revised: where p d ( x ) represents the data distribution.

Original: Different autoregressive models adopt different orderings of input dimensions and parameterize the conditional probability p θ ( x i | x 1 , · · · , x i − 1 ) , i = 1 , · · · , D in different ways.
Revised: Different autoregressive models adopt different orderings of input dimensions and parameterize the conditional probability p θ ( x i | x 1 , · · · , x i − 1 ) , i = 1 , · · · , D in different ways.

Original: Most architectures over images order the variables x 1 , · · · , x D of image x in raster scan order ( i.e . , left- right then top to bottom).
Revised: Most architectures over images order the variables x 1 , · · · , x D of image x in raster scan order ( i.e . , left-toright then top-to-bottom).

Original: Popular autoregressive architectures include MADE (Germain et al., 2015), PixelCNN (Oord et al., 2016b; van den Oord et al., 2016; Salimans et al., 2017) and Transformer (Vaswani et al., 2017), where they respectively use masked linear layers, convolutional  and self-attention blocks to ensure that the output corresponding to pθ ( i |
Revised: Popular autoregressive architectures include MADE (Germain et al., 2015), PixelCNN (Oord et al., 2016b; van den Oord et al., 2016; Salimans et al., 2017) and Transformer (Vaswani et al., 2017), where they respectively use masked linear layers, convolutional layers and self-attention blocks to ensure that the output corresponding to p

Original:  ) is oblivious of x i , x  +1 , · · · , 
Revised: θ ( x i | x 1 , · · · , x i − 1 ) is oblivious of x i , x i +1 , · · · , x D .

Original: Cost of Sampling During training, we can evaluate autoregressive models efﬁciently because x 1 , · · · , x D are provided by data and all conditionals p ( x i | x 1 , · · · , x i − 1 ) can be computed in parallel.
Revised: Cost of Sampling During training, we can evaluate autoregressive models efﬁciently because x 1 , · · · , x D are provided by data and all conditionals p ( x i | x 1 , · · · , x i − 1 ) can be computed in parallel.

Original: In contrast, sampling from autoregressive models is an inherently sequential process and cannot be easily accelerated by parallel computing: we ﬁrst need to sample x 1 , after which we sample x 2 from p θ ( x 2 | x 1 ) and so on—the i -th variable x
Revised: In contrast, sampling from autoregressive models is an inherently sequential process and cannot be easily accelerated by parallel computing: we ﬁrst need to sample x 1 , after which we sample x 2 from p θ ( x 2 | x 1 ) and so on—the i -th variable x

Original: i can only be obtained after we have already computed x 1 , · · · , x i − 1 .
Revised: i can only be obtained after we have already computed x 1 , · · · , x i − 1 .

Original: Thus, the run-time of autoregressive generation grows at least linearly with respect to the length of a sample.
Revised: Thus, the run-time of autoregressive generation grows at least linearly with respect to the length of a sample.

Original: In practice, the sample length D can be more than hundreds of thousands for real-world image and audio data.
Revised: In practice, the sample length D can be more than hundreds of thousands for real-world image and audio data.

Original: This poses a major challenge to fast autoregressive generation on a small computing budget.
Revised: This poses a major challenge to fast autoregressive generation on a small computing budget.

Original:  A NYTIME S AMPLING WITH O RDERED A UTOENCODERS
Revised:  A NYTIME S AMPLING WITH O RDERED A UTOENCODERS

Original: Our goal is to circumvent the non-interruption and linear time complexity of autoregressive models by pushing the task of autoregressive modeling from the original data space (e.g., pixel space) into an ordered representation space.
Revised: Our goal is to circumvent the non-interruption and linear time complexity of autoregressive models by pushing the task of autoregressive modeling from the original data space (e.g., pixel space) into an ordered representation space.

Original: In doing so, we develop a new class of autoregressive models where premature truncation of the autoregressive sampling process leads to the generation of a lower quality sample instead of an incomplete sample.
Revised: In doing so, we develop a new class of autoregressive models where premature truncation of the autoregressive sampling process leads to the generation of a lower quality sample instead of an incomplete sample.

Original: In this section, we shall ﬁrst describe the learning of the ordered representation space via the use of an ordered autoencoder .
Revised: In this section, we shall ﬁrst describe the learning of the ordered representation space via the use of an ordered autoencoder .

Original: We then describe how to achieve anytime sampling with ordered autoencoders.
Revised: We then describe how to achieve anytime sampling with ordered autoencoders.

Original: 3. O RDERED A UTOENCODERS
Revised: 3. O RDERED A UTOENCODERS

Original: Consider an autoencoder that encodes an input x ∈ R D to a code z ∈ R K .
Revised: Consider an autoencoder that encodes an input x ∈ R D to a code z ∈ R K .

Original: Let z
Revised: Let z

Original: =
Revised: =

Original: e θ ( x ) :
Revised: e θ ( x ) :

Original: R D →
Revised: R D →

Original: R K be the encoder parameterized by θ and x (cid:48)
Revised: R K be the encoder parameterized by θ and x (cid:48)

Original: =
Revised: =

Original: d φ ( z ) : R K → R D be the decoder parameterized by φ .
Revised: d φ ( z ) : R K → R D be the decoder parameterized by φ .

Original: We deﬁne e θ ( · ) ≤ i : x ∈ R D (cid:55)→ ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , which truncates the representation to the ﬁrst i dimensions of the encoding z = e θ ( x ) , masking out the remainder of the dimensions with a zero value.
Revised: We deﬁne e θ ( · ) ≤ i : x ∈ R D (cid:55)→ ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , which truncates the representation to the ﬁrst i dimensions of the encoding z = e θ ( x ) , masking out the remainder of the dimensions with a zero value.

Original: We deﬁne the ordered autoencoder objective as
Revised: We deﬁne the ordered autoencoder objective as

Original: [Equation]
Revised: [Equation]

Original: We note that Eq .
Revised: We note that Eq .

Original: (3) is equivalent to Rippel et al.
Revised: (3) is equivalent to Rippel et al.

Original: ’s
Revised: ’s

Original: nested dropout formulation using a uniform sampling of possible truncations.
Revised: nested dropout formulation using a uniform sampling of possible truncations.

Original: 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 .
Revised: 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 .

Original: (3) recovers PCA.
Revised: (3) recovers PCA.

Original: . T HEORETICAL A NALYSIS
Revised: 3.1. T HEORETICAL A NALYSIS

Original: Rippel et al.
Revised: Rippel et al.

Original: (2014)’s analysis of the ordered autoencoder is limited to linear/sigmoid encoder and a linear decoder.
Revised: (2014)’s analysis of the ordered autoencoder is limited to linear/sigmoid encoder and a linear decoder.

Original: In this section, we extend the analysis to general autoencoder architectures by employing an information-theoretic framework to analyze the importance of the i -th latent code to reconstruction for ordered autoencoders.
Revised: In this section, we extend the analysis to general autoencoder architectures by employing an information-theoretic framework to analyze the importance of the i -th latent code to reconstruction for ordered autoencoders.

Original: We ﬁrst reframe our problem from a probabilistic perspective.
Revised: We ﬁrst reframe our problem from a probabilistic perspective.

Original: In lieu of using deterministic autoencoders, we assume that both the encoder and decoder are stochastic functions.
Revised: In lieu of using deterministic autoencoders, we assume that both the encoder and decoder are stochastic functions.

Original: In particular, we let q e θ ( z | x ) be a probability distribution over z ∈ R K conditioned on input x , and similarly let p d φ ( x | z ) be the stochastic counterpart to d φ ( z ) .
Revised: In particular, we let q e θ ( z | x ) be a probability distribution over z ∈ R K conditioned on input x , and similarly let p d φ ( x | z ) be the stochastic counterpart to d φ ( z ) .

Original: We then use q e θ ( z | x ) ≤ i to denote the distribution of ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , where z ∼ q e θ ( z | x ) , and let p d φ ( x | z ) ≤ i represent the distribution of p d φ ( x | ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K ) .
Revised: We then use q e θ ( z | x ) ≤ i to denote the distribution of ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , where z ∼ q e θ ( z | x ) , and let p d φ ( x | z ) ≤ i represent the distribution of p d φ ( x | ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K ) .

Original: We can modify
Revised: We can modify

Original: Eq .
Revised: Eq .

Original: (3) to have the following form:
Revised: (3) to have the following form:

Original: [Equation]
Revised: [Equation]

Original: where U{ 1 , K } denotes a uniform distribution over { 1 , 2 , · · · , K } , and p d ( x ) represents the data distribution.
Revised: where U{ 1 , K } denotes a uniform distribution over { 1 , 2 , · · · , K } , and p d ( x ) represents the data distribution.

Original: We can choose both the encoder and decoder to be fully factorized Gaussian distributions with a ﬁxed variance σ 2 , then Eq .
Revised: We can choose both the encoder and decoder to be fully factorized Gaussian distributions with a ﬁxed variance σ 2 , then Eq .

Original: (13) can be simpliﬁed to
Revised: (13) can be simpliﬁed to

Original: [Equation]
Revised: [Equation]

Original: The stochastic encoder and decoder in this case will become deterministic when σ → 0 , and the above equation will yield the same encoder/decoder pair as Eq .
Revised: The stochastic encoder and decoder in this case will become deterministic when σ → 0 , and the above equation will yield the same encoder/decoder pair as Eq .

Original: (3) when σ → 0 and N → ∞ .
Revised: (3) when σ → 0 and N → ∞ .

Original: The optimal encoders and decoders that minimize Eq .
Revised: The optimal encoders and decoders that minimize Eq .

Original: (13) satisfy the following property.
Revised: (13) satisfy the following property.

Original: Theorem 1.
Revised: Theorem 1.

Original: Let x denote the input random variable.
Revised: Let x denote the input random variable.

Original: Assuming both the encoder and decoder are optimal in terms of minimizing Eq .
Revised: Assuming both the encoder and decoder are optimal in terms of minimizing Eq .

Original: (13), and ∀ i ∈ 3 , · · · , K, z i − 1 ⊥ z i | x , z ≤ i − 2 , we have[Equation] [Equation] [Equation]
Revised: (13), and ∀ i ∈ 3 , · · · , K, z i − 1 ⊥ z i | x , z ≤ i − 2 , we have

Original: |

Original: [Equation]
Revised: [Equation]

Original: ≤

Original: [Equation]

Original: −

Original: [Equation] ) ≤ [Equation] z 
Revised: where z ≤ i denotes ( z 1 , z 2 , · · · , z i ) .

Original: [Equation]
Revised: [Figure]

Original: −

Original: [Equation]

Original: x | z ≤ i − 2 ) , where z ≤ i denotes ( z 1 , z 2 , · · · , z i ) .

Original: We defer the proof to Appendix A.1.
Revised: We defer the proof to Appendix A.1.

Original: The assumption z i − 1 ⊥ z i | x , z ≤ i − 2 holds whenever the encoder q θ ( z | x ) is a factorized distribution, which is a common choice in variational autoencoders (Kingma & Welling, 2013), and we use I ( a ; b | c ) to denote the mutual information between random variables a and b conditioned on c .
Revised: The assumption z i − 1 ⊥ z i | x , z ≤ i − 2 holds whenever the encoder q θ ( z | x ) is a factorized distribution, which is a common choice in variational autoencoders (Kingma & Welling, 2013), and we use I ( a ; b | c ) to denote the mutual information between random variables a and b conditioned on c .

Original: Intuitively, the above theorem states that for optimal encoders and decoders that minimize Eq .
Revised: Intuitively, the above theorem states that for optimal encoders and decoders that minimize Eq .

Original: (13), one can extract less additional information about the raw input as the code gets longer.
Revised: (13), one can extract less additional information about the raw input as the code gets longer.

Original: Therefore, there exists a natural ordering among different dimensions of the code based on the additional information they can provide for reconstructing the inputs.
Revised: Therefore, there exists a natural ordering among different dimensions of the code based on the additional information they can provide for reconstructing the inputs.

Original: 3. A NYTIME S AMPLING
Revised: 3. A NYTIME S AMPLING

Original: 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).
Revised: 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).

Original: For each input x i in a dataset x 1 , · · · , x N , we feed it to the encoder to get z i =
Revised: For each input x i in a dataset x 1 , · · · , x N , we feed it to the encoder to get z i =

Original: e θ ( x i ) .
Revised: e θ ( x i ) .

Original: The resulting codes z 1 , z 2 , · · · , z N are used as training data.
Revised: The resulting codes z 1 , z 2 , · · · , z N are used as training data.

Original: After training both the ordered autoencoder and autoregressive model, we can perform anytime sampling on a large spectrum of computing budgets.
Revised: After training both the ordered autoencoder and autoregressive model, we can perform anytime sampling on a large spectrum of computing budgets.

Original: Suppose for example we can afford to generate T code dimensions from the autoregressive[Figure] model, denoted as z ≤ T ∈ R T .
Revised: Suppose for example we can afford to generate T code dimensions from the autoregressive model, denoted as z ≤ T ∈ R T .

Original: We can simply zero-pad it to get ( z 1 , z 2 , · · · , z T , 0 , · · · , 0) T ∈ R K and decode it to get a complete sample.
Revised: We can simply zero-pad it to get ( z 1 , z 2 , · · · , z T , 0 , · · · , 0) T ∈ R K and decode it to get a complete sample.

Original: Unlike the autoregressive part, the decoder has access to all dimensions of the latent code at the same time and can decode in parallel.
Revised: Unlike the autoregressive part, the decoder has access to all dimensions of the latent code at the same time and can decode in parallel.

Original: The framework is shown in Fig . 1(b).
Revised: The framework is shown in Fig . 1(b).

Original: 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.
Revised: 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.

Original: On modern
Revised: On modern

Original: GPUs, the code length has minimal effect on the run-time of decoding, as long as the decoder is not itself autoregressive (see empirical veriﬁcations in Section 5.2.2).
Revised: GPUs, the code length has minimal effect on the run-time of decoding, as long as the decoder is not itself autoregressive (see empirical veriﬁcations in Section 5.2.2).

Original:  O RDERED VQ-VAE
Revised:  O RDERED VQ-VAE

Original: In this section, we apply the ordered autoencoder framework to the vector quantized variational autoencoder (VQ-VAE) and its extension (van den Oord et al., 2017; Razavi et al., 2019).
Revised: In this section, we apply the ordered autoencoder framework to the vector quantized variational autoencoder (VQ-VAE) and its extension (van den Oord et al., 2017; Razavi et al., 2019).

Original: Since these models are quantized autoencoders paired with a latent autoregressive model, they admit a natural extension to ordered VQ-VAEs (OVQ-VAEs) under our framework—a new family of VQVAE models capable of anytime sampling.
Revised: Since these models are quantized autoencoders paired with a latent autoregressive model, they admit a natural extension to ordered VQ-VAEs (OVQ-VAEs) under our framework—a new family of VQVAE models capable of anytime sampling.

Original: Below, we begin by describing the VQ-VAE, and then highlight two key design choices (ordered discrete codes and channel-wise quantization) critical for
Revised: Below, we begin by describing the VQ-VAE, and then highlight two key design choices (ordered discrete codes and channel-wise quantization) critical for

Original: OVQ-VAEs.
Revised: OVQ-VAEs.

Original: We show that, with small changes of the original VQ-VAE, these two choices can be applied straightforwardly.
Revised: We show that, with small changes of the original VQ-VAE, these two choices can be applied straightforwardly.

Original: 4. VQ-VAE
Revised: 4. VQ-VAE

Original: To construct a VQ-VAE with code length of K discrete latent variables , the encoder must ﬁrst map the raw input x to a continuous representation z e =
Revised: To construct a VQ-VAE with code length of K discrete latent variables , the encoder must ﬁrst map the raw input x to a continuous representation z e =

Original: e θ ( x ) ∈ R K × D , before feeding it to a vector-valued quantization function q : R K × D → { 1 , 2 , · · · , C } K deﬁned as
Revised: e θ ( x ) ∈ R K × D , before feeding it to a vector-valued quantization function q : R K × D → { 1 , 2 , · · · , C } K deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where q ( z e ) j ∈ { 1 , 2 , · · · , C } denotes the j -th component of the vector-valued function q ( z e ) , z ej ∈ R D denotes the j -th row of z e , and e i denotes the i -th row of the embedding matrix E ∈ E C × D .
Revised: where q ( z e ) j ∈ { 1 , 2 , · · · , C } denotes the j -th component of the vector-valued function q ( z e ) , z ej ∈ R D denotes the j -th row of z e , and e i denotes the i -th row of the embedding matrix E ∈ E C × D .

Original: Next, we view q ( z e ) as a sequence of indices and use them to look up embedding vectors from the codebook E .
Revised: Next, we view q ( z e ) as a sequence of indices and use them to look up embedding vectors from the codebook E .

Original: This yields a latent representation z d ∈ R K × D , given by z dj =
Revised: This yields a latent representation z d ∈ R K × D , given by z dj =

Original: e q ( z e ) j , where z dj ∈ R D denotes the j -th row of z d .
Revised: e q ( z e ) j , where z dj ∈ R D denotes the j -th row of z d .

Original: Finally, we can decode z d to obtain the reconstruction d φ ( z d ) .
Revised: Finally, we can decode z d to obtain the reconstruction d φ ( z d ) .

Original: This procedure can be viewed as a regular autoencoder with a non-differentiable nonlinear function that maps each latent vector z ej to 1 -of- K embedding vectors e i .
Revised: This procedure can be viewed as a regular autoencoder with a non-differentiable nonlinear function that maps each latent vector z ej to 1 -of- K embedding vectors e i .

Original: During training, we use the straight-through gradient estimator (Bengio et al., 2013) to propagate gradients through the quantization function, i.e . , gradients are directly copied from the decoder input z d to the encoder output z e .
Revised: During training, we use the straight-through gradient estimator (Bengio et al., 2013) to propagate gradients through the quantization function, i.e . , gradients are directly copied from the decoder input z d to the encoder output z e .

Original: The loss function for training on a single data point x is given by
Revised: The loss function for training on a single data point x is given by

Original: [Equation]
Revised: [Equation]

Original: ) where sg stands for the stop_gradient operator, which is deﬁned as identity function at forward computation and has zero partial derivatives at backward propagation.
Revised: where sg stands for the stop_gradient operator, which is deﬁned as identity function at forward computation and has zero partial derivatives at backward propagation.

Original: β is a hyper-parameter ranging from 0 .
Revised: β is a hyper-parameter ranging from 0 .

Original: 1 to 2 . 0 .
Revised: 1 to 2 . 0 .

Original: The ﬁrst term of Eq .
Revised: The ﬁrst term of Eq .

Original: (5) is the standard reconstruction loss, the second term is 
Revised: (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).

Original: Samples from a VQ-VAE can be produced by ﬁrst training an autoregressive model on its latent space, followed by decoding samples from the autoregressive model into the raw data space.
Revised: Samples from a VQ-VAE can be produced by ﬁrst training an autoregressive model on its latent space, followed by decoding samples from the autoregressive model into the raw data space.

Original: 4. O RDERED D ISCRETE C ODES
Revised: 4. O RDERED D ISCRETE C ODES

Original: Since the VQ-VAE outputs a sequence of discrete latent codes q ( z e ) , we wish to impose an ordering that prioritizes the code dimensions based on importance to reconstruction.
Revised: Since the VQ-VAE outputs a sequence of discrete latent codes q ( z e ) , we wish to impose an ordering that prioritizes the code dimensions based on importance to reconstruction.

Original: In analogy to Eq .
Revised: In analogy to Eq .

Original: (3), we can modify the reconstruction error term in Eq .
Revised: (3), we can modify the reconstruction error term in Eq .

Original: (5) to learn ordered latent representations.
Revised: (5) to learn ordered latent representations.

Original: The modiﬁed loss function is an order-inducing objective given by
Revised: The modiﬁed loss function is an order-inducing objective given by

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where e θ ( x ) ≤ i and z d ≤ i denote the results of keeping the top i rows of e θ ( x ) and z d and then masking out the remainder rows with zero vectors.
Revised: where e θ ( x ) ≤ i and z d ≤ i denote the results of keeping the top i rows of e θ ( x ) and z d and then masking out the remainder rows with zero vectors.

Original: We uniformly sample the masking index i ∼ U{ 1 , K } to approximate the average in Eq .
Revised: We uniformly sample the masking index i ∼ U{ 1 , K } to approximate the average in Eq .

Original: (6) when K is large.
Revised: (6) when K is large.

Original: An alternative sampling distribution is the geometric distribution (Rippel et al., 2014).
Revised: An alternative sampling distribution is the geometric distribution (Rippel et al., 2014).

Original: In Appendix C.3, we show that OVQ-VAE is sensitive to the choice of the parameter in geometric distribution.
Revised: In Appendix C.3, we show that OVQ-VAE is sensitive to the choice of the parameter in geometric distribution.

Original: Because of the difﬁculty, Rippel et al.
Revised: Because of the difﬁculty, Rippel et al.

Original: (2014) applies additional tricks.
Revised: (2014) applies additional tricks.

Original: This results in the learning of ordered discrete latent variables, which can then be paired with a latent autoregressive model for anytime sampling.
Revised: This results in the learning of ordered discrete latent variables, which can then be paired with a latent autoregressive model for anytime sampling.

Original:  C HANNEL -W ISE Q UANTIZATION
Revised: 4. C HANNEL -W ISE Q UANTIZATION

Original: In Section 4.1, we assume the encoder output to be a K × D matrix ( i.e . , z e ∈ R K × D ).
Revised: In Section 4.1, we assume the encoder output to be a K × D matrix ( i.e . , z e ∈ R K × D ).

Original: In practice, the output can have various sizes depending on the encoder network, and we need to reshape it to a two-dimensional matrix.
Revised: In practice, the output can have various sizes depending on the encoder network, and we need to reshape it to a two-dimensional matrix.

Original: 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 × H × W .
Revised: 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 × H × W .

Original: Here L , H , and W stand for the channel, height, and width of the feature maps.
Revised: Here L , H , and W stand for the channel, height, and width of the feature maps.

Original: We discuss below how the reshaping procedure can signiﬁcantly impact the performance of anytime sampling and propose a reshaping procedure that facilitates high-performance anytime sampling.
Revised: We discuss below how the reshaping procedure can signiﬁcantly impact the performance of anytime sampling and propose a reshaping procedure that facilitates high-performance anytime sampling.

Original: Consider convolutional encoders on image data, where the output feature map has a size of L × H × W .
Revised: Consider convolutional encoders on image data, where the output feature map has a size of L × H × W .

Original: The most common way of reshaping this 3D feature map, as in van den
Revised: The most common way of reshaping this 3D feature map, as in van den

Original: Oord et al.
Revised: Oord et al.

Original: (2017), is to let H × W be the code length, and let the number of channels L be the size of embedding vectors, i.e . ,
Revised: (2017), is to let H × W be the code length, and let the number of channels L be the size of embedding vectors, i.e . ,

Original: K
Revised: K

Original: = H × W and D =
Revised: = H × W and D =

Original: L .
Revised: L .

Original: 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.
Revised: 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.

Original: Since the code dimensions correspond to spatial locations of the feature map, they encode local features due to a limited receptive ﬁeld.
Revised: Since the code dimensions correspond to spatial locations of the feature map, they encode local features due to a limited receptive ﬁeld.

Original: This is detrimental to anytime sampling, because early dimensions cannot capture the global information needed for reconstructing the entire image.
Revised: This is detrimental to anytime sampling, because early dimensions cannot capture the global information needed for reconstructing the entire image.

Original: We demonstrate this in Fig .
Revised: We demonstrate this in Fig .

Original: 2, which shows that OVQ-VAE with spatial-wise quantization is only able to reconstruct the top rows of an image with 1 / 4 of the code length.
Revised: 2, which shows that OVQ-VAE with spatial-wise quantization is only able to reconstruct the top rows of an image with 1 / 4 of the code length.

Original: 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 . 1(a) for visual comparison of spatial-wise and channel-wise quantization).
Revised: 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 . 1(a) for visual comparison of spatial-wise and channel-wise quantization).

Original: Speciﬁcally, the code length is L ( i.e . , K = L ), and the size of the embedding vectors is H × W ( i.e . , D
Revised: Speciﬁcally, the code length is L ( i.e . , K = L ), and the size of the embedding vectors is H × W ( i.e . , D

Original: = H × W ).
Revised: = H × W ).

Original: In this case, one code dimension includes all spatial locations in the feature map and can capture global information better.
Revised: In this case, one code dimension includes all spatial locations in the feature map and can capture global information better.

Original: As shown in the right panel of Fig . 2, channel-wise quantization clearly outperforms spatial-wise quantization for anytime sampling.
Revised: As shown in the right panel of Fig . 2, channel-wise quantization clearly outperforms spatial-wise quantization for anytime sampling.

Original: We use channel-wise quantization in all subsequent experiments.
Revised: We use channel-wise quantization in all subsequent experiments.

Original: Note that in practice we can easily apply the channel-wise quantization on VQ-VAE by changing the code length from H × W to L , as shown in Fig . 1(a).
Revised: Note that in practice we can easily apply the channel-wise quantization on VQ-VAE by changing the code length from H × W to L , as shown in Fig . 1(a).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In our experiments, we focus on anytime sampling for autoregressive models trained on the latent space of OVQ-VAEs, as shown in Fig . 1(b).
Revised: In our experiments, we focus on anytime sampling for autoregressive models trained on the latent space of OVQ-VAEs, as shown in Fig . 1(b).

Original: We ﬁrst verify that our learning objectives in Eq .
Revised: We ﬁrst verify that our learning objectives in Eq .

Original: Eq .
Revised: Eq .

Original: (6) are effective at inducing ordered representations.
Revised: (6) are effective at inducing ordered representations.

Original: Next, we demonstrate that our OVQ-VAE models achieve comparable sample quality to regular VQ-VAEs on several image and audio datasets, while additionally allowing a graceful trade-off between sample quality and computation time via anytime sampling.
Revised: Next, we demonstrate that our OVQ-VAE models achieve comparable sample quality to regular VQ-VAEs on several image and audio datasets, while additionally allowing a graceful trade-off between sample quality and computation time via anytime sampling.

Original: Due to limits on space, we defer the results of audio generation to Appendix B and provide additional experimental details, results and code links in Appendix E and G.
Revised: Due to limits on space, we defer the results of audio generation to Appendix B and provide additional experimental details, results and code links in Appendix E and G.

Original: 5. ORDERED V ERSUS U NORDERED C ODES
Revised: 5. O

Revised: [Equation] [Figure]

Original: Our proposed ordered autoencoder framework learns an ordered encoding that is in contrast to
Revised: Our proposed ordered autoencoder framework learns an ordered encoding that is in contrast to

Original: [Figure]
Revised: [Figure]

Original: the encoding learned by a standard autoencoder (which we shall refer to as unordered).
Revised: the encoding learned by a standard autoencoder (which we shall refer to as unordered).

Original: 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.
Revised: 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.

Original: In particular, we compare the importance of the i -th code—as measured by the reduction in reconstruction error ∆( i ) —for
Revised: In particular, we compare the importance of the i -th code—as measured by the reduction in reconstruction error ∆( i ) —for

Original: PCA, standard (unordered) VQ-VAE, and ordered VQ-VAE.
Revised: PCA, standard (unordered) VQ-VAE, and ordered VQ-VAE.

Original: For VQ-VAE and ordered VQ-VAE, we deﬁne the reduction in reconstruction error ∆( i ) for a data point x as∆ x (
Revised: For VQ-VAE and ordered VQ-VAE, we deﬁne the reduction in reconstruction error ∆( i ) for a data point x as

Original: [Table]
Revised: [Equation]

Original: Averaging ∆ x ( i ) over the entire dataset thus yields ∆( i ) .
Revised: Averaging ∆ x ( i ) over the entire dataset thus yields ∆( i ) .

Original: Similarly we deﬁne ∆( i ) as the reduction on reconstruction error of the entire dataset, when adding the i -th principal component for PCA.
Revised: Similarly we deﬁne ∆( i ) as the reduction on reconstruction error of the entire dataset, when adding the i -th principal component for PCA.

Original: Fig . 3(a) shows the ∆( i ) ’s of the three models on the CIFAR-10.
Revised: Fig . 3(a) shows the ∆( i ) ’s of the three models on the CIFAR-10.

Original: Since PCA and ordered VQ-VAE both learn an ordered encoding, their ∆( i ) ’s decay gradually as i increases.
Revised: Since PCA and ordered VQ-VAE both learn an ordered encoding, their ∆( i ) ’s decay gradually as i increases.

Original: In contrast, the standard
Revised: In contrast, the standard

Original: VQ-VAE with an unordered encoding exhibits a highly irregular ∆( i ) , indicating no meaningful ordering of the dimensions.
Revised: VQ-VAE with an unordered encoding exhibits a highly irregular ∆( i ) , indicating no meaningful ordering of the dimensions.

Original: Fig . 3(b) further shows how the reconstruction error decreases as a function of the truncated code length for the three models.
Revised: Fig . 3(b) further shows how the reconstruction error decreases as a function of the truncated code length for the three models.

Original: Although unordered VQ-VAE and ordered VQ-VAE achieve similar reconstruction errors for sufﬁciently large code lengths, it is evident that an ordered encoding achieves signiﬁcantly better reconstructions when the code length is aggressively truncated.
Revised: Although unordered VQ-VAE and ordered VQ-VAE achieve similar reconstruction errors for sufﬁciently large code lengths, it is evident that an ordered encoding achieves signiﬁcantly better reconstructions when the code length is aggressively truncated.

Original: When sufﬁciently truncated, we observe even PCA outperforms unordered VQ-VAE despite the latter being a more expressive model.
Revised: When sufﬁciently truncated, we observe even PCA outperforms unordered VQ-VAE despite the latter being a more expressive model.

Original: In contrast, ordered VQ-VAE achieves superior reconstructions compared to PCA and unordered VQ-VAE across all truncation lengths.
Revised: In contrast, ordered VQ-VAE achieves superior reconstructions compared to PCA and unordered VQ-VAE across all truncation lengths.

Original: We repeat the experiment on standard VAE to disentangle the speciﬁc implementations in VQ-VAE.
Revised: We repeat the experiment on standard VAE to disentangle the speciﬁc implementations in VQ-VAE.

Original: We observe that the order-inducing objective has consistent results on standard VAE (Appendix C.1).
Revised: We observe that the order-inducing objective has consistent results on standard VAE (Appendix C.1).

Original: 5. I MAGE G ENERATION
Revised: 5. I MAGE G ENERATION

Original: We test the performance of anytime sampling using OVQ-VAEs ( Anytime + ordered ) on several image datasets.
Revised: We test the performance of anytime sampling using OVQ-VAEs ( Anytime + ordered ) on several image datasets.

Original: We compare our approach to two baselines.
Revised: We compare our approach to two baselines.

Original: One is the original VQ-VAE model proposed by van den Oord et al.
Revised: One is the original VQ-VAE model proposed by van den Oord et al.

Original: (2017) without anytime sampling.
Revised: (2017) without anytime sampling.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: We evaluate the model performance on the MNIST, CIFAR-10 (Krizhevsky, 2009) and CelebA
Revised: We evaluate the model performance on the MNIST, CIFAR-10 (Krizhevsky, 2009) and CelebA

Original: (Liu et al., 2014) datasets.
Revised: (Liu et al., 2014) datasets.

Original: For CelebA, the images are resized to 64 × 64 .
Revised: For CelebA, the images are resized to 64 × 64 .

Original: All pixel values are scaled to the range [0 , 1] .
Revised: All pixel values are scaled to the range [0 , 1] .

Original: We borrow the model architectures and optimizers from van den Oord et al.
Revised: We borrow the model architectures and optimizers from van den Oord et al.

Original: 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.
Revised: 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.

Original: We train a Transformer (Vaswani et al., 2017) on our
Revised: We train a Transformer (Vaswani et al., 2017) on our

Original: VQ-VAEs, as opposed to the PixelCNN model used in van den Oord et al.
Revised: VQ-VAEs, as opposed to the PixelCNN model used in van den Oord et al.

Original: PixelCNNs use standard 2D convolutional layers to capture a bounded receptive ﬁeld and model the conditional dependence.
Revised: PixelCNNs use standard 2D convolutional layers to capture a bounded receptive ﬁeld and model the conditional dependence.

Original: Transformers apply attention mechanism and feed forward network to model the conditional dependence of 1D sequence.
Revised: Transformers apply attention mechanism and feed forward network to model the conditional dependence of 1D sequence.

Original: 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).
Revised: 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).

Original: Our experiments on different autoregressive models in Appendix C.2 further support the arguments.[Figure] [Figure] [Figure]
Revised: Our experiments on different autoregressive models in Appendix C.2 further support the arguments.

Original: [Figure]

Original: [Figure]
Revised: [Figure]

Original: 5.2. I MAGE Q UALITY
Revised: 5.2. I MAGE Q UALITY

Revised: In Fig . 4(a)

Original: [Figure]
Revised: [Figure]

Original: In Fig . 4(a), we report FID (Heusel et al., 2017) scores (lower is better) on CIFAR-10 and CelebA when performing anytime sampling for ordered versus unordered VQ-VAE.
Revised: we report FID (Heusel et al., 2017) scores (lower is better) on CIFAR-10 and CelebA when performing anytime sampling for ordered versus unordered VQ-VAE.

Original: FID (Fréchet Inception
Revised: FID (Fréchet Inception

Original: Distance) score is the Fréchet distance between two multivariate Gaussians, whose means and covariances are estimated from the 2048-dimensional activations of the Inception-v3 (Szegedy et al., 2016) network for real and generated samples respectively.
Revised: Distance) score is the Fréchet distance between two multivariate Gaussians, whose means and covariances are estimated from the 2048-dimensional activations of the Inception-v3 (Szegedy et al., 2016) network for real and generated samples respectively.

Original: As a reference, we also report the FID scores when using the original VQ-VAE model (with residual blocks and spatial-wise quantization) sampled at the full code length (van den Oord et al., 2017).
Revised: As a reference, we also report the FID scores when using the original VQ-VAE model (with residual blocks and spatial-wise quantization) sampled at the full code length (van den Oord et al., 2017).

Original: Our main ﬁnding is that OVQ-VAE achieves a better FID score than unordered VQ-VAE at all fractional code lengths (ranging from 20% to 100%); in other words, OVQ-VAE achieves strictly superior anytime sampling performance compared to unordered VQ-VAE on both CIFAR-10 and CelebA. On CIFAR-10 dataset, a better FID score is achieved by OVQ-VAE even when sampling full codes.
Revised: Our main ﬁnding is that OVQ-VAE achieves a better FID score than unordered VQ-VAE at all fractional code lengths (ranging from 20% to 100%); in other words, OVQ-VAE achieves strictly superior anytime sampling performance compared to unordered VQ-VAE on both CIFAR-10 and CelebA. On CIFAR-10 dataset, a better FID score is achieved by OVQ-VAE even when sampling full codes.

Original: In Appendix C.4, we show that the regularization effect of the ordered codes causes this phenomenon.
Revised: In Appendix C.4, we show that the regularization effect of the ordered codes causes this phenomenon.

Original: In Fig .
Revised: In Fig .

Original: 5 (more in Appendix G.1), we visualize the sample quality degradation as a function of fractional code length when sampling from the OVQ-VAE.
Revised: 5 (more in Appendix G.1), we visualize the sample quality degradation as a function of fractional code length when sampling from the OVQ-VAE.

Original: We observe a consistent increase in sample quality as we increased the fractional code length.
Revised: We observe a consistent increase in sample quality as we increased the fractional code length.

Original: In particular, we observe the model to initially generate a global structure of an image and then gradually ﬁll in local details.
Revised: In particular, we observe the model to initially generate a global structure of an image and then gradually ﬁll in local details.

Original: We further show in Appendix F that, samples sharing the highest priority latent code have similar global structure.
Revised: We further show in Appendix F that, samples sharing the highest priority latent code have similar global structure.

Original:  Although our method was inspired by PCA, we encountered limited success when training an autoregressive model on the PCA-represented data.
Revised: [Figure] Although our method was inspired by PCA, we encountered limited success when training an autoregressive model on the PCA-represented data.

Original: Please refer to Appendix D for more details.
Revised: Please refer to Appendix D for more details.

Original: 5.2. I NFERENCE S PEED
Revised: 5.2. I NFERENCE S PEED

Original: We compare the inference speed of our approach vs. the original VQ-VAE model by the wall-clock time needed for sampling.
Revised: We compare the inference speed of our approach vs. the original VQ-VAE model by the wall-clock time needed for sampling.

Original: We also include the decoding time in our approach.
Revised: We also include the decoding time in our approach.

Original: We respectively measure the time of generating 50000 and 100000 images on CIFAR-10 and CelebA datasets, with a batch size of 100.
Revised: We respectively measure the time of generating 50000 and 100000 images on CIFAR-10 and CelebA datasets, with a batch size of 100.

Original: All samples are produced on a single NVIDIA TITAN
Revised: All samples are produced on a single NVIDIA TITAN

Original: Xp GPU.
Revised: Xp GPU.

Original: Fig . 4(b) shows that the time needed for anytime sampling increases almost linearly with respect to the sampled code length.
Revised: Fig . 4(b) shows that the time needed for anytime sampling increases almost linearly with respect to the sampled code length.

Original: This supports our argument in Section 3.2 that the decoding time is negligible compared to the autoregressive component.
Revised: This supports our argument in Section 3.2 that the decoding time is negligible compared to the autoregressive component.

Original: Indeed, the decoder took around 24 seconds to generate all samples for CelebA, whereas the sampling time of the autoregressive model was around 610 seconds—over an order of magnitude larger.
Revised: Indeed, the decoder took around 24 seconds to generate all samples for CelebA, whereas the sampling time of the autoregressive model was around 610 seconds—over an order of magnitude larger.

Original: Moreover, since we can achieve roughly the highest sample quality with only 60% of the full code length on CelebA, anytime sampling can save around 40% run-time compared to naïve sampling without hurting sample quality.
Revised: Moreover, since we can achieve roughly the highest sample quality with only 60% of the full code length on CelebA, anytime sampling can save around 40% run-time compared to naïve sampling without hurting sample quality.

Original: In addition, our method is faster than the original VQ-VAE even when sampling the full code length, without compromising sample quality ( cf . , Section 5.2.1).
Revised: In addition, our method is faster than the original VQ-VAE even when sampling the full code length, without compromising sample quality ( cf . , Section 5.2.1).

Original: This is because the Transformer model we used is sufﬁciently shallower than the Gated PixelCNN (van den Oord et al., 2016) model in the original VQ-VAE paper.
Revised: This is because the Transformer model we used is sufﬁciently shallower than the Gated PixelCNN (van den Oord et al., 2016) model in the original VQ-VAE paper.

Original: Compared to PixelCNN++ (Salimans et al., 2017), an autoregressive model on the raw pixel space, the sampling speed of our method can be an order of magnitude faster since our autoregressive models are trained on the latent space with much lower dimensionality.
Revised: Compared to PixelCNN++ (Salimans et al., 2017), an autoregressive model on the raw pixel space, the sampling speed of our method can be an order of magnitude faster since our autoregressive models are trained on the latent space with much lower dimensionality.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Prior work has tackled the issue of slow autoregressive generation by improving implementations of the generation algorithm.
Revised: Prior work has tackled the issue of slow autoregressive generation by improving implementations of the generation algorithm.

Original: For example, the sampling speed of convolutional autoregressive models can be improved substantially by caching hidden state computation (Ramachandran et al., 2017).
Revised: For example, the sampling speed of convolutional autoregressive models can be improved substantially by caching hidden state computation (Ramachandran et al., 2017).

Original: While such approaches provide substantial speedups in generation time, they are still at best linear in the dimension of the sample space.
Revised: While such approaches provide substantial speedups in generation time, they are still at best linear in the dimension of the sample space.

Original: van den Oord et al.
Revised: van den Oord et al.

Original: (2018) improves the inference speed by allowing parallel computing.
Revised: (2018) improves the inference speed by allowing parallel computing.

Original: Compared to our approach, they do not have the test-time adaptivity to computational constraints.
Revised: Compared to our approach, they do not have the test-time adaptivity to computational constraints.

Original: In contrast, we design methods that allow trade-offs between generation speed and sample quality on-the-ﬂy based on computational constraints.
Revised: In contrast, we design methods that allow trade-offs between generation speed and sample quality on-the-ﬂy based on computational constraints.

Original: For example, running apps can accommodate to the real-time computational resources without model re-training.
Revised: For example, running apps can accommodate to the real-time computational resources without model re-training.

Original: In addition, they can be combined together with our method without sacriﬁcing sample quality.
Revised: In addition, they can be combined together with our method without sacriﬁcing sample quality.

Original: Speciﬁcally,
Revised: Speciﬁcally,

Original: Ramachandran et al.
Revised: Ramachandran et al.

Original: (2017) leverages caches to speed up autoregressive sampling, which can be directly applied to our autoregressive model on ordered codes without affecting sample quality.
Revised: (2017) leverages caches to speed up autoregressive sampling, which can be directly applied to our autoregressive model on ordered codes without affecting sample quality.

Original: van den Oord et al.
Revised: van den Oord et al.

Original: (2018) proposes probability density distillation to distill autoregressive models into fast implicit generators.
Revised: (2018) proposes probability density distillation to distill autoregressive models into fast implicit generators.

Original: We can apply the same technique on our latent autoregressive model to allow a similar speedup.
Revised: We can apply the same technique on our latent autoregressive model to allow a similar speedup.

Original: In order to enable anytime sampling, our method requires learning an ordered latent representation of data by training ordered autoencoders.
Revised: In order to enable anytime sampling, our method requires learning an ordered latent representation of data by training ordered autoencoders.

Original: Rippel et al.
Revised: Rippel et al.

Original: (2014) proposes a generalization of Principal
Revised: (2014) proposes a generalization of Principal

Original: Components Analysis to learn an ordered representation.
Revised: Components Analysis to learn an ordered representation.

Original: Instead of the uniform distribution over discrete codes in our method, they adopted a geometric distribution over continuous codes during the training of the ordered autoencoders.
Revised: Instead of the uniform distribution over discrete codes in our method, they adopted a geometric distribution over continuous codes during the training of the ordered autoencoders.

Original: Because of the difference they require additional tricks such as unit sweeping and adaptive regularization coefﬁcients to stabilize the training, while our method is more stable and scalable.
Revised: Because of the difference they require additional tricks such as unit sweeping and adaptive regularization coefﬁcients to stabilize the training, while our method is more stable and scalable.

Original: In addition, they only focus on fast retrieval and image compression.
Revised: In addition, they only focus on fast retrieval and image compression.

Original: By contrast, we further extend our approach to autoencoders with discrete latent codes ( e.g . , VQ-VAEs) and explore their applications in anytime sampling for autoregressive models.
Revised: By contrast, we further extend our approach to autoencoders with discrete latent codes ( e.g . , VQ-VAEs) and explore their applications in anytime sampling for autoregressive models.

Original: Another work related to our approach is hierarchical nonlinear PCA (Scholz & Vigário, 2002).
Revised: Another work related to our approach is hierarchical nonlinear PCA (Scholz & Vigário, 2002).

Original: We generalize their approach to latent spaces of arbitrary dimensionality, and leverage Monte Carlo estimations to improve the efﬁciency when learning very high dimensional latent representations.
Revised: We generalize their approach to latent spaces of arbitrary dimensionality, and leverage Monte Carlo estimations to improve the efﬁciency when learning very high dimensional latent representations.

Original: The denoising generative models proposed by Song & Ermon (2019); Ho et al.
Revised: The denoising generative models proposed by Song & Ermon (2019); Ho et al.

Original: (2020) progressively denoise images into better quality, instead of modeling images from coarse to ﬁne like our methods.
Revised: (2020) progressively denoise images into better quality, instead of modeling images from coarse to ﬁne like our methods.

Original: This means that interrupting the sampling procedure of diffusion models at an early time might lead to very noisy samples, but in our case it will lead to images with corrector coarse structures and no noise which is arguably more desirable.
Revised: This means that interrupting the sampling procedure of diffusion models at an early time might lead to very noisy samples, but in our case it will lead to images with corrector coarse structures and no noise which is arguably more desirable.

Original: A line of works draw connections between ordered latent codes and the linear autoencoders.
Revised: A line of works draw connections between ordered latent codes and the linear autoencoders.

Original: Kunin et al.
Revised: Kunin et al.

Original: (2019) proves that the principal directions can be deduced from the critical points of Lregularized linear autoencoders, and Bao et al.
Revised: (2019) proves that the principal directions can be deduced from the critical points of Lregularized linear autoencoders, and Bao et al.

Original: (2020) further shows that linear autoencoders can directly learn the ordered, axis-aligned principal components with non-uniform L 2 regularization.
Revised: (2020) further shows that linear autoencoders can directly learn the ordered, axis-aligned principal components with non-uniform L 2 regularization.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Sampling from autoregressive models is an expensive sequential process that can be intractable when on a tight computing budget.
Revised: Sampling from autoregressive models is an expensive sequential process that can be intractable when on a tight computing budget.

Original: To address this difﬁculty, we consider the novel task of adaptive autoregressive sampling that can naturally trade-off computation with sample quality.
Revised: To address this difﬁculty, we consider the novel task of adaptive autoregressive sampling that can naturally trade-off computation with sample quality.

Original: Inspired by
Revised: Inspired by

Original: PCA, we adopt ordered autoencoders, whose latent codes are prioritized based on their importance to reconstruction.
Revised: PCA, we adopt ordered autoencoders, whose latent codes are prioritized based on their importance to reconstruction.

Original: We show that it is possible to do anytime sampling for autoregressive models trained on these ordered latent codes—we may stop the sequential sampling process at any step and still obtain a complete sample of reasonable quality by decoding the partial codes.
Revised: We show that it is possible to do anytime sampling for autoregressive models trained on these ordered latent codes—we may stop the sequential sampling process at any step and still obtain a complete sample of reasonable quality by decoding the partial codes.

Original: With both theoretical arguments and empirical evidence, we show that ordered autoencoders can induce a valid ordering that facilitates anytime sampling.
Revised: With both theoretical arguments and empirical evidence, we show that ordered autoencoders can induce a valid ordering that facilitates anytime sampling.

Original: Experimentally, we test our approach on several image and audio datasets by pairing an ordered VQ-VAE (a powerful autoencoder architecture) and a Transformer (an expressive autoregressive model) on the latent space.
Revised: Experimentally, we test our approach on several image and audio datasets by pairing an ordered VQ-VAE (a powerful autoencoder architecture) and a Transformer (an expressive autoregressive model) on the latent space.

Original: We demonstrate that our samples suffer almost no loss of quality (as measured by FID scores) for images when using only 60% to 80% of all code dimensions, and the sample quality degrades gracefully as we gradually reduce the code length. 
Revised: We demonstrate that our samples suffer almost no loss of quality (as measured by FID scores) for images when using only 60% to 80% of all code dimensions, and the sample quality degrades gracefully as we gradually reduce the code length. A CKNOWLEDGEMENTS

Revised: We are grateful to Shengjia Zhao, Tommi Jaakkola and anonymous reviewers in ICLR for helpful discussion.

Revised: We would like to thank Tianying Wen for the lovely ﬁgure.

Revised: YX was supported by the

Revised: HDTV Grand Alliance Fellowship.

Original:  A PROOFS A. PROOF FOR THEOREM
Revised: SE was supported by NSF (#1651565, #1522054, #1733686), ONR (N00014- 19-1-2145), AFOSR (FA9550-19-1-0024), Amazon AWS, and FLI. YS was partially supported by the Apple PhD Fellowship in AI/ML. A PROOFS A. PROOF FOR THEOREM

Original: Proof.
Revised: Proof.

Original: For simplicity we denote the distribution of the stochastic part (ﬁrst i dimensions) of q e θ ( z | x ) ≤ i as q e θ ( z ≤ i | x ) , and similarly we denote p d φ ( x | z ) ≤
Revised: For simplicity we denote the distribution of the stochastic part (ﬁrst i dimensions) of q e θ ( z | x ) ≤ i as q e θ ( z ≤ i | x ) , and similarly we denote p d φ ( x | z ) ≤

Original: i as p d φ ( x | z ≤ i ) .
Revised: i as p d φ ( x | z ≤ i ) .

Original: We ﬁrst reformu-
Revised: We ﬁrst reformu-

Original: late the objective
Revised: late the objective

Original: [Equation]
Revised: [Equation]

Original: The inequality (8) holds because KL-divergences are non-negative.
Revised: The inequality (8) holds because KL-divergences are non-negative.

Original: Under the assumption of an optimal decoder φ , we can achieve the equality in (8)
Revised: Under the assumption of an optimal decoder φ , we can achieve the equality in (8)

Original: [Equation]
Revised: [Equation]

Original: in which case the objective equals
Revised: in which case the objective equals

Original: [Equation]
Revised: [Equation]

Original: We can deﬁne the following modiﬁed objective by adding the data entropy term, since it is a constant independent of θ .
Revised: We can deﬁne the following modiﬁed objective by adding the data entropy term, since it is a constant independent of θ .

Original: [Equation]
Revised: [Equation]

Original: If there exists an integer i ∈ { 2 , · · · , K } such that I ( z i ; x | z ≤ i − 1 ) >
Revised: If there exists an integer i ∈ { 2 , · · · , K } such that I ( z i ; x | z ≤ i − 1 ) >

Original: I ( z i − 1 ; x | z ≤ i − 2 ) , we can exchange the position of z i and z i − 1 to increase the value of objective (9).
Revised: I ( z i − 1 ; x | z ≤ i − 2 ) , we can exchange the position of z i and z i − 1 to increase the value of objective (9).

Original: By the chain rule of mutual information we have:
Revised: By the chain rule of mutual information we have:

Original: [Equation]
Revised: [Equation]

Original: When I ( z i ; x | z ≤ i − 1 ) >
Revised: When I ( z i ; x | z ≤ i − 1 ) >

Original: I ( z i − 1 ; x | z ≤ i − 2 ) , we can show that I ( x ; z i , z ≤ i − 2 ) >
Revised: I ( z i − 1 ; x | z ≤ i − 2 ) , we can show that I ( x ; z i , z ≤ i − 2 ) >

Original: I ( x ; z i − 1 , z ≤ i − 2 ) :
Revised: I ( x ; z i − 1 , z ≤ i − 2 ) :

Original: [Equation]
Revised: [Equation]

Original: ( x ; z i , z ≤ i −
Revised: ( x ; z i , z ≤ i −

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: ( x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: z i−
Revised: z i

Original: [Equation]
Revised: [Equation]

Original: ≤

Original: [Equation]

Original: [Equation]

Original: where Eq .
Revised: where Eq .

Original: (10) holds by the chain rule of mutual information:
Revised: (10) holds by the chain rule of mutual information:

Original: [Equation]
Revised: [Equation]

Original: 2 ) , and I ( z i ; z i − 1 | x , z ≤ i − 2 ) = 0 by the conditional in- dependence assumption zi ⊥
Revised: and I ( z i ; z i − 1 | x , z ≤ i − 2 ) = 0 by the conditional in- dependence assumption z

Original: z
Revised: i ⊥

Original:  i − 1 | x , z ≤ i − 2 .
Revised: z i − 1 | x , z ≤ i − 2 .

Original: Hence if we exchange the position of z i and z i − 1 of the latent vector z , we can show that the new objective value after the position switch is strictly smaller than the original one:
Revised: Hence if we exchange the position of z i and z i − 1 of the latent vector z , we can show that the new objective value after the position switch is strictly smaller than the original one:

Original: [Equation]
Revised: [Equation]

Original: Inequality (12) holds by I ( x ; z i , z ≤ i − 2 )
Revised: Inequality (12) holds by I ( x ; z i , z ≤ i − 2 )

Original: > I ( x ; z i [Equation] 1 , z ≤ i − 2 ) .
Revised: > I ( x ; z i − 1 , z ≤ i − 2 ) .

Original: From above we can conclude that if the encoder is optimal, then the following inequalities must hold:
Revised: From above we can conclude that if the encoder is optimal, then the following inequalities must hold:

Original: [Equation]
Revised: [Equation]

Original: Otherwise we can exchange the dimensions to make the objective Eq .
Revised: Otherwise we can exchange the dimensions to make the objective Eq .

Original: (9) smaller, which contradicts the optimally of encoder e θ .
Revised: (9) smaller, which contradicts the optimally of encoder e θ .

Original: B AUDIO GENERATION
Revised: B AUDIO GENERATION

Original: [Figure]
Revised: [Figure]

Original: Our method can also be applied to audio data.
Revised: Our method can also be applied to audio data.

Original: We evaluate anytime autoregressive models on the
Revised: We evaluate anytime autoregressive models on the

Original: VCTK dataset (Veaux et al., 2017), which consists of speech recordings from 109 different speakers.
Revised: VCTK dataset (Veaux et al., 2017), which consists of speech recordings from 109 different speakers.

Original: The original VQ-VAE uses an autoregressive decoder, which may cost more time than the latent autoregressive model and thus cannot be accelerated by anytime sampling.
Revised: The original VQ-VAE uses an autoregressive decoder, which may cost more time than the latent autoregressive model and thus cannot be accelerated by anytime sampling.

Original: Instead, we adopt a non-autoregressive decoder inspired by the generator of WaveGan (Donahue et al., 2018).
Revised: Instead, we adopt a non-autoregressive decoder inspired by the generator of WaveGan (Donahue et al., 2018).

Original: Same as images, we train a Transformer model on the latent codes.
Revised: Same as images, we train a Transformer model on the latent codes.

Original: We compare waveforms sampled from ordered vs. unordered VQ-VAEs in Fig . 6, and provide links to audio samples in Appendix G.2.
Revised: We compare waveforms sampled from ordered vs. unordered VQ-VAEs in Fig . 6, and provide links to audio samples in Appendix G.2.

Original: By inspecting the waveforms and audio samples, we observe that the generated waveform captures the correct global structure using as few as 6.25% of the full code length, and gradually reﬁnes itself as more code dimensions are sampled.
Revised: By inspecting the waveforms and audio samples, we observe that the generated waveform captures the correct global structure using as few as 6.25% of the full code length, and gradually reﬁnes itself as more code dimensions are sampled.

Original: In contrast, audio samples from the unordered VQ-VAE contain considerably more noise when using truncated codes.
Revised: In contrast, audio samples from the unordered VQ-VAE contain considerably more noise when using truncated codes.

Original: C ADDITIONAL EXPERIMENTAL RESULTS C. ORDERED VERSUS UNORDERED CODES ON VAE
Revised: C ADDITIONAL EXPERIMENTAL RESULTS C. ORDERED VERSUS UNORDERED CODES ON VAE

Original: To better understand the effect of order-inducing objective, we disentangle the order-inducing objective Eq .
Revised: To better understand the effect of order-inducing objective, we disentangle the order-inducing objective Eq .

Original: (6) with the speciﬁc implementation on VQ-VAE, such as stop_gradient operator and quantization.
Revised: (6) with the speciﬁc implementation on VQ-VAE, such as stop_gradient operator and quantization.

Original: We adopt a similar order-inducing objective on standard VAE:
Revised: We adopt a similar order-inducing objective on standard VAE:

Original: [Equation]
Revised: [Equation]

Original: where q e θ ( z | x ) ≤ i denotes the distribution of ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , z ∼ q e θ ( z | x ) , and p d φ ( x | z ) ≤ i represents the distribution of p d φ ( x | ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K ) .
Revised: where q e θ ( z | x ) ≤ i denotes the distribution of ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K , z ∼ q e θ ( z | x ) , and p d φ ( x | z ) ≤ i represents the distribution of p d φ ( x | ( z 1 , z 2 , · · · , z i , 0 , · · · , 0) T ∈ R K ) .

Original: We set the prior p ( z ) ≤ i as the i dimensional unit normal distribution.
Revised: We set the prior p ( z ) ≤ i as the i dimensional unit normal distribution.

Original: We repeat the experiments on CIFAR-10 and observe similar experimental results between VQ-VAE and standard VAE.
Revised: We repeat the experiments on CIFAR-10 and observe similar experimental results between VQ-VAE and standard VAE.

Original: Fig .
Revised: Fig .

Original: 7(a) shows that the standard VAE has irregular ∆( i ) , while the ordered VAE has much better ordering on latent codes.
Revised: 7(a) shows that the standard VAE has irregular ∆( i ) , while the ordered VAE has much better ordering on latent codes.

Original: Fig . 7(b) further shows that the ordered VAE has better reconstructions under different truncated code lengths.
Revised: Fig . 7(b) further shows that the ordered VAE has better reconstructions under different truncated code lengths.

Original: The experimental results are consistent with
Revised: The experimental results are consistent with

Original: Fig . 3(a) and Fig . 3(b)..
Revised: Fig . 3(a) and Fig . 3(b)..

Original: [Figure]
Revised: [Figure]

Original: C. ABLATION STUDY ON AUTOREGRESSIVE MODEL
Revised: C. ABLATION STUDY ON AUTOREGRESSIVE MODEL

Original: We study the performance of different autoregressive model on ordered codes.
Revised: We study the performance of different autoregressive model on ordered codes.

Original: We compare Transformer (Vaswani et al., 2017), PixelCNN (Oord et al., 2016b) and LSTM (Hochreiter & Schmidhuber, 1997).
Revised: We compare Transformer (Vaswani et al., 2017), PixelCNN (Oord et al., 2016b) and LSTM (Hochreiter & Schmidhuber, 1997).

Original: PixelCNN adopts standard 2D convolutional layers to model the conditional dependence.
Revised: PixelCNN adopts standard 2D convolutional layers to model the conditional dependence.

Original: In order to enlarge the receptive ﬁeld, PixelCNN stacks many convolutional layers.
Revised: In order to enlarge the receptive ﬁeld, PixelCNN stacks many convolutional layers.

Original: Transformer applies attention mechanism and feed forward network to model the conditional dependence of 1D sequences.
Revised: Transformer applies attention mechanism and feed forward network to model the conditional dependence of 1D sequences.

Original: LSTM uses four different types of gates to improve the long-term modeling power of the recurrent models on 1D sequences.
Revised: LSTM uses four different types of gates to improve the long-term modeling power of the recurrent models on 1D sequences.

Original: Transformer and LSTM can be naturally applied to the 1D channel-wise quantized codes.
Revised: Transformer and LSTM can be naturally applied to the 1D channel-wise quantized codes.

Original: Since
Revised: Since

Original: PixelCNN operates on 2D data, we reshape the 1D codes into 2D tensors.
Revised: PixelCNN operates on 2D data, we reshape the 1D codes into 2D tensors.

Original: More speciﬁcally, for 70-dimensional 1D codes on CIFAR-10, we ﬁrstly pad the codes into 81 dimensions then reshape it into 9 × 9 tensors.
Revised: More speciﬁcally, for 70-dimensional 1D codes on CIFAR-10, we ﬁrstly pad the codes into 81 dimensions then reshape it into 9 × 9 tensors.

Original: Fig . 8(a) shows the FID scores of different autoregressive model on CIFARdataset.
Revised: Fig . 8(a) shows the FID scores of different autoregressive model on CIFARdataset.

Original: The results show that PixelCNN has inferior performance in all cases except when used with 0.2 fractions of full code length.
Revised: The results show that PixelCNN has inferior performance in all cases except when used with 0.2 fractions of full code length.

Original: This is because PixelCNN works well only when the input has strong local spatial correlations, but there is no spatial correlation for channel-wise quantized codes.
Revised: This is because PixelCNN works well only when the input has strong local spatial correlations, but there is no spatial correlation for channel-wise quantized codes.

Original: In contrast, autoregressive models tailored for 1D sequences work better on channel-wise dequantized codes, as they have uniformly better FID scores when using 0.4/0.6/0.8/1.0 fractions of full code length.
Revised: In contrast, autoregressive models tailored for 1D sequences work better on channel-wise dequantized codes, as they have uniformly better FID scores when using 0.4/0.6/0.8/1.0 fractions of full code length.

Original: C. ABLATION STUDY ON SAMPLING DISTRIBUTION
Revised: C. ABLATION STUDY ON SAMPLING DISTRIBUTION

Original: We study the effect of different sampling distributions in order-inducing objective Eq .
Revised: We study the effect of different sampling distributions in order-inducing objective Eq .

Original: We compare the adopted uniform distribution with the geometric distribution used in Rippel et al.
Revised: We compare the adopted uniform distribution with the geometric distribution used in Rippel et al.

Original: (2014) on CIFAR-10, as shown in Fig .
Revised: (2014) on CIFAR-10, as shown in Fig .

Original: 8(b).
Revised: 8(b).

Original: We normalize the geometric distribution on ﬁnite indices with length K , i.e. Pr( i = k )
Revised: We normalize the geometric distribution on ﬁnite indices with length K , i.e. Pr( i = k )

Original: =
Revised: =

Original: (1 − p )
Revised: (1 − p )

Original: k − 1 p 1 − (1 − p ) K , i ∈ { 1 , 2 , . . .
Revised: k − 1 p 1 − (1 − p ) K , i ∈ { 1 , 2 , . . .

Original: , K } , and denote the normalized geometric distribution as Geo ( p ) .
Revised: , K } , and denote the normalized geometric distribution as Geo ( p ) .

Original: Note that when p → 0 , Geo ( p ) recovers the uniform distribution.
Revised: Note that when p → 0 , Geo ( p ) recovers the uniform distribution.

Original: Fig . 8(c) shows that OVQ-VAEs trained with all the different distributions can trade off the quality and computational budget.
Revised: Fig . 8(c) shows that OVQ-VAEs trained with all the different distributions can trade off the quality and computational budget.

Original: We ﬁnd that OVQ-VAE is sensitive to the parameter of the geometric distribution.
Revised: We ﬁnd that OVQ-VAE is sensitive to the parameter of the geometric distribution.

Original: The performance of OVQ-VAE with Geo(0.03) is marginally worse than the uniform distribution.
Revised: The performance of OVQ-VAE with Geo(0.03) is marginally worse than the uniform distribution.

Original: But when changing the distribution to Geo(0.1), the FID scores become much worse with large code length (0.6/0.8/1.0 fractions of the full code length).
Revised: But when changing the distribution to Geo(0.1), the FID scores become much worse with large code length (0.6/0.8/1.0 fractions of the full code length).

Original: Since for i ∼ Geo ( p ) , Pr( i ≥ t ) = (1 − p ) t − 1 (1 − (1 − p )
Revised: Since for i ∼ Geo ( p ) , Pr( i ≥ t ) = (1 − p ) t − 1 (1 − (1 − p )

Original: K − t +1 ) / (1 − (1 − p ) K ) ≤ (1 − p ) t − 1 , which indicates that the geometric distribution allocates exponentially smaller probability to code with higher index.
Revised: K − t +1 ) / (1 − (1 − p ) K ) ≤ (1 − p ) t − 1 , which indicates that the geometric distribution allocates exponentially smaller probability to code with higher index.

Original: [Figure]
Revised: [Figure]

Original: C. ON THE REGULARIZATION EFFECT OF ORDERED CODES
Revised: C. ON THE REGULARIZATION EFFECT OF ORDERED CODES

Original: Imposing an order on latent codes improves the inductive bias for the autoregressive model to learn the codes.
Revised: Imposing an order on latent codes improves the inductive bias for the autoregressive model to learn the codes.

Original: When using full codes on CIFAR-10 dataset, even though the OVQ-VAE has higher training error than unordered VQ-VAE, a better FID score is achieved by the ordered model.
Revised: When using full codes on CIFAR-10 dataset, even though the OVQ-VAE has higher training error than unordered VQ-VAE, a better FID score is achieved by the ordered model.

Original: This validates the intuition that it is easier to model an image from coarse to ﬁne.
Revised: This validates the intuition that it is easier to model an image from coarse to ﬁne.

Original: These results are corroborated by lower FID scores for the anytime model with full length codes under different number of training samples.
Revised: These results are corroborated by lower FID scores for the anytime model with full length codes under different number of training samples.

Original: As shown in Fig .
Revised: As shown in Fig .

Original: 9(a), the ordered model has an increasingly larger FID improvement over the unordered model when the dataset becomes increasingly smaller.
Revised: 9(a), the ordered model has an increasingly larger FID improvement over the unordered model when the dataset becomes increasingly smaller.

Original: These results indicate that training on ordered codes has a regularization effect.
Revised: These results indicate that training on ordered codes has a regularization effect.

Original: We hypothesize that ordered codes capture the inductive bias of coarse-to-ﬁne image modeling better.
Revised: We hypothesize that ordered codes capture the inductive bias of coarse-to-ﬁne image modeling better.

Original: C. COMPARISON TO TALORED VQ-VAE
Revised: C. COMPARISON TO TALORED VQ-VAE

Original: We compare the anytime sampling to the unordered VQ-VAE with tailored latent space (Tailored
Revised: We compare the anytime sampling to the unordered VQ-VAE with tailored latent space (Tailored

Original: VQ-VAE) on CIFAR-10.
Revised: VQ-VAE) on CIFAR-10.

Original: The Tailored VQ-VAE has a pre-speciﬁed latent size, using the same computational budget as truncated codes of anytime sampling.
Revised: The Tailored VQ-VAE has a pre-speciﬁed latent size, using the same computational budget as truncated codes of anytime sampling.

Original: For a fair comparison, we experiment with transformers on the latent space of Tailored VQ-VAE.
Revised: For a fair comparison, we experiment with transformers on the latent space of Tailored VQ-VAE.

Original: Fig . 9(b) shows that anytime sampling always has better FID scores than Tailored VQ-VAE, except when the code is very short.
Revised: Fig . 9(b) shows that anytime sampling always has better FID scores than Tailored VQ-VAE, except when the code is very short.

Original: We hypothesize that the learning signals from training on the full length codes with OVQ-VAE improves the quality when codes are shorter, thus demonstrating a FID improvement over the Tailored VQ-VAE.
Revised: We hypothesize that the learning signals from training on the full length codes with OVQ-VAE improves the quality when codes are shorter, thus demonstrating a FID improvement over the Tailored VQ-VAE.

Original: Moreover, OVQ-VAE has the additional beneﬁt of allowing anytime sampling when the computational budget is not knownin advance.
Revised: Moreover, OVQ-VAE has the additional beneﬁt of allowing anytime sampling when the computational budget is not known advance.

Original: D TRAIN AUTOREGRESSIVE MODEL ON PCA REPRESENTATION
Revised: D TRAIN AUTOREGRESSIVE MODEL ON PCA REPRESENTATION

Original: An alternative way to induce order on the latent space is by projecting data onto the PCA representation.
Revised: An alternative way to induce order on the latent space is by projecting data onto the PCA representation.

Original: However, we encounter limited success in training the autoregressive model on the top of PCA representation.
Revised: However, we encounter limited success in training the autoregressive model on the top of PCA representation.

Original: When training an autoregressive model on PCA-represented data, we observe inferior log-likelihoods.
Revised: When training an autoregressive model on PCA-represented data, we observe inferior log-likelihoods.

Original: We ﬁrst prepare the data by uniformly dequantizing the discrete pixel values to continuous ones.
Revised: We ﬁrst prepare the data by uniformly dequantizing the discrete pixel values to continuous ones.

Original: Then we project these continuous data into the PCA space by using the orthogonal projection matrix composed of singular vectors of the data covariance matrix.
Revised: Then we project these continuous data into the PCA space by using the orthogonal projection matrix composed of singular vectors of the data covariance matrix.

Original: Note that this projection preserves the volume of the original pixel space since the projection matrix’s determinant is 1, so log-likelihoods of a model trained on the raw continuous data space and the PCA projected space are comparable.
Revised: Note that this projection preserves the volume of the original pixel space since the projection matrix’s determinant is 1, so log-likelihoods of a model trained on the raw continuous data space and the PCA projected space are comparable.

Original: We report the bits/dim (lower is better), which is computed by dividing the negative log-likelihood (log base 2) by the dimension of data.
Revised: We report the bits/dim (lower is better), which is computed by dividing the negative log-likelihood (log base 2) by the dimension of data.

Original: We train transformer models on the projected space and the raw data space.
Revised: We train transformer models on the projected space and the raw data space.

Original: Surprisingly, on MNIST, the transformer model obtains 1.22 bits/dim on the projected space versus 0.80 bits/dim on the raw data, along with inferior sample quality.
Revised: Surprisingly, on MNIST, the transformer model obtains 1.22 bits/dim on the projected space versus 0.80 bits/dim on the raw data, along with inferior sample quality.

Original: We hypothesize two reasons.
Revised: We hypothesize two reasons.

Original: First, models have been tuned with respect to inputs that are unlike the PCA representation, but rather on inputs such as raw pixel data.
Revised: First, models have been tuned with respect to inputs that are unlike the PCA representation, but rather on inputs such as raw pixel data.

Original: Second, PCA does not capture multiple data modalities well, unlike
Revised: Second, PCA does not capture multiple data modalities well, unlike

Original: OVQ-VAE.
Revised: OVQ-VAE.

Original: Moreover, autoregressive models typically do not perform well on continuous data.
Revised: Moreover, autoregressive models typically do not perform well on continuous data.

Original: In contrast, our Transformer model operates on discrete latent codes of the VQ-VAE.
Revised: In contrast, our Transformer model operates on discrete latent codes of the VQ-VAE.

Original: [Figure]
Revised: [Figure]

Revised: scores with different computational budgets.

Original: E EXTRA IMPLEMENTATION DETAILS
Revised: E EXTRA IMPLEMENTATION DETAILS

Original: Our code is released via the anonymous link https://anonymous.4open.science/r/3946e9c8-8f98-4836abc1-0f711244476d/ and included in the supplementary material as well.
Revised: Our code is released via the anonymous link https://anonymous.4open.science/r/3946e9c8-8f98-4836abc1-0f711244476d/ and included in the supplementary material as well.

Original: Below we introduce more details on network architectures and the training processes.
Revised: Below we introduce more details on network architectures and the training processes.

Original: E. NETWORKS
Revised: E. NETWORKS

Original: For MNIST dataset, the encoder has 3 convolutional layers with ﬁlter size (4,4,3) and stride (2,2,1) respectively.
Revised: For MNIST dataset, the encoder has 3 convolutional layers with ﬁlter size (4,4,3) and stride (2,2,1) respectively.

Original: For CelebA and CIFAR-10 datasets, the encoder has 4 convolutional layers with ﬁlter size (4,4,4,3) and stride (2,2,2,1) respectively.
Revised: For CelebA and CIFAR-10 datasets, the encoder has 4 convolutional layers with ﬁlter size (4,4,4,3) and stride (2,2,2,1) respectively.

Original: The decoders for dataset above are the counterpart of the corresponding encoders.
Revised: The decoders for dataset above are the counterpart of the corresponding encoders.

Original: For VCTK dataset, we use a encoder that has 5 convolutional layers with a ﬁlter size of 25 and stride of 4.
Revised: For VCTK dataset, we use a encoder that has 5 convolutional layers with a ﬁlter size of 25 and stride of 4.

Original: The activation functions are chosen to be LeakyRelu-0.2.
Revised: The activation functions are chosen to be LeakyRelu-0.2.

Original: We adopt a decoder architecture which has 4 convolutional and upsampling layers.
Revised: We adopt a decoder architecture which has 4 convolutional and upsampling layers.

Original: The architecture is the same with the generator architecture in Donahue et al.
Revised: The architecture is the same with the generator architecture in Donahue et al.

Original: (2018), except for the number of layers.
Revised: (2018), except for the number of layers.

Original: For all the datasets, we use a 6-layer Transformer decoder with an embedding size of 512, latent size of 2048, and dropout rate of 0.1.
Revised: For all the datasets, we use a 6-layer Transformer decoder with an embedding size of 512, latent size of 2048, and dropout rate of 0.1.

Original: We use 8 heads in multi-head self-attention layers.
Revised: We use 8 heads in multi-head self-attention layers.

Original: E. IMAGE GENERATION
Revised: E. IMAGE GENERATION

Original: The FID scores are computed using the ofﬁcial code from TTUR (Heusel et al., 2017) 1 authors.
Revised: The FID scores are computed using the ofﬁcial code from TTUR (Heusel et al., 2017) 1 authors.

Original: We compute FID scores on CIFAR-10 and CelebA based on a total of 50000 samples andsamples respectively.
Revised: We compute FID scores on CIFAR-10 and CelebA based on a total of 50000 samples andsamples respectively.

Original: We pre-train the VQ-VAE models with full code lengths for 200 epochs.
Revised: We pre-train the VQ-VAE models with full code lengths for 200 epochs.

Original: Then we train the VQ-VAE models with the new objective Eq .
Revised: Then we train the VQ-VAE models with the new objective Eq .

Original: (6) for 200 more epochs.
Revised: (6) for 200 more epochs.

Original: We use the Adam optimizer with learning rate 1 . 

Revised: We use the Adam optimizer with learning rate 1 .

Original: × 10 − 3 for training.
Revised: × 10 − 3 for training.

Original: We train the autoregressive model for 50 epochs on both MNIST and
Revised: We train the autoregressive model for 50 epochs on both MNIST and

Original: CIFAR-10, and 100 epochs on CelebA.
Revised: CIFAR-10, and 100 epochs on CelebA.

Original: We use the Adam optimizer with a learning rate of 2 . 

Revised: We use the Adam optimizer with a learning rate of 2 .

Original: × 10 −for the Transformer decoder.
Revised: × 10 −for the Transformer decoder.

Original: We select the checkpoint with the smallest validation loss on every epoch.
Revised: We select the checkpoint with the smallest validation loss on every epoch.

Original: The batch size is ﬁxed to be 128 during all training processes.
Revised: The batch size is ﬁxed to be 128 during all training processes.

Original: E. AUDIO GENERATION
Revised: E. AUDIO GENERATION

Original: We randomly subsample all the data points in VCTK dataset to make all audios have the same length (15360).
Revised: We randomly subsample all the data points in VCTK dataset to make all audios have the same length (15360).

Original: The VQ-VAE models are pre-trained with full code length for 20 epochs, and then ﬁne-tuned with our objective Eq .
Revised: The VQ-VAE models are pre-trained with full code length for 20 epochs, and then ﬁne-tuned with our objective Eq .

Original: (6) for 20 more epochs.
Revised: (6) for 20 more epochs.

Original: We use the Adam optimizer with learning rate 2 . 

Revised: We use the Adam optimizer with learning rate 2 .

Original: × 10 − 4 for training the VQ-VAE model.
Revised: × 10 − 4 for training the VQ-VAE model.

Original: We train the Transformer for 50 epochs on VCTK, use the Adam optimizer with a learning rate of 2 . 

Revised: We train the Transformer for 50 epochs on VCTK, use the Adam optimizer with a learning rate of 2 .

Original: × 10 − 3 .
Revised: × 10 − 3 .

Original: We select the checkpoint with the smallest validation loss on every epoch.
Revised: We select the checkpoint with the smallest validation loss on every epoch.

Original: The batch size is ﬁxed to be 8 for the VQ-VAE model and 128 for the
Revised: The batch size is ﬁxed to be 8 for the VQ-VAE model and 128 for the

Original: Transformer during training.
Revised: Transformer during training.

Original: F SAMPLES ON THE SAME PRIORITY CODE
Revised: F SAMPLES ON THE SAME PRIORITY CODE

Original: As further illustration of the ordered encoding, we show in Fig .
Revised: As further illustration of the ordered encoding, we show in Fig .

Original: 10 the result of full code length sampling when the ﬁrst (highest priority) discrete latent code is ﬁxed.
Revised: 10 the result of full code length sampling when the ﬁrst (highest priority) discrete latent code is ﬁxed.

Original: The ﬁxing of the ﬁrst latent code causes anytime sampling to produce a wide variety of samples that share high-level global similarities.
Revised: The ﬁxing of the ﬁrst latent code causes anytime sampling to produce a wide variety of samples that share high-level global similarities.

Original: G
Revised: G

Original: EXTRA SAMPLES G. IMAGE SAMPLES
Revised: EXTRA SAMPLES G. IMAGE SAMPLES

Original: We show extended samples from ordered VQ-VAEs in Fig .
Revised: We show extended samples from ordered VQ-VAEs in Fig .

Original: 11, Fig .
Revised: 11, Fig .

Original: 12 and Fig .
Revised: 12 and Fig .

Original: For comparison, we also provide samples from unordered VQ-VAEs in Fig .
Revised: For comparison, we also provide samples from unordered VQ-VAEs in Fig .

Original: 14, Fig . 15 and Fig .
Revised: 14, Fig . 15 and Fig .

Original: G. AUDIO SAMPLES
Revised: G. AUDIO SAMPLES

Original: We include the audio samples that are sampled from our anytime sampler in the supplementary material.
Revised: We include the audio samples that are sampled from our anytime sampler in the supplementary material.

Original: The audio / audio_baseline directory contains 90 samples from ordered / unordered
Revised: The audio / audio_baseline directory contains 90 samples from ordered / unordered

Original: VQ-VAEs respectively.
Revised: VQ-VAEs respectively.

Original: The fractions of full code length (0.0625, 0.25 and 1.0) used for generation are included in the names of .wav ﬁles.
Revised: The fractions of full code length (0.0625, 0.25 and 1.0) used for generation are included in the names of .wav ﬁles.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]