跳转至

Variable-Length Generative Protein Design via Generalized Poisson Flow

会议: NeurIPS 2026
arXiv: 2607.09039
OpenReview: https://openreview.net/forum?id=znfwhqtvhg
代码: https://github.com/
领域: 计算生物
关键词: 蛋白质设计、变长生成模型、泊松流、流匹配、多模态生成

一句话总结

本文提出广义泊松流(GPFlow),将非齐次广义泊松过程与组分生成流耦合,建立了一个支持连续坐标、离散序列及黎曼流形等多模态的免模拟变长蛋白质生成框架,在无条件结构/序列设计、基序脚手架与多肽共设计中显著超越定长基线。

研究背景与动机

近年来,基于扩散模型与流匹配(Flow Matching)的生成方法在从头蛋白质设计、构象采样及骨架生成中取得了突破性进展。然而,主流蛋白质生成范式(如 FrameDiff、Proteina、RFdiffusion 等)在数学公式上几乎均严格限定在固定维度的状态空间中运行。这意味着在进行采样生成前,用户必须预先人为指定目标蛋白质的氨基酸链长。这种定长约束与实际蛋白质工程需求存在尖锐矛盾:在基序脚手架(Motif Scaffolding)以及功能多肽设计等核心下游任务中,实现特定结合构象或折叠稳定性的最优支架长度预先完全未知,且长度本身与结构的物理可设计性(Designability)深度耦合。若采用对候选长度区间进行离散网格扫描的暴力策略,不仅推断开销成倍暴增,且一旦预选长度与条件基序的空间拓扑约束不兼容,模型就会彻底错失可行的设计解空间。

变长生成的瓶颈在于如何构建一个严谨自洽且在多模态下易于学习的变维随机概率路径。已有探索要么仅能处理连续欧氏空间的跳跃扩散(如 TDDM,无法自然扩展至离散序列),要么仅适用于离散编辑操作(如 EditFlow,无法处理连续三维原子骨架),或者依赖需要逐步删除残基的自回归/反向删除机制(如 SCISOR)。蛋白质设计本质上是一个典型的混合几何问题,既包含离散氨基酸类别,又包含连续的欧氏主链三维坐标,甚至涉及旋转群 \(\mathrm{SO}(3)\) 与扭转角环面 \(\mathbb{T}\) 等黎曼流形结构,亟需一个统一的概率流表述。

本文的核心切入点是将蛋白质链长的动态增长刻画为非齐次广义泊松过程,并与残基组分内部的演化生成元(Within-length Generator)紧密耦合。核心 idea:将蛋白质生成表述为非齐次广义泊松过程与多模态流匹配的联合概率流,通过对辅助插入时间的解析边际化实现免模拟(simulation-free)的负对数似然训练,并在群体层面上严格保证联合多模态分布的恢复与 KL 散度上界。

方法详解

整体框架

GPFlow 将变长蛋白质状态表示为二元组 \(Y_t = (X_t, S_t) \in \mathsf{Y}\),其中 \(X_t \in \mathbb{N}\) 表示当前时间点 \(t \in [0, 1]\) 处的残基长度,而 \(S_t = (S_t^1, \dots, S_t^{X_t}) \in \mathsf{S}^{X_t}\) 表示具有 \(X_t\) 个残基的多模态特征。这里的组分空间 \(\mathsf{S}\) 具有通用性:它可以是 \(\mathrm{C}_\alpha\) 的连续笛卡尔坐标(欧氏空间)、离散氨基酸残基类别(计数测度)、主链刚体方向与侧链二面角(黎曼流形),或是上述空间的笛卡尔积。

在采样生成时,模型从无残基的空状态 \(Y_0 = (0, \emptyset)\)(在无条件设计下)或给定功能基序状态 \(Y_0 = (X_c, S_c)\)(在条件设计下)出发,随时间推进协同演化。在两个离散插入事件之间,已有残基的状态沿着连续或离散的组分生成元 \(\mathcal{B}_t\) 持续更新;与此同时,非齐次泊松过程以预测速率 \(\lambda_t\) 触发残基插入事件。一旦事件发生,模型从学习到的采样分布 \(\rho_t\) 中抽取新残基的特征值及其在序列中的插入位置,并无缝融入当前多模态表征中,直至 \(t=1\) 生成具有天然合理长度与物理稳定性的完整蛋白质。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始状态<br/>Y0: 空序列 或 条件基序残基"] --> B["长度泊松流机制<br/>非齐次泊松过程驱动链长演化"]
    B --> C["保序插入与多模态耦合<br/>按槽位速率预测插入位置与新残基特征"]
    C --> D["组分流匹配与生成元精化<br/>连续坐标/黎曼流形/离散类型协同演化"]
    D -->|时间积分 t 从 0 推进至 1| E["输出目标蛋白质<br/>Y1: 自适应长度的结构与序列"]

关键设计

1. 长度泊松流与免模拟似然学习:将链长演化解析边际化至闭式损失
针对固定长度约束与连续模拟训练开销巨大的痛点,GPFlow 建立了一个前向吸收的广义泊松过程。设 \(\kappa_t: [0, 1] \to [0, 1]\) 为单调非减调度器(满足 \(\kappa_0 = 0, \kappa_1 = 1\)),残基插入的瞬时风险函数定义为 \(h_t = \frac{\dot{\kappa}_t}{1 - \kappa_t}\)(当 \(t < \tau_\kappa\))。对于终点长度为 \(L\) 的目标蛋白,条件事件速率严格遵循线性衰减关系 \(\lambda_t(X_t \mid L) = (L - X_t) h_t\)。数学上可严格证明该过程在任意中间时刻服从二项分布 \(X_t \sim B(L, \kappa_t)\),并在饱和时刻 \(t \geq \tau_\kappa\) 几乎必然被吸收于真实目标长度 \(L\)。进一步地,利用随机点过程的积分平滑公式,作者证明了对所有辅助插入时间 \(\{t_k\}\) 取期望后,过程的负对数似然可完全转化为对单步时刻 \(t\) 与二项采样状态 \(X_t\) 的闭式期望: $\(\mathcal{L}_{\text{GP}} = \mathbb{E}_{t, X_1 \sim q(X)} \left[ \lambda_t - h_t (X_1 - X_t) \log \lambda_t \right]\)$ 该设计彻底摆脱了传统连续时间跳跃过程繁重的数值模拟轨迹依赖,使模型能够像标准流匹配一样以单步采样高效训练。

2. 保序插入与多模态耦合生成元:统一处理欧氏、黎曼与离散残基空间
蛋白质分子具有严格的从 N 端到 C 端的残基拓扑顺序,单纯的无序点集插入会破坏序列共价骨架的几何连续性。GPFlow 通过将插入变量拓展为槽位-取值对 \(m = (i, a)\),在当前包含 \(k\) 个残基的状态上定义了 \(k+1\) 个相互独立的保序插入槽位 \(\Omega_0, \dots, \Omega_k\)。每个槽位分配有独立的条件插入速率 \(\lambda_t^i(y \mid Y_1) = |\Omega_i| h_t\) 与局部采样分布 \(\rho_t^i\)。在残基特征演化层面,GPFlow 针对不同物理模态量身定制流匹配生成元:对于连续的 \(\mathrm{C}_\alpha\) 坐标采用最优传输向量场回归;对于旋转群 \(\mathrm{SO}(3)\) 和扭转角环面 \(\mathbb{T}\) 采用黎曼测地插值与切空间向量场匹配;对于离散氨基酸类别则通过单纯形软连续嵌入或连续时间马尔可夫链转移核进行统一训练。整体优化目标由泊松速率损失、新插入特征重构损失与组分流匹配损失加权组合而成: $\(\mathcal{L} = \mathcal{L}_{\text{GP}} + w_{\text{rec}} \mathcal{L}_{\text{rec}} + w_{\text{FM}} \mathcal{L}_{\text{FM}}\)$ 这种解耦而协同的设计保证了模型能够以相同机理自洽处理从主链骨架到全原子侧链的复杂生物大分子表征。

3. 训练动力学的生成元 KL 散度上界与分布恢复理论保证
以往生成模型对于跳跃核和速率损失函数的选取多停留在启发式设计层面。本文从马尔可夫过程无穷小生成元(Generator)的数学视角,严格建立了真实数据流与生成模型流之间的积分 KL 散度显式上界: $\(\mathcal{D}_{\text{KL}}(p_T \parallel \tilde{p}_T) \leq \int_0^T \mathbb{E}_{x \sim p_t} \left[ \frac{\|u_t(x) - \tilde{u}_t(x)\|_2^2}{2\sigma_t^2} + \mathcal{D}_{\text{KL}}(Q_t \parallel \tilde{Q}_t) \right] \mathrm{d}t\)$ 其中跳跃测度的广义 KL 项恰好精确分解为速率项的广义 KL 与速率加权下的采样分布 KL:\(\mathcal{D}_{\text{KL}}(Q_t \parallel \tilde{Q}_t) = \mathcal{D}_{\text{KL}}(\lambda_t \parallel \tilde{\lambda}_t) + \lambda_t \mathcal{D}_{\text{KL}}(\rho_t \parallel \tilde{\rho}_t)\)。作者证明了经验 NLL 速率损失 \(\mathcal{L}_{\text{GP}}\) 的参数梯度与广义 KL 散度的期望梯度在数学上完全等价。这不仅为使用广义 KL 而非 MSE 提供了坚实的信息论解释,而且从群体统计层面证明了最优参数能够精确恢复真实的目标联合多模态分布。

损失函数 / 训练策略

模型采用端到端多任务联合优化,总目标函数包含三项: $\(\mathcal{L} = \mathcal{L}_{\text{GP}} + w_{\text{rec}} \mathcal{L}_{\text{rec}} + w_{\text{FM}} \mathcal{L}_{\text{FM}}\)$ 其中速率损失 \(\mathcal{L}_{\text{GP}}\) 监督插入事件的发生概率;特征重构损失 \(\mathcal{L}_{\text{rec}} = \mathbb{E}_{Y_1, Y_t} \left[ -h_t \sum_{k=1}^{X_1 - X_t} \log \rho_t(Y_t, s_k) \right]\) 监督待插入残基的几何与类型分布;组分流匹配损失 \(\mathcal{L}_{\text{FM}}\) 监督已有残基在欧氏空间及黎曼切空间上的速度向量场。训练时 \(t \sim \mathcal{U}[0, 1]\),中间状态通过二项分布直接解析采样,无需任何常微分/随机微分方程的推进模拟。

实验关键数据

主实验

论文在无条件结构生成(PDB 与 AFDB 数据集,生成 500 个样本且不提供长度输入)以及 RFdiffusion 基序脚手架(16 个严苛基序任务,每任务采样 1000 次并通过 FoldSeek 聚类评估独特成功簇)上与主流方法展开全面评测。

数据集 / 任务 评估指标 GPFlow (本文) 对应基线 (Proteina) 业界代表 (RFdiffusion / FrameFlow) 增益 / 优势
PDB (无条件结构) 可设计性 (Designability, %) ↑ 96.1 77.6 (Proteina-60M) 94.4 (RFDiffusion) / 88.6 (FrameFlow) 相比基线显著提升 +18.5%
AFDB (无条件结构) 可设计性 (Designability, %) ↑ 94.5 90.6 (Proteina-60M) - 相比对应基线提升 +3.9%
PDB (无条件结构) 多样性 (Diversity) ↑ / 新颖性 (PDB-TM) ↓ 0.26 / 0.76 0.25 / 0.78 0.46 / 0.71 (RFDiffusion) 新颖性更优 (TM 越低越新)
基序脚手架 (16项) 独特成功簇数 (Unique Success Avg/Total) ↑ 117.4 / 1878 94.9 / 1519 41.5 / 664 (RFD) / 51.1 / 817 (FrameFlow) 总成功簇超出基线 +23.6%,大幅领先业界
基序脚手架 (16项) 排名第一的任务数 (Rank 1 Tasks) ↑ 10 / 16 4 / 16 1 / 16 (RFD) / 1 / 16 (FrameFlow) 覆盖 62.5% 的任务标杆

消融实验

在多肽序列-结构共设计任务(Peptide Co-Design)中,模型需在不给定多肽链长的情况下,依据受体蛋白口袋微环境联合生成全原子多肽结构与序列。下表展示了 GPFlow 在摆脱真值长度先验(GT-length Oracle)约束下的核心性能表现:

模型配置 氨基酸恢复率 (AAR, %) ↑ 骨架 RMSD (Å) ↓ 二级结构比率 (SSR, %) ↑ 可设计性 (Designability, %) ↑ 必须依赖真值长度?
ProteinGenerator [35] 28.18 4.36 58.89 72.66 ✓ (需要)
RFdiffusion + MPNN [49] 26.69 4.17 62.70 84.72 ✓ (需要)
PepFlow 基线 [29] 46.35 1.99 78.98 37.27 ✓ (需要)
GPFlow (完整模型) 53.71 1.39 91.08 48.19 ✗ (完全自适应)

关键发现

  • 链长边缘分布的自涌现精确恢复:在无条件结构与序列生成中,GPFlow 无需外部指定长度分布先验,采样过程完全由学得的速率函数驱动,生成的链长直方图与训练集真实分布呈现近乎完美的重合,证明了联合分布建模的有效性。
  • 基序脚手架解空间的爆炸性扩展:在诸如 1BCF、1PRW、3IXT 等基序脚手架任务上,固定长度模板基线(如 Proteina)往往坍缩至极少数甚至单一聚类簇中,而 GPFlow 通过自适应探索可变支架长度,独特成功簇数取得了 10 倍至数十倍的惊人飞跃(如 3IXT 从 Proteina 的 8 簇猛增至 398 簇),并且能够自然生成超过 200 个残基的高质量长支架蛋白。
  • 无需长度 Oracle 的多肽鲁棒设计:在全原子多肽共设计中,定长基线(PepFlow)因继承真值索引容易在非天然残基剔除处产生空洞断裂,而 GPFlow 能敏锐捕捉口袋几何拓扑生成物理连续的紧密多肽,在摆脱长度先验的情况下骨架 RMSD 压低至 1.39 Å,AAR 提升至 53.71%。

亮点与洞察

  • 解析边际化带来的免模拟高效训练:将非齐次泊松过程在时间维度的随机跳跃通过解析积分消除,使得变维随机动力学的似然优化可以像标准连续流匹配一样进行单步无偏梯度估计,兼具数学优雅性与工程实用性。
  • 跨越连续、离散与黎曼模态的大一统视角:打破了先前变维生成算法局限于单一模态(TDDM 局限于连续坐标、EditFlow 局限于离散 token)的技术壁垒,通过测度论生成元框架将笛卡尔坐标、氨基酸类别与刚体旋转角完美融合。
  • 理论指引的广义 KL 散度与生成元上界:首次从点过程对数似然严谨推导出广义 KL 散度在速率匹配中的必然性,建立了连接有限步生成动力学与真实分布逼近误差的严密数学理论。

局限与展望

  • 作者承认的局限:模型性能对前向路径调度器(Scheduler \(\kappa_t\))的选择较为敏感,特别是在多模态混合场景下,连续坐标收敛速度与离散残基插入速率之间的步调失谐可能会影响局部结合界面的致密性。
  • 实验与现实场景局限:在无条件生成大尺度样本时,虽然采用了基于局部概率路径的启发式改进,但在超长序列(如 >1000 残基)或超复杂多聚体界面上的连续碰撞惩罚依然缺乏显式的物理力场约束。
  • 未来改进方向:可进一步引入化学反应体系中的加速随机模拟算法(如 \(\tau\)-leaping 近似),以及探索结构-序列多尺度层次化跳跃过程,提升超大蛋白复合物变长设计的推理效率与刚体精度。

相关工作与启发

  • vs Proteina / RFdiffusion: 后者作为当今主流的骨架生成大模型,其数学核心受限于欧氏或刚体流匹配的定维假设,在脚手架任务中必须依赖繁琐的人工 contig 长度扫描;GPFlow 在保留其强大主干网络表征能力的同时,通过外挂轻量速率预测头实现了真正的端到端变长自适应探索。
  • vs TDDM / EditFlow: TDDM 针对连续变量采用跳跃扩散核且依赖模拟近似,无法处理离散氨基酸 token;EditFlow 仅适用于纯离散词表编辑操作。GPFlow 在理论上统一了二者的重叠特例,并首次将统一概率流推展至包含流形与离散类别的复杂生物物理大分子体系。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出多模态统一广义泊松流,建立了变长蛋白质生成的严密概率理论与免模拟学习范式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 横跨无条件结构、无条件序列、基序脚手架以及全原子多肽共设计四大严苛基准,证据扎实详尽。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽,从测度论、无穷小生成元到算法实现层层递进,结构清晰流畅。
  • 价值: ⭐⭐⭐⭐⭐ 彻底打破了蛋白质生成模型长期依赖预设长度的根本枷锁,为实用化生物计算与药物分子设计提供了关键工具。