Small Transformers Compute Universal Metric Embeddings

TL;DR

This paper demonstrates that small neural networks called probabilistic transformers can effectively embed data from arbitrary metric spaces into Gaussian mixture representations with low distortion, avoiding the curse of dimensionality.

Contribution

It provides theoretical guarantees for probabilistic transformers to embed metric space data with low distortion, including bi-Hölder and bi-Lipschitz bounds, applicable to various structured datasets.

Findings

Probabilistic transformers achieve low-distortion embeddings of metric data.

Embedding guarantees extend to Riemannian manifolds, metric trees, and graphs.

Transformers can embed into Gaussian mixtures with arbitrarily small distortion.

Abstract

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\"{o}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\"{o}lder embeddings with arbitrarily small distortion.