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