基于广义泊松流的变长生成式蛋白质设计

# 变长生成式蛋白质设计：基于广义泊松流  
来源：https://arxiv.org/html/2607.09039  

\\useunder \\ul  

\\newkeytheoremtheorem\[name=定理, numberwithin=section\]  
\\newkeytheoremproposition\[name=命题, numberlike=theorem\]  

周安然∗，张涵尼∗，曲妍如∗，陈宇新，郭瑞涵，范佳俊，刘戈  
西贝尔计算与数据科学学院  
伊利诺伊大学厄巴纳-香槟分校  

###### 摘要  

在蛋白质设计中，生成可变长度的蛋白质至关重要，因为最优长度通常未知且与可设计性紧密耦合。现有的扩散和基于流的生成模型通常需要在采样前指定蛋白质长度，限制了其在探索可行设计空间时的灵活性。为解决这一局限，我们提出了**广义泊松流 (GPFlow)**，这是一个变长生成框架，通过最小化负对数似然来学习非齐次广义泊松过程的速率函数。我们建立了恢复联合多模态分布的整体性保证，并推导了数据分布与生成分布之间KL散度的上界。我们综合评估了GPFlow在结构设计、序列设计、基序支架构建和多肽共设计中的表现，涵盖欧几里得、分类和黎曼模态，以全面验证其变长生成质量。在无条件设计中，与相应的定长基线相比，GPFlow改善了结构可设计性并在序列设计中实现了最佳分布拟合，同时完美恢复了长度分布。在条件性基序支架构建中，GPFlow在16个基于结构的设计任务中有10个排名第一，取得了显著更多的独特成功案例，并且在基于序列的设计中也获得了更多通过的任务。在多肽共设计中，即使没有天然长度预言机，GPFlow仍保持竞争力。  

††注释：通讯作者：[email protected]。  

![[未标注图片]](https://arxiv.org/html/2607.09039v1/x1.png)  

## 1 引言  

扩散模型 [22 (https://arxiv.org/html/2607.09039#bib.bib22),42 (https://arxiv.org/html/2607.09039#bib.bib42)] 和流匹配 [34 (https://arxiv.org/html/2607.09039#bib.bib34)] 已成为蛋白质生成建模的广泛使用框架，并在蛋白质设计中取得了强劲结果 [40 (https://arxiv.org/html/2607.09039#bib.bib40),18 (https://arxiv.org/html/2607.09039#bib.bib18),31 (https://arxiv.org/html/2607.09039#bib.bib31)]。然而，它们的标准公式在固定维度的状态空间上运行。这种约束对于像基序支架构建这样的任务而言是限制性的，因为最优长度往往是*先验*未知的，并且与可设计性紧密耦合。在预先指定的长度范围内进行采样扫描会增加推理成本，并且当所选长度与条件约束不兼容时可能遗漏可行的设计。为弥补这一差距，我们提出了**广义泊松流 (GPFlow)**，这是一个变长生成框架，其中长度在非齐次广义泊松过程下演化。GPFlow将此长度过程与长度内的生成器耦合，从而在同一结构中容纳连续、离散、黎曼以及混合的长度依赖模态。基于条件过程的边缘化定理，我们从随机过程的负对数似然 (NLL) 推导出可计算的目标，并进一步将构造扩展到多模态生成，且具有联合分布恢复的保证。此外，我们推导了一个基于生成器的KL上界，其项与加性损失对齐，从而在最优边缘生成器处补充了种群级恢复保证。虽然最终的速率目标与现有工作一致，但我们的公式为连续、离散、黎曼和混合模态提供了统一的概率流构造，具有更深入的理论联系。  

我们考虑五种不同的蛋白质设计场景：(a) 无条件结构设计 (第5.1节 (https://arxiv.org/html/2607.09039#S5.SS1))，(b) 无条件序列设计 (第5.2节 (https://arxiv.org/html/2607.09039#S5.SS2))，(c) 基于结构的基序支架构建 (第5.3节 (https://arxiv.org/html/2607.09039#S5.SS3))，(d) 基于序列的基序支架构建 (附录E.4 (https://arxiv.org/html/2607.09039#A5.SS4))，以及 (e) 肽结构-序列共设计 (第5.4节 (https://arxiv.org/html/2607.09039#S5.SS4))。这一任务谱系评估了相同的变长构造如何适应不同的模态和生成目标。通过保留相应的定长基础架构，我们可以直接与其变长扩展进行比较。在这些任务中，我们的实验提供了证据，表明动态长度建模可以在评估的设置中改进相应的定长系统，同时消除指定目标长度的需要。对于无条件生成，我们的小型结构模型显著提高了相应Proteina基础模型的可设计性，而642M序列模型更接近UniRef50统计量并保持了高于DPLM基础模型的分集度。在基序支架构建基准中，结构模型在16个任务中有10个排名第一，并在几个具有挑战性的目标上取得了超过10倍的独特成功案例增长，而序列模型也获得了更多通过的任务。在多肽共设计中，平移、旋转、残基类型和扭转角与插入共同学习，GPFlow即使在没有天然长度预言机的情况下也保持竞争力并超越了基础PepFlow模型。  

我们的贡献总结如下：  
- • **统一的变维概率流**。我们在GPFlow中用非齐次广义泊松概率路径制定了变长生成建模，并将其与连续、离散、黎曼或混合长度内动态耦合。单一边缘化原则涵盖了这些模态类型以实现多模态生成。  
- • **具有理论保证的无模拟似然学习**。我们从广义泊松轨迹似然推导出速率目标，同时解析地边缘化其辅助插入时间，因此训练完全无模拟。我们进一步证明了联合多模态分布的种群级恢复，并建立了基于生成器的KL上界。  
- • **跨五个蛋白质设计设置的变长生成**。我们实例化了该框架，用于蛋白质结构设计、序列设计、基序支架构建以及肽结构-序列共设计。在所有任务中，与定长基础模型的直接比较表明了GPFlow优越的性能和灵活性。我们实现了优越的结构可设计性、更好的序列拟合度、基序支架构建中更多的独特成功和通过任务，以及无需天然长度预言机的稳定肽设计。  

## 2 相关工作  

##### 蛋白质设计模型  
基于扩散和流匹配，最近的蛋白质设计生成模型取得了强劲的性能。FrameDiff [51 (https://arxiv.org/html/2607.09039#bib.bib51)] 采用黎曼扩散来学习每个氨基酸的位置和方向，FrameFlow [50 (https://arxiv.org/html/2607.09039#bib.bib50)] 进一步将其扩展为黎曼流匹配。Proteina [18 (https://arxiv.org/html/2607.09039#bib.bib18)] 是一个仅用CA的基于流的结构生成模型，实现了显著的缩放结果。RFDiffusion [49 (https://arxiv.org/html/2607.09039#bib.bib49),8 (https://arxiv.org/html/2607.09039#bib.bib8)] 是一系列扩散模型，将RoseTTAFold [3 (https://arxiv.org/html/2607.09039#bib.bib3)] 的结构预测能力融入生成建模。通常，这些结构生成模型也支持基序支架构建，其中功能重要的蛋白质结构片段，称为*基序*，作为条件提供给模型。还有专门用于基序支架构建的模型，例如Genie2 [31 (https://arxiv.org/html/2607.09039#bib.bib31)] 是一个扩散生成模型，具有专门的多基序框架，用于设计共现基序。这些模型要么采用修复风格的采样，通过固定基序残基和结构来施加硬约束，要么将基序信息作为显式但软的条件信号注入。额外的序列和肽设计模型见附录C (https://arxiv.org/html/2607.09039#A3)。  

##### 变长生成模型  
标准的扩散或流匹配模型需要预定的采样长度，限制了其对变长数据的适用性。非扩散方法主要依赖自回归建模 [47 (https://arxiv.org/html/2607.09039#bib.bib47)]。现有两项工作通过相关的随机过程构造融入了数据长度。TDDM [9 (https://arxiv.org/html/2607.09039#bib.bib9)] 用与连续扩散动态耦合的可学习跳转核来建模维度变化，但未处理分类的长度依赖模态。EditFlow [20 (https://arxiv.org/html/2607.09039#bib.bib20)] 则从离散流匹配 [17 (https://arxiv.org/html/2607.09039#bib.bib17)] 推导出插入速率；它适用于分类模态，但不适用于连续坐标。重要的是，这两个模型均未应用于蛋白质设计领域。最近，SCISOR [4 (https://arxiv.org/html/2607.09039#bib.bib4)] 通过基于删除的过程处理变长蛋白质序列生成，其中在前向加噪时插入随机残基，由学习的反向过程移除。GPFlow主要在范围和概率解释上有所不同。它使用统一的变维概率流公式处理连续、离散、黎曼和混合模态。我们的速率目标来自广义泊松轨迹似然，其中辅助插入时间被解析地边缘化，从而产生无模拟的目标。在其重叠范围内，GPFlow恢复了TDDM和EditFlow的仅插入、特定模态的实例，并得到相同的速率目标。因此，我们的贡献在于统一的构造和似然解释以及关于训练动态的新KL界。虽然原始的TDDM和EditFlow论文未评估蛋白质设计任务，但我们的实验在五个综合设置中实例化了该框架，用于无条件及条件性蛋白质设计。  

## 3 广义泊松流  

为了使用流匹配（或扩散）为随机状态$Y_t$构建概率路径$\{p_t\}_{t\in[0,1]}$，需要 (a) $p_1\approx q$，其中$q$是数据分布，(b) 从$p_t$进行无模拟采样以实现可行训练，以及 (c) 用于学习边缘流的可计算损失。我们首先推导仅长度状态$Y_t=X_t\in\mathbb{N}$的理论结果，然后将其扩展到多模态状态$Y_t=(X_t,S_t)\in\mathsf{Y}$，其中$S_t$可以包含任意离散或连续的长度依赖模态或其混合。我们展示了如何使用广义泊松过程构建可行的概率路径。  

### 3.1 仅长度的广义泊松流  

一个*广义泊松过程* $\{X_t\}_{0\leq t\leq T}$是一个连续时间马尔可夫过程，由*速率函数* $\lambda_t:[0,T]\times\mathbb{N}\to\mathbb{R}^{\geq 0}$定义，使得$\Pr(X_{t+h}=X_t+1)=\lambda_t(X_t)h+o(h)$，且$\Pr(X_{t+h}=X_t)=1-\lambda_t(X_t)h+o(h)$。与流匹配惯例一致，我们假设广义泊松过程定义在时间区间$[0,1]$上，且$X_0=0$。我们现在通过以下定理展示使用广义泊松过程构建概率路径$p_t(X)$的一种可行构造：  

\{theorem\} [store=forwardrate]  
设$\kappa:[0,1]\to[0,1]$是一个非递减的绝对连续调度器，满足$\kappa_0=0$和$\kappa_1=1$，并定义*完成时间* $\tau_\kappa:=\inf\{t:\kappa_t=1\}$。对于任意固定的$L\geq 1$，假设对于$t<\tau_\kappa$有$\kappa_t<1$，并定义插入*风险*和*条件速率*为  

$$h_t:=\begin{cases}\dfrac{\dot{\kappa}_t}{1-\kappa_t},&t<\tau_\kappa,\\ 0,&t\geq\tau_\kappa,\end{cases}\qquad\lambda_t(X_t|L):=(L-X_t)h_t.$$  

那么，所得过程满足$X_t\sim B(L,\kappa_t)$（对于$t<\tau_\kappa$），并在$t\geq\tau_\kappa$时吸收到$X_t=L$。  

证明见附录A.1 (https://arxiv.org/html/2607.09039#A1.SS1)。定理3.1 (https://arxiv.org/html/2607.09039#S3.SS1) 提供了一种无模拟的方法来采样$X_t$，并保证$X_1=L$几乎必然成立。与连续模态的标准流匹配类似，我们建立以下边缘化定理：  

\{theorem\} [store=marginal]  
对于上述定义的任意调度器$\kappa_t$，设$q$是目标长度分布，定义$p_t(k|y):=\Pr(X_t=k|X_1=y)$，并设$p_t(k):=\sum_y p_t(k|y)q(y)$。对于每个满足$p_t(k)>0$的$k$，后验概率为$p(X_1=y|X_t=k):=p_t(k|y)q(y)/p_t(k)$。那么*边缘速率*  

$$\lambda_t^*(k):=\mathbb{E}_{X_1\sim p(\cdot|X_t=k)}[\lambda_t(k|X_1)]$$  

生成边缘路径$p_t$，从而$p_1=q$。  

证明见附录A.2 (https://arxiv.org/html/2607.09039#A1.SS2)，主要依赖作为连续性方程的科尔莫戈罗夫前向方程。定理1 (https://arxiv.org/html/2607.09039#S3.E1) 保证了任何长度分布$q$都可以用广义泊松过程建模，并且确保了数据约束$p_1=q$，用于在离散长度空间上构建概率流$p_t$。因此，我们称我们的生成框架为*广义泊松流*。为了有效学习该流，在命题A.3 (https://arxiv.org/html/2607.09039#A1.SS3)中，我们证明广义泊松过程实现路径的负对数似然可以计算为  

$$\mathcal{L}_{\text{NLL}}=\int_0^1 \lambda_t \mathrm{d}t - \sum_{k=1}^L \log \lambda_{t_k},$$  

其中两项可分别视为生存项和事件项的NLL。插入时间是用于构建概率路径的辅助变量，通过随机点过程的积分公式 [7 (https://arxiv.org/html/2607.09039#bib.bib7)]，插入时间上的期望和可以边缘化为对真实边缘速率$\lambda_t^*$的积分。结合定理1 (https://arxiv.org/html/2607.09039#S3.E1) 中的边缘化结果和公式 (1) (https://arxiv.org/html/2607.09039#S3.E1) 中的条件速率，我们有：  

$$\mathbb{E}_{\{t_k