Anytime Sampling for Autoregressive Models via Ordered Autoencoding

1 INTRODUCTION . 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 ) . 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 . 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 . As a result , the sampling process of autoregressive models can be slow and does not allow interruptions . 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 . 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 . 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 fly to accommodate instantaneous resource availability at time-of-deployment . To address this difficulty , we consider the novel task of adaptive autoregressive generation under computational constraints . 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 . In particular , we take advantage of a generalization of Principal Components Analysis ( PCA ) proposed by Rippel et al . ( 2014 ) , which learns an ordered representations induced by a structured application of dropout to the representations learned by an autoencoder . Such a representation encodes raw data into a latent space where dimensions are sorted based on their importance for reconstruction . Autoregressive modeling is then applied in the ordered representation space instead . 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 first few dimensions of the representation and decode it to a sample whose quality degrades smoothly with truncation . 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 . Through experiments , we show that our autoregressive models are capable of trading off sample quality and inference speed . 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 CelebA . Compared to standard autoregressive models , our approach allows the sample quality to degrade gracefully as we reduce the computational budget for sampling . We also observe that on the VCTK audio dataset ( Veaux et al. , 2017 ) , our autoregressive model is able to generate the low frequency features first , then gradually refine the waveforms with higher frequency components as we increase the number of sampled latent dimensions . 2 BACKGROUND . Autoregressive Models Autoregressive models define a probability distribution over data points x ∈ RD by factorizing the joint probability distribution as a product of univariate conditional distributions with the chain rule . Using pθ to denote the distribution of the model , we have : pθ ( x ) = D∏ i=1 pθ ( xi | x1 , · · · , xi−1 ) ( 1 ) The model is trained by maximizing the likelihood : L = Epd ( x ) [ log pθ ( x ) ] , ( 2 ) where pd ( x ) represents the data distribution . Different autoregressive models adopt different orderings of input dimensions and parameterize the conditional probability pθ ( xi | x1 , · · · , xi−1 ) , i = 1 , · · · , D in different ways . Most architectures over images order the variables x1 , · · · , xD of image x in raster scan order ( i.e. , left-toright then top-to-bottom ) . 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θ ( xi | x1 , · · · , xi−1 ) is oblivious of xi , xi+1 , · · · , xD . Cost of Sampling During training , we can evaluate autoregressive models efficiently because x1 , · · · , xD are provided by data and all conditionals p ( xi | x1 , · · · , xi−1 ) can be computed in parallel . In contrast , sampling from autoregressive models is an inherently sequential process and can not be easily accelerated by parallel computing : we first need to sample x1 , after which we sample x2 from pθ ( x2 | x1 ) and so on—the i-th variable xi can only be obtained after we have already computed x1 , · · · , xi−1 . Thus , the run-time of autoregressive generation grows at least linearly with respect to the length of a sample . In practice , the sample length D can be more than hundreds of thousands for real-world image and audio data . This poses a major challenge to fast autoregressive generation on a small computing budget . 3 ANYTIME SAMPLING WITH ORDERED AUTOENCODERS . 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 . 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 . In this section , we shall first describe the learning of the ordered representation space via the use of an ordered autoencoder . We then describe how to achieve anytime sampling with ordered autoencoders . 3.1 ORDERED AUTOENCODERS . Consider an autoencoder that encodes an input x ∈ RD to a code z ∈ RK . Let z = eθ ( x ) : RD → RK be the encoder parameterized by θ and x′ = dφ ( z ) : RK → RD be the decoder parameterized by φ . We define eθ ( · ) ≤i : x ∈ RD 7→ ( z1 , z2 , · · · , zi , 0 , · · · , 0 ) T ∈ RK , which truncates the representation to the first i dimensions of the encoding z = eθ ( x ) , masking out the remainder of the dimensions with a zero value . We define the ordered autoencoder objective as 1 N N∑ i=1 1 K K∑ j=1 ‖xi − dφ ( eθ ( xi ) ≤j ) ‖22 . ( 3 ) We note that Eq . ( 3 ) is equivalent to Rippel et al . ( 2014 ) ’ s nested dropout formulation using a uniform sampling of possible truncations . Moreover , when the encoder/decoder pair is constrained to be a pair of orthogonal matrices up to a transpose , then the optimal solution in Eq . ( 3 ) recovers PCA . 3.1.1 THEORETICAL ANALYSIS . Rippel et al . ( 2014 ) ’ s analysis of the ordered autoencoder is limited to linear/sigmoid encoder and a linear decoder . 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 . We first reframe our problem from a probabilistic perspective . In lieu of using deterministic autoencoders , we assume that both the encoder and decoder are stochastic functions . In particular , we let qeθ ( z | x ) be a probability distribution over z ∈ RK conditioned on input x , and similarly let pdφ ( x | z ) be the stochastic counterpart to dφ ( z ) . We then use qeθ ( z | x ) ≤i to denote the distribution of ( z1 , z2 , · · · , zi , 0 , · · · , 0 ) T ∈ RK , where z ∼ qeθ ( z | x ) , and let pdφ ( x | z ) ≤i represent the distribution of pdφ ( x | ( z1 , z2 , · · · , zi , 0 , · · · , 0 ) T ∈ RK ) . We can modify Eq . ( 3 ) to have the following form : Ex∼pd ( x ) , i∼U { 1 , K } Ez∼qeθ ( z|x ) ≤i [ − log pdφ ( x|z ) ≤i ] , ( 4 ) where U { 1 , K } denotes a uniform distribution over { 1 , 2 , · · · , K } , and pd ( x ) represents the data distribution . We can choose both the encoder and decoder to be fully factorized Gaussian distributions with a fixed variance σ2 , then Eq . ( 13 ) can be simplified to Epd ( x ) [ 1 K K∑ i=1 Ez∼N ( eθ ( x ) ≤i ; σ2 ) [ 1 2σ2 ‖x− dφ ( z ) ≤i‖22 ] ] . 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 . ( 3 ) when σ → 0 and N →∞ . The optimal encoders and decoders that minimize Eq . ( 13 ) satisfy the following property . Theorem 1 . Let x denote the input random variable . Assuming both the encoder and decoder are optimal in terms of minimizing Eq . ( 13 ) , and ∀i ∈ 3 , · · · , K , zi−1 ⊥ zi | x , z≤i−2 , we have ∀i ∈ { 3 , · · · , K } : I ( zi ; x|z≤i−1 ) ≤ I ( zi−1 ; x|z≤i−2 ) , where z≤i denotes ( z1 , z2 , · · · , zi ) . We defer the proof to Appendix A.1 . The assumption zi−1 ⊥ zi | 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. Intuitively , the above theorem states that for optimal encoders and decoders that minimize Eq . ( 13 ) , one can extract less additional information about the raw input as the code gets longer . Therefore , there exists a natural ordering among different dimensions of the code based on the additional information they can provide for reconstructing the inputs .

Summary: The paper considers the problem of slow sampling in autoregressive generative models. Sampling in such models is sequential, so its computational cost scales with the data dimensionality. Existing work speeds up autoregressive sampling by caching activations or distilling into normalizing flows with fast sampling. Authors of this work instead propose a method that returns (approximate) samples given an arbitrary computational budget, a behaviour referred to as *anytime sampling*. The proposed model is based on VQ-VAE by van den Oord et al. (2017), where an autoregressive model is fit to a latent space of a trained discrete autoencoder, rather than to raw pixels. Authors adapt the *nested dropout* idea by Rippel et al. (2014) to encourage the discrete autoencoder to order latent dimensions by their "importance" for reconstruction. Experiments demonstrate that the ordered latent space allows to stop the autoregressive sampling process at an arbitrary latent dimension and still obtain "complete" samples. The quality of samples increases as more latent dimensions are sampled, which allows to trade sample quality for reduced computational cost.