跳转至

Flexible Flows for Biological Sequence Design

会议: NeurIPS 2026
arXiv: 2606.10543
OpenReview: https://openreview.net/forum?id=rxmkkL4trc
论文: https://arxiv.org/abs/2606.10543
领域: 计算生物
关键词: 离散流匹配, 生物序列设计, 结构化耦合, 隐空间引导, 变长编辑流

一句话总结

FlexFlow 针对生物序列设计提出具有生物先验结构化耦合与隐空间编辑速率参数化的离散流匹配模型,结合测试期 Dirichlet 操作预算缩放与连续隐空间分类器自由引导,在 DNA 增强子/启动子生成与 MHC-II 肽段设计上显著超越现有离散扩散与流匹配基线。

研究背景与动机

生物功能序列(蛋白质、RNA 和 DNA)的设计本质上是在遵循严格进化约束与生物物理规律的超高维离散空间中进行寻优。早期的自回归语言模型(如 ProtGPT2、ProGen2、DNABERT)通过大规模预训练验证了生成可行性,但其固有的序列自回归采样存在误差累积慢、采样吞吐受限以及对长程结构和全局模体协调性捕捉不足等问题。随后基于连续时间马尔可夫链(CTMC)的离散扩散模型与离散流匹配(Discrete Flow Matching, DFM)逐渐成为非自回归采样的主流,兼具快速并行解码与更优全局可控性的潜力。

然而现有的离散流生成方案在生物序列场景中面临三大约束瓶颈:首先,传统 DFM 普遍依赖均匀噪声先验或 Mask 耦合,完全忽略了生物分子演化过程中固有的突变偏好与生物化学相似性结构;其次,对于真实生物分子中普遍存在的变长序列生成(如不同长度的活性肽段或结合配体),现有模型缺乏在推理阶段对插入、删除和替换等编辑操作进行精细化微调的控制旋钮;最后,在引入转录谱或结合靶标等功能条件引导时,现有的离散无分类器引导(CFG)主要在离散速率空间按 token 进行几何插值,割裂了序列全局上下文依赖,极易破坏复杂调控模体或结合口袋的整体协同构象。

为了解决上述矛盾,本文的核心 idea 是将生物演化替代先验无缝注入离散流匹配的端点耦合分布,并引入全局隐变量参数化的变长编辑速率网络与连续隐空间无分类器引导,在完全保留流匹配闭式训练优势的同时实现全局结构协调与测试期编辑操作预算的可控调节。

方法详解

整体框架

FlexFlow 整体框架由三大核心部分构成:面向前向扩散过程的生物结构化耦合核、基于全局隐变量的变长编辑速率网络、以及基于连续隐空间的无分类器引导与 Dirichlet 测试期操作缩放机制。在训练阶段,模型将先验序列与真实数据序列通过生物演化替代矩阵(如 BLOSUM62 或 HKY85)构成的双随机转移核进行耦合采样,并在引入空白符的扩充对齐空间中使用 Bregman 散度目标进行闭式训练;在推理阶段,网络通过编码器抽取连续全局隐变量,联合调控全序列各位置的插入、删除与替换概率分布,并通过隐空间引导向量实现跨位置强协同的条件引导生成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入与条件<br/>噪声序列 xt 与条件标签 c"] --> B["结构化先验耦合<br/>BLOSUM62/HKY85 诱导演化先验偏置"]
    B --> C["连续隐空间编码器<br/>推断全局上下文高斯隐变量 r"]
    C --> D["隐空间无分类器引导<br/>条件/无条件隐表示连续插值"]
    D --> E["变长编辑速率参数化<br/>基于共享隐变量预测局部 ins/sub/del"]
    E --> F["Dirichlet 温度预算缩放<br/>测试期调整 ins/del/sub 混合权重"]
    F --> G["CTMC 逆向流生成<br/>生成高亲和力多肽或活性启动子 DNA"]

关键设计

1. 结构化先验耦合(Structured Coupling):注入生物演化替代偏置

传统离散流匹配使用无信息的均匀噪声作为起点 \(p_0(\mathbf{x})\),导致前向加噪过程抹平了生物残基间的真实突变距离。FlexFlow 构建了一个参数化转移核 \(\mathrm{K}_\gamma \in \mathbb{R}^{|\mathcal{T}| \times |\mathcal{T}|}\),其中 \(\gamma\) 索引针对不同模态的领域替代矩阵(例如蛋白质序列采用 BLOSUM62 对数几率矩阵转换得到的双随机矩阵,DNA/RNA 采用 JC69 或 HKY85 替代模型)。在前向加噪轨迹的端点采样中: $\((\tilde{\mathbf{x}}_0, \mathbf{x}_1) \sim \pi(\tilde{\mathbf{x}}_0, \mathbf{x}_1, \mathrm{K}_\gamma) = p(\tilde{\mathbf{x}}_0 | \mathrm{K}_\gamma) p_{\mathrm{data}}(\mathbf{x}_1), \quad p(\tilde{\mathbf{x}}_0 | \mathbf{x}_0, \mathrm{K}_\gamma) = \prod_{i=1}^N \mathrm{Cat}(\mathrm{K}_\gamma \mathbf{e}_{\mathbf{x}_0^i})\)$ 论文在理论上证明了当 \(\mathrm{K}_\gamma\) 退化为均匀全 1 矩阵时,该耦合严格等价于标准均匀独立耦合;而在生物替代核下,该图结构先验将离散单纯形噪声限制在生物进化偏好的局部邻域内,无需改动流匹配的目标函数与训练策略即可大幅降低生成路径在非物理空间漫游的无效探索。

2. 基于全局隐变量的编辑速率参数化与 Dirichlet 操作预算控制

为了原生支持不同长度的生物序列生成并克服直接参数化全局编辑转移矩阵在组合维度上的不可行性,FlexFlow 提出层次化编辑速率建模:给定当前带噪序列 \(\mathbf{x}_t\),先利用非线性编码器推断全局上下文连续隐变量 \(\mathbf{r} \sim \mathcal{N}(\Psi(\mathbf{x}_t), \sigma_r^2 \mathbf{I})\);各位置 \(i\) 上的编辑操作 \(\mathrm{op}^i \in \{\mathrm{ins}, \mathrm{sub}, \mathrm{del}\}\) 则条件独立地由局部状态与全局隐变量联合给出:\(p_\Phi(\mathrm{op}^i | x_t^i, \mathbf{r}) = \mathrm{Cat}(\Phi(x_t^i, \mathbf{r}))\)。通过在全局隐空间上进行边缘化,全序列编辑转移速率天然具备了全局长程上下文感知能力。进一步地,为了在测试期灵活应对不同长度约束(例如需要扩增变异或保守截短),FlexFlow 引入了 Dirichlet 先验温度调节向量 \(\boldsymbol{\beta}_i = [1/\tau^i_{\mathrm{ins}}, 1/\tau^i_{\mathrm{sub}}, 1/\tau^i_{\mathrm{del}}] \sim \mathrm{Dir}(\boldsymbol{\alpha})\),通过浓度参数 \(\boldsymbol{\alpha} = [\alpha_{\mathrm{ins}}, \alpha_{\mathrm{sub}}, \alpha_{\mathrm{del}}]\) 在测试期即插即用控制生成时插入与替换的全局预算平衡,免去昂贵的重训练。

3. 连续隐空间无分类器引导(Latent Classifier-Free Guidance)

在给定外部功能条件 \(c\)(例如转录起始信号或 MHC-II 等位基因靶标)时,传统的速率空间引导是在每个离散 token 上独立对条件与非条件速率做几何插值 \(\tilde{u}_t \propto (u_t^{\emptyset})^{1-w} (u_t^c)^w\),容易引发逐位震荡与局部模体崩塌。FlexFlow 另辟蹊径,在训练期以概率 \(p_{\mathrm{drop}}\) 随机丢弃条件 \(c\),并在测试期直接在连续隐空间内执行向量外插: $\(\mathbf{r}_{\mathrm{LG}} = \mathbf{r}_\emptyset + w (\mathbf{r}_c - \mathbf{r}_\emptyset)\)$ 导出的引导潜变量 \(\mathbf{r}_{\mathrm{LG}}\) 随后统一输入至下游 MLP 预测编辑分布。论文从理论上证明了该隐式引导严格等价于在连续隐空间对隐式分类器对数后验 \(\log p(c|\mathbf{r})\) 进行梯度上升,从而以单步连续平移的方式确保了对复杂结合基序的全序列全局协同增强,并自然保证了离散速率矩阵的非负性而无需后验重归一化。

损失函数 / 训练策略

为解决多条编辑路径对应相同序列导致目标函数难以处理的问题,模型引入扩充空白符 \(\varepsilon\),在增广空间 \(\mathcal{Z} = \mathcal{T} \cup \{\varepsilon\}\) 上通过 Bregman 散度构建可计算的流匹配目标: $\(\min_\theta \mathcal{L}(\theta) = \mathbb{E}_{t, \mathbf{x}_t, \mathbf{z}_t} \left[ \sum_{\mathbf{x} \neq \mathbf{x}_t} u_t^\theta(\mathbf{x}|\mathbf{x}_t) - \sum_{i=1}^N \mathbb{I}_{[z_1^i \neq z_t^i]} \frac{\dot{\kappa}_t}{1-\kappa_t} \log u_t^\theta(\mathbf{x}(\mathbf{z}_t, i, z_1^i) | \mathbf{x}_t) \right]\)$ 其中第一项显式抑制非真实转移的冗余输出速率,第二项则通过带噪序列向目标对齐序列的单步有效编辑操作执行交叉熵对齐,通过 Adam 优化器在单卡或多卡环境下进行端到端稳定训练。

实验关键数据

主实验

论文在三个代表性生物序列设计任务上进行了系统评估:(1) 人类黑色素瘤与果蝇脑细胞增强子 DNA 无条件生成(序列长度 500,以 Fréchet 亲和力距离 FBD 衡量数据流形重合度);(2) 人类启动子 DNA 条件设计(1024 bp,依据 FANTOM5 CAGE 信号条件引导,以 SEI 模型预测的调控活性均方误差 MSE 衡量);(3) 具有复杂等位基因特异性的 MHC-II 肽段条件设计(13–18 氨基酸,以 DeepMHCII 结合分类器评分与 ESM-2 嵌入算得的 Fréchet 蛋白质距离 FPD 联合评估)。

任务 / 评估集 指标 FlexFlow (本文) 强基线对比 (Dirichlet FM / Fisher FM) 提升 / 收益
增强子 DNA (Melanoma) FBD (\(\downarrow\)) 3.4 5.3 (Dirichlet FM) / 3.8 (Fisher FM) 相比最强基线降低 10.5%
增强子 DNA (Fly Brain) FBD (\(\downarrow\)) 4.2 15.2 (Dirichlet FM) / 27.5 (Fisher FM) 相比最强基线降低 72.4%
启动子条件设计 (Promoter) MSE (\(\downarrow\)) 0.022 (Latent) / 0.024 (Rate) 0.026 (Dirichlet FM) / 0.030 (Fisher FM) 调控活性预测误差相对下降 15.4%
MHC-II 肽段生成 (EL Test) DeepMHCII 分数 (\(\uparrow\)) 0.66 (Latent) / 0.58 (Rate) 0.21 (Dirichlet FM) / 0.31 (Dirichlet FDM) 呈递亲和力预测分数提升超 112%
MHC-II 肽段生成 (EL Test) FPD (\(\downarrow\)) 0.58 (Rate Guidance) 0.79 (Dirichlet FM) / 0.72 (Dirichlet FDM) 序列流形逼近度显著改善

消融实验

在合成 2D/1D 拓扑密度估计(衡量对尖锐边界与周期长程依赖的拟合)与先验耦合消融实验中,FlexFlow 均展现出稳定增益:

实验设置 / 拓扑数据集 指标类型 Baseline (D3PM / EditFlow) FlexFlow (完整方法) 机制说明
先验消融: \(p_{\mathrm{unif}}\) 均匀基线 FBD (MEL / FB) 3.7 / 5.2 — 缺乏生物先验的无偏基线
先验消融: 仅加频率先验 \(p_{\mathrm{freq}}\) FBD (MEL / FB) 3.5 / 4.7 — 频率统计带来初级增益
完整模型: 频率先验 + 结构化耦合 \(\mathrm{K}_\gamma\) FBD (MEL / FB) — 3.4 / 4.2 演化替代核大幅缩小流形偏离
玩具分布: Sine wave \(\mathcal{W}_2\) 距离 (\(\downarrow\)) 0.033 / 0.031 0.024 更好保留周期结构
玩具分布: Concentric Rings \(\mathcal{W}_2\) 距离 (\(\downarrow\)) 0.085 / 0.080 0.040 多峰对称流形拟合误差减半
玩具分布: Checkerboard \(\mathcal{W}_2\) 距离 (\(\downarrow\)) 0.067 / 0.070 0.041 局部尖锐边界跳变建模更优

关键发现

  • 结构化耦合的决定性影响:单纯改变初始分布的频率先验只能带来有限的指标改善,只有结合了领域演化相似度矩阵 \(\mathrm{K}_\gamma\) 的结构化核,才能在复杂染色质特征数据集(如 Fly Brain)中实现从 15.2 暴降至 4.2 的断层式领先。
  • 隐空间引导的质量-多样性权衡:在肽段设计中,隐空间引导(Latent Guidance)在 DeepMHCII 评分上达到了惊人的 0.66(远超所有基线的 0.2~0.3 水平),但在高引导强度下 FPD 会从 0.58 轻微回升至 0.84,证实隐空间分数外推能极强地聚焦于结合口袋核心基序,但会略微压缩非结合区域的序列覆盖多样性。
  • 推理效率持平:得益于连续隐变量的轻量前传,FlexFlow 在 100 步逆向采样下耗时约 4.0 秒(单卡 500 bp),与传统固定长度的 Dirichlet FM(3.8 秒)处于同一水平,未引入明显的计算瓶颈。

亮点与洞察

  • 将生物替代模型优雅收敛为耦合核:无需修改流匹配的训练理论框架,仅通过在端点构造双随机转移矩阵就为扩散逆向过程注入了进化保真偏置,数学优美且通用。
  • 隐空间向量外插替代离散比率插值:指出了离散序列直接在 token rate 空间做 CFG 的独立性局限,证明连续隐空间的高斯均值外推等价于隐式分类器的梯度上升,为离散生成的可控引导开辟了新范式。
  • 解耦变长编辑的后验控制:引入 Dirichlet 先验控制每种编辑类型(插入/删除/替换)的采样温度,使生物学家在下游推理时能够针对特定突变筛选策略动态调整操作预算。

局限与展望

  • 采样步数依然依赖百步迭代:当前模型基于标准的 Euler/Heun 逆向积分器需要约 100 步采样,面对千万量级的高通量突变筛选仍有延迟瓶颈,未来需探索基于一致性蒸馏(Consistency Distillation)的一步或少步采样拓展。
  • 三维构象信息的间接性:虽然通过 MHC-II 结合靶标等位基因进行了有效条件引导,但模型目前仍主要停留在序列一维离散空间,未能直接显式结合肽链-受体复合物的原子级全息 3D 坐标进行几何协同流匹配。

相关工作与启发

  • vs Dirichlet Flow Matching (Stark et al., 2024):Dirichlet FM 借助连续单纯形重参数化处理离散 DNA 序列,但无法原生处理变长多肽序列,且缺乏突变偏好耦合;FlexFlow 通过离散编辑算子直接建模变长跳变,生成保真度更优。
  • vs Edit Flows (Havasi et al., 2025):Edit Flows 建立了增广离散编辑空间的基本理论,但其依赖均匀先验且仅能进行非受控的均匀编辑采样;FlexFlow 扩展了隐空间层次化参数化、Dirichlet 预算控制与隐空间引导,真正适配了复杂的生物分子工程需求。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙结合演化替代核、连续隐空间 CFG 与变长编辑流,理论推导严密且高度切合生物问题。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 拓扑流形、无条件 DNA 增强子、条件 DNA 启动子及自建 MHC-II 多肽基准,对比基线全面。
  • 写作质量: ⭐⭐⭐⭐⭐ 结构清晰严谨,命题与证明完备,图表展现丰富。
  • 价值: ⭐⭐⭐⭐⭐ 为蛋白质多肽工程及基因调控元件合成提供了兼具灵活性、保真度与可控性的全新基础生成架构。