Small Transformers Compute Universal Metric Embeddings
This provides a theoretical foundation for neural network-based metric embeddings, potentially benefiting machine learning applications in domains like Riemannian manifolds and combinatorial graphs, though it appears incremental as it builds on existing transport metric frameworks.
The paper tackles the problem of embedding data from arbitrary metric spaces into Gaussian mixture spaces with low distortion, showing that small probabilistic transformers can achieve bi-Hölder embeddings for any n-point dataset with depth about n log(n) and width about n^2, avoiding the curse of dimensionality.
We study representations of data from an arbitrary metric space $\mathcal{X}$ in the space of univariate Gaussian mixtures with a transport metric (Delon and Desolneux 2020). We derive embedding guarantees for feature maps implemented by small neural networks called \emph{probabilistic transformers}. Our guarantees are of memorization type: we prove that a probabilistic transformer of depth about $n\log(n)$ and width about $n^2$ can bi-Hölder embed any $n$-point dataset from $\mathcal{X}$ with low metric distortion, thus avoiding the curse of dimensionality. We further derive probabilistic bi-Lipschitz guarantees, which trade off the amount of distortion and the probability that a randomly chosen pair of points embeds with that distortion. If $\mathcal{X}$'s geometry is sufficiently regular, we obtain stronger, bi-Lipschitz guarantees for all points in the dataset. As applications, we derive neural embedding guarantees for datasets from Riemannian manifolds, metric trees, and certain types of combinatorial graphs. When instead embedding into multivariate Gaussian mixtures, we show that probabilistic transformers can compute bi-Hölder embeddings with arbitrarily small distortion.