你有没有想过，一个蛋白质口袋里的小分子药物，和电视里一个人走路的视频片段，本质上其实是同一类问题？

听起来有点荒谬。一个是原子尺度的微观运动，一个是宏观世界里人眼能直接看到的画面。可是如果你把它们都剥成最简单的骨架，剩下的东西惊人地相似：一堆散落在空间里的点，每个点都带着一个时间戳，你需要预测它们接下来会跑到哪里去。

这篇来自香港中文大学和上海人工智能实验室团队的论文，就是抓住了这个共性,提出了一个叫 MinkowskiPE 的方法，试图用同一套数学语言,同时搞定分子动力学预测和视频预测这两件看起来毫不相干的事。

先说说这件事难在哪

现在做时空建模,学界基本分成两条路子,而这两条路子都有点别扭。

第一条路子讲究物理直觉。你见过的循环神经网络、卷积网络、用微分方程描述连续演化的模型,都属于这一类。它们的共同点是:提前规定好时间该怎么往前走,一步一步推演。这样做的好处是物理约束强,模型不容易学出离谱的结果。坏处也很明显,规则定死了,灵活性就被锁住了,模型很难从数据里自己发现一些意料之外的规律。

第二条路子讲究数据驱动。Transformer 这类模型属于这一类,靠学出来的注意力机制去捕捉时间和空间之间的关联,比如把空间注意力和时间注意力拆开算,或者搞多尺度、搞轨迹感知的注意力。

Transformer:一种完全基于注意力机制的神经网络架构,最早用于自然语言处理,后来被广泛用于视觉、视频等多种任务。

这类方法确实灵活,但有个隐患,时间和空间到底是怎么耦合在一起的,模型自己心里有数,人却看不透,这种耦合关系基本是隐式的,埋在一堆参数里说不清楚。

于是问题就摆在这里了:physics驱动的方法限制了表达能力,学习驱动的方法又缺乏明确的物理先验。能不能找一个两头都占的办法?

论文作者把目光投向了一个听起来有点吓人的词:闵可夫斯基几何。

闵可夫斯基几何*:爱因斯坦狭义相对论里描述时空的几何结构,核心特点是用一种特殊的度量区分"时间方向"和"空间方向",让它们在数学上地位不对等。

这东西原本是物理学里用来描述光速不变、时空统一的工具。但这篇论文的作者发现,抛开相对论的物理含义不谈,这套几何结构本身恰好提供了一种数学框架:它能把时间和空间放进同一个坐标系里,又不至于把两者混为一谈。这正是时空建模缺的那块拼图。

位置编码:到底编码了什么

要理解 MinkowskiPE 怎么工作,得先搞清楚 Transformer 里"位置编码"这件事到底在解决什么问题。

位置编码*:Transformer 本身对输入顺序不敏感(它看一句话里哪个词在前哪个词在后,全靠额外加的位置信息),位置编码就是给每个输入token打上"我在哪里"的标签,让模型能区分不同位置的token。

一种比较理想的位置编码方式,叫"相对位置编码"。它的核心思想是:两个token之间的注意力分数,应该只取决于它们之间的相对距离,而不是它们各自在绝对坐标系里的具体位置。

打个比方。你在地铁上和朋友聊天,你们俩之间的对话关系(谁先说谁后说、说话间隔多久),跟地铁车厢本身开到了哪一站没有任何关系。如果位置编码没做好相对性,那就相当于你们的对话内容会随着车厢位置变化而莫名其妙地改变,这显然不对。相对位置编码要保证的就是:把整个坐标系平移一下(比如车厢往前开了一站),两个token之间的关系应该纹丝不动。

论文里把这个性质写成一个公式:两个token的注意力分数,只能通过它们坐标之差来计算,不能直接依赖各自的绝对坐标。这句话听起来简单,真正做到却不容易,尤其是当坐标里既有时间又有空间的时候。

洛伦兹变换怎么被借用过来

闵可夫斯基几何里有一个核心工具,叫洛伦兹变换。

洛伦兹变换*:狭义相对论中,保持时空间隔不变的坐标变换。它和普通的旋转不太一样,因为时间和空间在这个变换里的"地位"是相反的,数学上体现为一个正负号相反的度量矩阵。

关键的技术突破在于:这篇论文根本没打算让分子或者视频里的物理过程真的遵守相对论。它做的事情,是把每个token的时空坐标拿来,当成一个参数,去生成一个洛伦兹变换,然后把这个变换作用在query和key特征上,而不是作用在原始坐标本身。

这个区分很重要。论文原话讲得很直白:事件坐标本身不被变换,它们只是用来参数化一个作用在特征上的变换。换句话说,闵可夫斯基几何在这里是一个表示层面的设计工具,不是对数据物理规律的假设。

那这么绕一圈图什么呢?图的是一个数学上的巧合:洛伦兹变换满足一个"群同态"性质,两个变换的组合,恰好等于坐标相加后对应的那个变换。只要你把每个token的query和key都按照自己的坐标做一次这样的变换,再做内积,最后算出来的结果就会自动变成只依赖两个token相对坐标差的形式。这正是前面说的相对位置编码想要的效果,而且是自动达成的,不需要额外再去设计什么补丁。

这里有个细节特别巧妙。如果你直接对query和key做同样的洛伦兹变换再求内积,数学上并不会自动消去绝对位置,因为洛伦兹变换不是正交变换,它的转置不等于它的逆。论文解决这个问题的办法是,在key的变换里额外插入一个度量矩阵(就是那个让时间和空间符号相反的矩阵),靠这个度量把"转置"硬是转换成了"求逆"。这一步插入看似是个数学技巧,实际上是整套方法能成立的命门。如果省掉这一步,后面所有"只依赖相对坐标"的漂亮性质都会崩掉。

这就好比你想证明两个人站在同一辆移动的扶梯上,相对距离不会变。如果你只是把两人的坐标各自做一次相同的"伸缩",他们之间的距离确实不变。但洛伦兹变换不是单纯的伸缩,它还带着一种扭曲(把时间和空间掺在一起)。如果你不往其中一边加一个修正,这种扭曲会让"相对距离不变"这件事落空。度量矩阵的插入,就是这个修正。

从单个维度到多维时空

真实世界的坐标往往不止一维。分子动力学里每个原子有三维空间坐标加一维时间,视频里每个图像块有二维空间坐标加一维时间。论文需要把刚才那套只处理一维标量参数的洛伦兹变换,扩展到能处理多维坐标的情形。

直接扩展会遇到麻烦。严格来说,如果想要多维坐标的每个方向都对应一个独立的洛伦兹变换,还要保证组合起来依然满足"相加等于相乘"这条规则,这些变换的生成元必须两两对易(专业说法叫"交换")。但物理上,沿不同方向的"推进"(boost)天然是不对易的,这是洛伦兹群自身的代数结构决定的,躲不开。

论文的解决办法是退一步求其次,只用一个一维的洛伦兹变换子群。具体做法是:先用一个可学习的线性投影,把多维时空坐标压缩成一个标量,然后用这个标量去生成一个标准的(1+1)维洛伦兹变换(说白了就是双曲余弦和双曲正弦构成的那种矩阵)。论文还证明了一件挺有意思的事:在二维的洛伦兹特征块里,任何光滑的、满足"相加对应相乘"性质的映射,最终都必然退化成这种"先线性投影成标量,再做一个参数化boost"的形式。这不是工程上的权宜之计,而是这类构造里唯一可行的通用解法。

对于更高维的特征空间(每个注意力头的维度往往远大于2),论文的做法是把特征切成若干个二维小块,每个小块配一个独立学到的投影方向,各算各的洛伦兹变换,互不干扰。这个设计有点像给每个二维小块发一把专属钥匙,钥匙对应的锁孔方向各不相同,但开锁的原理是一样的。如果不这样切块,而是强行用一个大维度的洛伦兹表示去统一处理,前面说的"生成元必须对易"这条硬性限制就会卡死整个设计,模型压根没法训练。

说到这里得补一句背景。论文附录里专门分析过伽利略对称性(非相对论物理里更自然的对称结构)能不能替代洛伦兹变换来做类似的位置编码。结果发现,如果坚持用"有限维度、保持内积不变(专业说法是幺正)、每个token单独变换"这种构造方式,伽利略对称性会被证明只能保留时间差和速度差的信息,却必然丢掉两个token之间最关键的相对空间位移。这是一个数学上证明过的障碍,不是没试过,而是试了发现此路不通。这也从侧面说明了为什么论文最终选择洛伦兹变换而不是看起来更贴近日常物理直觉的伽利略变换。

两个战场:分子和视频

讲了这么多数学构造,实际效果到底怎么样?论文选了两个几乎完全不同的任务来检验。

第一个战场是分子动力学预测,具体场景是蛋白质和配体(可以理解为潜在药物小分子)的结合动态。给模型看配体过去一段时间的运动轨迹,加上周围蛋白质的结构信息,让它预测配体接下来会漂到哪里去。

这里有个设定挺讲究,叫"半柔性"设置,蛋白质被当成固定不动的环境,只有配体在动。每个配体原子在每个观测时刻,都被当作一个独立的"事件token",坐标是(时间,x,y,z)。模型结构上先做一次配体原子对蛋白质残基的交叉注意力,再做一次配体内部跨时间的自注意力,MinkowskiPE就加在这两种注意力模块里。

实验数据用的是 MISATO 数据集,涵盖十个蛋白质配体体系的单轨迹预测,以及三个规模递增的多轨迹数据集(MISATO-100、MISATO-1000、MISATO-Full)。

单轨迹设置下,MinkowskiPE 在十个体系、三项指标共30项评估里,拿到17项最优,26项进入前两名。多轨迹设置下更是全面碾压,三个数据集共九项评估,全部拿到最优。论文用的三项评估指标分别是:坐标误差(MAE),衡量原子间距离匹配程度的 Matching 指标,以及衡量结构稳定性的 Stability 指标(预测的原子间距离和真实值偏差在0.5埃以内的比例)。

第二个战场是视频预测,用的是经典的 KTH 人体动作数据集,输入10帧黑白视频,预测接下来20帧。这里每一帧先被卷积编码器压缩成16乘16的特征网格,网格上每个特征点当作一个事件token,坐标是(时间,行坐标,列坐标)。

这个场景跟分子动力学的差别其实挺大。分子里空间坐标是连续的三维位置,视频里的空间坐标是规则网格上的离散位置。可论文想证明的恰恰是,只要把它们都看成"事件"(带时空坐标的token),同一套 MinkowskiPE 的构造就能套用,不需要为每个领域单独设计一套位置编码方案。

实验结果显示,在 MSE、MAE、PSNR、LPIPS 四项指标上 MinkowskiPE 都拿到最优,相比表现最好的基线方法,MSE降低了9.9%,而且只用了约十分之一的参数量(2.5M对比最优基线通常是十几M到近四十M)。

参数量差这么多却效果更好,这事说明了什么?往小了说,这是工程效率的胜利。往大了说,它提示我们:很多时候模型效果不好,未必是容量不够,而是归纳偏置(模型对问题结构的先验假设)没找对。MinkowskiPE相当于直接把"时空如何耦合"这条规律,用几何结构内置进了注意力计算里,模型不需要再费力气从海量数据里自己摸索这条规律,参数自然就能省下来。

下面是单轨迹预测里几个代表性体系的具体数据,方便直观感受差距:

方法 | 4K6W MAE | 1XP6 MAE | 4G3E MAE | 3B9S MAE

VerletMD | 15.428 | 13.753 | 5.111 | 19.473

GNN-MD | 3.695 | 2.378 | 2.709 | 2.578

NeuralMD-ODE | 3.656 | 1.924 | 2.191 | 3.039

BioMD | 3.362 | 2.273 | 2.535 | 4.243

MinkowskiPE | **2.477** | **1.911** | **1.818** | **2.486**

多轨迹预测的整体数据:

数据集 | 方法 | MAE | Matching | Stability

MISATO-1000 | DenoisingLD(此前最优) | 4.467 | 0.430 | 85.42

MISATO-1000 | MinkowskiPE | **4.367** | **0.415** | **86.01**

MISATO-Full | DenoisingLD(此前最优) | 4.400 | 0.412 | 85.64

MISATO-Full | MinkowskiPE | **4.315** | **0.392** | **86.29**

论文还专门做了一个对照实验,叫"控制变量消融"。就是把模型架构、训练流程全部不动,唯独把 MinkowskiPE 从注意力模块里摘掉,看效果掉多少。结果是:在 MISATO-1000 上,去掉位置编码后 MAE 从4.367涨到4.488,Stability从86.01降到84.96;在KTH上,MSE从35.78涨到37.93,LPIPS从0.18342涨到0.19203。所有指标全线下滑。这个实验的意义在于证明:前面那些亮眼的结果,不是模型架构本身厉害,位置编码白加了也一样好,而是实实在在来自这个几何设计的贡献。

这件事带来的启发

读完这篇论文,我最大的感触是:很多看起来"学科壁垒分明"的问题,拆到足够细的层次,可能共享着同一个数学骨架。分子和视频,一个属于化学信息学,一个属于计算机视觉,平时几乎不会出现在同一次学术会议的同一个分会场里。但作者愣是把它们都还原成"带时空坐标的token序列",然后用一套几何工具同时伺候。这种跨领域的统一视角,比具体的性能提升数字更让我觉得有价值。

还有一个细节值得单独说一说。论文在附录里花了大篇幅去证明伽利略对称性走不通,这其实是一种挺诚实的学术态度,很多论文只会展示"我的方法有效",很少有人愿意花篇幅写"我尝试过另一条路,证明它在数学上行不通"。这种负结果的呈现,对后来者少走弯路是有实际价值的,哪怕它不会出现在摘要的任何一句话里。

论文也坦白承认了局限:目前用的是固定的几何偏置,不同任务可能需要不同的几何结构才是最优解,而且现在这套一维洛伦兹子群的构造,表达能力本身也是被"生成元必须对易"这条数学约束卡死的,想要更丰富的多参数构造,得先解决这个代数上的死结。这是一条明确留着没填的坑。

Q&A

Q1：MinkowskiPE是什么？

A：MinkowskiPE是一种相对位置编码方法,它用时空坐标参数化洛伦兹变换,作用在Transformer的query和key特征上,让注意力分数只依赖两个token之间的相对时空位移,同时保持和FlashAttention等高效注意力实现兼容。

Q2：MinkowskiPE在哪些任务上做了验证？

A：论文在两类任务上做了验证,一类是蛋白质配体结合动力学预测(基于MISATO数据集),另一类是KTH人体动作视频预测。前者在全部九项多轨迹评估中都拿到最优,后者MSE比最优基线降低9.9%,且只用约十分之一参数量。

Q3：MinkowskiPE会不会假设数据真的遵循相对论物理规律？

A：不会。论文明确说明闵可夫斯基几何在这里只是一种表示层面的归纳偏置,事件坐标本身不被变换,只是用来生成作用在特征上的洛伦兹变换,模型并不假设分子运动或视频画面遵守相对论。