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 neural score-matching models achieve high precision but require retraining for every target distribution. We introduce DiScoFormer (Density and Score Transformer), an equivariant Transformer that maps i.i.d. samples to both density values and score vectors. Unlike score matching, which learns a fixed function R d → R d for a single distribution, DiScoFormer learns a sequence-tosequence operator that generalizes across distributions and sample sizes without retraining. 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, kernellike behaviors. The model outperforms KDE for density and score estimation, and provides a plugin score oracle for score-debiased KDE, Fisher information computation, and Fokker-Planck-type PDEs.