# TraDE: Transformers for Density Estimation

arXiv: https://arxiv.org/abs/2004.02441 · HTML: https://mlsanity.com/paper/2004.02441/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 6 Apr 2020.

## Built on by (most cited first, 18 of 18)

- [Simplicity Bias in Transformers and their Ability to Learn Sparse Boolean Functions](https://mlsanity.com/paper/2211.12316/) (2022)
- [Benchmarking Multimodal AutoML for Tabular Data with Text Fields](https://mlsanity.com/paper/2111.02705/) (2021)
- [Gradient-Free Structured Pruning with Unlabeled Data](https://mlsanity.com/paper/2303.04185/) (2023)
- [Arbitrary Conditional Distributions with Energy](https://mlsanity.com/paper/2102.04426/) (2021)
- [Deep Generative Models for Detector Signature Simulation: A Taxonomic Review](https://mlsanity.com/paper/2312.09597/) (2023)
- [Differentially Private Normalizing Flows for Privacy-Preserving Density Estimation](https://mlsanity.com/paper/2103.14068/) (2021)
- [qecGPT: decoding Quantum Error-correcting Codes with Generative Pre-trained Transformers](https://mlsanity.com/paper/2307.09025/) (2023)
- [Transformer Neural Autoregressive Flows](https://mlsanity.com/paper/2401.01855/) (2024)
- [Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens](https://mlsanity.com/paper/2410.14655/) (2024)
- [The DEformer: An Order-Agnostic Distribution Estimating Transformer](https://mlsanity.com/paper/2106.06989/) (2021)
- [A Flexible Nadaraya-Watson Head Can Offer Explainable and Calibrated Classification](https://mlsanity.com/paper/2212.03411/) (2022)
- [Transformers as Unsupervised Learning Algorithms: A study on Gaussian Mixtures](https://mlsanity.com/paper/2505.11918/) (2025)
- [Sum-Product-Attention Networks: Leveraging Self-Attention in Probabilistic Circuits](https://mlsanity.com/paper/2109.06587/) (2021)
- [TreeFlow: Going beyond Tree-based Gaussian Probabilistic Regression](https://mlsanity.com/paper/2206.04140/) (2022)
- [AQuaMaM: An Autoregressive, Quaternion Manifold Model for Rapidly Estimating Complex SO(3) Distributions](https://mlsanity.com/paper/2301.08838/) (2023)
- [Learning Invariant Representations with a Nonparametric Nadaraya-Watson Head](https://mlsanity.com/paper/2309.13377/) (2023)
- [Soft Learning Probabilistic Circuits](https://mlsanity.com/paper/2403.14504/) (2024)
- [Are Deep Sequence Classifiers Good at Non-Trivial Generalization?](https://mlsanity.com/paper/2210.13082/) (2022)

## Abstract (arXiv)

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.