通过多尺度结构生成的蛋白质自回归建模

Protein Autoregressive Modeling via Multiscale Structure Generation

arXiv: 2602.04883v1

论文信息

标题: Protein Autoregressive Modeling via Multiscale Structure Generation

作者: Yanru Qu, Cheng-Yen Hsieh, Zaixiang Zheng, et al.

发布日期: 2026-02-04

arXiv ID: 2602.04883v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题:这篇论文要解决的是如何将自回归(AR)模型的强大能力应用于蛋白质骨架结构生成。传统的 AR 模型在处理蛋白质结构时存在两大挑战:一是需要离散化连续的 3D 坐标,导致结构细节丢失;二是蛋白质残基间存在双向依赖关系,与 AR 模型的单向预测假设冲突。

  • 核心方法:论文提出了 PAR 框架,采用多尺度自回归策略。通过多尺度下采样将蛋白质结构分解为不同粒度的层次,利用 AR Transformer 进行下一尺度预测,生成条件嵌入,再由基于流动匹配(Flow Matching)的骨架解码器直接生成 Cα原子坐标,从而避免离散化并保留双向空间依赖。

  • 关键结果:在无条件生成基准测试上,PAR 在 PDB 数据集子集上微调后,实现了 96.6% 的可设计性(Designability),并且将 Fréchet 蛋白质结构距离(FPSD)降至 161.0(见论文表 1),展现出优异的分布建模能力和结构质量。

  • 主要局限:论文作者明确指出,在多尺度生成过程中,增加尺度层级(例如从 3 层增加到 4 层或 5 层)可能会导致错误累积和曝光偏差加剧,反而降低生成结构的可设计性(见论文表 4)。此外,AR Transformer 部分的规模扩展未能带来显著的性能提升,作者将其归因于更严重的曝光偏差问题(见论文第 8.7 节)。

  • 适合读者:对蛋白质设计、深度生成模型(特别是自回归和扩散模型)以及 AI 在科学领域应用感兴趣的研究人员和工程师。尤其适合那些关注如何将大规模语言模型的成功范式迁移到科学计算领域的读者。

论文背景和研究动机

蛋白质的三维结构决定了其生物学功能,因此,设计和生成具有特定结构的新型蛋白质是生物医药和纳米技术领域的核心目标。近年来,深度生成模型,特别是扩散模型及其变体(如 Flow Matching),已成为蛋白质骨架生成的主流方法。这些方法通常直接在 SE(3)帧表示或简化原子坐标(如 Cα原子)上进行操作。

然而,在自然语言处理和计算机视觉领域,自回归模型凭借其强大的可扩展性和出色的零样本泛化能力,已成为主导性范式。论文作者敏锐地指出,将自回归模型应用于蛋白质结构生成领域的研究几乎是空白。他们识别出阻碍这一交叉的两个关键障碍:

  1. 连续数据的离散化问题:标准 AR 模型处理连续数据的常见方法是先对其进行向量量化(VQ),将连续坐标转化为离散的令牌(token)。但这种离散化不可避免地会损失原子级别的结构细节,限制了生成结构的保真度。
  2. 单向性与双向依赖的冲突:蛋白质结构具有内在的双向依赖性,序列上相距甚远的残基在空间中可能紧密接触,形成氢键或疏水作用。传统的 AR 模型采用 “下一个令牌预测” 方式,其隐式的单向因果假设与蛋白质的这种全局、双向的生物物理关系相冲突。

因此,论文的核心动机是提出一个全新的框架,既能释放 AR 模型在可扩展性和泛化性上的潜力,又能优雅地克服其在处理蛋白质结构数据时的固有局限。

核心方法和技术细节

PAR 框架的设计哲学源自蛋白质结构本身的层级性质,即其结构可以跨越多尺度,从粗略的全局拓扑到精细的局部原子坐标。基于此,PAR 提出了一种类似 “雕塑” 的生成过程:先勾勒粗糙轮廓,再逐步精雕细琢。该框架包含三个核心组件:

1. 多尺度下采样 这是 PAR 实现多尺度 AR 建模的基础。对于一个长度为 LL 的蛋白质骨架 x∈RL×3\mathbf{x} \in \mathbb{R}^{L \times 3},此步骤将其分解为 nn 个不同粒度的表示 X={x1,…,xn}X = \{\mathbf{x}^1, \dots, \mathbf{x}^n\}。其通过一个确定性映射来实现:xi=Down(x,size(i))\mathbf{x}^i = \text{Down}(\mathbf{x}, \mathtt{size}(i))。这里的 Down\text{Down} 操作是沿序列维度进行插值,产生 size(i)\mathtt{size}(i) 个 3D 质心点,从而构建一个从粗粒度到细粒度的结构金字塔。论文默认采用按长度定义尺度序列,例如 S={64,128,256}\mathcal{S}=\{64, 128, 256\}。

2. 自回归 Transformer 与下一尺度预测 此组件定义了多尺度的生成顺序。与标准 AR 模型预测下一个 “令牌” 不同,PAR 进行 “下一尺度” 预测。一个标准的、非等变性的 Transformer Tθ\mathcal{T}_{\theta} 负责处理所有先前粗尺度的信息,并生成一个用于指导当前尺度生成的条件嵌入 zi\mathbf{z}^i。具体来说,它会将之前生成的粗结构 xi−1\mathbf{x}^{i-1} 上采样后的特征与一个可学习的 bos 嵌入拼接,然后预测出 zi\mathbf{z}^i。

3. 基于 Flow Matching 的骨架解码器 这是解决连续数据生成问题的关键。PAR 没有将结构离散化,而是使用一个基于 Flow Matching 的解码器网络 vθ\mathbf{v}_{\theta},直接对 Cα原子的连续坐标进行建模。这个解码器以 AR Transformer 预测的条件嵌入 zi\mathbf{z}^i 为输入,将一个标准正态分布通过普通/随机微分方程(ODE/SDE)演化到目标尺度 xi\mathbf{x}^i 的数据分布。其训练目标是最小化网络预测的速度场与真实速度场之间的差距:

L(θ)=Ex∼pD[1n∑i=1n1size(i)Eti∼p(ti),ϵi∼N(0,I)∥vθ(xtii,ti,zi)−(xi−ϵi)∥2].\mathcal{L}(\theta) = \mathbb{E}_{\mathbf{x}\sim p_{\mathcal{D}}}\left[\frac{1}{n}\sum_{i=1}^{n}\frac{1}{\mathtt{size}(i)}\mathbb{E}_{t^i\sim p(t^i),\boldsymbol{\epsilon}^i\sim\mathcal{N}(0,I)}\left\| \mathbf{v}_{\theta}(\mathbf{x}^i_{t^i}, t^i, \mathbf{z}^i) - (\mathbf{x}^i - \boldsymbol{\epsilon}^i) \right\|^2\right].

此外,为缓解 AR 模型中普遍存在的 “曝光偏差” 问题(即训练时基于真实数据,推理时基于模型自己预测的数据,两者不匹配导致误差累积),PAR 引入了两项针对性技术:

  • 噪声上下文学习:在训练时,向提供给 AR Transformer 的历史结构上下文添加不同程度的噪声,迫使模型学会从 “不完美” 的历史信息中预测下一尺度,增强其鲁棒性。
  • 计划采样:在训练过程中,以一定概率(如 0.5)用模型当前预测的结构替代真实的下一尺度结构作为后续步骤的输入,使训练过程更贴近推理场景。

在推理时,整个过程是完全自回归的:从可学习的 bos 嵌入和噪声开始,AR Transformer 生成 z1\mathbf{z}^1,解码器通过 ODE/SDE 采样生成最粗尺度的结构 x1\mathbf{x}^1;然后将 x1\mathbf{x}^1 上采样后反馈给 AR Transformer 以生成 z2\mathbf{z}^2,解码器再生成 x2\mathbf{x}^2。此 “预测-生成-上采样” 的循环迭代 nn 次,直至生成完整分辨率的骨架结构 xn=x\mathbf{x}^n=\mathbf{x}。

创新点和贡献

PAR 的贡献不仅仅是提出一个新模型,更在于开创性地证明了自回归范式在蛋白质骨架生成领域巨大潜力。

  1. 首个多尺度、非离散的自回归框架:PAR 是首个成功将多尺度 AR 建模应用于蛋白质骨架生成的工作。它巧妙地通过从粗到细的 “下一尺度预测” 策略,解决了传统 AR 模型与蛋白质双向空间依赖之间的矛盾,同时又利用 Flow Matching 解码器,规避了离散化带来的结构信息损失。

  2. 卓越的零样本泛化能力:这是 PAR 作为 AR 模型最引人注目的特性之一。它无需任何微调,就能灵活应对多种条件生成任务。例如,用户只需指定少量 3D 点(如 16 个)作为粗略的全局结构提示(Prompt),PAR 就能生成符合该拓扑结构的完整蛋白质骨架(见论文图 3)。同样,在 “Motif Scaffolding” 任务中,只需在生成过程中将特定功能区域(motif)的真实坐标 “强制” 输入到每一尺度的 AR Transformer 中,模型就能生成保留该 Motif 且结构多样的新骨架(见论文图 4)。

  3. 高效且可解释的生成过程:PAR 的逐尺度生成过程具有高度的可解释性,如同雕塑家由粗到精地创作。通过对 AR Transformer 各尺度的注意力图进行可视化分析,论文揭示了模型在每个尺度上主要关注其前一尺度的信息,同时保留了对更早期尺度的非零注意力,体现了其有效整合多尺度信息的能力(见论文图 6)。此外,利用多尺度的优势,PAR 可以在粗尺度上使用随机性强的 SDE 采样以确保全局拓扑的多样性,而在细尺度上使用确定性强的 ODE 采样以快速完成局部细节的细化,从而在不损失质量的前提下,将采样速度相比单尺度方法提升2.5 倍(见论文表 2)。

实验结果分析

论文在无条件生成、条件生成、可扩展性和消融实验等多个维度对 PAR 进行了全面评估。

  • 无条件生成能力:在与 FrameDiff、RFDiffusion、Proteina 等主流方法的对比中,PAR 在衡量分布匹配度的关键指标 FPSD 上表现突出。在 PDB 数据集上微调后,PAR 的 FPSD 得分(161.0)远优于其他基线模型(见论文表 1),表明其生成的样本分布与真实蛋白质结构分布最为接近。同时,其生成结构的可设计性(96.6%)也达到了顶尖水平,证明其生成的结构不仅多样,而且具备物理可行性。

  • 可扩展性分析:类似于对大型语言模型的研究,论文探索了 PAR 的规模效应。结果表明,单纯增加训练步数带来的边际收益递减,但扩大 Flow Matching 解码器的模型参数量(从 60M 到 400M)能够持续且显著地改善 FPSD 和结构质量(见论文图 5)。值得注意的是,单独扩大 AR Transformer 的规模并未带来明显提升,作者分析这可能是由于 AR Transformer 更容易在有限数据下因曝光偏差而过拟合。

  • 曝光偏差缓解措施的有效性:消融实验清晰地展示了噪声上下文学习和计划采样的重要性。在采用这些技术后,模型在自一致性 RMSD(sc-RMSD,即生成的结构与其设计序列反折叠结构之间的差异)指标上从 2.20 Å 大幅降低到 1.48 Å(见论文表 3),生成质量得到实质性提升。

  • 采样效率:PAR 通过在不同尺度混合 SDE 和 ODE 采样,实现了计算效率的重大飞跃。在初始尺度使用 400 步 SDE 采样以构建可靠的全局拓扑,而在后续两个更精细、序列更长的尺度上,采样步数可以从 400 步骤降至2 步。这相比所有尺度都跑 400 步的单尺度模型,推理速度最高提升了4.7 倍,同时保持了97% 的高设计性(见论文表 2)。

实践建议

尽管 PAR 是一个前沿的研究框架,但其展现的特性为有意涉足蛋白质设计的从业者提供了明确的实践方向:

  1. 利用零样本能力进行快速原型设计:对于需要进行特定拓扑结构设计的研究人员,可以直接利用 PAR 的 “提示生成” 功能,通过制定极少量(如 16 个)3D 控制点来探索广阔的蛋白质结构空间,无需为每个特定任务训练或微调模型,极大地降低了使用门槛。

  2. 选择合适的基础模型进行微调:论文表明,更大的 Flow Matching 解码器能带来更好的生成质量。如果资源允许,在预训练的 400M 参数 PAR 模型上进行面向特定家族或功能蛋白质的微调,可能会比从头训练一个更小的模型获得更优效果。但需注意,AR Transformer 部分的尺度扩展收益有限,可优先将计算资源分配给骨架解码器。

  3. 在效率和多样性/质量间权衡的采样策略:在生成大规模库或长序列蛋白质时,可以借鉴论文中的多尺度 SDE/ODE 编排策略。在初始较短的粗尺度上投入更多的 SDE 采样步骤来探索多样的全局构象,而在后续计算量大的细尺度上,利用 ODE 快速完成局部弛豫。这对于需要生成数百万个结构以进行虚拟筛选的场景,是一个非常有价值的效率优化技巧。

  4. 处理时序或动态问题的潜力:论文在讨论部分猜测,PAR 的 “下采样-再上采样” 流程可能隐含地对局部构象动力学进行零样本建模。对于研究蛋白质动态和别构效应的学者,这可能是一个值得探索的方向,即通过观察从同一粗尺度骨架出发能否生成合理的多种构象,来间接研究蛋白质的柔性。