[论文解读] Small Transformers Compute Universal Metric Embeddings

Anastasis Kratsios, Valentin Debarnot+1 more|arXiv (Cornell University)|Sep 14, 2022

Medical Imaging and Analysis被引用 5

一句话总结

该论文表明，小型概率变压器能够以约 $ n\log n $ 的深度和约 $ n^2 $ 的宽度，在传输度量下将任意有限度量空间嵌入到单变量高斯混合分布空间中，实现低失真。此外，论文建立了概率性的双-Lipschitz保证，并表明通过多变量高斯分布进行的嵌入可实现任意小的失真。

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ö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.

研究动机与目标

• 开发一种基于神经网络的通用表示框架，用于任意度量空间，并具备理论保证。
• 克服欧几里得空间与双曲空间在表示现实世界数据中复杂且非均匀几何结构时的局限性。
• 为一般度量空间上的小型神经架构建立嵌入保证，特别是双-Hölder 和概率性双-Lipschitz 边界。
• 证明概率变压器能够以可控失真将数据从黎曼流形、度量树和组合图中嵌入。
• 证明通过小型变压器，利用多变量高斯混合分布空间可实现任意小的失真。

提出的方法

• 使用概率变压器——一种小型、近期的具有注意力机制的神经网络——学习从一般度量空间 $\mathcal{X}$ 到单变量高斯混合分布空间的特征映射。
• 在高斯混合分布上采用 Wasserstein 传输度量作为目标表示空间，以确保几何保真度。
• 推导出记忆型嵌入保证：深度约为 $\sim n\log n$、宽度约为 $\sim n^2$ 的变压器可对任意 $ n $-点度量空间实现低失真双-Hölder 嵌入。
• 引入概率性双-Lipschitz 保证，以失真水平与实现该失真的概率之间进行权衡。
• 应用度量嵌入理论中的结果，包括 Linial 等人（2002）和 Naor 等人（2006）提出的马尔可夫类型与失真下界，以分析嵌入的可行性。
• 算法化构造偏置项，以确保正确嵌入到高斯混合分布的 Wasserstein 空间中，从而实现单射且失真可控的映射。

实验结果

研究问题

• RQ1 小型概率变压器能否以低失真实现对任意有限度量空间的通用度量嵌入？
• RQ2 为使变压器实现度量空间中任意 $ n $-点数据集的双-Hölder 嵌入，所需的架构深度与宽度分别是多少？
• RQ3 如何推导出概率性双-Lipschitz 保证，以平衡失真与嵌入成功率？
• RQ4 通过小型变压器将嵌入到多变量高斯混合分布中，能否实现任意小的失真？
• RQ5 该框架在嵌入分层结构或黎曼几何数据时的理论极限是什么？

主要发现

• 深度约为 $\sim n\log n$、宽度约为 $\sim n^2$ 的概率变压器可对任意 $ n $-点度量空间实现低失真的双-Hölder 嵌入，避免了维度灾难。
• 建立了概率性双-Lipschitz 保证：对于任意失真水平 $ D $，存在一个概率 $ p(D) $，使得随机点对的嵌入失真不超过 $ D $。
• 当底层度量空间 $\mathcal{X}$ 具有足够规则的几何结构（例如具有截面负曲率的黎曼流形）时，对所有点都存在更强的统一双-Lipschitz 保证。
• 通过小型变压器将嵌入到多变量高斯混合分布中，可实现任意小的失真，表明该表示空间具有高度表达能力。
• 该框架可实现对来自黎曼流形、度量树和组合图的数据的通用嵌入，并具备理论失真边界。
• 来自度量嵌入理论的理论结果——特别是马尔可夫类型与失真下界——证实此类嵌入在失真缩放方面是最优的。