# Small Transformers Compute Universal Metric Embeddings

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 153 references that are in the map.
- First posted on arXiv 14 Sep 2022.

## Builds on (most cited first, 30 of 53)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [Error bounds for approximations with deep ReLU networks](https://mlsanity.com/paper/1610.01145/) (2016)
- [What Neural Networks Memorize and Why: Discovering the Long Tail via Influence Estimation](https://mlsanity.com/paper/2008.03703/) (2020)
- [The generalization error of random features regression: Precise asymptotics and double descent curve](https://mlsanity.com/paper/1908.05355/) (2019)
- [Hyperbolic Neural Networks](https://mlsanity.com/paper/1805.09112/) (2018)
- [Nearly-tight VC-dimension and pseudodimension bounds for piecewise linear neural networks](https://mlsanity.com/paper/1703.02930/) (2017)
- [Universal Approximation with Deep Narrow Networks](https://mlsanity.com/paper/1905.08539/) (2019)
- [Universality of Deep Convolutional Neural Networks](https://mlsanity.com/paper/1805.10769/) (2018)
- [Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: optimal rate and curse of dimensionality](https://mlsanity.com/paper/1810.08033/) (2018)
- [Path-SGD: Path-Normalized Optimization in Deep Neural Networks](https://mlsanity.com/paper/1506.02617/) (2015)
- [Geoopt: Riemannian Optimization in PyTorch](https://mlsanity.com/paper/2005.02819/) (2020)
- [Generalization error of random features and kernel methods: hypercontractivity and kernel matrix concentration](https://mlsanity.com/paper/2101.10588/) (2021)
- [Minimum Width for Universal Approximation](https://mlsanity.com/paper/2006.08859/) (2020)
- [Embedding Text in Hyperbolic Spaces](https://mlsanity.com/paper/1806.04313/) (2018)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [geomstats: a Python Package for Riemannian Geometry in Machine Learning](https://mlsanity.com/paper/1805.08308/) (2018)
- [From Trees to Continuous Embeddings and Back: Hyperbolic Hierarchical Clustering](https://mlsanity.com/paper/2010.00402/) (2020)
- [Representation formulas and pointwise properties for Barron functions](https://mlsanity.com/paper/2006.05982/) (2020)
- [Multilevel Clustering via Wasserstein Means](https://mlsanity.com/paper/1706.03883/) (2017)
- [Inferring Concept Hierarchies from Text Corpora via Hyperbolic Embeddings](https://mlsanity.com/paper/1902.00913/) (2019)
- [HoroPCA: Hyperbolic Dimensionality Reduction via Horospherical Projections](https://mlsanity.com/paper/2106.03306/) (2021)
- [Universal Approximation Theorems for Differentiable Geometric Deep Learning](https://mlsanity.com/paper/2101.05390/) (2021)
- [On the Optimal Memorization Power of ReLU Neural Networks](https://mlsanity.com/paper/2110.03187/) (2021)
- [Wasserstein Regression](https://mlsanity.com/paper/2006.09660/) (2020)
- [Provable Memorization via Deep Neural Networks using Sub-linear Parameters](https://mlsanity.com/paper/2010.13363/) (2020)
- [Universal Approximation Under Constraints is Possible with Transformers](https://mlsanity.com/paper/2110.03303/) (2021)

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

- [Transformers are Universal In-context Learners](https://mlsanity.com/paper/2408.01367/) (2024)
- [Designing Universal Causal Deep Learning Models: The Case of Infinite-Dimensional Dynamical Systems from Stochastic Analysis](https://mlsanity.com/paper/2210.13300/) (2022)
- [A Brief Survey on the Approximation Theory for Sequence Modelling](https://mlsanity.com/paper/2302.13752/) (2023)
- [Neural Snowflakes: Universal Latent Graph Inference via Trainable Latent Geometries](https://mlsanity.com/paper/2310.15003/) (2023)
- [Transformers through the lens of support-preserving maps between measures](https://mlsanity.com/paper/2509.25611/) (2025)
- [Neural Spacetimes for DAG Representation Learning](https://mlsanity.com/paper/2408.13885/) (2024)
- [The Curved Spacetime of Transformer Architectures](https://mlsanity.com/paper/2511.03060/) (2025)
- [Capacity Bounds for Hyperbolic Neural Network Representations of Latent Tree Structures](https://mlsanity.com/paper/2308.09250/) (2023)
- [Transformers Can Solve Non-Linear and Non-Markovian Filtering Problems in Continuous Time For Conditionally Gaussian Signals](https://mlsanity.com/paper/2310.19603/) (2023)
- [Beyond Universal Approximation Theorems: Algorithmic Uniform Approximation by Neural Networks Trained with Noisy Data](https://mlsanity.com/paper/2509.00924/) (2025)
- [Tighter Learning Guarantees on Digital Computers via Concentration of Measure on Finite Spaces](https://mlsanity.com/paper/2402.05576/) (2024)
- [Classification Fields: Arbitrarily Fine Recursive Hierarchical Clustering From Few Examples](https://mlsanity.com/paper/2605.07119/) (2026)

## Abstract (arXiv)

We study representations of data from an arbitrary metric space 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 probabilistic transformers. Our guarantees are of memorization type: we prove that a probabilistic transformer of depth about nlog(n) and width about n2 can bi-H"older embed any n-point dataset from 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 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"older embeddings with arbitrarily small distortion.