Original: A DA C AT : Adaptive Categorical Discretization for Autoregressive Models
Revised: A DA C AT : Adaptive Categorical Discretization for Autoregressive Models

Original: Abstract
Revised: Abstract

Original: Autoregressive generative models can estimate complex continuous data distributions, like trajectory rollouts in an RL environment, image intensities, and audio.
Revised: Autoregressive generative models can estimate complex continuous data distributions, like trajectory rollouts in an RL environment, image intensities, and audio.

Original: Most state-of-the-art models discretize continuous data into several bins and use categorical distributions over the bins to approximate the continuous data distribution.
Revised: Most state-of-the-art models discretize continuous data into several bins and use categorical distributions over the bins to approximate the continuous data distribution.

Original: The advantage is that the categorical distribution can easily express multiple modes and are straightforward to optimize.
Revised: The advantage is that the categorical distribution can easily express multiple modes and are straightforward to optimize.

Original: However, such approximation cannot express sharp changes in density without using significantly more bins, which makes it parameter inefficient.
Revised: However, such approximation cannot express sharp changes in density without using significantly more bins, making it parameter inefficient.

Original: We propose an efficient, expressive, multimodal parameterization called Adaptive
Revised: We propose an efficient, expressive, multimodal parameterization called Adaptive

Original: Categorical Discretization (A DA C AT ).
Revised: Categorical Discretization (A DA C AT ).

Original: A DA C AT discretizes each dimension of an autoregressive model adaptively, which allows the model to allocate density to fine intervals of interest, improving parameter efficiency.
Revised: A DA C AT discretizes each dimension of an autoregressive model adaptively, which allows the model to allocate density to fine intervals of interest, improving parameter efficiency.

Original: A DA C AT generalizes both categoricals and quantile-based regression.
Revised: A DA C AT generalizes both categoricals and quantile-based regression.

Original: A DAC AT is a simple add-on to any discretization-based distribution estimator.
Revised: A DAC AT is a simple add-on to any discretization-based distribution estimator.

Original: In experiments, A DA C AT improves density estimation for real-world tabular data, images, audio, and trajectories, and improves planning in model-based offline RL.
Revised: In experiments, A DA C AT improves density estimation for real-world tabular data, images, audio, and trajectories, and improves planning in model-based offline RL.

Original:  INTRODUCTION
Revised:  INTRODUCTION

Original: Deep generative models estimate complex, high dimensional distributions from samples.
Revised: Deep generative models estimate complex, highdimensional distributions from samples.

Original: Autoregressive models like NADE [Larochelle and Murray, 2011,Uria et al., 2016],
Revised: Autoregressive models like NADE [Larochelle and Murray, 2011,

Original:  PixelRNN [Van Oord et al., 2016] and GPT [Radford et al., 2018] express a joint distribution by decomposing it into a product of simpler one-dimensional conditionals.
Revised: Uria et al., 2016], PixelRNN [Van Oord et al., 2016] and GPT [Radford et al., 2018] express a joint distribution by decomposing it into a product of simpler one-dimensional conditionals.

Original: Each of these conditionals p ( x t | x 1 , x 2 . . .
Revised: Each of these conditionals p ( x t | x 1 , x 2 . . .

Original: x t − 1 ) is parameterized by a neural networkmapping from a subset of observed vari-
Revised: x t − 1 ) is parameterized by a neural network

Original: [Figure]
Revised: [Table]

Original: , an autoregressive model parameterizes p ( x 1 ) , the marginal distribution over the first dimension, and p ( x 2 | x 1 ) a conditional distribution over the second. Uniform discretization (middle) divides their 1-D support into 16 equal-sized intervals and parameterizes each conditional with a categorical. However, it poorly fits the continuous samples. In contrast, parameterizing p ( x t | x <t ) with A DA C AT closely approximates the target distribution with the same number of bins.
Revised: [Figure]

Original: ables to logits over the next dimension.
Revised: mapping from a subset of observed variables to logits over the next dimension.

Original: For discrete data like language tokens, the conditional takes the form of a categorical.
Revised: For discrete data like language tokens, the conditional takes the form of a categorical distribution.

Original: Categorical distributions are relatively easy to optimize, flexible and can easily express multimodal distributions as each bin’s logit is independently predicted.
Revised: Categorical distributions are relatively easy to optimize, flexible and can easily express multimodal distributions as each bin’s logit is independently predicted.

Original: Ordinal and continuous data such as image intensities ranging from 0 to 1 have a natural ordering between possible values of each dimension x t .
Revised: Ordinal and continuous data such as image intensities ranging from 0 to 1 have a natural ordering between possible values of each dimension x t .

Original: The categorical does not exploit this ordering, instead separately predicting each bin.
Revised: The categorical does not exploit this ordering, instead separately predicting each bin.

Original: Categorical distributions also scale poorly when encoding highly precise data like agent trajectories, tabular datasets and audio [Oord et al., 2016].
Revised: Categorical distributions also scale poorly when encoding highly precise data like agent trajectories, tabular datasets and audio [Oord et al., 2016].

Original: Auditory quality degrades if the waveform is quantized to less than 8-16 bits (25665k intensity levels).
Revised: Auditory quality degrades if the waveform is quantized to less than 8-16 bits (25665k intensity levels).

Original: Control applications often need high precision as well.
Revised: Control applications often need high precision as well.

Original: Unfortunately, categorical likelihood degrades rapidly at high quantization levels.
Revised: Unfortunately, categorical likelihood degrades rapidly at high quantization levels.

Original: The uniformly discretized model in Figure 1 has a negative log likelihood − 0 .
Revised: The uniformly discretized model in Figure 1 has a negative log-likelihood − 0 .

Original: with 16 bins, while our adaptively discretized approach achieves NLL − 1 .
Revised: with 16 bins, while our adaptively discretized approach achieves NLL − 1 .

Original: 02 with the same architecture andnumber of bins (lower is better). Halving the width of a
Revised: 02 with the same architecture and

Original: [Figure]
Revised: [Figure]

Revised: number of bins (lower is better).

Revised: We note that the negative log-likelihood can be negative as it is computed on the continuous density by treating the discretized distribution as a mixture of uniform distributions.

Original:  particular bin in the categorical distribution would require double the parameters in the final layer of the network.
Revised: Halving the width of a particular bin in the categorical distribution would require double the parameters in the final layer of the network.

Original: Past work tries to improve the efficiency of categoricals with hierarchical softmax [Morin and Bengio, 2005] or quantilebased discretization [Janner et al., 2021].
Revised: Past work tries to improve the efficiency of categoricals with hierarchical softmax [Morin and Bengio, 2005] or quantilebased discretization [Janner et al., 2021].

Original: Heuristic, handengineered discretizations like the µ -law
Revised: Heuristic, handengineered discretizations like the µ -law

Original: [Oord et al., 2016] reduce quantization error and improve perceptual quality.
Revised: [Oord et al., 2016] reduce quantization error and improve perceptual quality.

Original: As an alternative, a single gaussian, gaussian mixtures [Bishop, 1994] or logistic mixtures [Salimans et al., 2017] are frequently used for parameter efficient conditionals, but can be hard to optimize, especially as the number of mixture components increases.
Revised: As an alternative, a single Gaussian, Gaussian mixtures [Bishop, 1994] or logistic mixtures [Salimans et al., 2017] are frequently used for parameter efficient conditionals, but can be hard to optimize, especially as the number of mixture components increases.

Original: In this work, we propose a parameterization of 1-D conditionals that is parameter efficient, expressive and multimodal.
Revised: In this work, we propose a parameterization of 1-D conditionals that is parameter efficient, expressive and multimodal.

Original: We propose Adaptive Categorical Discretization (A DA C AT ).
Revised: We propose Adaptive Categorical Discretization (A DA C AT ).

Original: Based on the observation that high precision is often only required in a small subset of a distribution’s support, A DA C AT is a distribution parameterized by a vector of interval masses and interval widths.
Revised: Based on the observation that high precision is often only required in a small subset of a distribution’s support, A DA C AT is a distribution parameterized by a vector of interval masses and interval widths.

Original: A DA C AT is depicted in
Revised: A DA C AT is depicted in

Original: Figure 2.
Revised: Figure 2.

Original: In contrast to categoricals with equal bin widths, variable bin widths allow the network to localize mass precisely without increasing precision elsewhere.
Revised: In contrast to categoricals with equal bin widths, variable bin widths allow the network to localize mass precisely without increasing precision elsewhere.

Original: Compared to non-uniform but fixed discretizations like quantiles, A DAC AT parameters are adaptive: they are predicted by a neural network conditioned on prior dimensions, which is important as the best discretization for a particular conditional differs from the best for the marginal.
Revised: Compared to non-uniform but fixed discretizations like quantiles, A DAC AT parameters are adaptive: they are predicted by a neural network conditioned on prior dimensions, which is important as the best discretization for a particular conditional differs from the best for the marginal.

Original: We also propose an analytic target smoothing strategy to ease optimization, and draw connections between target  and dequantization [Uria et al., 2016] and score matching [Vincent, 2011].
Revised: We also propose an analytic target smoothing strategy to ease optimization, and draw connections between target smoothing and dequantization [Uria et al., 2016] and score matching [Vincent, 2011].

Original: In experiments, A DA C AT with target smoothing scales better to few parameters or is competitive with strong baselines on image density estimation, offline reinforcement learning, tabular data and audio.
Revised: In experiments, A DA C AT with target smoothing scales better to few parameters or is competitive with strong baselines on image density estimation, offline reinforcement learning, tabular data and audio.

Original:  ADAPTIVE CATEGORICAL DISCRETIZATION 2. ADACAT DISTRIBUTION
Revised:  ADAPTIVE CATEGORICAL DISCRETIZATION 2. ADACAT DISTRIBUTION

Original: The A DA C AT distribution is a particular subfamily of mixtures of uniform distributions where each mixture component has non-overlapping support.
Revised: The A DA C AT distribution is a particular subfamily of mixtures of uniform distributions where each mixture component has non-overlapping support.

Original: A standard A DA C AT distribution A DA C AT k ( w, h ) has k components with a support over [0 , 1) .
Revised: A standard A DA C AT distribution A DA C AT k ( w, h ) has k components with a support over [0 , 1) .

Original: It is parameterized by two vectors in the k -dimensional simplex: w, h ∈ ∆ k − 1 .
Revised: It is parameterized by two vectors in the k -dimensional simplex: w, h ∈ ∆ k − 1 .

Original: Thus, w and h are normalized.
Revised: Thus, w and h are normalized.

Original: w is additionally constrained to be non-zero in all of its elements.
Revised: w is additionally constrained to be non-zero in all of its elements.

Original: The probability density function (PDF) of an A DA C AT distribution is defined as,
Revised: The probability density function (PDF) of an A DA C AT distribution is defined as,

Original: [Equation]
Revised: [Equation]

Original: where c i = (cid:80) i − 1 j =1 w j is the prefix sum of the dimensions of parameter w and I [ · ] is the indicator function.
Revised: where c i = (cid:80) i − 1 j =1 w j is the prefix sum of the dimensions of parameter w and I [ · ] is the indicator function.

Original: Intuitively, w i captures the size of each discretized bin (support of each mixture component), h i captures the probability mass assigned to each bin, and h i w i is the density contributed by each bin.
Revised: Intuitively, w i captures the size of each discretized bin (support of each mixture component), h i captures the probability mass assigned to each bin, and h i w i is the density contributed by each bin.

Original: 2.2 RELATIONSHIP WITH UNIFORM AND QUANTILE DISCRETIZATION
Revised: 2.2 RELATIONSHIP WITH UNIFORM AND QUANTILE DISCRETIZATION

Original: Connection to Uniform Discretization Generative models over ordinal data like PixelRNNs [Van Oord et al., 2016] commonly divide the support of 1-D distributions into equalwidth bins, and share the same bins across all dimensions of the data.
Revised: Connection to Uniform Discretization Generative models over ordinal data like PixelRNNs [Van Oord et al., 2016] commonly divide the support of 1-D distributions into equalwidth bins, and share the same bins across all dimensions of the data.

Original: This allows neural networks to parameterize the distribution with a simple classification head that predicts a categorical distribution over bins.
Revised: This allows neural networks to parameterize the distribution with a simple classification head that predicts a categorical distribution over bins.

Original: A DA C AT generalizes 1-D categorical distributions with uniformly discretized support.
Revised: A DA C AT generalizes 1-D categorical distributions with uniformly discretized support.

Original: If w is set to be w i = 1 k , ∀ i , the distribution is reduced to a categorical distribution over { 0 , 1 k , 2 k , · · · , k − 1 k } augmented with a uniform noise of magnitude 1 k .
Revised: If w is set to be w i = 1 k , ∀ i , the distribution is reduced to a categorical distribution over { 0 , 1 k , 2 k , · · · , k − 1 k } augmented with a uniform noise of magnitude 1 k .

Original: Figure 2 shows how
Revised: Figure 2 shows how

Original: A DA C AT is more expressive than a uniformly discretized categorical, allowing bin widths to vary and more closely approximating the modes of a mixture of two Gaussians.
Revised: A DA C AT is more expressive than a uniformly discretized categorical, allowing bin widths to vary and more closely approximating the modes of a mixture of two Gaussians.

Original: Connection to Quantile-based Discretization
Revised: Connection to Quantile-based Discretization

Original: A DA C AT also generalizes quantile-based discretization, which discretizes a distribution’s support by binning data into groups with equal numbers of observed data points.
Revised: A DA C AT also generalizes quantile-based discretization, which discretizes a distribution’s support by binning data into groups with equal numbers of observed data points.

Original: If h is set to be h i = 1 k , ∀ i , i.e. the same mass in every bin, the vector w can be interpreted as the k -quantile of the distribution.
Revised: If h is set to be h i = 1 k , ∀ i , i.e. the same mass in every bin, the vector w can be interpreted as the k -quantile of the distribution.

Original: This strategy is employed by generative models like the Trajectory Transformer [Janner et al., 2021], which pre-computes and fixes the bin widths w separately for each dimension to achieve equal mass 1 k per bin of the marginal distributions of the training set, then predicts mass h with a neural network based on observed dimensions.
Revised: This strategy is employed by generative models like the Trajectory Transformer [Janner et al., 2021], which pre-computes and fixes the bin widths w separately for each dimension to achieve equal mass 1 k per bin of the marginal distributions of the training set, then predicts mass h with a neural network based on observed dimensions.

Original: 2.3 AUTOREGRESSIVE PARAMETERIZATION
Revised: 2.3 AUTOREGRESSIVE PARAMETERIZATION

Original: In problems with dimension greater than 1, we use deep autoregressive models to factorize the joint density f ( x ) into multiple 1-D conditional A DA C AT distributions:
Revised: In problems with dimension greater than 1, we use deep autoregressive models to factorize the joint density f ( x ) into multiple 1-D conditional A DA C AT distributions:

Original: [Equation]
Revised: [Equation]

Original: This factorization can express any target distribution by imposing an order over data dimensions, e.g. the natural ordering of the data in a sequential problem.

Original: For each dimension, conditioned on observed or generated values of prior dimensions, the neural net g θ outputs two unconstrained parameters { ϕ t , ψ t } =
Revised: For each dimension, conditioned on observed or generated values of prior dimensions, the neural net g θ outputs two unconstrained parameters { ϕ t , ψ t } =

Original: g θ ( x <t ) , where ϕ t , ψ t ∈ R k .
Revised: g θ ( x <t ) , where ϕ t , ψ t ∈ R k .

Original: The predicted ϕ and ψ represent the unnormalized log values for h and w for each dimension.
Revised: The predicted ϕ and ψ represent the unnormalized log values for h and w for each dimension.

Original: These parameters are normalized independently using a softmax to satisfy the normalization and positivity constraints:
Revised: These parameters are normalized independently using a softmax to satisfy the normalization and positivity constraints:

Original: [Equation]
Revised: [Equation]

Original: where ϕ t =
Revised: where ϕ t =

Original: ϕ
Revised: ϕ

Original: θ ( x <t ) , ψ t = ψ θ ( x <t ) .
Revised: θ ( x <t ) , ψ t = ψ θ ( x <t ) .

Original: Unlike uniform, quantile-based, or heuristic discretization strategies, our autoregressive model can adaptively choose how to discretize each dimension’s conditional distribution based on observations.
Revised: Unlike uniform, quantile-based, or heuristic discretization strategies, our autoregressive model can adaptively choose how to discretize each dimension’s conditional distribution based on observations.

Original: Adaptivity improve expressiveness, since density can be precisely localized in regions of interest, and the discretization can vary across data dimensions.
Revised: Adaptivity improves expressiveness, since density can be precisely localized in regions of interest, and the discretization can vary across data dimensions.

Original: This is especially important for problems where the optimal discretization is not known a priori .
Revised: This is especially important for problems where the optimal discretization is not known a priori .

Original: In the 2-D dataset shown in Figure 1, fixed discretizations poorly express the inherent multimodality in the data, while A DA C AT ’s adaptivity allows the network to shift modes of p θ ( x 2 | x 1 ) for different values of the first dimension, x 1 . 
Revised: In the 2-D dataset shown in Figure 1, fixed discretizations poorly express the inherent multimodality in the data, while A DA C AT ’s adaptivity allows the network to shift modes of p θ ( x 2 | x 1 ) for different values of the first dimension, x 1 . smoothed objective shrinks size of most bins they are

Original:  OPTIMIZING WITH ANALYTIC TARGET SMOOTHING
Revised:  OPTIMIZING WITH ANALYTIC TARGET SMOOTHING

Original: Autoregressive models with A DA C AT conditionals can be estimated by minimizing the KL-divergence between the target distribution and the learned density D KL ( p data ( x ) ∥
Revised: Autoregressive models with A DA C AT conditionals can be estimated by minimizing the Kullback–Leibler (KL)- divergence between the target distribution and the learned density D KL ( p data ( x ) ∥

Original: p θ ( x )) .
Revised: p θ ( x )) .

Original: The KL reduces to the empirical log likelihood objective below, where x 1 , · · · , x n are sampled from the data distribution p data :
Revised: The KL reduces to the empirical log-likelihood objective below, where x 1 , · · · , x n are sampled from the data distribution p data :

Original: [Equation]
Revised: [Equation]

Original: [Equation] However, due to the discontinuous nature of the A DA C AT density function, we observed that naïvely maximizing the empirical log likelihood encourages the model to get trapped in poor local optima.
Revised: However, due to the discontinuous nature of the A DA C AT density function, we observed that naïvely maximizing the empirical log-likelihood encourages the model to get trapped in poor local optima.

Original: This phenomenon is best illustrated in 1-D, as in Figure 3.
Revised: This phenomenon is best illustrated in 1-D, as in Figure 3.

Original: The density in red is estimated with maximum likelihood ˆ L ll (3), but the bin widths shrink over the course of optimization and reach small values.
Revised: The density in red is estimated with maximum likelihood ˆ L ll (3), but the bin widths shrink over the course of optimization and reach small values.

Original: Density is overestimated in between modes and underestimated in regions where a single uniform mixture component needs to cover a large interval.
Revised: Density is overestimated in between modes and underestimated in regions where a single uniform mixture component needs to cover a large interval.

Original: We provide one possible explanation for bin collapse.
Revised: We provide one possible explanation for bin collapse.

Original: Rewriting the log likelihood with A DA C AT ’s PDF (1) based on a summation over bins i = 1 to k ,
Revised: Rewriting the log-likelihood with A DA C AT ’s PDF (1) based on a summation over bins i = 1 to k ,

Original: [Equation]
Revised: [Equation]

Original: Due to the constraint that mixture components are nonoverlapping, only a single term of the inner summation is non-zero in (4).
Revised: Due to the constraint that mixture components are nonoverlapping, only a single term of the inner summation is non-zero in (4).

Original: The loss separates into terms for h and w ,
Revised: The loss separates into terms for h and w ,

Original: [Equation]
Revised: [Equation]

Original: Maximizing the loss for data point x d pushes for higher density h ti w ti when x td lies in for bin i of the conditional p θ ( ·| x <td ) .
Revised: Maximizing the loss for data point x d pushes for higher density h ti w ti when x td lies in for bin i of the conditional p θ ( ·| x <t d ) .

Original: This density is increased by either increasing log mass log h ti or decreasing log bin width log w i .
Revised: This density is increased by either increasing log mass log h ti or decreasing log bin width log w i .

Original: For uniform and heuristic discretizers, log w ti is fixed.
Revised: For uniform and heuristic discretizers, log w ti is fixed.

Original: However, updating the bin width w ti with finite step sizes can make the data point x td move out of the current bin discontinuously.
Revised: However, updating the bin width w ti with finite step sizes can make the data point x td move out of the current bin discontinuously, which can result in biased gradient estimates (see Appendix A).

Original: The gradient ddw ti ˆ L ll is also constant for any value of a sample within bin i as the density is piecewise constant, so the gradient encouraging bin collapse does not attenuate as the sample approaches a bin boundary.
Revised: The gradient ddw ti ˆ L ll is also constant for any value of a sample within bin i as the density is piece-wise constant, so the gradient encouraging bin collapse does not attenuate as the sample approaches a bin boundary.

Original: Once a bin is updated to exclude a particular data point, only the normalization of w t encourages the bin to grow to include the data point again, but we empirically find that this is not enough to prevent collapse.
Revised: Once a bin is updated to exclude a particular data point, only the normalization of w t encourages the bin to grow to include the data point again, but we empirically find that this is not enough to prevent collapse.

Original: Instead, optimization could shrink the new bin w ti +1 or w ti − 1 , repeating until a majority of the mixture components collapse to support a small fraction of the overall interval.
Revised: Instead, optimization could shrink the new bin w ti +1 or w ti − 1 , repeating until a majority of the mixture components collapse to support a small fraction of the overall interval.

Original: Luckily, this issue can be largely alleviated by using a smoothed objective:
Revised: Luckily, this issue can be largely alleviated by using a smoothed objective:

Original: [Equation]
Revised: [Equation]

Original: where ζ (˜ x | x ) is any smoothing density function that is centered around x .
Revised: where ζ (˜ x | x ) is any smoothing density function that is centered around x .

Revised: This smoothed objective can be interpreted as the NLL objective under a smoothed data distribution (by applying the smoothing function on top of the data).

Original: We  find that both Uniform and Gaussian distributions with mean x effectively prevent the bins from collapsing, and use ζ ( ·| x ) =
Revised: We discuss this in more details in Appendix B. In practice, we find that both Uniform and Gaussian distributions with mean x effectively prevent the bins from collapsing, and use ζ ( ·| x ) =

Original: Unif [ x − λ 2 , x + λ 2 )  in all experiments with λ ∈ { 0 . 00001 , 0 . 0001 , 0 . 001 } for simplicity, truncating on the boundaries of the support of x ∈ [0 , 1) .
Revised: Unif [ x − λ 2 , x + λ 2 ) or ζ ( ·| x ) = N ( x, λ 2 ) in all experiments, truncating on the boundaries of the support of x ∈ [0 , 1) .

Original: By optimizing ˆ L s with uniform target smoothing, the density in blue in Figure 3 converges to a close approximation of the target density.
Revised: By optimizing ˆ L s with uniform target smoothing, the density in blue in Figure 3 converges to a close approximation of the target density.

Original: The smoothed objective might seem intractable with an integration in the inner summation.
Revised: The smoothed objective might seem intractable with an integral in the inner summation.

Original: Fortunately, the form of the conditional log p θ (˜ x | x <ti ) with A DA C AT ’s simple density function allows us to evaluate the integration analytically as long as the smoothing density has an analytic cumulative density function (CDF).
Revised: Fortunately, the form of the conditional log p θ (˜ x | x <ti ) with A DA C AT ’s simple density function allows us to evaluate the integral analytically as long as the smoothing density has an analytic cumulative density function (CDF).

Original: If F ( · ) is the CDF of ζ , then the integration can be analytically computed as: 
Revised: If F ( · ) is the CDF of ζ , then the integral can be analytically computed as: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  where we recall that c is the prefix sum of w as defined previously.
Revised: 8) where we recall that c is the prefix sum of w as defined previously.

Original: Only the bins that intersect with the support of the smoothing density function contribute to this objective.
Revised: Only the bins that intersect with the support of the smoothing density function contribute to this objective.

Original: 3.1 RELATIONSHIP WITH DENOISING SCORE MATCHING
Revised: 3.1 RELATIONSHIP WITH DENOISING SCORE MATCHING

Original: Energy based models and denoising autoencoders trained by denoising score matching (DSM, Vincent [2011]) minimize a reconstruction objective:
Revised: Energy based models and denoising autoencoders trained by denoising score matching (DSM, Vincent [2011]) minimize a reconstruction objective:

Original: [Equation]
Revised: [Equation]

Original: Assuming the observation model p θ ( ·| ˜ x ) =
Revised: Assuming the observation model p θ ( ·| ˜ x ) =

Original: N (ˆ x θ (˜ x ) , I ) is a standard Gaussian,[Equation] [Equation] |
Revised: N (ˆ x θ (˜ x ) , I ) is a standard Gaussian,

Original: [Equation]
Revised: [Equation]

Original: x )] resembling (6).
Revised: resembling (6).

Original: However, our target smoothed loss is designed to regularize the output conditional distribution, so perturbations are employed on the output space, not the input space, and our generative model is conditioned on clean, unperturbed observations.
Revised: However, our target smoothed loss is designed to regularize the output conditional distribution, so perturbations are employed on the output space, not the input space, and our generative model is conditioned on clean, unperturbed observations.

Original: Recent works introduce multi-scale perturbations [Song and Ermon, 2019], and denoising diffusion probabilistic models reweight a related variational bound for this class of models to improve sample quality [Ho et al., 2020].
Revised: Recent works introduce multi-scale perturbations [Song and Ermon, 2019], and denoising diffusion probabilistic models reweight a related variational bound for this class of models to improve sample quality [Ho et al., 2020].

Original: 3.2 RELATIONSHIP WITH DEQUANTIZATION
Revised: 3.2 RELATIONSHIP WITH DEQUANTIZATION

Original: Other continuous density estimators like normalizing flows trained on discrete data suffer from degenerate solutions if trained naïvely via maximum likelihood (Ho et al.
Revised: Other continuous density estimators like normalizing flows trained on discrete data suffer from degenerate solutions if trained naïvely via maximum likelihood (Ho et al.

Original: Sec. 3.1).
Revised: Sec. 3.1).

Original: Flows suffer from a different failure case than non-smoothed A DA C AT .
Revised: Flows suffer from a different failure case than non-smoothed A DA C AT .

Original: Continuous estimators trained on discrete data e.g. a mixture of dirac δ functions, can arbitrarily increase density on  training data as the empirical distribution is supported on a set with measure zero.
Revised: Continuous estimators, e.g., a mixture of Dirac δ functions, can arbitrarily increase density on discrete training data as the empirical distribution is supported on a set with measure zero.

Original: Dequantization avoids the problem by adding continuous noise to observed discrete samples [Theis et al., 2015].
Revised: Dequantization avoids the problem by adding continuous noise to observed discrete samples [Theis et al., 2015, Hoogeboom et al., 2021].

Original: For uniform noise, the  dequantized estimator maximizes the following log likelihood of discrete data sample x ,
Revised: As an example, the following dequantized objective gives a lower bound of the log-likelihood of discrete data sample x :

Original: [Equation]
Revised: [Equation]

Original: a variational lower bound on the continuous density defined by the discrete log likelihood log P θ ( x  ) .
Revised: where λ is chosen such that [ x, x + λ ) with different discrete sample x do not overlap with each other.

Original: While (9) closely resembles (7), it differs subtly in that dequantization perturbs all dimensions of the data x , not just the 1-D target, and that the integration is done via a stochastic sample from ζ rather than analytically.
Revised: While (9) closely resembles (7), it differs subtly in that dequantization perturbs all dimensions of the data x , not just the 1-D target, and that the integral is done via a stochastic sample from ζ rather than analytically.

Original: We observed bin collapse even on continuous data like the mixtures in Figures 2, 3, and find that single-sample estimates of the expectation do not prevent collapse.
Revised: We observe bin collapse even on continuous data like the mixtures in Figures 2, 3, and find that single-sample estimates of the expectation do not prevent collapse.

Original: These findings suggest that analytic target smoothing helps with the discontinuity in the model conditional rather than a property of the data.
Revised: These findings suggest that analytic target smoothing helps with the discontinuity in the model conditional rather than a property of the data.

Original:  EVALUATION
Revised:  EVALUATION

Original: In experiments, we evaluate the performance of autoregressive density estimators with adaptive categorical conditional distributions for several data modalities.
Revised: In experiments, we evaluate the performance of autoregressive density estimators with adaptive categorical conditional distributions for several data modalities.

Original: We evaluate on standard benchmarks for real-world tabular data (Section 4.1), image generation (4.2), speech synthesis (4.3) and offline reinforcement learning (4.4).
Revised: We evaluate on standard benchmarks for real-world tabular data (Section 4.1), image generation (4.2), speech synthesis (4.3) and offline reinforcement learning (4.4).

Original: A DA C AT outperforms uniform discretization strategies in all settings, and is competitive with hand-engineered conditional distributions.
Revised: A DA C AT outperforms uniform discretization strategies in all settings, and is competitive with hand-engineered conditional distributions.

Original: Beyond density estimation, our results suggest that A DA C AT can improve downstream task performance including speech quality and control.
Revised: Beyond density estimation, our results suggest that A DA C AT can improve downstream task performance, including speech quality and control.

Original: 4.1 TABULAR DATA MODELING
Revised: 4.1 TABULAR DATA MODELING

Original: We compare the performance of autoregressive models with
Revised: We compare the performance of autoregressive models with

Original: A DA C AT and uniform parameterizations on real-world tabular density estimation benchmarks, the UCI datasets of
Revised: A DA C AT and uniform parameterizations on real-world tabular density estimation benchmarks, the UCI datasets of

Original: Dua and Graff [2017].
Revised: Dua and Graff [2017].

Original: The state-of-the-art performances on these benchmarks are also included for reference.
Revised: The state-of-the-art performances on these benchmarks are also included for reference.

Original: We use a 4-layer feed-forward network to predict the A DA C AT parameters for each dimension of the data ( e.g. , we use
Revised: We use a 4-layer feed-forward network to predict the A DA C AT parameters for each dimension of the data ( e.g. , we use

Original: MLPs for the POWER dataset since it has 6-dimensional data).
Revised: MLPs for the POWER dataset since it has 6-dimensional data).

Original: Each network has 500 hidden units for all datasets except for GAS where we use 1000 hidden units.
Revised: Each network has 500 hidden units for all datasets except for GAS, where we use 1000 hidden units.

Original: We also use a Fourier encoding of the input inspired by
Revised: We also use a Fourier encoding of the input inspired by

Original: Tancik et al.
Revised: Tancik et al.

Original: [2020], Kingma et al.
Revised: [2020], Kingma et al.

Original: [2021] to allow a shallow model to capture high-frequency variations in the input.
Revised: [2021] to allow a shallow model to capture high-frequency variations in the input.

Original: Specifically, we augment eachthe input element x t with b pairs of additional features: { sin(2 j x t ) , cos(2 j x t ) } b − 1 j =0 .
Revised: Specifically, we augment each input element x t with b pairs of additional features: { sin(2 j x t ) , cos(2 j x t ) } b − 1 j =0 .

Original: We choose the feature count b = 32 for GAS and POWER, b = 8 for MINIBOONE, and b = 4 for HEPMASS.
Revised: We choose the feature count b = 32 for GAS and POWER, b = 8 for MINIBOONE, and b = 4 for HEPMASS.

Original: The uniform baseline uses the same architecture except the widths of the bins are forced to be uniform.
Revised: The uniform baseline uses the same architecture except the widths of the bins are forced to be uniform.

Original: We search over the number of bins in { 100 , 200 , 300 , 500 , 1000 } for both the uniform baseline and A DA C AT and select the best to report in the table.
Revised: We search over the number of bins in { 100 , 200 , 300 , 500 , 1000 } for both the uniform baseline and A DA C AT and select the best to report in the table.

Original: All models are trained for 400 epochs using Adam [Kingma and Ba, 2014] with a learning rate of 0 .
Revised: All models are trained for 400 epochs using Adam [Kingma and Ba, 2014] with a learning rate of 0 .

Original: 0003 and the learning rate halves every 100 epochs.
Revised: 0003 and the learning rate halves every 100 epochs.

Original: We use uniform target smoothing for A DA C AT with λ = 0 . 00001 for POWER and λ = 0 .  for all other datasets. 
Revised: We use truncated Gaussian target smoothing for A DA C AT with λ = 0 . 00001 for POWER, and λ = 0 . 0001 for all other datasets. See more details in Appendix C.1.

Original: Table 1 reports results.
Revised: Table 1 reports results.

Original: Overall, A DA C AT consistently outperforms the uniform baseline across all datasets, reducing the NLL by 1.9, 4.8, 3.1 and 4.0 nats on POWER, GAS,
Revised: Overall, A DA C AT consistently outperforms the uniform baseline across all datasets, reducing the NLL by 1.9, 4.8, 3.1 and 4.0 nats on POWER, GAS,

Original: HEPMASS and MINIBOONE respectively.
Revised: HEPMASS and MINIBOONE, respectively.

Original: Our approach also obtains competitive performance with the state-of-the-
Revised: Our approach also obtains competitive performance with the state-of-theart normalizing flow models on GAS.

Original: 4.2 IMAGE DENSITY ESTIMATION
Revised: 4.2 IMAGE DENSITY ESTIMATION

Original: Table 2 compares the performance of A DA C AT against existing parameterizations on the grayscale MNIST [LeCun and
Revised: Table 2 compares the performance of A DA C AT against existing parameterizations on the grayscale MNIST [LeCun and

Original: Cortes, 2010] image generation task in terms of negative log likelihood measured in bits/dimension.
Revised: Cortes, 2010] image generation task in terms of negative loglikelihood measured in bits/dimension.

Original: The autoregressive architecture we use for this task is a GPT-like Transformer decoder Vaswani et al.
Revised: The autoregressive architecture we use for this task is a GPT-like Transformer decoder Vaswani et al.

Original: [2017] with 4 layers and 4 heads.
Revised: [2017] with 4 layers and 4 heads.

Original: Since the image data is discrete, instead of dequantizing and smoothing the target, we directly minimize the cross entropy loss in the original discrete space.
Revised: Since the image data is discrete, instead of dequantizing and smoothing the target, we directly minimize the cross entropy loss in the original discrete space.

Original: We compute the log probability of the i th discrete target by analytically computing the total probability mass assigned to (cid:2) i 256 , i +1256 (cid:3) in our continuous distribution, i.e. density integrated over the interval.
Revised: We compute the log probability of the i th discrete target by analytically com- [Table] [Table] puting the total probability mass assigned to (cid:2) i 256 , i +1256 (cid:3) in our continuous distribution, i.e. density integrated over the interval.

Original: This corresponds to mapping our continuous distribution from [0 , 1] to the 256 discrete values uniformly such that the i th discrete value is mapped from (cid:2) i 256 , i +1 256 (cid:3) .
Revised: This corresponds to mapping our continuous distribution from [0 , 1] to the 256 discrete values uniformly such that the i th discrete value is mapped from (cid:2) i 256 , i +1256 (cid:3) .

Original: The results are grouped according to the number of parameters used to express the intensity distribution of each pixel, allowing us to examine the parameter efficiency of each approach.
Revised: The results are grouped according to the number of parameters used to express the intensity distribution of each pixel, allowing us to examine the parameter efficiency of each approach.

Original: For the uniform baseline, we evenly divideup the [0 , 1] intensity interval into k bins and use k parameters to model the probability assigned to each bin (with unnormalized log probability).
Revised: For the uniform baseline, we evenly divide the [0 , 1] intensity interval into k bins and use k parameters to model the probability assigned to each bin (with unnormalized log probability).

Original: DMoL uses 3 k parameters for a k -component mixture model ( e.g. , a 256 parmaeter count budget for DMoL corresponds to a 86 -component mixture model).
Revised: DMoL uses 3 k parameters for a k -component mixture model ( e.g. , a 256 parameter count budget for DMoL corresponds to a 86 -component mixture model).

Original: A DA C AT uses 2 k parameters for a k -component mixture model, using k bins of variable size.We examine [Table] [Table]
Revised: A DA C AT uses 2 k parameters for a k -component mixture model, using k bins of variable size.

Original:  parameter counts ranging from 32 to 512 .
Revised: We examine parameter counts ranging from 32 to 512 .

Original: Overall, A DA C AT has better performance on most parameter counts.
Revised: Overall, A DA C AT has better performance on most parameter counts.

Original: It only underperforms the uniform discretization at 256 parameters.
Revised: It only underperforms the uniform discretization at 256 parameters.

Original: We note that the MNIST dataset is discrete with 256 classes, which means that the uniform discretization has the optimal bin division.
Revised: We note that the MNIST dataset is discrete with 256 classes, which means that the uniform discretization has the optimal bin division.

Original: Therefore, we don’t expect A DA C AT to be able to outperform uniform because A DA C AT has effectively half of the bins available.
Revised: Therefore, we do not expect A DA C AT to be able to outperform uniform because A DA C AT has effectively half of the bins available.

Original: DMoL achieves the best performance when there are 32 parameters, but scales poorly with the number of components and underperforms A DA C AT for most settings with more than 32 parameters.
Revised: DMoL achieves the best performance when there are 32 parameters, but scales poorly with the number of components and underperforms A DA C AT for most settings with more than 32 parameters.

Original: We also experimented with an adaptive quantile baseline where we keep the probability mass assigned to each bin to be the same rather than the width.
Revised: We also experiment with an adaptive quantile baseline where we keep the probability mass assigned to each bin to be the same rather than the width.

Original: However, we found that the adaptive quantile baseline is very unstable to train.
Revised: However, we found that the adaptive quantile baseline is very unstable to train.

Original: We only report its result on a parameter count of 32 because the model at a higher parameter count often diverges early in training which results in inconsistent performances across runs.
Revised: We only report its result on a parameter count of 32 because the model at a higher parameter count often diverges early in training which results in inconsistent performances across runs.

Original: Adaptive quantiles outperform uniform withparameters, yet is still much less expressive than A DA C AT .
Revised: The adaptive quantile baseline outperforms uniform with 32 parameters, yet is still much less expressive than A DA C AT .

Original: 4.3 AUDIO DENSITY ESTIMATION AND VOCODING
Revised: 4.3 AUDIO DENSITY ESTIMATION AND VOCODING

Original: Neural vocoders synthesize human-like speech expressed as a waveform, conditioned on phonetic or frequency spectrum based features.
Revised: Neural vocoders synthesize human-like speech expressed as a waveform, conditioned on phonetic or frequency spectrum based features.

Original: Speech waveforms are long sequences as audio must be sampled at a high rate for fidelity, typically 1624 kHz.
Revised: Speech waveforms are long sequences as audio must be sampled at a high rate for fidelity, typically 1624 kHz.

Original: Thus, efficient generative models are essential for practical applications, and the output layer of the network can be a significant fraction of the compute.
Revised: Thus, efficient generative models are essential for practical applications, and the output layer of the network can be a significant fraction of the compute.

Original: WaveNet
Revised: WaveNet

Original: [Oord et al., 2016] is a popular autoregressive vocoder that synthesizes waveforms conditioned on a mel-spectrogram representation of the amplitude of audio frequencies across multiple bands.
Revised: [Oord et al., 2016] is a popular autoregressive vocoder that synthesizes waveforms conditioned on a Mel-spectrogram representation of the amplitude of audio frequencies across multiple bands.

Original: Despite the conditioning information, vocoding is challenging as WaveNet needs to reconstruct the phase of the audio frequencies.
Revised: Despite the conditioning information, vocoding is challenging as WaveNet needs to reconstruct the phase of the audio frequencies.

Original: This is done by estimating the distribution of audio in a dataset via maximum likelihood. 
Revised: This is done by estimating the distribution of audio in a dataset via maximum likelihood. [Table]

Revised: dim) and waveform vocoding MCD error for

Revised: WaveNet with different parameterizations of conditional distributions on the LJSpeech dataset.

Revised: † Discretized Mixture of

Revised: Logistics (DMoL) requires 3 parameters per mixture component (weight, mean and log scale), so output parameters are approximately matched to baselines.

Original: We train WaveNet with the standard dilated CNN architecture on the open LJSpeech dataset [Ito and Johnson, 2017] of wav files using an open-source implementation.
Revised: We train WaveNet with the standard dilated CNN architecture on the open LJSpeech dataset [Ito and Johnson, 2017] of wav files using an open-source implementation.

Original: A groundtruth mel spectrogram is extracted and used for conditioning
Revised: A groundtruth Mel-spectrogram is extracted and used for conditioning

Original: WaveNet.
Revised: WaveNet.

Original: For baselines, we use different parameterizations of the conditional distribution: a uniformly discretized categorical, a categorical discretized by a hand-engineered µ -law strategy that sets bin widths logarithmically with intensity, and a discretized mixture of logistics.
Revised: For baselines, we use different parameterizations of the conditional distribution: a uniformly discretized categorical, a categorical discretized by a hand-engineered µ -law strategy that sets bin widths logarithmically with intensity, and a discretized mixture of logistics.

Original: All models use 24 convolutional layers and are optimized for 500k iterations with Adam.
Revised: All models use 24 convolutional layers and are optimized for 500k iterations with Adam.

Original: The learning rate is initially 0.001 and is decayed by half every 200k iterations, with batch size 8.
Revised: The learning rate is initially 0.001 and is decayed by half every 200k iterations, with batch size 8.

Original: An exponential moving average of model parameters is used for testing.
Revised: An exponential moving average of model parameters is used for testing.

Original: Audio is sampled at 22,050 Hz and windowed into blocks of 1024 samples.
Revised: Audio is sampled at 22,050 Hz and windowed into blocks of 1024 samples.

Original: We evaluate the continuous negative log likelihood (NLL) of [Table]
Revised: We evaluate the continuous negative log-likelihood (NLL) of test waveforms.

Original: ) and waveform vocoding MCD error for

Original: WaveNet with different parameterizations of conditional distributions on the LJSpeech dataset.

Original: † Discretized Mixture of

Original: Logistics (DMoL) requires 3 parameters per mixture component (weight, mean and log scale), so output parameters are approximately matched to baselines.

Original: test waveforms.

Original: NLL is measured in bits/dimension with the waveform scaled to a [ − 1 , 1] amplitude.
Revised: NLL is measured in bits/dimension with the waveform scaled to a [ − 1 , 1] amplitude.

Original: We also measure the NLL of µ -law transformed data with µ = 256 , which is also scaled to [ − 1 , 1] .
Revised: We also measure the NLL of µ -law transformed data with µ = 256 , which is also scaled to [ − 1 , 1] .

Original: Following Chen et al.
Revised: Following Chen et al.

Original: [2020], the Mel
Revised: [2020], the Mel

Original: Cepstral Distance (MCD) objectively quantifies the perceptual similarity of our synthesized audio and reference audio based on an aligned mean squared error metric
Revised: Cepstral Distance (MCD) objectively quantifies the perceptual similarity of our synthesized audio and reference audio based on an aligned mean squared error metric

Original: [Kubichek, 1993].
Revised: [Kubichek, 1993].

Original: Samples from the uniformly discretized model led to numerical instabilities in the open-source library used to compute the MCD metric and have clear auditory artifacts, so the MCD metric is omitted.
Revised: Samples from the uniformly discretized model led to numerical instabilities in the open-source library used to compute the MCD metric and have clear auditory artifacts, so the MCD metric is omitted.

Original: Table 3 shows results, grouped by the number of parameters output by WaveNet for each conditional.
Revised: Table 3 shows results grouped by the number of parameters output by WaveNet for each conditional.

Original: Note that A DAC AT has half the bins of categorical baselines at the same parameter count due to using two parameter vectors, w and h .
Revised: Note that A DAC AT has half the bins of categorical baselines at the same parameter count due to using two parameter vectors, w and h .

Original: Across all settings, A DA C AT achieves significantly better negative log likelihood than uniform discretization: 4 .bpd lower NLL with 30 output parameters, 2 . 

Revised: Across all settings, A DA C AT achieves significantly better negative log-likelihood than uniform discretization: 4 .bpd lower NLL with 30 output parameters, 2 .

Original: bpd with 256 parameters, and 1 .
Revised: bpd with 256 parameters, and 1 .

Original: 41 bpd with 512 parameters.
Revised: 41 bpd with 512 parameters.

Original: A DA C AT is also competitive with hand-engineered quantization in the µ -law intensity space, despite not having prior knowledege about humans’ logarithmic perception of sound intensity, and without DMoL’s instabilities.
Revised: A DA C AT is also competitive with hand-engineered quantization in the µ -law intensity space, despite not having prior knowledge about humans’ logarithmic perception of sound intensity, and without DMoL’s instabilities.

Original: Still, well-tuned µ -law and DMoL strategies perform well, so the main advantage of A DA C AT in the audio domain is capturing most of their performance without human-provided inductive bias.
Revised: Still, well-tuned µ -law and DMoL strategies perform well, so the main advantage of A DA C AT in the audio domain is capturing most of their performance without human-provided inductive bias.

Original: A DA C AT and the µ -law transform are complementary, and could be used in concert by learning an adaptive discretization of the heuristically transformed interval.
Revised: A DA C AT and the µ -law transform are complementary, and could be used in concert by learning an adaptive discretization of the heuristically transformed interval.

Original: 4.4 MODEL-BASED OFFLINE REINFORCEMENT LEARNING
Revised: 4.4 MODEL-BASED OFFLINE REINFORCEMENT LEARNING

Original: [Equation]

Original:  AT ’s parameterization can also  adopted in the dynamics model of a model-based planner for reinforcement learning tasks.
Revised: A DA C AT ’s parameterization can also be adopted in the dynamics model of a model-based planner for reinforcement learning tasks.

Original: We tested our parameterization with Trajectory Transformer, a recent work in model-based offline RL that uses a Transformer-based architecture to learn the dynamics model of an environment from a dataset offline, and then uses the model to plan online to produce actions for RL agents.
Revised: We tested our parameterization with Trajectory Transformer [Janner et al., 2021], a recent work in model-based offline RL that uses a Transformer-based architecture to learn the dynamics model of an environment from a dataset offline, and then uses the model to plan online to produce actions for RL agents.

Original: The original Trajectory Transformer architecture discretizes each dimension of states and actions into tokens and uses a one-hot embedding to encode them similar to how a language model handles vocabulary.
Revised: The original Trajectory Transformer architecture discretizes each dimension of states and actions into tokens and uses a one-hot embedding to encode them, similar to how a language model handles vocabulary.

Original: Since our discretization is done adaptively with context dependency, continuous inputs are more informative.
Revised: Since our discretization is done adaptively with context dependency, continuous inputs are more informative.

Original: Thus, we minimally modify the architecture by replacing the one-hot embedding layer with a linear layer that takes  in a scalar input and outputs its embedding.
Revised: Thus, we minimally modify the architecture by replacing the one-hot embedding layer with a linear layer that takes [Table] in a scalar input and outputs its embedding.

Original: This modification arguably loses some capacity since it has many fewer parameters than the original architecture.
Revised: This modification arguably loses some capacity since it has many fewer parameters than the original architecture.

Original: Yet, as we show in our experiments, the gain from the flexibility of our parameterization out-weighs the potential capacity reduction.
Revised: Yet, as we show in our experiments, the gain from the flexibility of our parameterization outweighs the potential capacity reduction.

Original: We also reduce the number of bins by a factor of 2 to match the parameter size of the output layer since A DA C AT requires 2 × more parameters than uniform and non-adaptive quantile-based discretization.
Revised: We also reduce the number of bins by a factor of 2 to match the parameter size of the output layer since A DA C AT requires 2 × more parameters than uniform and non-adaptive quantile-based discretization.

Original: We use  a smoothing coefficient of λ = 0 .
Revised: We use uniform smoothing with a smoothing coefficient of λ = 0 .

Original: 001 for the target smoothing objective.
Revised: 001 for the target smoothing objective.

Original: We also keep the planning hyperparameters the same as the original work for a fair comparison.
Revised: We also keep the planning hyperparameters the same as the original work for a fair comparison (except for one hyperparameter on action sampling).

Original: [Table]
Revised: See more details on other minor differences between our training and planning procedure compared to the original training and planning procedure in Appendix C.3.

Original: Table 4 shows the performance of the RL agent on three
Revised: Table 4 shows the performance of the RL agent on three

Original: D4RL datasets [Fu et al., 2020] under A DA C AT ’s parameterization and two discrete parameterizations used by Janner et al.
Revised: D4RL datasets [Fu et al., 2020] under A DA C AT ’s parameterization and two discrete parameterizations used by Janner et al.

Original: A DA C AT performs better than or on par with the uniform and quantile parameterizationsthat were used in the original paper, improving return by 2.9% and 5.2% on average, respectively.
Revised: A DA C AT performs better than or on par with the uniform and quantile parameterizations used in the original paper, improving return by 2.9% and 5.2% on average, respectively.

Original: This demonstrates its effectiveness in accurately modeling continuous data, and the downstream benefits of more expressive discretization.
Revised: This demonstrates its effectiveness in accurately modeling continuous data, and the downstream benefits of more expressive discretization.

Original:  RELATED WORK
Revised:  RELATED WORK

Original: Adaptive Discretization The idea of adaptive discretization has found tremondous applications in different fields such as reinforcement learning [Chow and Tsitsiklis, 1991], finite element analysis [Liszka and Orkisz, 1980] and computer graphics [Jevans and Wyvill, 1988].
Revised: Adaptive Discretization The idea of adaptive discretization has found tremendous applications in different fields such as reinforcement learning [Chow and Tsitsiklis, 1991], finite element analysis [Liszka and Orkisz, 1980] and computer graphics [Jevans and Wyvill, 1988].

Original: We bring this powerful idea into density modeling by introducing A DAC AT .
Revised: We bring this powerful idea into density modeling by introducing A DAC AT .

Original: Unlike most existing works on discretization that rely on heuristics and prior knowledge of the data domain [Tang and Agrawal, 2020, Ghasemipour et al., 2021], A DA C AT can be jointly optimized with the rest of the network param-
Revised: Unlike most existing works on discretization that rely on heuristics and prior knowledge of the data domain [Tang and Agrawal, 2020, Ghasemipour et al., 2021], A DA C AT can be jointly optimized with the rest of the network parameters and learns to adaptively discretize. Bhat et al. [2021] (AdaBins) is closest to our work. Though the idea is similar,

Revised: AdaBins is different from A DA C AT in several important aspects.

Revised: AdaBins parameterizes the bin widths directly while

Original:  Efficient softmax Language models often have vocabulary sizes of 10k-100k tokens.
Revised: A DA C AT parameterizes the unnormalized log bin widths. AdaBins uses additional regularization loss that encourages the bin centers to be close to the data while A DA C AT does not need any additional regularization. AdaBins is primarily used in depth prediction in vision, whereas our work focuses on generative modeling across multiple domains. Efficient softmax Language models often have vocabulary sizes of 10k-100k tokens.

Original: Computing the softmax normalizer for such a large vocabular can be expensive, motivating more efficient softmax variants, surveyed by Ruder [2016].
Revised: Computing the softmax normalizer for such a large vocabulary can be expensive, motivating more efficient softmax variants, surveyed by Ruder [2016].

Original: Hierarchical softmax [Morin and Bengio, 2005] is one variant that groups tokens in a tree structure, and does not rely on data being ordinal.
Revised: Hierarchical softmax [Morin and Bengio, 2005] is one variant that groups tokens in a tree structure, and does not rely on data being ordinal.

Original: However, unlike A DA C AT , hierarchical softmax groups a fixed, discrete vocabulary, rather than supporting continuous intervals.
Revised: However, unlike A DA C AT , hierarchical softmax groups a fixed, discrete vocabulary, rather than supporting continuous intervals.

Original: Image density estimation Until recently [Kingma et al., 2021], autoregressive models were the-state-of-the-art on image density estimation benchmarks, and are still widely used.
Revised: Image density estimation Until recently [Kingma et al., 2021], autoregressive models were the-state-of-the-art on image density estimation benchmarks, and are still widely used.

Original: Order agnostic models Uria et al.
Revised: Order agnostic models Uria et al.

Original: [2014] improve the flexibility of autoregressive models in downstream tasks like inpainting and outpainting [Jain et al., 2020], but do not change the form of the conditional distribution.
Revised: [2014] improve the flexibility of autoregressive models in downstream tasks like inpainting and outpainting [Jain et al., 2020], but do not change the form of the conditional distribution.

Original: Audio synthesis Likelihood based models are popular in text-to-speech.
Revised: Audio synthesis Likelihood-based models are popular in text-to-speech.

Original: Efficiency is a key concern, motivating parameter-efficient conditional distributions.
Revised: Efficiency is a key concern, motivating parameter-efficient conditional distributions.

Original: Paine et al.
Revised: Paine et al.

Original: [2016] cache intermediate WaveNet activations for improved speed, and van den Oord et al.
Revised: [2016] cache intermediate WaveNet activations for improved speed, and van den Oord et al.

Original: [2017] distill WaveNet into a parallel flow.
Revised: [2017] distill WaveNet into a parallel flow.

Original: Tacotron [Wang et al., 2017] performed end-to-end speech synthesis from text with a WaveNet vocoder, so our work can be applied to text-to-speech (TTS) system.
Revised: Tacotron [Wang et al., 2017] performed end-to-end speech synthesis from text with a WaveNet vocoder, so our work can be applied to text-to-speech (TTS) systems.

Original: Other approaches include Flow-based models, diffusion [Kong et al., 2021, Chen et al., 2020] and GANs [Kumar et al., 2019].
Revised: Other approaches include Flow-based models, diffusion [Kong et al., 2021, Chen et al., 2020] and GANs [Kumar et al., 2019].

Original:  LIMITATIONS
Revised:  LIMITATIONS

Original: One fundamental limitation of A DA C AT is that it can never model more modes than the number of the bins it has.
Revised: One fundamental limitation of A DA C AT is that it can never model more modes than the number of bins.

Original: One possible solution is to add flow steps [Dinh et al., 2016, Rezende and Mohamed, 2015] to transform the space and run the autoregressive model with A DA C AT in the transformed space.
Revised: One possible solution is to add flow steps [Dinh et al., 2016, Rezende and Mohamed, 2015] to transform the space and run the autoregressive model with A DA C AT in the transformed space.

Original: It is possible that using A DA C AT in a transformed or latent space can lead to a more parameter efficient way of modeling continuous distributions.
Revised: It is possible that using A DA C AT in a transformed or latent space can lead to a more parameter-efficient way of modeling continuous distributions.

Original: Another limitation of A DA C AT is that the parameterization is discontinuous, with more discontinuities as the number of bins increases.
Revised: Another limitation of A DA C AT is that the parameterization is discontinuous, with more discontinuities as the number of bins increases.

Original: We largely resolved the issue using target smoothing.
Revised: We largely resolved the issue using target smoothing.

Original: It is possible to use non-uniform density within each bin like spline parameterizations [Durkan et al., 2019], but our preliminary experiments suggested that these are challenging to optimize in practice.
Revised: It is possible to use non-uniform density within each bin like spline parameterizations [Durkan et al., 2019], but our preliminary experiments suggested that these are challenging to optimize in practice.

Original:  CONCLUSION
Revised:  CONCLUSION

Original: Likelihood-based autoregressive generative models can estimate complex distributions of high-dimensional real-world data, but often struggle with efficiency and can be difficult to expressively parameterize on continuous data.
Revised: Likelihood-based autoregressive generative models can estimate complex distributions of high-dimensional real-world data, but often struggle with efficiency and can be difficult to expressively parameterize on continuous data.

Original: In this paper, we presented A DA C AT , a flexible, efficient, multimodal parameterization of 1-D conditionals with applications in autoregressive models.
Revised: In this paper, we presented A DA C AT , a flexible, efficient, multimodal parameterization of 1-D conditionals with applications in autoregressive models.

Original: We demonstrated the effectiveness of A DA C AT on a diverse set of tasks: image generation, tabular density estimation, audio vocoding and model-based planning.
Revised: We demonstrated the effectiveness of A DA C AT on a diverse set of tasks: image generation, tabular density estimation, audio vocoding and model-based planning.

Original: A DA C AT improves density estimation and downstream task performance over uniform discretization, and is competitive with other hand-engineered discretizers. 
Revised: A DA C AT improves density estimation and downstream task performance over uniform discretization, and is competitive with other hand-engineered discretizers. Acknowledgements We would like to thank Michael Janner, Xinyang Geng and

Revised: Hao Liu for helpful discussions at the early stage of this work.

Revised: We also thank Sergey Levine for insightful feedback on the offline RL experiments.

Revised: We thank Katie Kang, Dibya

Revised: Ghosh, and other members of the RAIL lab for feedback on paper drafts.

Revised: Qiyang Li is supported by Berkeley Fellowship.

Revised: A WHY IS TRAINING WITH THE NON-SMOOTHED LOSS UNSTABLE?

Revised: We hypothesize that the training instability comes from the biased gradients from the non-smoothed loss.

Revised: We demonstrate where this bias could come from in a simple 1-D example below by analyzing the gradient of the negative log-likelihood (NLL) and the gradient of the Monte-Carlo approximation of the NLL:

Revised: [Equation]

Revised: The first gradient ∇ θ L ll is the correct gradient that we hope to approximate with ∇ θ ˆ L ll during training.

Revised: However, due to the non-differentiability of p θ ( x d ) with respect to θ at the boundary of each uniform mixture component, the following statement is not generally true:

Revised: [Equation]

Revised: We provide a simple example where the gradient approximation is biased.

Revised: Let p data be a uniform distribution in [0 ., 1 . ) and p θ be parameterized by the AdaCat distribution with w 1 = w 2 = 0 .

Original: {+5 , h 
+}
Revised: {+5 , h 
+}

Revised: =

Revised: h

Revised: =

Revised: 5 (2 components).

Revised: Recall that θ =

Revised: [ ψ 1 , ψ 2 , ϕ 1 , ϕ 2 ] , which produces w 1 , w 2 , h 1 , hthrough softmax.

Revised: The correct gradient and the expectation of the approximated gradient can be computed analytically as follows:

Revised: [Equation]

Revised: + w

Revised: [Equation]

Revised: log (cid:18) hw(cid:19)(cid:21) = − (cid:34) log (cid:18)

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: (cid:19) ∇

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: log (cid:18)

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: (cid:19) ∇ θ w

Revised: [Equation]

Revised: + w

Revised: [Equation]

Revised: ∇

Revised: [Equation]

Revised: log (cid:18)

Revised: hw

Revised: [Equation]

Revised: (cid:19)

Revised: + w

Revised: [Equation]

Revised: θ log (cid:18) hw

Revised: [Equation]

Revised: (cid:19) (cid:35)

Revised: E

Revised: [Equation]

Revised: ] =

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: (cid:19)

Revised: +

Revised: I {

Revised: [Equation]

Revised: [Equation]

Revised: (cid:18)

Revised: [Equation]

Revised: w

Revised: [Equation]

Revised: (cid:19) (cid:35) =

Revised: − E x (cid:34)

Revised: I { x < 0 .

Revised: 5 } (cid:35) ∇

Revised: [Equation]

Revised: log (cid:18) hw(cid:19)

Revised: +

Revised: E x (cid:34)

Revised: I { x ≥ 0 .

Revised: 5 } (cid:35) ∇

Revised: [Equation]

Revised: log (cid:18) hw(cid:19) = − w 1 ∇ θ log (cid:18) hw(cid:19) − w

Revised: [Equation]

Revised: ∇ θ log (cid:18) hw(cid:19)

Revised: The correct gradient (first equation) contains two more terms (that involves ∇ θ w 1 and ∇ θ w 2 ) than the expectation of the sample gradient (second equation).

Revised: B INTERPRETATION OF THE SMOOTHED OBJECTIVE

Revised: The gradient of our smoothed objective is an unbiased gradient estimator of the negative log-likelihood of the smoothed data distribution under the model.

Revised: We state the connection in more precise terms (with m = 1 for readability):

Revised: Claim B

Revised: [Equation]

Revised: [Equation]

Revised: where the smoothed data distribution is defined as

Revised: [Equation]

Revised: Before we move on to the proof

Revised: [Equation]

Revised: we first check that the smoothed data distribution ˜ p data (˜ x ) is in fact a valid distribution:

Revised: [Equation]

Revised: Proof.

Revised: [Equation]

Revised: We would like to emphasize that the reason we could move the ∇ θ operator outside the expectation is because the integral (cid:82) ˜ x ζ (˜ x | x ) log p θ (˜ x ) d ˜ x is differentiable with respect to θ : we analytically evaluate the integral using the cumulative density function (Equation 8).

Revised: This is in contrast to the non-smoothed objective, which uses the discontinuous probability density function as analyzed in the previous section.

Revised: Furthermore, as long as the smoothing kernel ζ has small bandwidth ( λ ), ˜ p data would be close to p data .

Revised: As a result, the gradient of the smoothed objective would closely approximate the true objective that needs to be optimized E x ∼ p data [log p θ ( x )] .

Revised: C EXPERIMENT DETAILS

Revised: To complete all the experiments we did in this paper, we approximatedly used 200 hours of machine time on

Revised: NVIDIA DGX-1 machine with Tesla V100.

Revised: The machine has 8 Tesla V100 GPUs and 2x 20-core Intel(R) Xeon(R)

Revised: CPU E5-2698 v4 @

Revised: 2.20GHz (see more detailed specs here https://images.nvidia.com/content/pdf/ dgx1-v100-system-architecture-whitepaper. pdf ).

Revised: We also provide additional details of our hyperparameters for each task below.

Revised: C.1 TABULAR DATASETS

Revised: We provide the parameters used for A DA C AT and the uniform baseline in Table 5 and 6.

Revised: Note that the parameters are different because they are individually tuned for each method via a grid search.

Revised: Other shared hyperparameters are provided in Table 7.

Revised: [Table]

Revised: H is the number of hidden neurons used in each four-layer network.

Revised: b is the number of Fourier feature pairs used (see

Revised: Section 4.1), λ is the width of the smoothing distribution, k is the number of mixture components used, and n is the batch size.

Revised: [Table]

Revised: C.2 MNIST

Revised: For MNIST image generation experiments, we mostly borrow the default training hyperparameters from the minGPT repo ( github.com/karpathy/minGPT ) except for the batch size and the size of the network (number of layers and number of heads).

Revised: See Table 8 below.

Revised: [Table]

Revised: [Table]

Revised: C.3 OFFLINE RL

Revised: We used the same hyperparameters as used in the official trajectory transformer codebase ( github.com/jannerm/ trajectory-transformer ).

Revised: The only differences aside from the embedding layer change lie in the implementation details of trajectory sampling:

Revised: Mid-Point Sampling vs. Uniform

Revised: Sampling The original trajectory transformer codebase implements mid-point sampling where a bin is first sampled with the corresponding predicted probability.

Revised: Then, the mid-point of the bin is used as the sampling result.

Revised: This is different from our implementation where we use a uniform sampling over the selected bin since we no longer treat states and actions as discretized tokens.

Revised: Action Sampler and Observation Sampler

Revised: The original trajectory transformer codebase employs different sampling strategies for actions and observations.

Revised: In particular, the observation is decoded greedily by always choosing the most probable bin and the action is decoded by sampling from the top 40% bins ( cdf_act=0.6 ).

Revised: We follow the same strategies with the only difference that we use uniform sampling within the bin whereas the original code uses mid-point sampling (see the discussion in the previous paragraph).

Revised: In addition, we removed the 40% cut-off ( cdf_act=0 ) for the Hopper-Medium task as we found it to improve the planning performance.

Revised: Exponential Moving Average We also planned with the exponential moving average (with a coefficient of 0 . 995 ) of the parameters for the HalfCheetah-Medium and Walker2d-Medium tasks to help reduce the variance of the planning results.

Revised: We did not use it for the

Revised: Hopper-Medium task because we did not find it to help.

Revised: C.4 WAVENET

Revised: Our WaveNet experiment code closely follows the implementation of Yamamoto [2019], available at github.

Revised: com/r9y9/wavenet_vocoder .

Revised: The model is a

Revised: WaveNet audio decoder based on dilated convolutions and conditioned on ground-truth Mel-spectrograms extracted from the audio files of the LJSpeech 1.1 dataset.

Revised: Audio is sampled at 22050 Hz.

Revised: Our µ -law baseline uses a categorical output distribution with an n -bit µ -law quantized waveform as input (8-bit for 256 bins).

Revised: Waveform samples are represented through one-hot vectors.

Revised: For Mixture of Logistics (MoL) and Gaussian conditionals, the input waveform is represented as 16-bit raw audio.

Revised: WaveNet is optimized with

Revised: Adam at a learning rate of 0.001 that halves every 200K steps.

Revised: Checkpoints are averaged during training with an exponential moving average, with a decay rate of 0.9999.

Revised: The

Revised: EMA model is used for evaluation.

Revised: For AdaCat conditionals, the model is trained with the uniformly smoothed objective with λ = 0 . 0001 , 0 .

Revised: 00005 or 0 .