DiScoFormer: Plug-In Density and Score Estimation with Transformers

Estimating probability density and its score from samples remains a core problem in generative modeling, Bayesian inference, and kinetic theory. Existing methods are bifurcated: classical kernel density estimators (KDE) generalize across distributions but suffer from the curse of dimensionality, while modern neural score models achieve high precision but require retraining for every target distribution. We introduce DiScoFormer (Density and Score Transformer), a ``train-once, infer-anywhere" equivariant Transformer that maps i.i.d. samples to both density values and score vectors, generalizing across distributions and sample sizes. Analytically, we prove that self-attention can recover normalized KDE, establishing it as a functional generalization of kernel methods; empirically, individual attention heads learn multi-scale, kernel-like behaviors. The model converges faster and achieves higher precision than KDE for density estimation, and provides a high-fidelity plug-in score oracle for score-debiased KDE, Fisher information computation, and Fokker-Planck-type PDEs.

Imagine a demographer, collecting several statistics about humans, for example, age, height, and weight. The demographer collected these numbers from a few hundred humans and now would like to make conclusions about all humans: what fraction of humans have a certain combination of age, height and weight characteristics? A statistician would say that to answer this question the demographer should perform density estimation -- estimate the density (proportion) of humans with certain characteristics. If the number of characteristics is small, it is not a hard problem, and a well-studied method in statistics called Kernel Density Estimation can solve it. However, if the number of characteristics is large, it is a very hard problem, and Kernel Density Estimation does not perform well. In this paper, we propose using a type of Neural Network called Transformer to solve this hard problem of density estimation with many characteristics. This is the same type of Neural Network that powers the modern AI like ChatGPT. In our paper we mathematically prove that the Transformer Neural Network is particularly well-suited for this specific problem, and perform many experiments showing that in practice our Transformer does density estimation much better than Kernel Density Estimation. We also show how to use this to solve the equations that describe the motion of plasma -- hot charged gas that powers nuclear fusion, the technology that promises cheap abundant clean energy.