从非欧几何视角解释和改造Transformer

在人工智能的浩瀚宇宙中，Transformer 架构无疑是当今最耀眼的恒星。从 ChatGPT 的惊艳对话到 Midjourney 的梦幻绘图，它都是这一切奇迹背后的核心引擎。

然而，对于大多数非巨头的研究者来说，现实往往是骨感的：手头没有成堆的 H100，也没有那烧得起整座发电厂的预算，也想玩转 Transformer 怎么办？

硬（硬件）的不行，要不咱就来点软（数学）的？

是的，既然拼不过算力的暴力美学，那我们就试试数学的精巧逻辑。这就好比武林剑客，力量上拼不过人家，就得在剑术的精妙上下点功夫。

不过呢，稍微复杂一点的数学搬上去未必立刻这就 SOTA，但至少，它能从更抽象的第一性原理层面，为我们提供一种全新的解题思路。

在改进 Transformer 的道路上，通常有两派人马：一派在做减法，比如搞线性注意力（Linear Attention），试图让模型跑得更快、吞吐量更大；而另一派偏偏反其道而行之，他们在做加法。

咱就属于后者，主张把模型整复杂，先弄一点点黎曼几何和李群李代数等知识试试水，试图进一步压榨架构的潜力。

虽然这听起来似乎有点像是为了数学而数学，但其背后的野心在于：与其用海量数据去暴力拟合，不如给 AI 装上一个更加符合几何直觉的大脑。

至于性能优化？交给工程师嘛，咱只负责优雅。

好了，让我们系好安全带，看看咱是如何用微分几何的视角，把 Transformer 扔进一个弯曲的世界里魔改的。

1. 注意力流形

让我们回到 Transformer。它的强大毋庸置疑，然而，你是否想过，这个强大的引擎里头其实是一个路痴？这就不得不说说它的一个出厂设定缺陷：置换不变性（Permutation Invariance）。

痛苦的根源：当 AI 分不清谁是赢家

对于标准的自注意力机制（Self-Attention）来说，它看一句话就像是看一堆散落在地上的麻将牌。它只知道有什么字，却完全不知道字在哪。如果我们将这个概念放在文言文的语境下，这种缺陷足以引发一场历史级的混乱。

请看这个经典的跨朝代乱斗案例：

【原句：武圣之威】

关公 战 败 秦琼（注：关羽跨越时空，把唐朝的秦琼打趴下了。）

【乱序：瓦岗逆袭】

秦琼 战 败 关公（注：秦叔宝逆天改命，武圣颜面扫地。）

在人类眼里，这是两个截然相反的结局，胜负关系完全取决于关公和秦琼谁站在战败的前面。

然而，对于一个没有位置编码的 Transformer 来说，它收到的输入仅仅是一个词袋（Bag of Words）：

关公，秦琼，战，败

在它看来，左边的武圣发威和右边的瓦岗逆袭生成的特征向量是一猫一样的。这就是所谓的薛定谔的战役：只要不观测位置，你永远不知道到底是谁赢了。

为了解决这个问题，传统的做法是给每个词脸上贴个编号（1, 2, 3...），这就是位置编码（Positional Encoding）。但这就像是在一副写意水墨画上，生硬地用圆珠笔标上了 坐标。虽然让 AI 分清了胜负，但这种做法总感觉不够优雅和本质，也不符合道法自然的几何直觉。

那么有没有优雅一些的做法呢？

有一天，俺在刷锅的时候突然想到：与其强行标号，不如让这战场本身就是弯曲的（非欧的）如何？

黎曼流形：让注意力在曲面上流淌

我们认为，Token（单词或图像块）不应该待在平坦的欧几里得空间里，而应该住在一个黎曼流形上。这是一个弯曲的空间，就像球面或者起伏的山峦。

在这个框架下，不仅知道 Token 是什么（内容），还能通过几何结构感知它在哪里。

为了让形象更具辨识度，咱这里改成西游记中的三兄弟。

如果在平地上，你可以随意平移一个箭头（向量），它的方向和大小不变。但在球面上，你带着一个指向北极的箭头走到赤道，再沿着赤道走，箭头指的方向相对于当地经纬线就会发生变化。因此，我们需要想办法来比较不同地点的箭头。

我们将 Transformer 的自注意力机制重构为注意力流形（Attention Manifold）上的交互。在此框架下，Token 是流形上的点，特征嵌入是切空间中的向量。流形的几何由两部分定义：黎曼度量（调节局部交互与语义显著性）和平行移动算子（保证不同位置向量比较的一致性）。这为设计具有强归纳偏置的注意力机制提供了基于第一性原理的视角。

我们的核心理念发生了根本转变：从外在几何转向内蕴几何。重要的不是 Token 在环境空间中的绝对坐标，而是它们之间如何通过变换建立联系。我们不再试图参数化绝对位置，而是专注于定义一种协议，使切向量（特征嵌入）能在离散点的切空间之间进行转换，即平行移动。

尽管平行移动提供了切空间间的线性同构，但在离散设置中，我们缺乏显式的坐标、路径及联络定义。因此，我们摒弃显式位置编码，转而令模型直接学习切空间之间的线性映射。

这相当于隐式地构建了沿 Token 序列的平行移动，从而在无需具体位置信息的情况下，实现了几何上合法的内积运算。也就是说，我们无心构建整个连续的流形，而是转向学习一个离散联络。

2. 黎曼注意力

物理学中有一句名言：物质告诉时空如何弯曲，时空告诉物质如何运动。这话不是爱因斯坦自己说的，是别人对其理论的注解。

我们不妨借用一下这个思想，但此时需要忽略“时空”中的时间，因为我们只考虑空间。

在传统的 Attention 中，空间是均匀的。但在黎曼注意力中，Token 的内容（Embedding）决定了局部的几何结构。于是，我们引入了所谓的动态度量（Dynamic Metric）。如果某个 Token 很重要（比如图像中的主体），它的质量就大，会让周围的空间收缩或弯曲，形成一个类似重力井的东西。

这种由嵌入决定度量的方式具体可以通过一个精心设计的 mlp 来实现。比如，下式这种较为简单的投影操作表现也还可以

这种由嵌入决定动态度量的机制使得自注意力不再是无脑的全局平权扫描，而是体现一定的社会学马太效应或心理学光环效应。重要的信息会自动拉近与其他信息的几何距离，从而获得更多的关注。

我们不妨将其称为黎曼注意力（Riemannian Attention）。请看它完整的公式，

这里我们暂时还保留 Softmax，以后再来替换。

传统的 Attention 机制通过 Softmax 强制归一化，某种程度上是在维持一种众生平等的假象。然而，黎曼注意力打破了这种平衡，引入动态度量 后，模型表现出了鼓励马太效应的强者恒强。当一个 Token 被判定为语义上的富人（具有高重要性 ）时，它不仅自身携带的信息量大，更会通过度量乘积效应，成倍地放大它与周围 Token 的引力连接。

这种机制让模型不再温吞地平均分配注意力，而是敢于让关键特征形成垄断地位，迅速从噪声中通过资本积累脱颖而出。

万有引力 vs 黎曼注意力

依稀记得初遇 Transformer 之时，看着注意力机制将散落的 向量一一捕获、汇聚，一种万有引力的既视感便油然而生。那时未曾深究，只觉得它虽有引力之形，却无动力学之实，看起来更像是统计学中冰冷的聚类。

然则，现在终于有意把这一机制安放在黎曼几何的基座之上，通往爱因斯坦广义相对论的大门便已悄然敞开。至于如何赋予其随时间演化的完整动力学，那是后话；此刻，我们不妨先迈出关键的一步：打破统计学的束缚，用几何学的语言，将注意力机制魔改为万有引力场。

让我们回到高中物理课堂。牛老师告诉我们，万有引力公式是这样的：

这简直就是万有引力公式的 1:1 复刻嘛！

通过这一改动，黎曼注意力实际上引入了一个动态能量系统：有些 token 影响力有限，让其逐渐失去存在感（能量），而有些影响力巨大，让它们拥有相对更大的影响力。这不仅是几何学，也带那么点物理，虽然目前还不多。

当然，这里另外一个思路就是把 Softmax 改掉 。。。

3. 离散联络