[论文解读] A mathematical perspective on Transformers

Borjan Geshkovski, Cyril Letrouit+2 more|arXiv (Cornell University)|Dec 17, 2023

Topic Modeling被引用 6

一句话总结

本文将Transformer建模为在单位球面上演化的相互作用粒子系统，证明自注意力机制的动力学会导致令牌在长时间尺度上聚类为单一点，其过程可通过Wasserstein梯度流描述。关键结果表明，在高维空间或耦合参数较小时，所有令牌均会收敛至同步状态，从而通过亚稳态聚类动力学解决了在概率性下一个词预测中出现确定性输出的悖论。

ABSTRACT

Transformers play a central role in the inner workings of large language models. We develop a mathematical framework for analyzing Transformers based on their interpretation as interacting particle systems, which reveals that clusters emerge in long time. Our study explores the underlying theory and offers new perspectives for mathematicians as well as computer scientists.

研究动机与目标

• 通过将Transformer建模为在单位球面上的连续时间相互作用粒子系统，建立理解其行为的严格数学框架。
• 分析自注意力动力学的长时间行为，特别是高维空间中聚类现象的出现。
• 通过亚稳态聚类态解释在下一个词预测中确定性输出与语言建模概率本质之间的明显矛盾。
• 建立Transformer动力学与Wasserstein梯度流、球面上最优点构型以及非线性输运方程等经典数学概念之间的联系。

提出的方法

• 将Transformer的离散层建模为连续时间动力系统，将层索引视为时间变量。
• 将每个令牌表示为单位球面 $\mathbb{S}^{d-1}$ 上的粒子，受层归一化约束，并通过所有粒子的经验测度建模自注意力的非线性耦合。
• 利用连续性方程表述经验测度的演化，将其与非线性输运PDE联系起来。
• 引入一种代理自注意力机制，使其成为与球面上最优点构型相关的能量泛函的Wasserstein梯度流。
• 使用微分几何与优化工具（包括Hessian分析与球面流形上临界点的分类）分析该系统。
• 通过研究 $\beta \to 0$ 的极限并利用Hessian特征值关于 $\beta$ 的连续性，证明仅有同步构型是稳定的，从而获得收敛性结果。

实验结果

研究问题

• RQ1 Transformer中的自注意力机制如何随时间诱导令牌聚类？
• RQ2 在Transformer中由令牌构成的粒子系统在长时间尺度下的行为如何，特别是在高维设置下？
• RQ3 尽管建模的是概率分布，为何Transformer在下一个词预测中会产生确定性输出？
• RQ4 在何种条件下系统会收敛至单一同步状态，这与耦合参数 $\beta$ 的关系如何？
• RQ5 球面流形的几何与动力学特性如何影响注意力机制的稳定性和收敛性？

主要发现

• 在高维空间（$d \geq n$）中，$n$ 个初始随机分布在 $\mathbb{S}^{d-1}$ 上的粒子在 $t \to \infty$ 时收敛至单一点，表明全局聚类。
• 当 $\beta$ 较小时，能量泛函 $\mathsf{E}_{\beta}$ 的所有临界点（除同步构型外）均为严格鞍点，表明同步状态具有稳定性。
• 系统表现出双时间尺度动力学：快速聚类为少数几组，随后缓慢合并为单一簇。
• 代理自注意力机制是某一在球面上最优构型理论中已有深入研究的能量泛函的Wasserstein梯度流。
• 能量泛函在临界点处的Hessian与黎曼度量无关，确保了鞍点性质在不同梯度流公式下保持不变。
• 当 $\beta \leq c n^{-1}$ 且 $c > 0$ 与 $n$ 无关时，所有非同步临界点均为严格鞍点，证实了在长时间极限下同步状态的主导性。