TraDE: A Simple Self-Attention-Based Density Estimator
Authors: Rasool Fakoor
Pratik Anil Chaudhari
Jonas Mueller
Alex Smola
Venue:  (2021)
Keywords: density estimation
self-attention

Abstract: We present TraDE, a self-attention-based architecture for auto-regressive density estimation with continuous and discrete valued  data. Our model is trained using a penalized maximum likelihood objective, which  ensures that samples from the density estimate resemble the training data distribution. The use of self-attention means that the model need not retain conditional sufficient statistics during the auto-regressive process beyond what is needed for each covariate. On standard tabular and image data  benchmarks, TraDE produces significantly better density estimates than existing approaches such as normalizing flow estimators and recurrent auto-regressive models. However log-likelihood on held-out data only partially reflects how useful these estimates are in real-world applications. In order to systematically evaluate density estimators, we present a suite of tasks such as regression using generated samples, out-of-distribution detection, and robustness to noise in the training data and demonstrate that TraDE works well in these scenarios. 

Under review as a conference paper at ICLR 2021
TRADE: A SIMPLE SELF-ATTENTION-BASEED DENSITY ESTIMATOR
Anonymous authors Paper under double-blind review
ABSTRACT
We present TraDE, atttenttonnnbsee architecture for auto-regressive density estimation with continuous and discrete valued data. Our model is trained using penalized maximum likelihood objective, which ensures that samples from the density estimate resemble the training data distribution. The use of self-attention means that the model need not retain conditional sufficient statistics during the auto- regressive process beyond what is needed for each covariate. On standard tabular and image data benchmarks, TraDE produces significantly better density estimates than existing approaches such as normalizing flow estimators and recurrent auto- regressive models. However fog-likelihood on held-out data only partially reflects how useful these estimates are in real-world applications. In order to systematically evaluate density estimators, we present a suite of tasks such as regression using generated samples, -istr--distribution detection, and robustness to noise in the training data and demonstrate that TraDE works well in these scenarios.
INTRODUCTION
Density estimation involves estimating probability density p(x), given independent, identically distributed (iid) samples from it. This is versatile and important problem as it allows one to generate synthetic data or perform novelty and outlier detection. It is also an important subroutine in applications of graphical models. Deep neural networks are powerful function class and learning complex distributions with them is promising. This has resulted in a resurgence of interest in the classical problem of density estimation.
One of the more popular techniques for den- sity estimation is to sample data from a simple reference distribution and then to learn a (se- quence of) invertible transformations that allow us to adapt to a target distribution. Flow-based methods (Durkan et al., 2019b) employ this with great success. A more classical approach is to decompose p(a:) in iterative manner via ditional probabilities p(#:41|21...) and fit this distribution using the data (Murphy, 2013). One may even employ implicit generative models to sample from p(r) directly, perhaps without the ability to compute density estimates. This is the case with Generative Adversarial Networks (GANs) that reign supreme for image synthesis via sampling (Goodfellow et al., 2014; Karras et al., 2017).
Figure 1: TraDE is well suited to density estimation of Transformers. Left: Bumblebee (true density), Right: density estimated from data.
Implementing these above methods however requires special care, e.g., the normalizing transform requires the network to be invertible with efficiently computable Jacobian. Auto-regressive models using recurrent networks are difficult scale nigh-dimensional data due to the need store potentially high-dimensional conditional sufficient statistic (and also due to vanishing gradients). Generative models can be difficult to train and GANs lack a closed density model. Much of the current work is devoted to mitigating these issues. The main contributions of this paper include:
1. We introduce TraDE, a simple but novel auto-regressive density estimator that uses self-attention along with recurrent neural network (RNN)- input embedding to approximate arbitrary continuous and discrete conditional densities. It more flexible than contemporary architectures such those proposed by Durkan et al. (2019b;a); Kingma et al. (2016); De Cao al. (2019) for

---

Under review as a conference paper at ICLR 2021
this problem. yet it remains capable of approximating any density function. To our knowledge, this is the first adaptation of Transformer-lik architectures for continuous-valud density estimation. held-out data is the prevalent metric evaluate density estimators. However, this only provides a partial view of their performance in real-world applications. We propose a suite of experiments to systematically evaluate the performance of density estimators in downstream tasks such as classification and regression using generated samples, detection of out-of outoo-diistrbuution samples, and robustness noise in the training data.
2. Log- se-likkelhhodd
3. We provide extensive empirical evidence that TraDE substantially outperforms other density estimators on standard and additional benchmarks, along with thorough ablation experiments to dissect the empirical gains. The main feature of this work is the simplicity of our proposed method along with strong (systematically evaluated) empirical performance.
BACKGROUND AND RELATED WORK
Given a dataset {r",...,,z*" x"}
where each sample s' E IR is drawn iid from a distribution p(r). the maximun-likelihode formulation of density estimation finds a 0-parameterized distribution q with
argmax log ;
(1)
The candidate distribution q can be parameterized in a variety of ways as we discuss next.
Normalizing flows write z from which one can draw samples easily (Papamakarios et al., 2019). If this mapping is fo : 2 - 2, two distributions can be related using the determinant of the Jacobian as
transformation of samples from some base distribution P=
d dz
g( 0) :- Ps (z)
practical limitation of flow-hased models is that so must be a diffeomorphism, i.e., it is invertible and both fo and fo! are differentiable. Good performance using normalizing flows imposes nontrivial restrictions on how one can parametrize fe: it must be flexible yet invertible with Jacobian that can be computed efficiently. There are a number of techniques to achieve this, e-g., linear mappings, planar/radial flows (Rezende & Mohamed, 2015; Tabak & Tumnrr, 2013), Sylvester flows (Berg et al., 2018), coupling (Dinh al., 2014) and rego-segeee models (Larochelle & Murray, 2011). One may also compose the transformations, e.g., using monotonic mappings So in each layer (Huang et al., 2018; De Cao et al., 2019).
Auto-regressive models factorize the joint distribution product of univariate conditional distri- butions (x;; 0)= ;= II, (((x:(/1,..... 1;-1; 0). The auto-regressive approach to density estimation is straightforward and flexible there is no restriction on how each conditional distribution is modeled. Often, a single recurrent neural network (RNN) is used to sequentially estimate all conditionals with shared set of parameters (Oliva et al., 2018; Kingma al., 2016). For high- nigh-dimensional data, the challenge lies in handling the increasingly large state space æ1,- ........ I;-1 required to properly infer .z In recurrent auto-regressive models, these conditioned-upoon variables' values are stored in some representation h; which is updated via function h++1 = g(h,, This overcomes the problem of high-dimensional estimation, albeit at the expense of loss in fidelity. Techniques like masking the computational paths in a feed-forward network are popular to alleviate these problems further (Uria et al., 2016; Germain et al., 2015; Papamakarios al., 2017).
Many existing ruto-regressive algorithms are highly sensitive to the variable ordering chosen for factorizing g, and some methods must train complex ensembles over multiple orderings achieve good performance (Germain et al., 2015; Uria et al., 2014). While autoregressive models are commonly applied to natural language and time series data, this setting only involves variables that already naturally ordered (Chelba et al., 2013). In contrast, consider continuous (and discrete) density estimation of vector valued data, e.g. tabular data, where the underlying ordering and dependencies between variables is often unknown.
Generative models focus on drawing samples from the estimated distribution that look resemble the true distribution of data. There is rich history of learning explicit models from variational inference (Jordan et al., 1999) that allow both drawing samples and estimating the log-likelihood or implicit models such as Generative Adversarial Networks (GANs, see (Goodfellow et al., 2014))
2

---

Under review as a conference paper at ICLR 2021
where one may only draw samples. These have been shown to work well for natural images (Kingma & Welling, 2013) but have not obtained similar performance for tabular data. We also note the existence of many classical techniques that are less popular in deep learning, such as kernel density estimation (Silverman, 2018) and Chow-Liu trees (Chow & Liu, 1968; Choi et al., 2011).
TOOLS OF THE TRADE
Consider the S-dimensional Markov Random Field shown here, where the underlying graphical model is unknown in practice. Consider the following two orders in which to factorize the autoregressive model: (1, 3, 4, 5, 7.8) and (1,8,2, 7,3,6, 6,4, In the latter case the model becomes a simple sequence where e.g. ((,,2(11882277) = =p(#ala77 due to conditional independence. A latent variable auto-regressive model only needs to preserve the most recently encountered state in this latter ordering. In the first ordering, p(za (122) can be simplified further to p(za (T2), but we still need to carry the precise value of æ along until the end since p(rgg/a (I...))
p(1gg (21.)). This is fundamental weakness in models employing RNNs such as (Oliva et al., 2018). In practice, we may be unable to select favorable ordering for columns table (unlike for language where words are inherently ordered), especially as the underlying distribution is unknown.
3.1 VERTEX ORDERING AND SUFFICIENT STATISTICS
The above problem is also seen in sequence modeling, and Transformers were introduced to better model such long-range dependencies through self-attention (Vaswani al., 2017). A recurrent network can, in principle, absorb this information into its hidden state. In fact, Long-Short- Song-Shor--Teme Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) units were engincered specifically to store long-range dependencies until needed. Nonetheless, storing information costs parameter space. For an auto-regressive factorization where the true conditionals require one to store many variables' values for many time steps, the RNN/LSTM hidden state must be undesirably large to achieve low approximation error, which means these models must have many parameters (Collins et al., 2017). The following simple lemma formalizes this.
Lemma 1. Denote by G the graph of an undirected graphical model over random variables Ii...., Zd- Depending on the order vertices are traversed in our factorization the largest num- ber of latent variables a recurrent regreseeges model needs to store is bounded from above and below by the minimum and the maximum number of variables with cut edge of the graph G.
Proof. Given a subset of known variables S S f1... d} we want to estimate the conditional distribution of the variables on the complement C : = (1...dd d)s.s. For this we need to decompose into the Markov blanket M of C' and its remainder. By definition M consists of the variables with cut edge. Since p(=c(=s) = ((= we are done.
This problem with -ong-dependencies in auto-regressive models has been noted before, and is exacerbated in density estimation with continuous-valuee data, where values of conditioned-upon variables may need to be precisely retrieved. For instance, recent auto-regressive models employ masking to eliminate the sequential operations of recurrent models (Papamakarios et al., 2017). There are also models like Pixel RNN (Oord et al., 2016) which explicitly design a multi-scale masking mechanism suited for natural images (and also require discretization of continuous-valued data). Note that while there is a natural ordering of random variables in text/image data, variables tabular data do not follow any canonical ordering.
Applicable to both continuous and discrete valued data, our proposed TraDE model circumvents this issue by utilizing self-attention to retrieve feature values relevant for the conditioning, which does not require more parameters to retrieve the values of relevant features that happened to appear early in the auto regrereeree factorization order (Yun et al., 2019). Thus major benefit of self-attention here is its effectiveness at maintaining an accurate representation of the feature values xj for j < i when inferring x; (irrespective of the distance between and j).
3

---

Under review as a conference paper at ICLR 2021
3.2
THE ARCHITECTURE OF TRADE
TraDE is an auto-regressive density estimator and factorizes the distribution as
d
(x;;0) 0) :- II q (x.. 2.;1::.1....11::)
(2)
Here the hth univariate conditional qi conditions the feature z4 upon the features preceding it, and may be easier to model than a high-dimensional joint distribution. Parameters 0 are shared amongst conditionals. Our main observation is that auto-regressive conditionals can be accurately modeled using the attention mechanism in Transformer architecture (Vaswani et al., 2017).
Self-attention. The Transformer is a neural sequence transduction model and consists of a multi layer encoder/decoder pair. We only need the encoder for building TraDE. The encoder takes the input sequence (I1, .............. ...) and predicts the con-conniionnal Zi (r.].21..... ... ;;=1; For discrete data, each conditional is parameterized as categorical distribution. For continuous data, each conditional distribution is parametrized as a mixture of Gaussians, where the mcan/ reaam//ariance/poooottoo of each
mixture component depend on F1,..,,
(1.1/21..... x1-1;0) = Σ The N (2x R..))
(3)
with the mixture proportions k= Th, 1. All three of Tk.i Pk, and Th,s are predicted by the model as separate outputs at each feature They are parametrized by and depend on ,,,...,,;; -;
The crucial property of the Transformer's encoder is the self-attention module outputs a representation that captures correlations in different parts of its input.
nutshell, the self-attention map outputs 2 = Σj= a)p((z; where pzz,) is an embedding of the feature 2j and normalized weights aj (p(z;, p(z,)) compute the similarity between the embedding of x; and that of 2j. Self- attention therefore amounts a linear combination of the embedding of each feature with features that are more similar to getting a larger weight. We also use multi-heades attention like (Vaswani et al., 2017) which computes self- aff-attention independently for different embeddings. Self-attention is the crucial property that allows TraDE to handle long- congerange and complex correlations in the input features; effectively this eliminates the vanishing gradient problem RNNs by allowing direct connections between far away input neurons (Vaswani et al., 2017). Self-attention also enables permutation equivariance and naturally enables TraDE to be agnostic to the ordering of the features.
Masking is used to prevent Fi, Ii+ ...,, from taking part in the computation of the output qi and thereby preserve the auto- rugr-segreseive property of our density estimator. We keep residual connections, layer normalization and dropout in the encoder unchanged from the original architecture of (Vaswani et al., 2017). The final output layer of our TraDE model, is a position-wii wiie fully fully-connected layer where, for continuous data, the ith position outputs the mixing proportions Tk, means jAk.a and standard deviations oii that together specify a Gaussian mixture model. For discrete data, this final layer instead outputs the parameters of categorical distribution.
Positional encoding involves encoding the position & of feature as an additional input, and required to ensure that Transformers can identify which values correspond which inputs information that is otherwise lost in self-attention. To circumvent this issue, Vaswani et al. (2017) append Fourier position features to each input. Picking hyper- parameertssssss for the frequencies is however difficult and it does not work well either for density estimation (see Sec. 4). An alternative that we propose is to use a simple recurrent network the input to embed the input values each position. Here the time-steps of the RNN implicitly encode the positional information, and we use Gated Recurrent Unit (GRU) model to better handle long-range dependencies (Cho et al., 2014). This parallels recent findings from language modeling where (Wang al., 2019) also used an initial RNN embedding to generate inputs to the transformer. Observe that the GRU does not slow down TraDE at inference time since sampling is performed in regresegres fashion and remains Odd. Complexity of training is marginally higher but this is more than made up for by the superior performance.
Remark (Compared to other methods, TraDE is simple and effective). Our architecture for TraDE can be summarized simply using the encoder of Transformer (Vaswani et al., 2017) with appropriate masking to achieve rego-segressive dependencies with an output layer consisting of mixture of multi-variate Gaussians and an input embedding layer built using RNN. And
4

---

Under review as a conference paper at ICLR 2021
yet it is more flexible than architectures for normalizing flows without restrictive constraints on the input-output map (c.g. invertible functions, Jacobian computational costs, etc.) As compared to other regressive models, TraDE can handle long-range dependencies and does not need to permute input features during training/inferene like Germain et al. (2015); Uria et al. (2014). Finally, the objective/architecture of TraDE is general enough handle both continuous and discrete data, unlike many existing density estimators. As experiments in Sec. 4 shows these properties make TraDE very well-suited for auto-regressive density estimation.
Unlike discrete language data, tabular datasets contain many numerical values and their categorical features do not share common vocabulary. Thus existing applications of Transformers to such data remain limited. More generally, successful applications of Transformer models to antinuous-valuee data as considered here have remained rare to date (without coarse discretization as done for images in Parmar et al. (201 8)), limited to few applications in speech (Ren et al., 2019) and time-series (Li et al., 2019; Benidis et al., 2020; Wu et al., 2020).
Beyond Transformers' insensitivity to feature order when conditioning variables, they are good fit for tabular data because their lower layers model lower-order feature interactions (starting with pairwise interactions in the first layer and building up in each additional layer). This relates them to ANOVA models (Wahba et al., 1995) which only progressively blend features unlike in the fully-connected layers of feedforward neural networks.
Lemma 3. TraDE can approximate any continuous density p(z) on a compact domain x E X C R4. Proof. Consider any auto-regressive factorization of p(x). Here the ith univariate dis- tribution p(r; Ti,...,Zi-I) can be approximated arbitrarily well by a Gaussian mixture Et=1 Thi N(x;; Pk,.; k.,), assuming sufficient number of mixture components (Sorenson & Alspach, 1971). By the universal approximation capabilities of RNNs (Siegelmann & Sontag, 1995), the GRU input layer of TraDE can generate nearly any desired positional encodings that input to the initial self-attention layer. Finally, Yun al. (2019) establish that Transformer architecture with resitional-encoding proximate any continuous sequence-to-sequee -o-euueece function over compact domain. Their constructive proof of this result (Theorem 3 in Yun et al. (2019)) guarantees that, even with our use of masking to regr-eeggyyyvvely restrict information flow, the remaining pieces of the TraDE model can approximate (TTK, Mk ;)))
self-attention, and position wiise fully connected layers can ap-
the mapping {Z1,...,, 11=1))
3.3 THE LOSS FUNCTION OF TRADE
The MLE objective in (1) does not have a regularization term to deal with limited data. Furthermore, MLE-training only drives our network consider how it is modeling each conditional distribution in (2), rather than how it models the full joint distribution p(x) and the quality of samples drawn
from its joint distribution estimate g(x; 0). To encourage both, we combine the MLE objective with regularization penalty based on the Maximum Mean Discrepancy (MMD) to get the loss function of TraDE. The former ensures consistency of the estimate while the MMD term is effective in detecting obvious discrepancies when the samples drawn from the model do not resemble samples in the training dataset (details in Appendix A). The objective minimized in TraDE is a penalized likelihood:
L(0) nd
log q. Bepelll............. + MMD q(x)).
(4)
where hyper-parameter À 2 0 controls the degree of regularization and p denotes the empirical data distribution. This objective is minimized using stochastic gradient methods, where the gradient of the log- ligeelilioood term can be computed using standard back pockkpropagttion Computing the gradient of the MMD term involves differentiating the samples from qo with respect to the parameters O. For continuous- valued data, this easily done using the reparametrization trick (Kingma Welling, 2013) since we model each conditional calculate the gradient using the Gumbel softmax trick (Maddison et al., 2016). The TraDE is thus general enough to handle both continuous and discrete data, unlike many existing density estimators.
mixture of Gaussians. For categorical features,
We also show experiments in Appendix B where MMD regularization helps density estimation with limited data (beyond using other standard regularizers like weight decay and dropout); this is common in biological applications, where experiments are often too costly to be extensively replicated (Krishnaswamy et al., 2014; Chen al., 2020). Finally we note that there exist large
5

---

Under review as a conference paper at ICLR 2021
number of classical techniques, such maximum entropy and approximate moment matching techniques (Phillips et al., 2004; Altun & Smola, 2006) for regularization in density estimation; MMD is conceptually close to moment matching.
4 EXPERIMENTS
We first evaluate TraDE both qualitatively (Fig. 2 and Fig. S1 in the Appendix) and quantitatively on standard benchmark datasets (Sec. 4.1). We then present three additional ways to evaluate the performance of density estimators in downstream tasks (Sec. 4.3) along with some ablation studies (Sec. 4.2 and Appendix A.2) Details for all the experiments in this section, including pper-ppaameters,
provided in the Appendix D.
Remark 4 (Reporting log-likelihood for density estimators). The objective in (I) is a maximum likelihood objective and therefore the maximizer need not be close to p(æ). This is often ignored in the current literature and algorithms are compared based on their likelihood on held-out test data. However, the log-likes lilkeiiiihodd may be insufficient to ascertain the real- world performance of these models, e.g., in terms of verisimil- itude of the data they generate. This major motivation for to develop a complementary evaluation methodologies in Sec. 4.3. We emphasize these evaluation methods are an- other novel contribution of our paper. However, we also report log- likeliooddd results for comparison with others.
Figure 2: Qualitative evaluation on 2-dimensional datasets. We train TraDE on samples from six 2- dimensional densities and evaluate model likelihood over the entire domain by sampling a fine grid; original den- sities are shown in shown in rows 4. This setup is similar to (Nash Durkan, 2019). These distributions are highly multi-modal with complex corre- lations but TraDE learns accurate es- timate of the true density across large number examples.
4.1 RESULTS ON BENCHMARK DATASETS
& 3 and es-
timated densities
We follow the experimental setup of (Papamakarios et al., 2017) to ensure the same /aaiiing/valiaaiionnees dataset splits in our evaluation. In particular, preprocessing of all datasets is kept the same as that of (Papamakarios et al., 2017). The datasets named POWER GAS (Vergara al., 2012), HEPMASS, MINI- BOONE and BSDS300 (Martin et al., 2001) were taken from the UCI machine learning repository (Dua & Graff, 2017).
The MNIST dataset (LeCun et al., 1990) is used to evaluate TraDE on high-dimensional image-based data. We follow the variational inference literature, e.g., (Oord et al., 2016), and use the binarized version of MNIST. The datasets for anomaly detection tasks, namely Pendigits, ForestCover and Satimage-2 are from the Outlier Detection DataSets (OODS) library (Rayana, 2016). We normalized the OODS data by subtracting the per-feature mean and dividing by the standard deviation. We show the results on benchmark datasets in Ta- ble 2. There is wide diversity in the algorithms for density
Table Negative average test log-
likelihood in nats (smaller is better) on binarized MNIST.
LOG-LIKELHIOOD
VAE
82.14 + 0.07
PLANAR FLOWS
81.91 : 0.22
IAF
80.79 + 0.12
SYLVESTER
80.22 0.03
BLOCK NAF
80.71
PIXELRNN
79.20
estimation but we make effort provide complete compar- ison of known results irrespective of the specific methodology. Some methods like Neural Spline Flows (NSF) by (Durkan al., 2019b) are quite complex imple- ment; others like Masked Autoregressive Flows (MAF) (Papamakarios et al., 2017) use ensembles to estimate the density; some others like Autoregressive Energy Machines (ARM) of (Nash Durkan, 2019) average the log- likellooooo over large number of importance samples. As the table shows, TraDE obtains performance improvements over all these methods in terms of the log- log-likelihood. This performance is persistent across all datasets except MINIBOONE where TraDE is competitive although not the best. The improvement is drastic for the POWER, HEPMASS and BSDS300.
TRADE (OURS)
78.92 + 0.00
We also evaluate TraDE on the MNIST dataset in terms of the log- ligeelood test data. As Table 1 shows TraDE obtains high log- likelihood even compared to sophisticated models such as Pixel- RNN (Oord al., 2016), VAE (Kingma & Welling, 2013), Planar Flows (Rezende & Mohamed,

---

Under review as a conference paper at ICLR 2021
Table Average test log-likelihood in nats (higher better) for benchmark datasets. Entries marked with evaluate standard deviation across 3 independent runs of the algorithm; all others are mean + standard error across samples in the test datasct. TraDE achieves significantly better log-likelihood than other algorithms on all datasets except MINIBOONE.
POWER
GAS
HEPMASS
MINIBOONE
BSDS300
REAL NVP (DINH ET AL., 2016)
0.17 = 0.01
1.33:.0.1
-18.71 =: 0.02
-1.88 1 0.52
153.28 1.78
MADE MOG (GERMAIN ET AL., 2015)
0.4
8.47 0.02
+15.15
-12.27 0.47 0.44
153.71
0.28
MAF MOG (PAPAMAKARIOS BY Al..,
0.3++.001
9.59 + 0.02
-17.39
-11.68 0.44
156.36 0.28
2017)
FEJORD (GRATHWOHL AL.,, 2018)
0.46
8.59 0.00
14.92
-10.43 0.00
0.00
NAF (HUANG ET AL., 2018)
0.62
11.96 1 0.33
-15.09 0.4
.8.86 1 0.15
157.43 1 0.3
TAN (OLIVA AL., 2018)
0.6
12.06 + 0.02
-10.78
-11.01 0.48
159.8 0.07
BNAF' (DE CAO BT AL., 2019)
0.61 0.01
12.06 + 0.09
0,387
-8.95 : 0.07
157.36 0.03
NSF' (DURKAN AL., 201
0.06
13.09 0.002
0.000
157.31 0.42 0.28
AEM (NASH & DURKAN, 2019)
0.70
13,03 1 0.01
-12.85
0.00
-10.17 1: 0.26
158.71 0.14
TRADE (OURS)
0.70
13.27 0.01
-12.01 + 0.03
-9.49 + 0.13
160.01 0.02
2015), IAF (Kingma et al., 2016), Sylvester (Berg et al., 2018), and Block NAF (De Cao et al.. 2019). This is difficult dataset for density estimation because of its high dimensionality. We also show the quality of the samples generated by our model in Fig. $1.
4.2 ABLATION STUDY
Table 3: Average test log-likelihood in nats (higher is better) on benchmark datasets for ablation experiments.
We begin with an RNN trained as an auto-regressive density estimator; the performance of this basic model in Ta- ble is quite poor for all datasets. Next, we use Transformer net- work trained in auto-regressive fash- ion without position encoding; this
POWER
GAS
HEPMAASS
MINIBOONE
DSDS300
0.51
6.26
-1587
-13.13
157.20
Trandormer (without position
0.11
12.95
-1580
-22.00
134.71
encoding)
Transtormer (with position on- coding)
0.73
-1389
-2225
147.94
TraDE without MMD
0.72
13,26
12222
99.4
159,97
leads to significant improvements compared to the RNN but not on all datasets. Compare this to the third row which shows the results for Transformer with position encoding (instead of the GRU-based embedding in TraDE); this performs much better than a Transformer architecture without position encoding. This suggests that incorporating the information about the position is critical for auto-regressive models (also see Sec. 3). A large performance gain on all datasets is observed upon using a GRU-based embedding that replaces the position embedding of the Transformer.
Note that although the MMD regularization helps produce sizeable improvement for HEPMASS, this effect is not consistent across all datasets. This is akin to any other regularizer; regularizers do not generally improve performance on all datasets and models. H.MD-penalization of the MLE objective additionally helps obtain higher-fideliyy samples that are similar to those in the training dataset; this helps when few data are available for density estimation which we study in Appendix B. This is also evident in the regression experiment (Table 4) where the MSE of TraDE is significantly better than MLE-based RNNs and Transformers trained without the MMD term
4.3 SYSTEMATIC EVALUATION OF DENSITY ESTIMATORS
We propose evaluating density estimation in four canonical ways, regression using the generated samples, a two-sample test to check the quality of generated samples, out-of-distributio detection, and robustness of the density estimator to noise in the training data. This section shows that TraDE performs well on these tasks which demonstrates that it not only obtains high log- ligliiillhood on held-out data but can also be readily used for downstream tasks.
1. Regression using generated samples. We first characterize the quality of samples generated by the model, based on a regression task where Ig is regressed using data from the others (æ1,..., 1)
The procedure is as follows: first we use the training set of the HEPMASS (d 21) to fit the density estimator, and then create a synthetic dataset with both inputs & = (#).... xz-1) 1) and targets y = za sampled from the model. Two random forest regressors are fitted, one on the real data and another
7

---

Under review as a conference paper at ICLR 2021
on this synthetic data. These regressors are tested on real test data from HEPMASS. If the model synthesizes good samples, we expect that the test performance of the regressor fitted on synthetic data would be comparable to that of the regressor fitted on real data. Table 4 shows the results. Observe that the classifier trained on data synthesized by TraDE performs very similarly to the one trained on the original data. The MSE of RNN-based auto-regressive density estimator, which is higher, is provided for comparison. The MSE of standard Transformer with position embedding is much worse at 1.37.
2. Two-sample test on the generated data. A standard way to ascertain the quality of the generated samples is to perform a two-sample test (Wasserman, 2006). The idea is similar to a two-sample test (Lopez-Paz & Oquab, 2016) in the discrimina- tor of GAN: if the samples generated by the auto-regresse model are good, the discriminator should have an accuracy of 50%. We train density estimator on the training data; then fit a random forest classifier to differentiate between real validation data and synthesized validation data; then compute the accuracy of the classifier on the real test data and the synthesized test data. The accuracy using TraDE is 51 = 1% while the ac- curacy using RNN is 55 + 4% and that of standard Transformer with position embedding is much worse at 88 = 15% Standard deviation computed across different subsets of features (.1,,..,,.,, as inputs to the classifier. This suggests that samples generated by TraDE are much closer those in the real dataset than those generated by the RNN model.
Table 4: Mean squared error of regression on HEPMASS.
Real data
0.773
TraDE 0.780
RNN Transformer
0.803
1.37
2 (x1,22)) 22), ..... (I'
3. Out-of-distribution detection. This is classical ap- plication of density estimation techniques where we seek to discover anomalous samples in given dataset. We follow the setup of (Oliva et al., 2018): we call a datum out-of-distribution if the likelihood of datum under the model (x;; 0) < for chosen threshold t 2 O. We can compute the average precision of detecting out-of- distribution samples by sweeping across different values of f. The results are shown in Table 5. Observe that TraDE obtains extremely good performance, of more than 0.95 average precision, on the three datasets.
Table 5: Average precision for out- of-distribution detection The results for NADE, NICE and TAN were (carefully) eye- balled from the plots of (Oliva et al., 2018).
NADE
NICE
TAN
TraDE
Pendigits
0.91
0.92
0.97
0.98
ForestCover
0.87
0.80
0.94
0.95
Satimuge-2
0.98
0.975
0.98
1.0
4. TraDE builds robustness to noisy data. Real data may contain noise and order to be useful on downstream tasks. Density estima- tion must be insensitive to such noise but the maximum likeliooo objective is sensitive to noise in the training data. Methods such as NSF (Durkan al., 2019h) MAF (Papamakarios et al., 2017) indirectly mitigate this sensitivity using permutations of the input data or masking within hidden layers but these operations are not designed to be robust to noisy data. We study how TraDE deals with this scenario. We add noise to 10% of the entries in the training data; we then fit both TraDE and NSF on this noisy data; both models are evaluated on clean test data. As Table shows, the degradation of both TraDE and NSF is about the same; the former obtains a higher log-likelihood as noted in Table 2.
Table 6: Average test log- likelihood (nats) for HEP- MASS dataset with and without additive noise in the training data.
Clean Data
Noisy Data
NSF TraDE
-14.51
-14.98
-11.98
-12.43
5 DISCUSSION
This paper demonstrates that self-attente is naturally suited to building auto-regressive models with strong performance in continuous and discrete valued density estimation tasks. Our proposed method is a universal density estimator that is simpler and more flexible (without restrictions of invertibility and tractable Jacobian computation) than architectures for normalizing flows, and also handles long-range dependencies better than other auto-regressive models based on recurrent structures. We contribute a suite of downstream tasks such as regression, out of distribution detection, and robustness to noisy data, which evaluate how useful the density estimates in real-world applications. TraDE demonstrates state-of-tt of-the---the-rr empirical results that are better than many competing approaches across several extensive benchmarks.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Yasemin Altun and Alex Smola. Unifying divergence minimization and statistical inference via convex duality, In International Conference on Computational Learning Theory, pp. 139- 153. Springer, 2006.
Konstantinos Benidis, Syama Sundar Rangapuram. Valentin Flunkert, Bernie Wang, Danielle Maddix, Caner Turkmen, Jan Gasthaus, Michael Bohlke- Seikke-Schneidrr, David Salinas, Lorenzo Stella, et al. Neural forecasting: Introduction and literature overview. arXiv preprint 2004, 10240, 2020.
Rianne van den Berg. Leonard Hasenclever, Jakub M Tomezak, and Max Welling. Sylvester normalizing flows for variational inference. arXiv preprint arXiv: 1803.080339) 2018.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv:1312.3005 2013.
Mengjie Chen, Qi Zhan, Zepeng Mu, Lili Wang, Zhaohui Zheng, Jinlin Miao, Ping Zhu. and Yang I Li. Alignment of single-cell maa-seq samples without over-correction using kernel density matching. bioRxiv, 2020. doi: 10.1101/2020.01 16.8899336.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rmn encoder-decoder for statistical machine translation. arXiv preprint arXiv: |406. 1078, 2014.
Myung Jin Choi, Vincent YF Tan, Animashree Anandkumar, and Alan S Willsky. Leaminng latent tree graphical models. Journal of Machine Learning Research, 171M11):7711-181,, 2011.
C. Chow and Cong Liu. Approximating discrete probability distributions with dependence trees. IEEE transac- tions on Information Theory, 14(3)462-4677 1968.
Jasmine Collins, Jascha Sohl-Dickstein, and David Sussillo. Capacity and trainability in recurrent neural networks. In International Conference on Learning Representations, 2017.
H. Daniels. The asymptotic efficiency a maximum likelihood estimator. Proceedings of the Fourth Berkeley Symposium on Mathematical Statistics and Probability, Volume I: Contributions to the Theory of Statistics, pp. 151-163, Berkeley, Calif., 1961. University of California Press.
Nicola De Cao, Ivan Titov, and Wilker Aziz. Block neural autoregressive flow, arXiv preprint 2019.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv: 2014.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. arXiv preprint arXiv: arXiv:
2016.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.
Richard M Dudley. Real analyxis and probubility. Chapman and Hal//RCC 2018.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Cubic-spline flows. arXiv preprint arXiv: 1et.//9060.25555 2019a.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, pp. 7509-7520, 2019b.
Robert Fortet and Edith Mourier. Convergence de répartition empirique vers répartition théorique. In Annales scientifiques de l'École Normale Supérieure, volume 70, pp. 267-285, 1953.
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. MADE: Masked autoencoder for distribution estimation. In International Conference on Machine Learning, pp. 881-889, 2015.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair. Aaron
Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems, pp. 1672-26800 2014.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamies for scalable reversible generative models. arXiv preprint Xrvi: Xrvi 1810.01366 2018.
Arthur Gretton, Kenji Fukumizu, Zaid Harchaoui, and Bharath K Sriperumbudur. A fast, consistent kernel two-sample test. In Advances neural information processing systems, pp. 673-681, 2009.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal Machine Learning Research, 13 132M3337733-77, 2012.
2012.
Sepp Hochreiter and Jürgen Schmidhuher. Long short-term memory. Neural computation, 133:1735 1997.

---

Under review as a conference paper at ICLR 2021
Chin-Wet (ein Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows, arXiv preprint arXiv:/804.0077 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatioo with gumbel-softmax, 2016.
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv: 1710.10196, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv:1412.69900 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprini arXiv:1312.6114, 2013.
Diederik P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparamcterization trick. In NIPS, pp. 2575-2583, 2015.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743-4751, 2016.
Smita Krishnaswamy, Matthew H. Spitzer, Michael Mingueneau, Sean C. Bendall, Oren Litvin, Erica Stone, Dana Pe' er, and Garry P. Nolan. Conditional density-based analysis of cell signaling in single-cell data. Science, 346(6213) 13)), 2014. ISSN 0036-8075. doi: 10. 1061cc6/cce 1250689.
Hugo Larochelle lain Murray. The neural autoregressive distribution estimator. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pp. 29-37, 2011.
Yann LeCun, Bernhard Boser. John S Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence D Jackel. Handwritten digit recognition with a baco-proongation network. In Advances in neural information processing systems, pp. 396-404, 1990.
Chun-Liang Li, Wei-Cheng Chang. Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. In Advances in Newral Information Processing Systems, pp. 2203-2 2013-2213,20 2017.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, Fox, and R. Garnett Advances Neural Information Processing Systems 32, pp. 5243-5253. Curran Associates, Inc., 2019.
Po-Ling Loh and Martin Wainwright. Regularized m-estimators with nonconvexity: Statistical algorithmic theory for local optima. In C. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 26, pp. 476-484, Curran Associates, Inc., 2013.
David Lopez-Paz and Maxime Oquab. Revisiting classifier two-sample tests. arXiv preprint 111v::6610005555. 2016.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint 16X1v:1611.00712. 2016.
D. Martin, C. Fowlkes, D. Tal, and J. Malik. A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics. In Proc. 8th Int'l Conf. Computer Vision, volume pp. 416-423, July 2001.
Alfred Müller. Integral probability metries and their generating classes of functions. Advances in Applied Probahility, 29(2):4294443, 1997.
Kevin P. Murphy. Machine learning probabilistic perspective. MIT Press, Cambridge, Mass. [u.a.], 2013.
Charlie Nash and Conor Durkan. Autoregressive energy machines. arXiv preprint arXiv:1904.05622. 2019,
Junier B Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider. Transformation autoregressive networks. arXiv preprint arXiv:1801.000 199119 2018
Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. arXiv preprint 10010066.,. 2016.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advonces in Neural Information Processing Systems, pp. 2338-2347, 2017.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayaaan Normalizing flows for probabilistic modeling and inference. arXiv preprint 49Xiv: 2019.
Niki Parmar, Ashish Vaswani. Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In ICML. 2018
10

---

Under review as a conference paper at ICLR 2021
Steven Phillips, Miroslav Dudik, and Robert Schapire. A maximum entropy approach species distribution modeling. Proceedings of the twenty-first international conference on Machine learning, pp. 83, 2004.
Shebuti Rayana. ODDS library, 2016. http://odds.c http://odds.....oonnyyyye cs.stonyyrooo.eeuu
Yi Ren, Yangjun Ruan, Xu Tan, Tao Qin, Sheng Zhao, Zhou Zhao, and Tie-Yan Liu. Fastspeech: Fast, robust and controllable text speech. In H. Wallach, H. Larochelle, Beygelzimer, E. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 3171-3180, Curran Associates, Inc., 2019.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. arXiv preprint TrXiv:1505.05777 2015.
Hava T Siegelmann and Eduardo D Sontag. On the computational power of neural nets, Journal of computer and system sciences, 10(1):332-150 1995.
Beraard W Silverman. Density estimation for statistics and data analysis. Routledge, 2018.
Harold Sorenson and Daniel L. Alspach. Recursive Bayesian estimation using Gaussian sums, Automatica, (4):65 479, 1971.
Bharath Sriperumbudur
the optimal estimation of probability measures weak strong topologies. Bernoulli, 223991883,, 2016.
Ingo Steinwart. influence the kernel on the consistency support vector machines. Journal of machine learning research, 2(N09): 2001.
Esteban Tabak and Cristina Turner. family nonparametric density estimation algorithms. Communica- tions on Pure and Applied Mathematics, 66(2):145-164, 2013.
Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator, 2014.
Benigno Uria, Maxe-Alexandre Côté, Karol Gregor, lain Murra,, and Hugo Larochelle. Neural autoregressive distribution estimation. The Journal of Machine Learning Reseurch, 17(1):7184-7220, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998 6008, 2017.
Alexander Vergara, Shankar Vembu, Tuba Ayhan, Margaret A Ryan, Margie I. Homer, and Ramón Huerta. Chemical gas sensor drift compensation using classifier ensembles. Sensors and Actuators B: Chemical, 166: 320. 2012.
Grace Wahha, Yuedong Wang. Chong Gu, Ronald Klein, Barbara Klein, et al. Smoothing spline anova for exponential families, with application the wisconsin epidemiological study of diabetic retinopathy: the 1994 neyman memorial lecture. Annals of Statistics, 23(6):1855-1888 1995.
Chenguang Wang, Mu Li, and Alexander Smola. Language models with transformers. arXiv preprint arXiv:1904. 09408, 2019.
Larry Wasserman. All of ponparametric statistics. Springer Science & Business Media, 2006.
Nco Wu, Bradley Green, Xue Ben, and Shawn O' O'Banion. Deep transformer models for time series forecasting:
The influenza prevalence case, 2020.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are Transformers
universal approximators of sequence-to-sequeee functions? International Conference on Learning Representations, 2019. 2019.

---

Under review as a conference paper at ICLR 2021
Appendix
TraDE: A Simple Self-Attention-Basel Density Estimator
A REGULARIZATION USING MAXIMUM MEAN DISCREPANCY (MMD)
An alternative to maximum likelihood estimation, and in some cases a dual to it, is to perform non-pan non-parametric moment matching (Altun & Smola, 2006). One can combine log-likelihood loss and two-sample discrepancy loss to ensure high fidelity, i.e., the samples resemble those from the original dataset.
We can test whether two distributions and supported on space X' are different using samples
drawn from each of them by finding smooth function that is large on samples drawn from p and small on samples drawn from q. If x ~ p and y q, then p = q if and only if E,g [fx) = E, (f() for all bounded continuous functions/ f on X (Lemma 9.3.2 in (Dudley, 2018)). We can exploit this result computationally by restricting the test functions to some class f E F and finding the worst test function. This leads to the Maximum Mean Discrepancy metric defined next (Fortet & Mourier, 1953; Müller, 1997; Gretton et al., 2012: Sriperumbudur et al., 2016). For a class F of functions f: X' + R., the MMD between distributions p. is
MMD[F, p. q) sup fer
(Ezep [f(x)] Eyng (f())).
(5)
is cumbersome to find the supremum over a general class of functions F to compute the MMD. We can however restrict F to be the unit ball in universal Reproducing Kernel Hilbert Space (RKHS) (Gretton et al., 2012)) with kernel k.
The MMD is a metric in this case and is given by
MMD [k: p.
E [k x')] 2 E [k(2, 8) + E [k(y, v/)l
(6)
y-y"~p
With a universal kernel (e.g. Gaussian, Laplace), MMD will capture any difference between distri- butions (Steinwart, 2001). We can easily obtain an empirical estimate of the MMD above using samples (Gretton et al., 2012).
A.1 GRADIENT OF MMD TERM
Evaluating the gradient of the MMD term (6) involves differentiating the samples from go with respect to the parameters 0. When TraDE handles continuous variables, it samples from mixture of Gaussians for each conditional hence gradient of the MMD term is done using the reparametrization trick (Kingma et al., 2015). On the other hand, gradient of the MMD term (6) is computed using the Gumbel softmax trick (Maddison et al., 2016; Jang et al., 2016) for discrete variables (i.e., binarized MNIST). The objective of TraDE is thus general enough to easily handle both continuous and discrete data distributions which is not usually the case with other methods.
A.2 WHAT IF ONLY MMD IS USED AS A LOSS FUNCTION?
In theory, MMD with a universal kernel would also produce consistent estimates (Gretton al., 2009), but maximizing likelihood ensures our estimate is statistically efficient (Daniels, 1961; Loh & Wainwright, 2013). Although two-sample discrepancy loss like MMD can help produce samples that resemble those from the original dataset, using MMD as the only objective function will not result in good density estimator given limited data. To test this hypothesis, we run experiments in which only MMD is utilized as a loss function without the maximum-likellhon term.
Note that using our MMD term with fixed kernel (without solving min-max problem to optimize the kernel) only implies that the distributions match up to the chosen kernel. In other words, the global minima of our MMD objective with a fixed kernel are not necessarily the global minima the maximum likelihood term, an additional reason for the poor NLLs as shown in Table SI. While the MLE and MMD objectives should both result in a consistent estimator of p(x) in the limit of infinite data and model model-capacity, in finite settings these objectives favor estimators with different properties. In practice a combination of both objectives yields superior results, and makes performance less dependent on the MMD kernel bandwidth.
12

---

Under review as a conference paper at ICLR 2021
Table S1: Test likelihood (higher better) using ONLY the MMD objective for training vs. TraDE.
Dataset
Only MMD
TraDE
POWER
-5.00 + 0.24
0.73 0.00
GAS
100.85 + 0.16
13.27 + 0.01
HEPMASS
-28.18 + 0.12
-12.01 0.03
MINIBOONE
-68.06
-9,49 + 0.13
BSDS300
70.70 + 0.32
160.01 0.02
B DENSITY ESTIMATION WITH FEW DATA
MMD regularization especially helps when we have few training data. Test likelihoods (higher is better) after training on sub-sampled MINIBOONE dataset are: -55.62 vs. -55.53 (100 samples) and -49.07 vs. -48.90 (500 samples) for TraDE without and with MMDD-eenliization, respectively. Density estimation from few samples common in biological applications. where experiments are often too costly to be extensively replicated (Krishnaswamy al., 2014; Chen et al., 2020).
C SAMPLES FROM TRADE TRAINED ON MNIST
We also evaluate TraDE on the MNIST dataset in terms of the log-likelihood on test data. As Table 1 shows TraDE obtains high log- li-liikli even compared to sophisticated models such as Pixel- RNN (Oord et al., 2016). This is difficult dataset for density estimation because of the high dimensionalityy. Fig. SI shows the quality of the samples generated by TraDE.
1
8
1
2
0 3
3
0
8
3
0
2 4 3
3
5
2, 4
& 8
9
4
8
5
3
8
4 3
6
1855
4 1 7
8
6
22 9 9 5 4r3044004
0
8
5
9
4 8
3
2
6
5
3
/
2
9
8
0
3
9
4
5
7
9
3
Figure $1: Samples from TraDE fitted binary MNIST
D HYPER-PARAMETERS
All models are trained for 1000 epochs with the Adam optimizer (Kingma & Ba, 2014). The MMD kernel a mixture of 5 Gaussians for all datasets, i.e. k( v) = Σ k,(2, y) where each (r, = e -113///1/// with bandwidths o, E (1,2,4,6 16) (Li et al., 2017). Each layer of TraDE consists of GRU followed by self-attention block which consists of
and fully connected feed-forward layer. A layer normalization and a residual connection is used around each of these components (i.e. sett---aad-sll--tenttiooo and fully connected) self-attention block (Vaswani et al., 2017). Multiple layers of self-attention block can be stacked together in TraDE, Finally, the output layer of TraDE outputs T Pk,i, and ok which specify a univariate Gaussian mixture model approximation of each conditional distribution in the auto-regressive factorization. Table S2 shows TraDE pprem-eareeeerss. We used a coarse random search based on validation NLL to select suitable prper-paramtter.. Note that dataset used in this work have very different number of samples and dimensions as shown in Table S3
13

---

Under review as a conference paper at ICLR 2021
Table S2: Hyper-parameters for benchmark datasets.
POWER
GAS
HEPMASS
MINIBOONE
BSDS300
MNIST
MMD coefficient A
0.2
0.1
0.1
0.4
0.2
0.1
Gaussian mixture components
150
100
100
20
100
Number of layers
8
5
Multi-head attention head
16
8
2
Hidden neurons
512
400
64
128
256
Dropout
0.2
0.3
0.1
Learning rate
3B-4
3E-4
5E-4
5E-4
5E-4
5E-4
Mini-batch size
512
512
512
64
512
16
Weight decay
IE-6
1E-6
IE-6
0
IE-6
IE-6
Gradient clipping norm
5
Gumbel softmax temperature
n/a
Table S3: Dataset information. In order to ensure that results comparable and setups and data splits are used as previous works, we closely follow the experimental setup of (Papamakarios al., 2017) for dataset splits. In particular, preprocessing of all the datasets kept the same as that of (Papamakarios et al., 2017). Note that, these setups were used on all other baselines as well.
Dataset
dimension
Training size
Validation Size
Test size
POWER
6
1659917
184435
204928
GAS
8
352174
94685
105206
HEPMASS
21
315123
35013
174987
MINIBOONE
43
29556
3284
3648
BSDS300
63
1000000
50000
250000
MNIST
784
50000
10000
10000
FovestCover
10
252499
28055
5494
Pendigits
16
5903
655
312
Satimage-2
36
5095
566
142
14


Reviews:
[{"review_id": "ztwBdPebNYP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary\nThis paper uses the Transformer architecture for density estimation. It performs well on several non-trivial synthetic datasets and standard benchmark datasets. The authors also tested the model on other tasks that rely on density estimation.\n\n## Pros\n1. Addresses the issue of variable ordering on learning auto-regressive models and long-range dependencies with a new model architecture\n1. Attempts to develop new evaluations other than log-likelihoods on practical uses of learnt density models\n1. Impressive empirical results\n1. Excellent written quality and comprehensive review of related literature.\n\n## Cons\n1. There is not much novelty. Simply borrowing the Transformer architecture does not seem sufficient for an academic venue\n1. The additional evaluation tasks are not new, and the particular instances implemented have issues to be addressed\n1. No mentioning of computational costs for training compared with benchmark methods.\n1. Signs of heavy hyper-parameter tuning.\n\n\n# Recommendation\nReject due mainly to a lack of novelty and some inadequate evaluations. I may raise the score slightly if the latter is addressed. \n\n# Issues and questions:\n1. Lemma 1: it's not clear why \"variables with a cut edge\" means. Do you mean \"variables the removal of which disconnects a graph\"?\n1. How can the model take the entire sequence but only return conditional distributions $p(x_i|x_{1:i-1})$? I think I may be missing something here. \n1. The equation of $alpha_j=\\dots$ on page 4 is missing a parenthesis \")\"\n1. If the model is modelling each conditional correctly, why would it not model the joint well? This is mentioned in a strange place right after Lemma 3 which says it can model any joint.\n1. All synthetic densities have a colour shift. A colour bar would help.\n1. Introducing the MMD component does not seem to improve the log-likelihood much, so seems it is not essential for log-likelihood results. Can the author simply mention that adding this improves the model on small datasets?\n\n## Questions regarding additional three empirical performance measures:\n1. Why not conduct regression on all datasets? \n1. Two-sample testing is not novel, and the results can be reported at the same place as log-likelihoods. There are also a few statistical testing paradigms which are more rigorous than simply reporting a classification acc (e.g. Liu et al 20).\n1. OOD: how is the threshold swept? A better quantity to report would be the area-under-the-curve (AUC).\n1. Robustness to noise: adding noise should definitely affect test likelihood. Why do the authors believe that a better model should be more robust to changes in the training dataset?\n1. An overarching question is: are all the benefit of the proposed method a result of introducing the MMD regulariser?\n ", "extra_scores": {}, "extra_text": {}}, {"review_id": "szzzvBb3C2d", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes TraDE, a transformer-based density estimator that is capable of learning a density of real-valued tabular data. Compared to previously proposed transformers, there are three main differences in TraDE model: 1) the output is modeled as a mixture of Gaussians, 2) maximum mean discrepancy (MMD) is added to the loss, and 3) Gated Recurrent Unit (GRU) is used to provide positional encoding. Tested on a suite of benchmark tasks, the proposed method shows promising results over baselines.\n\nStrengths:\n\nThe paper is the first to apply transformers on continuous-valued tabular data. The motivation of using transformers for auto-regressive modeling of continuous data stated in Section 3 is persuasive.\n\nIn addition to reporting the test likelihood, the authors conduct supplementary experiments to validate the effectiveness of the proposed model. The experiments include training a regression model with generated samples, two-sample testing using a classifier, detecting out-of-distribution samples, and learning on noise-corrupted data. These analyses provide a rich view of how the proposed model behaves and confirms the effectiveness of the proposed approach.\n\nWeakness:\n\nI reckon the contribution of the paper as a density estimation method is marginal. Compared to the previously proposed transformers, [1,2], The key architectural difference is the use of a mixture of Gaussian as an output distribution instead of a discrete distribution generated from a softmax function. This change of output parametrization seems trivial and straightforward, compared to architectural improvements presented in [1, 2]. Another difference is the use of maximum mean discrepancy (MMD) as a regularizer (or an auxiliary objective function). However, as shown in Table 3, the gain from the use of MMD is not consistent.\n\nThe analyses conducted to evaluate generative models (contribution no. 2 on page 1) are valid, but the paper is not the first to perform such experiments and therefore it is not adequate to claim those experiments as a core contribution. Examining the predictive performance of a model trained on generated samples is used to evaluated generative adversarial networks [3, 4]. Also, measuring out-of-distribution detection performance is used widely in generative modeling literature [5, 6].\n\nRemark 4 on page 6 is wrong. The test likelihood is an estimator of cross entropy between the model distribution and the data distribution and therefore is connected to Kullback-Leibler (KL) divergence between them. As long as KL divergence is a meaningful measure of discrepancy, test likelihood gives meaningful information. Also, the maximizer of the objective in (1) is indeed p(x), given infinite data and a correctly specified model.\n\nThere is no explanation of how the outlier classes in Pendigits, ForestCover, and Satimage-2 datasets are defined.\n\nThere is no description of what transformers are used as the baseline in Table 3. Also, what is the \"standard Transformer\" used in Table 4? How exactly do these transformers differ from TraDE?\n\nMinor comments:\n\n- It would be more appropriate to use Proposition instead of Lemma for Lemma 1 and Lemma 3.\n- Currently, the numberings of lemmas and remarks are confusing. I suggest to number them separately.\n- It would be nice to mention that the datasets used in the experiments are tabular data, just in case if a reader is not familiar to the datasets.\n\n[1] Katharopoulos, Angelos, et al. \"Transformers are rnns: Fast autoregressive transformers with linear attention.\" arXiv preprint arXiv:2006.16236 (2020).  \n[2] Child, Rewon, et al. \"Generating long sequences with sparse transformers.\" arXiv preprint arXiv:1904.10509 (2019).  \n[3] Ye, Yuancheng, et al. \"GAN Quality Index (GQI) By GAN-induced Classifier.\" (2018).  \n[4] Borji, Ali. \"Pros and cons of gan evaluation measures.\" Computer Vision and Image Understanding 179 (2019): 41-65.  \n[5] Du, Yilun, and Igor Mordatch. \"Implicit generation and modeling with energy based models.\" Advances in Neural Information Processing Systems. 2019.  \n[6] Grathwohl, Will, et al. \"Your classifier is secretly an energy based model and you should treat it like one.\" arXiv preprint arXiv:1912.03263 (2019).", "extra_scores": {}, "extra_text": {}}, {"review_id": "ApOA1mut5yo", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThis work proposes a new auto-regressive density estimator built using self-attention module from the popular Transformer network. TraDE can be seen as an extension of decoder-only Transformer network where an input embeddings are given by a simple RNN-based encoder. Like Transformer, TraDE leverages multiple layers of self-attention module to implicitly model long-range dependencies. This effectively eliminates the need for explicit vertex ordering and hence useful on data with no known canonical ordering. The proposed model is general and can be applied to both continuous as well as discrete data. Along with the MLE objective, TraDE is additionally regularised using MMD penalty which can be easily back-propagated using reparametrization / Gumbel softmax trick.\n\nOn standard benchmark dataset, TraDE produces better density estimates compared to recently established state of the art baselines. To further evaluate the qualitative performance of density estimators, it proposes suite of various tasks on which TraDE is shown to work well.\n\n**Quality**\nThe paper is very well written and easy to follow. The experimental evaluation followed are standard (for density evaluation) and additional tasks depicts the usefulness of the sampled samples.\n\n**Originality**\nAs summarized above, TraDE is a simple extension of decoder of Transformer. Minor modification like RNN-based inputs (inspired by Wang et. al) and MMD loss led to dramatic improvement for density estimation tasks. However, overall the work lacks novelty and I feel it is only simple integration of different building blocks that produces better density estimator.\n\nMoreover, the tasks used for qualitative evaluation are also not completely novel. As pointed by the author, some of them (regression, two sample test, OOD) are already employed by the prior work.\n\n**Significance**\nDespite lack of novelty this work demonstrates,\n1. better empirical results as well as qualitative evaluation using various real world tasks, and\n2. carefully integrated self-attention module can perform better than many complex density estimators. \n\n**Clarity**\n1. How do one sample data using TraDE. I do not find any mention of methodology employed for sampling.\n2. Please include citations in Table 1.\n3. From Table 3, I note that inclusion of MMD loss has very minor effect on performance. What is the training time tradeoff for including MMD loss ? Also compare quantitative results of (TraDE - MMD) model for various tasks. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work explores an auto-regressive density estimator based on transformer networks. The model is trained via MLE with an additional MMD regularization term. 
Various experiments are performed on small benchmarks and show good results on density estimation. It is great to see that such a simple model is indeed very effective for density estimation on various small benchmarks (such as 2D density estimation and MNIST). 

The ablation experiments are informative and justify some of the model choices (such as the use of RNN to encode "positions"). Experiments are nicely chosen and paint a broad picture of the behaviour of the studied model.

The paper and author responses, however, excessively exaggerate the extent to which these results are relevant to the bigger picture in comparison to existing literature (e.g. flows and existing auto-regressive models).

As it has been extensively discussed with the reviewers, the proposed model is a straightforward application of a transformer network to auto-regressive modelling, this is specially so in light of existing work on auto-regressive models with transformers [e.g 1, 6, 8], self-attention [e.g 2]. BERT [7] itself can be used for auto-regressive modelling almost out-of-the-box (with the appropriate choice of masks during training).

At various points in the paper and author responses, it refers to flow models as "complicated/expensive" counterparts. These arguments are unfounded: auto-regressive models are particular cases of flows [3], and there are no obstructions to using transformer networks inside flows (in fact they have been already used, to achieve permutation equivariance and long-range correlations [e.g. 4]). 
The paper leaves comparisons to spline-flows out, arguing they are "hard to implement". This is quite conspicuous, as not only spline-flows are straightforward to implement, they produce results entirely on-par with the presented model (as an example, look at Fig 2 from [9] in comparison to Fig 1 from this paper).
Finally, the paper also misses an important discussion about the computational complexity of the proposed method. Auto-regressive models are considerably slower to sample from in relation to other types of directed models. Even more so with transformer networks as conditioners. For instance, flows [3, 5] allow for substantially faster sampling of large-dimensional data relative to auto-regressive models (by exploiting parallel sampling).


Extra comments:

The paper says "... Self-attention also enables
permutation equivariance and naturally enables TraDE to be agnostic to the ordering of the features ... "
This is true only for a *single* conditional $p(x_i | \text{Transformer}(x_{0 \ldots (i-1)}))$, not for the *joint* density. It is actually not straightforward to build auto-regressive models that are permutation invariant or that incorporate other forms of domain knowledge in general.
As an example, see [4] for how transformers and spline-flows can be used to produce exact permutation-invariant densities.


[1] Chen, M., Radford, A., Child, R., Wu, J., Jun, H., Luan, D. and Sutskever, I., 2020, November. Generative pretraining from pixels. In International Conference on Machine Learning (pp. 1691-1703). PMLR.

[2] Parmar, N., Vaswani, A., Uszkoreit, J., Kaiser, Ł., Shazeer, N., Ku, A. and Tran, D., 2018. Image transformer. arXiv preprint arXiv:1802.05751.

[3] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762.

[4] Wirnsberger, P., Ballard, A.J., Papamakarios, G., Abercrombie, S., Racanière, S., Pritzel, A., Rezende, D.J. and Blundell, C., 2020. Targeted free energy estimation via learned mappings. arXiv preprint arXiv:2002.04913.

[5] Huang, C.W., Krueger, D., Lacoste, A. and Courville, A., 2018. Neural autoregressive flows. arXiv preprint arXiv:1804.00779.

[6] Sun, C., Myers, A., Vondrick, C., Murphy, K. and Schmid, C., 2019. Videobert: A joint model for video and language representation learning. In Proceedings of the IEEE International Conference on Computer Vision (pp. 7464-7473).

[7] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova; ACL 2019.

[8] Child, R., Gray, S., Radford, A. and Sutskever, I., 2019. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509.

[9] Durkan, C., Bekasov, A., Murray, I. and Papamakarios, G., 2019. Neural spline flows. In Advances in Neural Information Processing Systems (pp. 7511-7522).
Decision: Reject