SE(3)-MeanFlow：李群上的少步蛋白质骨架生成

# SE(3)-MeanFlow：李群上的少步蛋白质骨架生成  
来源：https://arxiv.org/html/2607.27431  
Binghang Lu²，Yikai Liu³，Elaheh Akbari⁴，Soheil Kolouri⁴，Linxuan Wang⁵，Ping He⁴，Shuchan Wang⁶，Ruqi Zhang⁷¹¹footnotemark:1，Guang Lin¹,2,5,7¹¹footnotemark:1  
1普渡大学数学系  
2普渡大学电子与计算机工程系  
3普渡大学生物化学系  
4范德堡大学计算机科学系  
5普渡大学统计系  
6柏林自由大学数学与计算机科学系  
7普渡大学计算机科学系  
通讯作者：[email protected], [email protected].

###### 摘要

蛋白质骨架的生成式建模有望实现具有预定结构和功能特性的蛋白质的从头设计。现有的扩散模型和流匹配模型可以在SE(3)N\mathrm{SE}(3)^{N}上生成高质量骨架，但推理需要对ODE进行数百次网络评估的数值积分，每一步都涉及李群指数映射——这成为高通量设计管线中的瓶颈。我们引入SE(3)-MeanFlow，这是一个少步生成框架，将MeanFlow从欧几里得空间扩展到蛋白质坐标系的李群几何。在李代数so(3)\mathfrak{so}(3)和R3\mathbb{R}^{3}中天然地工作，我们推导了旋转和平移的闭式平均速度恒等式，从而得到无模拟训练目标。我们进一步引入SE(3)\mathrm{SE}(3)α\alpha-Flow目标，它消除了旋转分支中的雅可比-向量积，并作为预训练阶段；之后训练切换到小tt稳定的MeanFlow损失，该损失用于预训练的剩余部分以及基于矫正（rectification）的后训练。在蛋白质骨架生成中，SE(3)-MeanFlow匹配或超过使用数倍采样步数的流匹配基线，并且在少步区域内优势扩大，其中矫正使其在每一个匹配预算下都领先——代价是多样性略有降低。

## 1 引言

蛋白质是生命的基本组成部分之一。它们复杂的几何结构促成了特定的分子间相互作用，从而实现关键的生物学功能——作为化学反应中的催化剂、分子的转运体以及免疫反应的介导者。随着计算技术的出现，合理地设计具有所需结构的新型蛋白质以编程其功能成为可能，为应对长期存在的全球健康挑战开辟了道路，包括流感（Strauch et al., 2017 (https://arxiv.org/html/2607.27431#bib.bib1)）、COVID-19（Cao et al., 2020 (https://arxiv.org/html/2607.27431#bib.bib4)）和癌症免疫治疗（Silva et al., 2019 (https://arxiv.org/html/2607.27431#bib.bib6)）。

蛋白质骨架可以建模为NN个刚体的序列，每个残基一个刚体，每个刚体在保持方向的刚性变换下关联一个坐标系——即特殊欧几里得群SE(3)\mathrm{SE}(3)（Jumper et al., 2021 (https://arxiv.org/html/2607.27431#bib.bib7)）。因此，整个骨架由乘积群SE(3)N\mathrm{SE}(3)^{N}描述，而蛋白质从头设计问题就归结为从这个空间上的学习分布中采样。

近期工作在SE(3)N\mathrm{SE}(3)^{N}上的生成建模方面取得了重大进展。基于扩散的方法如FrameDiff（Yim et al., 2023b (https://arxiv.org/html/2607.27431#bib.bib8)）和RFDiffusion（Watson et al., 2023 (https://arxiv.org/html/2607.27431#bib.bib5)）实现了强可设计性，而流匹配方法——FoldFlow（Bose et al., 2024 (https://arxiv.org/html/2607.27431#bib.bib9)）和FrameFlow（Yim et al., 2023a (https://arxiv.org/html/2607.27431#bib.bib37)）——通过以无模拟方式学习SE(3)\mathrm{SE}(3)上的时变向量场，进一步提高了训练稳定性和灵活性。尽管取得了这些进展，所有现有的SE(3)\mathrm{SE}(3)生成模型都有一个共同的瓶颈：推理仍然需要对ODE进行多步数值积分，通常需要100–500次网络函数评估（NFE）。在SE(3)\mathrm{SE}(3)上，每一步都涉及网络评估和应用李群指数映射，使得推理比欧几里得空间中的推理昂贵得多。这限制了在高通量药物发现管线中的实际部署，因为每次活动必须生成数百万个候选结构。

减少推理步数已在欧几里得生成建模中得到广泛研究。一致性模型（Song et al., 2023 (https://arxiv.org/html/2607.27431#bib.bib11); Song and Dhariwal, 2024 (https://arxiv.org/html/2607.27431#bib.bib12)）和渐进蒸馏（Salimans and Ho, 2022 (https://arxiv.org/html/2607.27431#bib.bib13)）将推理压缩为一步或几步，但需要预训练的教师模型或精心分阶段的训练课程。MeanFlow（Geng et al., 2026a (https://arxiv.org/html/2607.27431#bib.bib35)）提供了一种更有原则性且自包含的替代方案。与标准流匹配（Lipman et al., 2022 (https://arxiv.org/html/2607.27431#bib.bib14); Tong et al., 2023 (https://arxiv.org/html/2607.27431#bib.bib20)）中对瞬时速度v(t,xt)v(t,x_{t})建模不同，MeanFlow引入了*平均速度*u(s,t,xt)u(s,t,x_{t})的概念——流轨迹在区间[s,t][s,t]上的平均速度。仅通过乘积法则和微积分基本定理，就从平均速度的定义推导出uu与vv之间的良定义恒等式。该恒等式产生了一个易于处理、无模拟的训练目标，仅使用瞬时速度作为监督。在推理时，整个流路径通过单次网络评估uθ(0,1,x1)u_{\theta}(0,1,x_{1})近似实现，从而无需蒸馏或预训练即可进行1-NFE生成。然而，MeanFlow在欧几里得空间Rd\mathbb{R}^{d}中表述，没有考虑到SE(3)\mathrm{SE}(3)的非平凡几何。将MeanFlow朴素地提升到弯曲流形是不适定的：轨迹上不同点的速度位于不同的切空间，其平均值需要沿路径进行平行移动。黎曼MeanFlow（记为RMF-PT）Zhong et al. (2026 (https://arxiv.org/html/2607.27431#bib.bib15))通过将MeanFlow扩展到一般黎曼流形来解决这个问题，利用平行移动定义平均速度，并推导出相应的黎曼MeanFlow恒等式。虽然具有原则性，但这个通用框架没有利用SE(3)\mathrm{SE}(3)作为李群的特定代数结构。另一项并行工作，黎曼MeanFlow（RMF）Woo et al. (2026 (https://arxiv.org/html/2607.27431#bib.bib23))在李代数中定义平均速度；我们在理论上的比较见附录E.6 (https://arxiv.org/html/2607.27431#A5.SS6)，经验上的比较见第4节 (https://arxiv.org/html/2607.27431#S4)。

在本工作中，我们提出SE(3)-MeanFlow，一个基于SE(3)\mathrm{SE}(3)李群结构的少步生成模型，用于蛋白质骨架的从头设计。我们的核心见解是分解SE(3)≅SO(3)×R3\mathrm{SE}(3)\cong\mathrm{SO}(3)\times\mathbb{R}^{3}使得MeanFlow恒等式可以在李代数so(3)\mathfrak{so}(3)和R3\mathbb{R}^{3}中分别推导，两者都允许闭式、无模拟的训练目标，无需平行移动。我们的主要贡献包括：

- •**基于积分的SE(3)-MeanFlow及其理论。**我们提出了一个基于积分的SE(3)-MeanFlow公式，以当前状态(Rt,xt)(R_{t},x_{t})为条件，并在李代数so(3)\mathfrak{so}(3)和R3\mathbb{R}^{3}中推导了闭式平均速度恒等式，无需平行移动。我们提供了理论验证，表明所得训练目标良定义，并且相应损失与SE(3)\mathrm{SE}(3)上的预期MeanFlow目标一致。
- •**针对蛋白质的稳定训练算法。**我们引入了SE(3)\mathrm{SE}(3)α\alpha-Flow目标以及实用的数值稳定技术，使蛋白质骨架的训练和少步生成可靠。
- •**蛋白质设计结果。**在SCOPe基准上，我们的方法匹配或超过使用数倍采样步数的流匹配基线，在少步区域内增益最大，在预训练和后训练设置中都是如此，而多样性损失适中。

## 2 背景与记号

### 2.1 SE(3)\mathrm{SE}(3)上的李群记号

在本工作中，SE(3)\mathrm{SE}(3)空间是刚体运动的集合
SE(3)≅SO(3)×R3,\mathrm{SE}(3)\cong\mathrm{SO}(3)\times\mathbb{R}^{3},
其中SO(3)\mathrm{SO}(3)表示3×33\times 3旋转矩阵的空间，构成一个李群：
SO(3)={R∈R3×3:RR⊤=I,det(R)=1}.\mathrm{SO}(3)=\{R\in\mathbb{R}^{3\times 3}:RR^{\top}=I,\ \det(R)=1\}.
我们遵循Bose et al. (2024 (https://arxiv.org/html/2607.27431#bib.bib9))，在定义度量和损失时使用解耦的SO(3)×R3\mathrm{SO}(3)\times\mathbb{R}^{3}几何。

SO(3)\mathrm{SO}(3)上的（连续时间）流可以定义为满足左平凡化ODE的曲线Rt∈SO(3)R_{t}\in\mathrm{SO}(3)
Ṙt=RtΩt,Ωt∈so(3),\dot{R}_{t}=R_{t}\,\Omega_{t},\qquad\Omega_{t}\in\mathfrak{so}(3),(1)
其中so(3)\mathfrak{so}(3)表示相应的李代数
so(3)={Ω∈R3×3:Ω⊤=−Ω}.\mathfrak{so}(3)=\{\Omega\in\mathbb{R}^{3\times 3}:\Omega^{\top}=-\Omega\}.
so(3)\mathfrak{so}(3)与R3\mathbb{R}^{3}空间同构，我们使用hat/vee映射(⋅)∧,(⋅)∨(\cdot)^{\wedge},(\cdot)^{\vee}来标识R3↔so(3)\mathbb{R}^{3}\leftrightarrow\mathfrak{so}(3)（例如，ω∈R3↦ω∧∈so(3)\omega\in\mathbb{R}^{3}\mapsto\omega^{\wedge}\in\mathfrak{so}(3)）。给定沿轨迹的角速度场Ωt:[0,1]→so(3)\Omega_{t}:[0,1]\to\mathfrak{so}(3)，方程 (1 (https://arxiv.org/html/2607.27431#S2.E1)) 的解可以写为
Rt=R0Texp(∫0tΩτdτ),R_{t}=R_{0}\,\mathcal{T}\exp\Big(\int_{0}^{t}\Omega_{\tau}\,d\tau\Big),
其中Texp\mathcal{T}\exp沿时间复合无穷小旋转。这里exp\exp表示从李代数到李群的（矩阵）指数映射。特别地，如果Ωτ≡Ω\Omega_{\tau}\equiv\Omega是常数，则
Rt=R0exp(tΩ).R_{t}=R_{0}\exp(t\Omega).
关于SO(3)\mathrm{SO}(3)/SE(3)\mathrm{SE}(3)流的详细讨论（包括常见度量下的闭式/测地线表达式），参见附录B (https://arxiv.org/html/2607.27431#A2)¹¹11本论文中引用的所有附录均包含在技术补充材料中。

### 2.2 蛋白质骨架参数化

遵循Bose et al. (2024 (https://arxiv.org/html/2607.27431#bib.bib9))中的设置，我们将蛋白质骨架参数化为NN个刚体的序列。每个残基i∈[N]i\in[N]关联一个坐标系Ti=(Ri,xi)∈SE(3)T_{i}=(R_{i},x_{i})\in\mathrm{SE}(3)，其中Ri∈SO(3)R_{i}\in\mathrm{SO}(3)表示方向，xi∈R3x_{i}\in\mathbb{R}^{3}表示CαC_{\alpha}原子的位置。骨架原子{N,Cα,C,O}i\{N,C_{\alpha},C,O\}_{i}的3D坐标通过对一组理想化坐标{N∗,Cα∗,C∗,O∗}\{N^{*},C_{\alpha}^{*},C^{*},O^{*}\}应用刚性变换TiT_{i}来恢复：
[N,Cα,C,O]i=Ti∘[N∗,Cα∗,C∗,O∗],[N,C_{\alpha},C,O]_{i}=T_{i}\circ[N^{*},C_{\alpha}^{*},C^{*},O^{*}],(2)
其中Cα∗=(0,0,0)C_{\alpha}^{*}=(0,0,0)固定在原点。关于SO(3)\mathrm{SO}(3)/SE(3)\mathrm{SE}(3)几何和记号的自包含回顾（包括(⋅)∧(\cdot)^{\wedge}和(⋅)∨(\cdot)^{\vee}的定义），参见附录B (https://arxiv.org/html/2607.27431#A2)–C (https://arxiv.org/html/2607.27431#A3)。

### 2.3 欧几里得空间中的MeanFlow

上述流匹配方法的一个局限性是，为了保证精度，我们通常要求步长1T\frac{1}{T}足够小，等价于需要大的推理步数TT。为了解决这个问题，在欧几里得空间中，Geng et al. (2026a (https://arxiv.org/html/2607.27431#bib.bib35),b (https://arxiv.org/html/2607.27431#bib.bib27))提出了MeanFlow方法。具体来说，给定由速度vtv_{t}生成的概率路径ptp_{t}，它定义平均速度：
(t−s)vavg(s,t,xt)=∫stv(τ,xτ)dτ.(t-s)v^{\text{avg}}(s,t,x_{t})=\int_{s}^{t}v(\tau,x_{\tau})d\tau.
它产生训练损失：
E[‖vθavg(s,t,xt)+(t−s)sg(ddtv^avg(s,t,xt))−v(t,xt)‖2].\mathbb{E}\left[\left\|v^{\text{avg}}_{\theta}(s,t,x_{t})+(t-s)\,\text{sg}\!\left(\frac{d}{dt}\hat{v}^{\text{avg}}(s,t,x_{t})\right)-v(t,x_{t})\right\|^{2}\right].(3)
其中随机性由t,s∼Unif([0,1]),s≤tt,s\sim\text{Unif}([0,1]),s\leq t和(x0,x1)∼q(x_{0},x_{1})\sim q给出，qq是某个联合分布。默认情况下，qq是pdatap_{\text{data}}与ppriorp_{\text{prior}}之间的独立耦合，
ddtv^avg(s,t,xt)=∂∂tvθavg(s,t,xt)+∇xvθavg(s,t,xt)⋅vt.\frac{d}{dt}\hat{v}^{\text{avg}}(s,t,x_{t})=\frac{\partial}{\partial t}v^{\text{avg}}_{\theta}(s,t,x_{t})+\nabla_{x}v^{\text{avg}}_{\theta}(s,t,x_{t})\cdot v_{t}.(4)
此外，在Geng et al. (2026b (https://arxiv.org/html/2607.27431#bib.bib27))中，vtv_{t}可以被vθavg(t,t,xt)v^{\text{avg}}_{\theta}(t,t,x_{t})替代。这里xt=(1−t)x0+tx1x_{t}=(1-t)x_{0}+tx_{1}，v(t,xt)=x1−x0v(t,x_{t})=x_{1}-x_{0}，sg是停止梯度算子。直观地说，积分在训练过程中由参数化模型估计；因此，在推理步骤中，它产生T步生成，其中T∈NT\in\mathbb{N}：
x^s=xt−1Tvθavg(s,t,xt)\displaystyle\hat{x}_{s}=x_{t}-\frac{1}{T}v^{\text{avg}}_{\theta}(s,t,x_{t})
t=1−i/T,s=t−1/T,i=0,1,...(T−1).\displaystyle t=1-i/T,s=t-1/T,i=0,1,...(T-1).(5)
当设定T=1T=1时，它产生一步生成。

参见图1的说明：

图1：SE(3)-MeanFlow的少步推理。双时间网络fθf_{\theta}（参见附录N (https://arxiv.org/html/2607.27431#A14)）预测[s,t][s,t]上的平均速度而不是tt处的瞬时速度，因此一次评估将状态推进整个步长——在SO(3)\mathrm{SO}(3)上是李群指数，在R3\mathbb{R}^{3}上是欧几里得更新——得到NFE=T\mathrm{NFE}=T。使这一点准确的训练目标是命题1 (https://arxiv.org/html/2607.27431#Thmproposition1)及其无JVP的α\alpha-Flow替代（第3.3节 (https://arxiv.org/html/2607.27431#S3.SS3)）。

## 3 方法：SE(3)\mathrm{SE}(3)空间中的MeanFlow

由于SE(3)=SO(3)×R3\mathrm{SE}(3)=\mathrm{SO}(3)\times\mathbb{R}^{3}，R3\mathbb{R}^{3}分量退化为常规的欧几里得MeanFlow；因此，下面我们专注于介绍SO(3)\mathrm{SO}(3)上的MeanFlow。

设分布路径{pt∈P(SO(3)):t∈[0,1]}\{p_{t}\in\mathcal{P}(\mathrm{SO}(3)):t\in[0,1]\}由……生成。