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.

💡 Research Summary

The paper introduces DiScoFormer, a universal transformer‑based operator that simultaneously estimates probability density f and its score ∇ log f from a set of i.i.d. samples. Traditional kernel density estimators (KDE) are distribution‑agnostic but suffer from the curse of dimensionality, while modern neural score‑matching models achieve high accuracy but must be retrained for each target distribution. DiScoFormer bridges this gap by learning a single “train‑once, infer‑anywhere” model that respects two essential symmetries: permutation equivariance (handled by the transformer’s set‑wise architecture without positional encodings) and affine equivariance (enforced through a whitening layer and random affine augmentations during training).

The authors provide a rigorous theoretical link between self‑attention and KDE. Proposition 3.2 proves that, with appropriate query, key, and value projections (Q = K = h·X, V = Iₙ), a single attention head computes exactly the normalized Gaussian kernel weights used in KDE. Consequently, the transformer can be viewed as a learnable, multi‑scale generalization of kernel smoothing, which is empirically confirmed by visualizing attention maps: different heads specialize in short‑range, long‑range, and directional interactions, effectively learning a mixture of kernels.

Training data consist of synthetic Gaussian mixture models (GMMs) with 1–10 components, sampled on the fly. Because GMMs admit closed‑form densities and scores, the model can be supervised with a loss that combines mean‑squared error on f and on ∇ log f (weighted by a hyperparameter α). The architecture comprises four transformer encoder layers (128‑dimensional hidden size, eight heads) and two output heads for density and score, totaling roughly 800 k parameters.

Empirical results span dimensions d = 1 to 10 and sample sizes n = 256 to 8192. Across all regimes DiScoFormer outperforms standard KDE and the recent score‑debiased KDE (SD‑KDE) in terms of mean‑squared error for both density and score. Notably, in high dimensions (d = 10) where KDE’s performance collapses, DiScoFormer maintains low error, demonstrating superior scalability. The model also generalizes to out‑of‑distribution GMMs with up to 19 components, despite being trained only on up to 10 components, indicating robust extrapolation.

Beyond density estimation, the learned score oracle is directly usable in downstream tasks: (i) plugging into SD‑KDE yields more accurate density estimates; (ii) computing Fisher information and differential entropy becomes straightforward; (iii) the oracle can drive deterministic solvers for Fokker‑Planck‑type PDEs (e.g., the Landau equation), eliminating the need for per‑distribution score retraining. Runtime comparisons show that inference scales linearly with n and d, and is competitive with fast KDE implementations.

In summary, DiScoFormer contributes (1) a theoretically grounded demonstration that transformers can emulate and extend kernel methods; (2) a practical, permutation‑ and affine‑equivariant architecture that jointly learns density and score; (3) extensive empirical evidence of superior accuracy, scalability, and generalization; and (4) a plug‑in component for a variety of statistical and physical applications. Future work may explore extensions to non‑synthetic data (images, time series), incorporation of sparse or low‑rank attention for further efficiency, and integration with diffusion‑based generative modeling pipelines.