11affiliationtext:  University of California Berkeley
Berkeley, CA, USA

#  AdaCat: Adaptive Categorical Discretization for Autoregressive Models

Qiyang Li Ajay Jain Pieter Abbeel

###### Abstract

Autoregressive generative models can estimate complex continuous data
distributions, like trajectory rollouts in an RL environment, image
intensities, and audio. 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. The advantage is that the
categorical distribution can easily express multiple modes and are
straightforward to optimize. However, such approximation cannot express sharp
changes in density without using significantly more bins, making it parameter
inefficient. We propose an efficient, expressive, multimodal parameterization
called Adaptive Categorical Discretization (AdaCat). AdaCat discretizes each
dimension of an autoregressive model adaptively, which allows the model to
allocate density to fine intervals of interest, improving parameter
efficiency. AdaCat generalizes both categoricals and quantile-based
regression. AdaCat is a simple add-on to any discretization-based distribution
estimator. In experiments, AdaCat improves density estimation for real-world
tabular data, images, audio, and trajectories, and improves planning in model-
based offline RL.

##  1 Introduction

Figure 1: In the 2-D two-spirals dataset, 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  AdaCat closely approximates the target distribution
with the same number of bins. Figure 2: AdaCat learns how to discretize the
support of continuous distributions for expressive, parameter efficient
density estimation, and generalizes other discretization strategies like
categoricals with equal bin widths (top). The flexibility afforded by adaptive
discretization allows closer approximations of target densities, such a
mixture of two Gaussians with different scales (bottom).

Deep generative models estimate complex, high-dimensional distributions from
samples. Autoregressive models like NADE [Larochelle and Murray, 2011, 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. Each of these conditionals
$p(x^{t}|x^{1},x^{2}\ldots x^{t-1})$ is parameterized by a neural network
mapping from a subset of observed variables to logits over the next dimension.
For discrete data like language tokens, the conditional takes the form of a
categorical distribution. Categorical distributions are relatively easy to
optimize, flexible and can easily express multimodal distributions as each
bin’s logit is independently predicted.

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}$. The
categorical does not exploit this ordering, instead separately predicting each
bin. Categorical distributions also scale poorly when encoding highly precise
data like agent trajectories, tabular datasets and audio [Oord et al., 2016].
Auditory quality degrades if the waveform is quantized to less than 8-16 bits
(256-65k intensity levels). Control applications often need high precision as
well. Unfortunately, categorical likelihood degrades rapidly at high
quantization levels. The uniformly discretized model in Figure 1 has a
negative log-likelihood $-0.85$ with 16 bins, while our adaptively discretized
approach achieves NLL $-1.02$ with the same architecture and number of bins
(lower is better). 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. Halving the width of a
particular bin in the categorical distribution would require double the
parameters in the final layer of the network.

Past work tries to improve the efficiency of categoricals with hierarchical
softmax [Morin and Bengio, 2005] or quantile-based discretization [Janner et
al., 2021]. Heuristic, hand-engineered discretizations like the $\mu$-law
[Oord et al., 2016] reduce quantization error and improve perceptual quality.
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.

In this work, we propose a parameterization of 1-D conditionals that is
parameter efficient, expressive and multimodal. We propose Adaptive
Categorical Discretization (AdaCat). Based on the observation that high
precision is often only required in a small subset of a distribution’s
support, AdaCat is a distribution parameterized by a vector of interval masses
and interval widths. AdaCat is depicted in Figure 2. In contrast to
categoricals with equal bin widths, variable bin widths allow the network to
localize mass precisely without increasing precision elsewhere. Compared to
non-uniform but fixed discretizations like quantiles, AdaCat 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.

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]. In experiments, AdaCat 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.111The code for reproducing the experiments in this paper is
available at github.com/ColinQiyangLi/AdaCat. Website:
colinqiyangli.github.io/adacat.

##  2 Adaptive Categorical Discretization

###  2.1 AdaCat Distribution

The AdaCat distribution is a particular subfamily of mixtures of uniform
distributions where each mixture component has non-overlapping support. A
standard AdaCat distribution $\textsc{AdaCat}_{k}(w,h)$ has $k$ components
with a support over $[0,1)$. It is parameterized by two vectors in the
$k$-dimensional simplex: $w,h\in\Delta^{k-1}$. Thus, $w$ and $h$ are
normalized. $w$ is additionally constrained to be non-zero in all of its
elements. The probability density function (PDF) of an AdaCat distribution is
defined as,

$\displaystyle f_{w,h,k}(x)=\sum_{i=1}^{k}\left\\{\mathbb{I}_{\left[c_{i}\leq
x<c_{i}+w_{i}\right]}\frac{h_{i}}{w_{i}}\right\\}$ (1)

where $c_{i}=\sum_{j=1}^{i-1}w_{j}$ is the prefix sum of the dimensions of
parameter $w$ and $\mathbb{I}_{[\cdot]}$ is the indicator function.

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 $\frac{h_{i}}{w_{i}}$ is the density contributed by each bin.

###  2.2 Relationship with Uniform and Quantile Discretization

#### 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 equal-width bins, and
share the same bins across all dimensions of the data. This allows neural
networks to parameterize the distribution with a simple classification head
that predicts a categorical distribution over bins. AdaCat generalizes 1-D
categorical distributions with uniformly discretized support. If $w$ is set to
be $w_{i}=\frac{1}{k},\forall i$, the distribution is reduced to a categorical
distribution over $\\{0,\frac{1}{k},\frac{2}{k},\cdots,\frac{k-1}{k}\\}$
augmented with a uniform noise of magnitude $\frac{1}{k}$. Figure 2 shows how
AdaCat 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.

#### Connection to Quantile-based Discretization

AdaCat also generalizes quantile-based discretization, which discretizes a
distribution’s support by binning data into groups with equal numbers of
observed data points. If $h$ is set to be $h_{i}=\frac{1}{k},\forall i$, i.e.
the same mass in every bin, the vector $w$ can be interpreted as the
$k$-quantile of the distribution. 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 $\frac{1}{k}$ per bin of the marginal distributions of the training
set, then predicts mass $h$ with a neural network based on observed
dimensions.

###  2.3 Autoregressive Parameterization

In problems with dimension greater than 1, we use deep autoregressive models
to factorize the joint density $f(x)$ into multiple 1-D conditional AdaCat
distributions:

$\displaystyle p_{\theta}(x)$
$\displaystyle=\prod_{t=1}^{m}p_{\theta}(x^{t}|x^{<t})=\prod_{t=1}^{m}f_{w^{t},h^{t},k}(x^{t})$

For each dimension, conditioned on observed or generated values of prior
dimensions, the neural net $g_{\theta}$ outputs two unconstrained parameters
$\\{\phi^{t},\psi^{t}\\}=g_{\theta}(x^{<t})$, where
$\phi^{t},\psi^{t}\in\mathbb{R}^{k}$. The predicted $\phi$ and $\psi$
represent the unnormalized log values for $h$ and $w$ for each dimension.
These parameters are normalized independently using a softmax to satisfy the
normalization and positivity constraints:

$\displaystyle
w^{t}_{i}=\frac{\exp(\psi^{t}_{i})}{\sum_{j=1}^{k}\left[\exp(\psi^{t}_{j})\right]},\quad
h^{t}_{i}=\frac{\exp(\phi^{t}_{i})}{\sum_{j=1}^{k}\left[\exp(\phi^{t}_{j})\right]}$
(2)

where $\phi^{t}=\phi_{\theta}(x^{<t}),\psi^{t}=\psi_{\theta}(x^{<t})$.

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. Adaptivity improves
expressiveness, since density can be precisely localized in regions of
interest, and the discretization can vary across data dimensions. This is
especially important for problems where the optimal discretization is not
known a priori. In the 2-D dataset shown in Figure 1, fixed discretizations
poorly express the inherent multimodality in the data, while AdaCat’s
adaptivity allows the network to shift modes of $p_{\theta}(x^{2}|x^{1})$ for
different values of the first dimension, $x^{1}$.

Figure 3: 1-D Toy Density Modeling: AdaCat optimized with the non-smoothed
objective (red) suffers from bin collapse. The non-smoothed objective shrinks
the size of most bins until they are degenerate with small support in order to
increase density at the modes. In contrast, with the smoothed objective
(blue), AdaCat closely approximates the target. The transparent curves show an
evolution of the learned density at different optimization iterations, with
more transparent ones being earlier in the optimization.

##  3 Optimizing with Analytic Target Smoothing

Autoregressive models with AdaCat conditionals can be estimated by minimizing
the Kullback–Leibler (KL)-divergence between the target distribution and the
learned density $D_{\text{KL}}(p_{\text{data}}(x)\;\|\;p_{\theta}(x))$. The KL
reduces to the empirical log-likelihood objective below, where
$x_{1},\cdots,x_{n}$ are sampled from the data distribution $p_{\text{data}}$:

$\displaystyle\hat{\mathcal{L}}_{\text{ll}}$
$\displaystyle=\frac{1}{n}\sum_{d=1}^{n}\sum_{t=1}^{m}\log
p_{\theta}(x^{t}_{d}|x^{<t}_{d})$ (3)

However, due to the discontinuous nature of the AdaCat density function, we
observed that naïvely maximizing the empirical log-likelihood encourages the
model to get trapped in poor local optima. This phenomenon is best illustrated
in 1-D, as in Figure 3. The density in red is estimated with maximum
likelihood $\hat{\mathcal{L}}_{\text{ll}}$ (3), but the bin widths shrink over
the course of optimization and reach small values. Density is overestimated in
between modes and underestimated in regions where a single uniform mixture
component needs to cover a large interval.

We provide one possible explanation for bin collapse. Rewriting the log-
likelihood with AdaCat’s PDF (1) based on a summation over bins $i=1$ to $k$,

$\displaystyle\hat{\mathcal{L}}_{\text{ll}}$
$\displaystyle=\frac{1}{n}\sum_{d=1}^{n}\sum_{t=1}^{m}\log\underbrace{\sum_{i=1}^{k}\left\\{\mathbb{I}_{\left[c_{i}^{t}\leq
x<c_{i}^{t}+w_{i}^{t}\right]}\frac{h_{i}^{t}}{w_{i}^{t}}\right\\}}_{f_{w^{t},h^{t},k}(x^{t})}$
(4)

Due to the constraint that mixture components are non-overlapping, only a
single term of the inner summation is non-zero in (4). The loss separates into
terms for $h$ and $w$,

$\displaystyle\hat{\mathcal{L}}_{\text{ll}}$
$\displaystyle=\frac{1}{n}\sum_{d=1}^{n}\sum_{t=1}^{m}\sum_{i=1}^{k}\mathbb{I}_{\left[c_{i}^{t}\leq
x<c_{i}^{t}+w_{i}^{t}\right]}\left\\{\log h_{i}^{t}-\log{w_{i}^{t}}\right\\}$
(5)

Maximizing the loss for data point $x_{d}$ pushes for higher density
$\frac{h_{i}^{t}}{w_{i}^{t}}$ when $x_{d}^{t}$ lies in for bin $i$ of the
conditional $p_{\theta}(\cdot|x_{d}^{<t})$. This density is increased by
either increasing log mass $\log h_{i}^{t}$ or decreasing log bin width $\log
w_{i}$. For uniform and heuristic discretizers, $\log w_{i}^{t}$ is fixed.
However, updating the bin width $w_{i}^{t}$ with finite step sizes can make
the data point $x_{d}^{t}$ move out of the current bin discontinuously, which
can result in biased gradient estimates (see Appendix A).

The gradient $\frac{d}{dw_{i}^{t}}\hat{\mathcal{L}}_{\text{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. 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. Instead, optimization could shrink the new bin
$w_{i+1}^{t}$ or $w_{i-1}^{t}$, repeating until a majority of the mixture
components collapse to support a small fraction of the overall interval.

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

$\displaystyle\hat{\mathcal{L}}_{s}$
$\displaystyle=\frac{1}{n}\sum_{d=1}^{n}\sum_{t=1}^{m}\mathbb{E}_{\zeta(\tilde{x}|x_{d}^{t})}\left[\log
p_{\theta}(\tilde{x}|x^{<t}_{d})\right]$ (6)
$\displaystyle=\frac{1}{n}\sum_{d=1}^{n}\sum_{t=1}^{m}\left[\int_{\tilde{x}}\zeta(\tilde{x}|x^{t}_{d})\log
p_{\theta}(\tilde{x}|x^{<t}_{d})d\tilde{x}\right]$ (7)

where $\zeta(\tilde{x}|x)$ is any smoothing density function that is centered
around $x$. 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). 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
$\zeta(\cdot|x)=\text{Unif}[x-\frac{\lambda}{2},x+\frac{\lambda}{2})$ or
$\zeta(\cdot|x)=\mathcal{N}(x,\lambda^{2})$ in all experiments, truncating on
the boundaries of the support of $x\in[0,1)$. By optimizing
$\hat{\mathcal{L}}_{s}$ with uniform target smoothing, the density in blue in
Figure 3 converges to a close approximation of the target density.

The smoothed objective might seem intractable with an integral in the inner
summation. Fortunately, the form of the conditional $\log
p_{\theta}(\tilde{x}|x^{<t}_{i})$ with AdaCat’s simple density function allows
us to evaluate the integral analytically as long as the smoothing density has
an analytic cumulative density function (CDF). If $F(\cdot)$ is the CDF of
$\zeta$, then the integral can be analytically computed as:

$\displaystyle\int_{x}$ $\displaystyle\zeta(x)\log f_{w,h,k}(x)dx$ (8)
$\displaystyle=\sum_{j=1}^{k}\left[(F(c_{j}+w_{j})-F(c_{j}))(\log h_{j}-\log
w_{j})\right]$

where we recall that $c$ is the prefix sum of $w$ as defined previously. Only
the bins that intersect with the support of the smoothing density function
contribute to this objective.

###  3.1 Relationship with denoising score matching

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

$\mathcal{L}_{\text{DSM}}(x)=\mathbb{E}_{\zeta(\tilde{x}|x)}\left\|x-\hat{x}_{\theta}(\tilde{x})\right\|^{2}_{2}$

Assuming the observation model
$p_{\theta}(\cdot|\tilde{x})=\mathcal{N}(\hat{x}_{\theta}(\tilde{x}),I)$ is a
standard Gaussian,

$\mathcal{L}_{\text{DSM}}(x)=-\mathbb{E}_{\zeta(\tilde{x}|x)}\left[\log
p_{\theta}(x|\tilde{x})\right]$

resembling (6). 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. 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].

###  3.2 Relationship with dequantization

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. [2019], Sec. 3.1). Flows suffer from a different failure
case than non-smoothed AdaCat. Continuous estimators, e.g., a mixture of Dirac
$\delta$ functions, can arbitrarily increase density on discrete training data
as the empirical distribution is supported on a set with measure zero.
Dequantization avoids the problem by adding continuous noise to observed
discrete samples [Theis et al., 2015, Hoogeboom et al., 2021]. As an example,
the following dequantized objective gives a lower bound of the log-likelihood
of discrete data sample $x$:

$\displaystyle\log p^{\text{DQ}}_{\theta}(x)$
$\displaystyle=\mathbb{E}_{\zeta(\tilde{x}|x)}\log p_{\theta}(\tilde{x})$ (9)
$\displaystyle=\int_{x}^{x+\lambda}\zeta(\tilde{x}|x)\log
p_{\theta}(\tilde{x})d\tilde{x}$
$\displaystyle\leq\log\int_{x}^{x+\lambda}\zeta(\tilde{x}|x)p_{\theta}(\tilde{x})d\tilde{x}=\log
P_{\theta}(x),$

where $\lambda$ is chosen such that $[x,x+\lambda)$ with different discrete
sample $x$ do not overlap with each other. 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 $\zeta$ rather than analytically. 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. These findings suggest
that analytic target smoothing helps with the discontinuity in the model
conditional rather than a property of the data.

##  4 Evaluation

In experiments, we evaluate the performance of autoregressive density
estimators with adaptive categorical conditional distributions for several
data modalities. 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). AdaCat outperforms uniform discretization
strategies in all settings, and is competitive with hand-engineered
conditional distributions. Beyond density estimation, our results suggest that
AdaCat can improve downstream task performance, including speech quality and
control.

###  4.1 Tabular data modeling

We compare the performance of autoregressive models with AdaCat and uniform
parameterizations on real-world tabular density estimation benchmarks, the UCI
datasets of Dua and Graff [2017]. The state-of-the-art performances on these
benchmarks are also included for reference. We use a 4-layer feed-forward
network to predict the AdaCat parameters for each dimension of the data (e.g.,
we use 6 MLPs for the POWER dataset since it has 6-dimensional data). Each
network has $500$ hidden units for all datasets except for GAS, where we use
$1000$ hidden units.

We also use a Fourier encoding of the input inspired by Tancik et al. [2020],
Kingma et al. [2021] to allow a shallow model to capture high-frequency
variations in the input. Specifically, we augment each input element $x^{t}$
with $b$ pairs of additional features:
$\\{\sin(2^{j}x^{t}),\cos(2^{j}x^{t})\\}_{j=0}^{b-1}$. We choose the feature
count $b=32$ for GAS and POWER, $b=8$ for MINIBOONE, and $b=4$ for HEPMASS.

The uniform baseline uses the same architecture except the widths of the bins
are forced to be uniform. We search over the number of bins in
$\\{100,200,300,500,1000\\}$ for both the uniform baseline and AdaCat and
select the best to report in the table. All models are trained for 400 epochs
using Adam [Kingma and Ba, 2014] with a learning rate of $0.0003$ and the
learning rate halves every 100 epochs. We use truncated Gaussian target
smoothing for AdaCat with $\lambda=0.00001$ for POWER, and $\lambda=0.0001$
for all other datasets. See more details in Appendix C.1.

Table 1 reports results. Overall, AdaCat 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, HEPMASS and MINIBOONE, respectively. Our approach also obtains
competitive performance with the state-of-the-art normalizing flow models on
GAS.

Method | POWER (m=6) | GAS (m=8) | HEPMASS (m=21) | MINIBOONE (m=43)
---|---|---|---|---
MADE [Germain et al., 2015] | $3.08$ | $-3.56$ | $20.98$ | $15.59$
MAF [Papamakarios et al., 2017] | $-0.24$ | $-10.08$ | $17.70$ | $11.75$
NAF-DDSF [Huang et al., 2018] | $-0.62$ | $-11.96$ | $15.09$ | $8.86$
TAN [Oliva et al., 2018] | $-0.48$ | $-11.19$ | $15.12$ | $11.01$
FFJORD [Grathwohl et al., 2019] | $-0.46$ | $-8.59$ | $14.92$ | $10.43$
Block NAF [De Cao et al., 2020] | $-0.61$ | $-12.06$ | $14.71$ | $8.95$
DDEs [Bigdeli et al., 2020] | $-0.97$ | $-9.73$ | $11.3$ | $6.94$
nMDMA [Gilboa et al., 2021] | $-1.78$ | $-8.43$ | $18.0$ | $18.6$
Uniform Discretization | $1.34$ | $-6.29$ | $21.37$ | $16.93$
AdaCat | $-0.56$ | $-11.27$ | $18.17$ | $14.14$
Table 1:  (Tabular Data) Test negative log-likelihood for density estimation
on UCI datasets [Dua and Graff, 2017]. We followed the same data pre-
processing pipeline as in Papamakarios et al. [2017]. AdaCat achieves
competitive performance on GAS and consistently outperforms the uniform
baseline.

###  4.2 Image density estimation

Figure 4: Test bits per dimension (bpd) on MNIST image generation task with
different output parameter count. The parameter count is in log scale. The
MNIST data is discrete with $2^{8}=256$ possible values for the intensity at
each pixel.

Parameters |  Uniform |  Adaptive Quantile | DMoL | AdaCat
---|---|---|---|---
512 | N/A | $\times$ | $0.761$ | $\mathbf{0.561}$
256 | $\mathbf{0.561}$ | $\times$ | $0.698$ | $0.573$
216 | $0.838$ | $\times$ | $0.704$ | $\mathbf{0.615}$
180 | $1.061$ | $\times$ | $0.684$ | $\mathbf{0.629}$
152 | $1.299$ | $\times$ | $0.776$ | $\mathbf{0.612}$
128 | $1.490$ | $\times$ | $0.700$ | $\mathbf{0.608}$
64 | $2.453$ | $\times$ | $0.720$ | $\mathbf{0.695}$
32 | $3.392$ | $1.276$ | $\mathbf{0.715}$ | $0.793$
Best | $\mathbf{0.561}$ | $1.276$ | $0.715$ | $\mathbf{0.561}$

Table 2:  (Image Generation) Test negative log-likelihood in bits per
dimension (bpd) on MNIST image generation task with different output parameter
count. AdaCat outperforms other baselines on most parameter counts. The
adaptive quantile baseline diverges with a parameter count higher than 32,
indicated by $\times$.

Table 4.2 compares the performance of AdaCat against existing
parameterizations on the grayscale MNIST [LeCun and Cortes, 2010] image
generation task in terms of negative log-likelihood measured in
bits/dimension. The autoregressive architecture we use for this task is a GPT-
like Transformer decoder Vaswani et al. [2017] with 4 layers and 4 heads.222We
use the implementation, training pipeline, and the default training
hyperparameters from github.com/karpathy/minGPT, and treat an image as a token
sequence with a vocabulary size of 256. We also use a smaller batch size of
$20$. See more details in Appendix C.2. 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. We compute the log
probability of the $i^{\text{th}}$ discrete target by analytically computing
the total probability mass assigned to
$\left[\frac{i}{256},\frac{i+1}{256}\right]$ in our continuous distribution,
i.e. density integrated over the interval. This corresponds to mapping our
continuous distribution from $[0,1]$ to the 256 discrete values uniformly such
that the $i^{\text{th}}$ discrete value is mapped from
$\left[\frac{i}{256},\frac{i+1}{256}\right]$.

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. 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). DMoL
uses $3k$ parameters for a $k$-component mixture model (e.g., a 256 parameter
count budget for DMoL corresponds to a $86$-component mixture model). AdaCat
uses $2k$ parameters for a $k$-component mixture model, using $k$ bins of
variable size. We examine parameter counts ranging from $32$ to $512$.

Overall, AdaCat has better performance on most parameter counts. It only
underperforms the uniform discretization at 256 parameters. We note that the
MNIST dataset is discrete with 256 classes, which means that the uniform
discretization has the optimal bin division. Therefore, we do not expect
AdaCat to be able to outperform uniform because AdaCat has effectively half of
the bins available. DMoL achieves the best performance when there are 32
parameters, but scales poorly with the number of components and underperforms
AdaCat for most settings with more than 32 parameters.

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.
However, we found that the adaptive quantile baseline is very unstable to
train. 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. The adaptive quantile baseline
outperforms uniform with 32 parameters, yet is still much less expressive than
AdaCat.

###  4.3 Audio density estimation and vocoding

Conditional dist. | Transform | Parameters | NLL (raw) $\downarrow$ | NLL ($\mu$-law $256$) $\downarrow$ | MCD $\downarrow$
---|---|---|---|---|---
Gaussian | – | 2 |  $-8.386\,890\,751\,457\,226$ $\pm$ $0.112\,123\,957\,441\,028\,8$ |  $-4.779\,327\,632\,510\,675$ $\pm$ $0.084\,373\,770\,595\,703\,72$ |  $3.083\,703\,872\,341\,475$ $\pm$ $0.015\,790\,094\,901\,568\,28$
Uniform Categorical | – | 30 |  $-3.783\,735\,063\,553\,031\,4$ $\pm$ $0.010\,705\,543\,565\,271\,626$ |  $0.627\,237\,099\,846\,973\,1$ $\pm$ $0.006\,552\,882\,891\,747\,916$ | –
Uniform Categorical |  $\mu$-law | 30 |  $-7.709\,024\,366\,355\,303$ $\pm$ $0.068\,005\,526\,933\,857\,46$ |  $-3.322\,211\,777\,553\,805$ $\pm$ $0.031\,624\,705\,329\,819\,584$ |  $17.003\,057\,930\,554\,17$ $\pm$ $0.490\,289\,184\,385\,245\,8$
DMoL, 10 components | – | 30 |  $-8.451\,060\,023\,837\,371$ $\pm$ $0.113\,886\,179\,788\,334\,97$ |  $-4.843\,496\,629\,718\,581$ $\pm$ $0.086\,227\,910\,057\,372\,83$ |  $2.999\,637\,888\,644\,914\,4$ $\pm$ $0.011\,966\,555\,044\,538\,363$
Adaptive Cat. (AdaCat) | – | 30 |  $-8.304\,251\,780\,649\,5$ $\pm$ $0.157\,773\,444\,248\,523$ |  $-3.878\,025\,790\,163\,652\,3$ $\pm$ $0.096\,287\,238\,807\,957\,18$ |  $4.871\,829\,278\,014\,099$ $\pm$ $0.043\,595\,760\,582\,669\,44$
Uniform Categorical | – | 256 |  $-6.280\,967\,525\,604\,935$ $\pm$ $0.039\,588\,908\,655\,366\,94$ |  $-0.845\,202\,480\,244\,617$ $\pm$ $0.036\,051\,270\,268\,624\,87$ | –
Uniform Categorical |  $\mu$-law | 256 |  $-8.762\,136\,326\,514\,831$ $\pm$ $0.103\,675\,197\,761\,540\,16$ |  $-4.375\,323\,221\,765\,382$ $\pm$ $0.068\,522\,433\,033\,192\,18$ |  $3.251\,568\,668\,452\,676\,6$ $\pm$ $0.012\,450\,773\,812\,721\,505$
DMoL, 85 components | – |  255† |  $-8.459\,562\,158\,165\,426$ $\pm$ $0.114\,082\,117\,922\,200\,34$ |  $-4.851\,999\,211\,201\,526$ $\pm$ $0.086\,493\,294\,653\,366\,26$ |  $3.006\,010\,470\,652\,975$ $\pm$ $0.012\,368\,450\,954\,111\,971$
Adaptive Cat. (AdaCat) | – | 256 |  $-8.372\,055\,390\,407\,677$ $\pm$ $0.101\,700\,324\,936\,775\,68$ |  $-3.985\,242\,836\,002\,713$ $\pm$ $0.065\,965\,731\,861\,305\,96$ |  $3.024\,945\,938\,650\,357\,7$ $\pm$ $0.033\,011\,902\,425\,867\,996$
Uniform Categorical | – | 512 |  $-6.921\,092\,730\,733\,378$ $\pm$ $0.051\,468\,628\,163\,777\,556$ |  $-1.443\,548\,559\,136\,500\,7$ $\pm$ $0.049\,167\,290\,753\,894\,61$ | –
Uniform Categorical |  $\mu$-law | 512 |  $-8.121\,395\,148\,324\,059$ $\pm$ $0.103\,804\,475\,798\,699\,4$ |  $-3.734\,582\,249\,953\,789$ $\pm$ $0.068\,095\,533\,491\,702\,03$ |  $1.989\,004\,025\,419\,469\,9$ $\pm$ $0.011\,102\,717\,415\,837\,168$
DMoL, 171 components | – |  513† |  $-8.455\,742\,079\,524\,459$ $\pm$ $0.114\,200\,415\,583\,688\,86$ |  $-4.848\,179\,029\,370\,971\,4$ $\pm$ $0.086\,594\,954\,553\,621\,5$ |  $3.074\,262\,611\,283\,590\,5$ $\pm$ $0.013\,312\,985\,026\,852\,051$
Adaptive Cat. (AdaCat) | – | 512 |  $-8.325\,184\,339\,430\,495$ $\pm$ $0.111\,415\,885\,906\,968\,05$ |  $-3.938\,371\,441\,060\,225\,7$ $\pm$ $0.075\,442\,567\,726\,441\,31$ |  $2.285\,745\,259\,392\,688$ $\pm$ $0.016\,842\,716\,007\,925\,206$
Table 3:  (Audio vocoding) Continuous negative log-likelihood (NLL, in
bits/dim) and waveform vocoding MCD error for WaveNet with different
parameterizations of conditional distributions on the LJSpeech dataset.
†Discretized Mixture of Logistics (DMoL) requires 3 parameters per mixture
component (weight, mean and log scale), so output parameters are approximately
matched to baselines.

Neural vocoders synthesize human-like speech expressed as a waveform,
conditioned on phonetic or frequency spectrum based features. Speech waveforms
are long sequences as audio must be sampled at a high rate for fidelity,
typically 16-24 kHz. Thus, efficient generative models are essential for
practical applications, and the output layer of the network can be a
significant fraction of the compute. WaveNet [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. Despite the conditioning information, vocoding is challenging
as WaveNet needs to reconstruct the phase of the audio frequencies. This is
done by estimating the distribution of audio in a dataset via maximum
likelihood.

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. A ground-truth Mel-spectrogram is extracted and used for
conditioning WaveNet. For baselines, we use different parameterizations of the
conditional distribution: a uniformly discretized categorical, a categorical
discretized by a hand-engineered $\mu$-law strategy that sets bin widths
logarithmically with intensity, and a discretized mixture of logistics. All
models use 24 convolutional layers and are optimized for 500k iterations with
Adam. The learning rate is initially 0.001 and is decayed by half every 200k
iterations, with batch size 8. An exponential moving average of model
parameters is used for testing. Audio is sampled at 22,050 Hz and windowed
into blocks of 1024 samples.

We evaluate the continuous negative log-likelihood (NLL) of test waveforms.
NLL is measured in bits/dimension with the waveform scaled to a $[-1,1]$
amplitude. We also measure the NLL of $\mu$-law transformed data with
$\mu=256$, which is also scaled to $[-1,1]$.

Following Chen et al. [2021], the Mel Cepstral Distance (MCD) objectively
quantifies the perceptual similarity of our synthesized audio and reference
audio based on an aligned mean squared error metric [Kubichek, 1993]. 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.

Table 3 shows results grouped by the number of parameters output by WaveNet
for each conditional. Note that AdaCat has half the bins of categorical
baselines at the same parameter count due to using two parameter vectors, $w$
and $h$. Across all settings, AdaCat achieves significantly better negative
log-likelihood than uniform discretization: $4.52$ bpd lower NLL with 30
output parameters, $2.09$ bpd with 256 parameters, and $1.41$ bpd with 512
parameters.

AdaCat is also competitive with hand-engineered quantization in the $\mu$-law
intensity space, despite not having prior knowledge about humans’ logarithmic
perception of sound intensity, and without DMoL’s instabilities. Still, well-
tuned $\mu$-law and DMoL strategies perform well, so the main advantage of
AdaCat in the audio domain is capturing most of their performance without
human-provided inductive bias. AdaCat and the $\mu$-law transform are
complementary, and could be used in concert by learning an adaptive
discretization of the heuristically transformed interval.

###  4.4 Model-based Offline Reinforcement Learning

AdaCat’s parameterization can also be adopted in the dynamics model of a
model-based planner for reinforcement learning tasks. 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. 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. Since our discretization is done
adaptively with context dependency, continuous inputs are more informative.

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. This modification arguably loses some capacity since it has many
fewer parameters than the original architecture. Yet, as we show in our
experiments, the gain from the flexibility of our parameterization outweighs
the potential capacity reduction. We also reduce the number of bins by a
factor of 2 to match the parameter size of the output layer since AdaCat
requires $2\times$ more parameters than uniform and non-adaptive quantile-
based discretization. We use uniform smoothing with a smoothing coefficient of
$\lambda=0.001$ for the target smoothing objective. We also keep the planning
hyperparameters the same as the original work for a fair comparison (except
for one hyperparameter on action sampling). 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.

Table 4 shows the performance of the RL agent on three D4RL datasets [Fu et
al., 2020] under AdaCat’s parameterization and two discrete parameterizations
used by Janner et al. [2021]. AdaCat 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. This demonstrates its
effectiveness in accurately modeling continuous data, and the downstream
benefits of more expressive discretization.

Dataset | Uniform | Quantile | AdaCat
---|---|---|---
HalfCheetah-Medium |  $44.0$ $\pm 0.31$ |  $46.9$ $\pm 0.4$ |  $\mathbf{47.8}$ $\pm\mathbf{0.22}$
Hopper-Medium |  $67.4$ $\pm 2.9$ |  $61.1$ $\pm 3.6$ |  $\mathbf{69.2}$ $\pm\mathbf{4.5}$
Walker2d-Medium |  $\mathbf{81.3}$ $\pm\mathbf{2.1}$ |  $79.0$ $\pm 2.8$ |  $79.3$ $\pm 0.8$
Table 4:  (Offline reinforcement learning) Normalized scores on three D4RL
locomotion (v2) tasks [Fu et al., 2020] using Trajectory Transformer [Janner
et al., 2021] with three different discretization methods. AdaCat
parameterization performs on par with or better than the uniform and quantile
methods used in the original paper. Both mean and standard error over 15
random seeds (5 independently trained Transformers and 3 trajectories per
Transformer) are reported, following the protocol in the original paper.

##  5 Related work

#### 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]. We bring this powerful idea into density modeling
by introducing AdaCat. 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], AdaCat 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, AdaBins is
different from AdaCat in several important aspects. AdaBins parameterizes the
bin widths directly while AdaCat parameterizes the unnormalized log bin
widths. AdaBins uses additional regularization loss that encourages the bin
centers to be close to the data while AdaCat 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. Computing the
softmax normalizer for such a large vocabulary can be expensive, motivating
more efficient softmax variants, surveyed by Ruder [2016]. Hierarchical
softmax [Morin and Bengio, 2005] is one variant that groups tokens in a tree
structure, and does not rely on data being ordinal. However, unlike AdaCat,
hierarchical softmax groups a fixed, discrete vocabulary, rather than
supporting continuous intervals.

#### 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.
Order agnostic models Uria et al. [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.

#### Audio synthesis

Likelihood-based models are popular in text-to-speech. Efficiency is a key
concern, motivating parameter-efficient conditional distributions. Paine et
al. [2016] cache intermediate WaveNet activations for improved speed, and Oord
et al. [2018] distill WaveNet into a parallel flow. 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. Other approaches
include Flow-based models, diffusion [Kong et al., 2021, Chen et al., 2021]
and GANs [Kumar et al., 2019].

##  6 Limitations

One fundamental limitation of AdaCat is that it can never model more modes
than the number of bins. 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 AdaCat in the transformed space. It is possible that
using AdaCat in a transformed or latent space can lead to a more parameter-
efficient way of modeling continuous distributions.

Another limitation of AdaCat is that the parameterization is discontinuous,
with more discontinuities as the number of bins increases. We largely resolved
the issue using target smoothing. 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.

##  7 Conclusion

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. In this paper, we presented AdaCat, a flexible, efficient, multimodal
parameterization of 1-D conditionals with applications in autoregressive
models. We demonstrated the effectiveness of AdaCat on a diverse set of tasks:
image generation, tabular density estimation, audio vocoding and model-based
planning. AdaCat 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 Hao Liu for helpful
discussions at the early stage of this work. We also thank Sergey Levine for
insightful feedback on the offline RL experiments. We thank Katie Kang, Dibya
Ghosh, and other members of the RAIL lab for feedback on paper drafts. Qiyang
Li is supported by Berkeley Fellowship.

## References

  * Bhat et al. [2021] Shariq Farooq Bhat, Ibraheem Alhashim, and Peter Wonka.  Adabins: Depth estimation using adaptive bins.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 4009–4018, 2021.
  * Bigdeli et al. [2020] Siavash A Bigdeli, Geng Lin, Tiziano Portenier, L Andrea Dunbar, and Matthias Zwicker.  Learning generative models using denoising density estimators.  _arXiv preprint arXiv:2001.02728_ , 2020.
  * Bishop [1994] Christopher M Bishop.  Mixture density networks.  1994\.
  * Chen et al. [2021] Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan.  Wavegrad: Estimating gradients for waveform generation.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=NsMLjcFaO8O.
  * Chow and Tsitsiklis [1991] CHEE-S Chow and John N Tsitsiklis.  An optimal one-way multigrid algorithm for discrete-time stochastic control.  _IEEE transactions on automatic control_ , 36(8):898–914, 1991.
  * De Cao et al. [2020] Nicola De Cao, Wilker Aziz, and Ivan Titov.  Block neural autoregressive flow.  In _Uncertainty in artificial intelligence_ , pages 1263–1273. PMLR, 2020.
  * Dinh et al. [2016] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real NVP.  _CoRR_ , abs/1605.08803, 2016.  URL http://arxiv.org/abs/1605.08803.
  * Dua and Graff [2017] Dheeru Dua and Casey Graff.  UCI machine learning repository, 2017.  URL http://archive.ics.uci.edu/ml.
  * Durkan et al. [2019] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  Neural spline flows.  _Advances in neural information processing systems_ , 32, 2019.
  * Fu et al. [2020] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning, 2020.
  * Germain et al. [2015] Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle.  Made: Masked autoencoder for distribution estimation.  In _International Conference on Machine Learning_ , pages 881–889. PMLR, 2015.
  * Ghasemipour et al. [2021] Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu.  Emaq: Expected-max q-learning operator for simple yet effective offline and online rl.  In _International Conference on Machine Learning_ , pages 3682–3691. PMLR, 2021.
  * Gilboa et al. [2021] Dar Gilboa, Ari Pakman, and Thibault Vatter.  Marginalizable density models.  _arXiv preprint arXiv:2106.04741_ , 2021.
  * Grathwohl et al. [2019] Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud.  Scalable reversible generative models with free-form continuous dynamics.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=rJxgknCcK7.
  * Ho et al. [2019] Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel.  Flow++: Improving flow-based generative models with variational dequantization and architecture design.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 2722–2730. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/ho19a.html.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Hoogeboom et al. [2021] Emiel Hoogeboom, Taco Cohen, and Jakub Mikolaj Tomczak.  Learning discrete distributions by dequantization.  In _Third Symposium on Advances in Approximate Bayesian Inference_ , 2021.  URL https://openreview.net/forum?id=a0EpGhKt_R.
  * Huang et al. [2018] Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville.  Neural autoregressive flows.  In _International Conference on Machine Learning_ , pages 2078–2087. PMLR, 2018.
  * Ito and Johnson [2017] Keith Ito and Linda Johnson.  The lj speech dataset.  https://keithito.com/LJ-Speech-Dataset/, 2017.
  * Jain et al. [2020] Ajay Jain, Pieter Abbeel, and Deepak Pathak.  Locally masked convolution for autoregressive models.  In _Conference on Uncertainty in Artificial Intelligence_ , pages 1358–1367. PMLR, 2020.
  * Janner et al. [2021] Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Jevans and Wyvill [1988] David Jevans and Brian Wyvill.  Adaptive voxel subdivision for ray tracing.  1988\.
  * Kingma et al. [2021] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  _Advances in neural information processing systems_ , 34:21696–21707, 2021.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kong et al. [2021] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro.  Diffwave: A versatile diffusion model for audio synthesis.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=a-xFK8Ymz5J.
  * Kubichek [1993] R. Kubichek.  Mel-cepstral distance measure for objective speech quality assessment.  In _Proceedings of IEEE Pacific Rim Conference on Communications Computers and Signal Processing_ , volume 1, pages 125–128 vol.1, 1993.  10.1109/PACRIM.1993.407206.
  * Kumar et al. [2019] Kundan Kumar, Rithesh Kumar, Thibault de Boissiere, Lucas Gestin, Wei Zhen Teoh, Jose Sotelo, Alexandre de Brébisson, Yoshua Bengio, and Aaron C Courville.  Melgan: Generative adversarial networks for conditional waveform synthesis.  _Advances in neural information processing systems_ , 32, 2019.
  * Larochelle and Murray [2011] Hugo Larochelle and Iain Murray.  The neural autoregressive distribution estimator.  In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, _Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics_ , volume 15 of _Proceedings of Machine Learning Research_ , pages 29–37, Fort Lauderdale, FL, USA, 11–13 Apr 2011. PMLR.  URL https://proceedings.mlr.press/v15/larochelle11a.html.
  * LeCun and Cortes [2010] Yann LeCun and Corinna Cortes.  MNIST handwritten digit database.  2010\.  URL http://yann.lecun.com/exdb/mnist/.
  * Liszka and Orkisz [1980] Tadeusz Liszka and Janusz Orkisz.  The finite difference method at arbitrary irregular grids and its application in applied mechanics.  _Computers & Structures_, 11(1-2):83–95, 1980\.
  * Morin and Bengio [2005] Frederic Morin and Yoshua Bengio.  Hierarchical probabilistic neural network language model.  In _International workshop on artificial intelligence and statistics_ , pages 246–252. PMLR, 2005.
  * Oliva et al. [2018] Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider.  Transformation autoregressive networks.  In _International Conference on Machine Learning_ , pages 3898–3907. PMLR, 2018.
  * Oord et al. [2018] Aaron Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George Driessche, Edward Lockhart, Luis Cobo, Florian Stimberg, et al.  Parallel wavenet: Fast high-fidelity speech synthesis.  In _International conference on machine learning_ , pages 3918–3926. PMLR, 2018.
  * Oord et al. [2016] Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu.  Wavenet: A generative model for raw audio.  _arXiv preprint arXiv:1609.03499_ , 2016.
  * Paine et al. [2016] Tom Le Paine, Pooya Khorrami, Shiyu Chang, Yang Zhang, Prajit Ramachandran, Mark A. Hasegawa-Johnson, and Thomas S. Huang.  Fast wavenet generation algorithm, 2016.
  * Papamakarios et al. [2017] George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  _Advances in neural information processing systems_ , 30, 2017.
  * Radford et al. [2018] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving language understanding by generative pre-training.  2018\.
  * Rezende and Mohamed [2015] Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International conference on machine learning_ , pages 1530–1538. PMLR, 2015.
  * Ruder [2016] Sebastian Ruder.  On word embeddings - Part 2: Approximating the Softmax.  http://ruder.io/word-embeddings-softmax, 2016.
  * Salimans et al. [2017] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma.  Pixelcnn++: A pixelcnn implementation with discretized logistic mixture likelihood and other modifications.  In _ICLR_ , 2017.
  * Song and Ermon [2019] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _CoRR_ , abs/1907.05600, 2019.  URL http://arxiv.org/abs/1907.05600.
  * Tancik et al. [2020] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng.  Fourier features let networks learn high frequency functions in low dimensional domains.  _Advances in Neural Information Processing Systems_ , 33:7537–7547, 2020.
  * Tang and Agrawal [2020] Yunhao Tang and Shipra Agrawal.  Discretizing continuous action space for on-policy optimization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 5981–5988, 2020.
  * Theis et al. [2015] Lucas Theis, Aäron van den Oord, and Matthias Bethge.  A note on the evaluation of generative models.  _arXiv preprint arXiv:1511.01844_ , 2015.
  * Uria et al. [2014] Benigno Uria, Iain Murray, and Hugo Larochelle.  A deep and tractable density estimator.  In _International Conference on Machine Learning_ , pages 467–475. PMLR, 2014.
  * Uria et al. [2016] Benigno Uria, Marc-Alexandre Côté, Karol Gregor, Iain Murray, and Hugo Larochelle.  Neural autoregressive distribution estimation.  _The Journal of Machine Learning Research_ , 17(1):7184–7220, 2016.
  * Van Oord et al. [2016] Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu.  Pixel recurrent neural networks.  In _International Conference on Machine Learning_ , pages 1747–1756. PMLR, 2016.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _CoRR_ , abs/1706.03762, 2017.  URL http://arxiv.org/abs/1706.03762.
  * Vincent [2011] Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011\.
  * Wang et al. [2017] Yuxuan Wang, RJ Skerry-Ryan, Daisy Stanton, Yonghui Wu, Ron J. Weiss, Navdeep Jaitly, Zongheng Yang, Ying Xiao, Zhifeng Chen, Samy Bengio, Quoc Le, Yannis Agiomyrgiannakis, Rob Clark, and Rif A. Saurous.  Tacotron: Towards end-to-end speech synthesis, 2017.
  * Yamamoto [2019] Ryuichi Yamamoto.  Wavenet vocoder.  https://github.com/r9y9/wavenet_vocoder, 2019.

##  Appendix A Why is training with the non-smoothed loss unstable?

We hypothesize that the training instability comes from the biased gradients
from the non-smoothed loss. 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:

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{ll}}$
$\displaystyle=-\nabla_{\theta}\int_{x}p_{\mathrm{data}}(x)\log
p_{\theta}(x)dx$ $\displaystyle\nabla_{\theta}\hat{\mathcal{L}}_{\text{ll}}$
$\displaystyle=-\frac{1}{n}\sum_{d=1}^{n}\nabla_{\theta}\log
p_{\theta}(x_{d})$

The first gradient $\nabla_{\theta}\mathcal{L}_{\text{ll}}$ is the correct
gradient that we hope to approximate with
$\nabla_{\theta}\hat{\mathcal{L}}_{\text{ll}}$ during training. However, due
to the non-differentiability of $p_{\theta}(x_{d})$ with respect to $\theta$
at the boundary of each uniform mixture component, the following statement is
_not_ generally true:

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{ll}}=\mathbb{E}_{x_{1},\cdots,x_{n}\stackrel{{\scriptstyle
i.i.d.}}{{\sim}}p_{\mathrm{data}}}\left[\nabla_{\theta}\hat{\mathcal{L}}_{\text{ll}}\right]$

We provide a simple example where the gradient approximation is biased. Let
$p_{\mathrm{data}}$ be a uniform distribution in $[0.,1.)$ and $p_{\theta}$ be
parameterized by the AdaCat distribution with
$w_{1}=w_{2}=0.5,h_{1}=h_{2}=0.5$ (2 components). Recall that
$\theta=[\psi_{1},\psi_{2},\phi_{1},\phi_{2}]$, which produces
$w_{1},w_{2},h_{1},h_{2}$ through softmax. The correct gradient and the
expectation of the approximated gradient can be computed analytically as
follows:

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{ll}}$
$\displaystyle=-\nabla_{\theta}\left[\int_{x}p_{\text{data}}(x)\log
p_{\theta}(x)dx\right]$
$\displaystyle=-\nabla_{\theta}\left[w_{1}\log\left(\frac{h_{1}}{w_{1}}\right)+w_{2}\log\left(\frac{h_{2}}{w_{2}}\right)\right]$
$\displaystyle=-\Bigg{[}\log\left(\frac{h_{1}}{w_{1}}\right)\nabla_{\theta}w_{1}+\log\left(\frac{h_{2}}{w_{2}}\right)\nabla_{\theta}w_{2}+$
$\displaystyle
w_{1}\nabla_{\theta}\log\left(\frac{h_{1}}{w_{1}}\right)+w_{2}\nabla_{\theta}\log\left(\frac{h_{2}}{w_{2}}\right)\Bigg{]}$
$\displaystyle\mathbb{E}_{x_{1:n}}\left[\nabla_{\theta}\hat{\mathcal{L}}_{\text{ll}}\right]$
$\displaystyle=-\mathbb{E}_{x}\left[\nabla_{\theta}\log p_{\theta}(x)\right]$
$\displaystyle=-\mathbb{E}_{x}\Bigg{[}\mathbb{I}\left\\{x<0.5\right\\}\nabla_{\theta}\log\left(\frac{h_{1}}{w_{1}}\right)+$
$\displaystyle\mathbb{I}\left\\{x\geq
0.5\right\\}\nabla_{\theta}\log\left(\frac{h_{2}}{w_{2}}\right)\Bigg{]}$
$\displaystyle=-\mathbb{E}_{x}\Bigg{[}\mathbb{I}\left\\{x<0.5\right\\}\Bigg{]}\nabla_{\theta}\log\left(\frac{h_{1}}{w_{1}}\right)+$
$\displaystyle\mathbb{E}_{x}\Bigg{[}\mathbb{I}\left\\{x\geq
0.5\right\\}\Bigg{]}\nabla_{\theta}\log\left(\frac{h_{2}}{w_{2}}\right)$
$\displaystyle=-w_{1}\nabla_{\theta}\log\left(\frac{h_{1}}{w_{1}}\right)-w_{2}\nabla_{\theta}\log\left(\frac{h_{2}}{w_{2}}\right)$

The correct gradient (first equation) contains two more terms (that involves
$\nabla_{\theta}w_{1}$ and $\nabla_{\theta}w_{2}$) than the expectation of the
sample gradient (second equation).

##  Appendix B Interpretation of the Smoothed Objective

The gradient of our smoothed objective is an _unbiased_ gradient estimator of
the negative log-likelihood of the _smoothed_ data distribution under the
model. We state the connection in more precise terms (with $m=1$ for
readability):

######  Claim B.1.

$\displaystyle\mathbb{E}_{x_{1},\cdots,x_{n}\stackrel{{\scriptstyle
i.i.d.}}{{\sim}}p_{\mathrm{data}}}\left[\nabla_{\theta}\hat{\mathcal{L}}_{s}\right]=\nabla_{\theta}\mathbb{E}_{\tilde{x}\sim\tilde{p}_{\mathrm{data}}}\left[\log
p_{\theta}(\tilde{x})\right]$

where the smoothed data distribution is defined as

$\displaystyle\tilde{p}_{\mathrm{data}}(\tilde{x})=\int_{x}\zeta(\tilde{x}|x)p_{\mathrm{data}}(x)dx$

Before we move on to the proof, we first check that the smoothed data
distribution $\tilde{p}_{\mathrm{data}}(\tilde{x})$ is in fact a valid
distribution:

$\displaystyle\int_{\tilde{x}}\tilde{p}_{\mathrm{data}}(\tilde{x})d\tilde{x}$
$\displaystyle=\int_{\tilde{x}}\int_{x}\zeta(\tilde{x}|x)p_{\mathrm{data}}(x)dxd\tilde{x}$
$\displaystyle=\int_{x}\left(\int_{\tilde{x}}\zeta(\tilde{x}|x)d\tilde{x}\right)p_{\mathrm{data}}(x)dx$
$\displaystyle=\int_{x}p_{\mathrm{data}}(x)dx$ $\displaystyle=1.$

###### Proof.

$\displaystyle\mathbb{E}_{x_{1:n}\stackrel{{\scriptstyle
i.i.d}}{{\sim}}p_{\mathrm{data}}(x)}\left[\nabla_{\theta}\hat{\mathcal{L}}_{s}\right]$
$\displaystyle\qquad=\mathbb{E}_{x_{1:n}}\left[\nabla_{\theta}\left[\frac{1}{n}\sum_{d=1}^{n}\int_{\tilde{x}_{d}}\zeta(\tilde{x}_{d}|x_{d})\log
p_{\theta}(\tilde{x}_{d})d\tilde{x}_{d}\right]\right]$
$\displaystyle\qquad=\nabla_{\theta}\mathbb{E}_{x_{1:n}}\left[\frac{1}{n}\sum_{d=1}^{n}\int_{\tilde{x}_{d}}\zeta(\tilde{x}_{d}|x_{d})\log
p_{\theta}(\tilde{x}_{d})d\tilde{x}_{d}\right]$
$\displaystyle\qquad=\nabla_{\theta}\mathbb{E}_{x\sim
p_{\mathrm{data}}(x)}\left[\int_{\tilde{x}}\zeta(\tilde{x}|x)\log
p_{\theta}(\tilde{x})d\tilde{x}\right]$
$\displaystyle\qquad=\nabla_{\theta}\int_{x}\int_{\tilde{x}}p_{\mathrm{data}}(x)\zeta(\tilde{x}|x)\log
p_{\theta}(\tilde{x})d\tilde{x}dx$
$\displaystyle\qquad=\nabla_{\theta}\int_{\tilde{x}}\left[\int_{x}p_{\mathrm{data}}(x)\zeta(\tilde{x}|x)dx\right]\log
p_{\theta}(\tilde{x})d\tilde{x}$
$\displaystyle\qquad=\nabla_{\theta}\int_{\tilde{x}}\tilde{p}_{\mathrm{data}}(\tilde{x})\log
p_{\theta}(\tilde{x})d\tilde{x}$
$\displaystyle\qquad=\nabla_{\theta}\mathbb{E}_{\tilde{x}\sim\tilde{p}_{\mathrm{data}}}\left[\log
p_{\theta}(\tilde{x})\right]$

∎

We would like to emphasize that the reason we could move the $\nabla_{\theta}$
operator outside the expectation is because the integral
$\int_{\tilde{x}}\zeta(\tilde{x}|x)\log p_{\theta}(\tilde{x})d\tilde{x}$ is
differentiable with respect to $\theta$: we analytically evaluate the integral
using the cumulative density function (Equation 8). This is in contrast to the
non-smoothed objective, which uses the discontinuous probability density
function as analyzed in the previous section. Furthermore, as long as the
smoothing kernel $\zeta$ has small bandwidth ($\lambda$),
$\tilde{p}_{\mathrm{data}}$ would be close to $p_{\mathrm{data}}$. As a
result, the gradient of the smoothed objective would closely approximate the
true objective that needs to be optimized $\mathbb{E}_{x\sim
p_{\mathrm{data}}}[\log p_{\theta}(x)]$.

##  Appendix C Experiment Details

To complete all the experiments we did in this paper, we approximatedly used
200 hours of machine time on NVIDIA DGX-1 machine with Tesla V100. The machine
has 8 Tesla V100 GPUs and 2x 20-core Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz
(see more detailed specs here
https://images.nvidia.com/content/pdf/dgx1-v100-system-architecture-
whitepaper.pdf). We also provide additional details of our hyperparameters for
each task below.

###  C.1 Tabular Datasets

We provide the parameters used for AdaCat and the uniform baseline in Table 5
and 6. Note that the parameters are different because they are individually
tuned for each method via a grid search. Other shared hyperparameters are
provided in Table 7.

Dataset | $b$ | $H$ | $\lambda$ | $k$ | $n$
---|---|---|---|---|---
POWER | $32$ | $500$ | $10^{-5}$ | $100$ | $10000$
GAS | $32$ | $1000$ | $10^{-4}$ | $1000$ | $10000$
HEPMASS | $4$ | $500$ | $10^{-4}$ | $100$ | $10000$
MINIBOONE | $32$ | $500$ | $10^{-4}$ | $100$ | $1000$
Table 5: Hyperparameters of AdaCat for UCI Datasets. $H$ is the number of hidden neurons used in each four-layer network. $b$ is the number of Fourier feature pairs used, $\lambda$ is the width of the smoothing distribution, $k$ is the number of mixture components used, and $n$ is the batch size. Dataset | $b$ | $H$ | $\lambda$ | $k$ | $n$
---|---|---|---|---|---
POWER | $32$ | $500$ | $10^{-5}$ | $100$ | $10000$
GAS | $32$ | $1000$ | $10^{-4}$ | $200$ | $10000$
HEPMASS | $4$ | $500$ | $10^{-4}$ | $300$ | $10000$
MINIBOONE | $32$ | $500$ | $10^{-4}$ | $100$ | $1000$
Table 6: Hyperparameters of the Uniform Baseline for UCI Datasets. $H$ is the number of hidden neurons used in each four-layer network. $b$ is the number of Fourier feature pairs used, $\lambda$ is the width of the smoothing distribution, $k$ is the number of mixture components used, and $n$ is the batch size. Hyperparameter | Value
---|---
Weight Decay | $10^{-4}$
Learning Rate |  $10^{-4}$ (Halved every 100 gradient steps)
# of Epochs | $400$
# of Layers |  $4$ (3 hidden layers)
Adam Optimizer | $(\beta_{1},\beta_{2},\epsilon)=(0.9,0.999,10^{-8})$
Table 7: Other Shared Hyperparameters for UCI Datasets

###  C.2 MNIST

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). See Table 8 below.

Hyperparameter | Value
---|---
Dropout Rate | $0.1$
Learning Rate | $3\times 10^{-4}$
Weight Decay | $0.1$
Gradient Clipping | $0.1$
# of Layers | $0.1$
Adam Optimizer | $(\beta_{1},\beta_{2},\epsilon)=(0.9,0.95,10^{-8})$
# of Epochs | $100$
Embedding Size | $768$
# of Layers | $4$
# of Att. Heads | $4$
Table 8: Hyperparameters for MNIST Experiment

###  C.3 Offline RL

We used the same hyperparameters as used in the official trajectory
transformer codebase (github.com/jannerm/trajectory-transformer). The only
differences aside from the embedding layer change lie in the implementation
details of trajectory sampling:

#### Mid-Point Sampling vs. Uniform Sampling

The original trajectory transformer codebase implements mid-point sampling
where a bin is first sampled with the corresponding predicted probability.
Then, the mid-point of the bin is used as the sampling result. 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.

#### Action Sampler and Observation Sampler

The original trajectory transformer codebase employs different sampling
strategies for actions and observations. 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). 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). 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.

#### 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. We did not use it
for the Hopper-Medium task because we did not find it to help.

###  C.4 WaveNet

Our WaveNet experiment code closely follows the implementation of Yamamoto
[2019], available at github.com/r9y9/wavenet_vocoder. The model is a 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.
Audio is sampled at 22050 Hz. Our $\mu$-law baseline uses a categorical output
distribution with an $n$-bit $\mu$-law quantized waveform as input (8-bit for
256 bins). Waveform samples are represented through one-hot vectors. For
Mixture of Logistics (MoL) and Gaussian conditionals, the input waveform is
represented as 16-bit raw audio. WaveNet is optimized with Adam at a learning
rate of 0.001 that halves every 200K steps. Checkpoints are averaged during
training with an exponential moving average, with a decay rate of 0.9999. The
EMA model is used for evaluation. For AdaCat conditionals, the model is
trained with the uniformly smoothed objective with $\lambda=0.0001$, $0.00005$
or $0.00001$.