TraDE: A Simple Self-Attention-Based Density Estimator

1 INTRODUCTION . Density estimation involves estimating a probability density p ( x ) , given independent , identically distributed ( iid ) samples from it . This is a 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 a 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 density estimation is to sample data from a simple reference distribution and then to learn a ( sequence of ) invertible transformations that allow us to adapt it to a target distribution . Flow-based methods ( Durkan et al. , 2019b ) employ this with great success . A more classical approach is to decompose p ( x ) in an iterative manner via conditional probabilities p ( xi+1|x1 ... i ) and fit this distribution using the data ( Murphy , 2013 ) . One may even employ implicit generative models to sample from p ( x ) 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 ) . Implementing these above methods however requires special care , e.g. , the normalizing transform requires the network to be invertible with an efficiently computable Jacobian . Auto-regressive models using recurrent networks are difficult to scale to high-dimensional data due to the need to store a 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 a recurrent neural network ( RNN ) -based input embedding to approximate arbitrary continuous and discrete conditional densities . It is more flexible than contemporary architectures such as those proposed by Durkan et al . ( 2019b ; a ) ; Kingma et al . ( 2016 ) ; De Cao et al . ( 2019 ) for this problem , yet it remains capable of approximating any density function . To our knowledge , this is the first adaptation of Transformer-like architectures for continuous-valued density estimation . 2 . Log-likelihood on held-out data is the prevalent metric to 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-distribution samples , and robustness to noise in the training data . 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 its strong ( systematically evaluated ) empirical performance . 2 BACKGROUND AND RELATED WORK . Given a dataset { x1 , . . . , xn } where each sample xl ∈ Rd is drawn iid from a distribution p ( x ) , the maximum-likelihood formulation of density estimation finds a θ-parameterized distribution q with θ̂ = argmax θ 1 n n∑ l=1 log q ( xl ; θ ) . ( 1 ) The candidate distribution q can be parameterized in a variety of ways as we discuss next . Normalizing flows write x ∼ q as a transformation of samples z from some base distribution pz from which one can draw samples easily ( Papamakarios et al. , 2019 ) . If this mapping is fθ : z → x , two distributions can be related using the determinant of the Jacobian as q ( x ; θ ) : = pz ( z ) ∣∣∣dfθ dz ∣∣∣−1 . A practical limitation of flow-based models is that fθ must be a diffeomorphism , i.e. , it is invertible and both fθ and f−1θ are differentiable . Good performance using normalizing flows imposes nontrivial restrictions on how one can parametrize fθ : it must be flexible yet invertible with a 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 & Turner , 2013 ) , Sylvester flows ( Berg et al. , 2018 ) , coupling ( Dinh et al. , 2014 ) and auto-regressive models ( Larochelle & Murray , 2011 ) . One may also compose the transformations , e.g. , using monotonic mappings fθ in each layer ( Huang et al. , 2018 ; De Cao et al. , 2019 ) . Auto-regressive models factorize the joint distribution as a product of univariate conditional distributions q ( x ; θ ) : = ∏ i qi ( xi|x1 , . . . xi−1 ; θ ) . The auto-regressive approach to density estimation is straightforward and flexible as 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 a shared set of parameters ( Oliva et al. , 2018 ; Kingma et al. , 2016 ) . For high-dimensional data , the challenge lies in handling the increasingly large state space x1 , . . . , xi−1 required to properly infer xi . In recurrent auto-regressive models , these conditioned-upon variables ’ values are stored in some representation hi which is updated via a function hi+1 = g ( hi , xi ) . 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 et al. , 2017 ) . Many existing auto-regressive algorithms are highly sensitive to the variable ordering chosen for factorizing q , and some methods must train complex ensembles over multiple orderings to 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 are already naturally ordered ( Chelba et al. , 2013 ) . In contrast , we 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 a 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 ) ) 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 ) . 3 TOOLS OF THE TRADE x1 x2 x3 x4 x5 x6 x7 x8 x1 x8 x2 x7 x3 x6 x4 x5 Consider the 8-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 , 2 , 3 , 4 , 5 , 6 , 7 , 8 ) and ( 1 , 8 , 2 , 7 , 3 , 6 , 4 , 5 ) . In the latter case the model becomes a simple sequence where e.g . p ( x3|x1,8,2,7 ) = p ( x3|x7 ) 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 ( x3|x1,2 ) can be simplified further to p ( x3|x2 ) , but we still need to carry the precise value of x1 along until the end since p ( x8|x1 ... 7 ) = p ( x8|x1,2 ) . This is a fundamental weakness in models employing RNNs such as ( Oliva et al. , 2018 ) . In practice , we may be unable to select a favorable ordering for columns in a 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 et al. , 2017 ) . A recurrent network can , in principle , absorb this information into its hidden state . In fact , Long-Short-Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) units were engineered 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 x1 , . . . , xd . Depending on the order vertices are traversed in our factorization the largest number of latent variables a recurrent auto-regressive model needs to store is bounded from above and below by the minimum and the maximum number of variables with a cut edge of the graph G. Proof . Given a subset of known variables S ⊆ { 1 , . . . d } we want to estimate the conditional distribution of the variables on the complement C : = { 1 , . . . d } \S . For this we need to decompose S into the Markov blanket M of C and its remainder . By definition M consists of the variables with a cut edge . Since p ( xC |xS ) = p ( xC |xM ) we are done . This problem with long-dependencies in auto-regressive models has been noted before , and is exacerbated in density estimation with continuous-valued 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 in 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-regressive factorization order ( Yun et al. , 2019 ) . Thus a major benefit of self-attention here is its effectiveness at maintaining an accurate representation of the feature values xj for j < i when inferring xi ( irrespective of the distance between i and j ) .

Summary: 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.