蒸馏式蛋白质主链生成
SPID: Distilled Protein Backbone Generation
AI总结:
针对扩散类蛋白质主链生成模型采样速度慢的瓶颈,研究通过适配SiD分数蒸馏策略,结合多步生成与推理噪声调制,训练出少步骤生成器,采样速度提升超20倍且性能接近Proteina教师模型,助力大规模蛋白质设计。
摘要:
基于扩散和流的生成模型近期在蛋白质主链生成任务中展现出优异性能,为从头蛋白质设计提供了前所未有的能力。然而,这类模型虽在生成质量上表现突出,却受限于生成速度,其反向扩散过程通常需要数百次迭代步骤。这一计算瓶颈限制了它们在大规模蛋白质发现中的实际应用——此类场景需要数千至数百万个候选结构。为解决这一挑战,我们探索了分数蒸馏技术,该技术已在视觉领域成功减少采样步骤数同时保持高生成质量。但直接套用这些方法会导致可设计性低到无法接受的程度。通过深入研究,我们找到了适配最新分数蒸馏策略——分数恒等蒸馏(Score identity Distillation,SiD)的恰当方式,用以训练少步骤蛋白质主链生成器,可大幅缩短采样时间,同时保持与预训练教师模型相当的性能。具体而言,多步生成结合推理时噪声调制是成功的关键。我们证明,所提出的蒸馏少步骤生成器采样速度提升超20倍,同时在可设计性、多样性和新颖性上达到与Proteina教师模型相近的水平。推理成本的降低实现了大规模计算机模拟蛋白质设计,从而让基于扩散的模型更贴近实际蛋白质工程应用。其PyTorch实现可在https://github.com/LY-Xie/SiD_Protein获取。
英文摘要:
Diffusion- and flow-based generative models have recently demonstrated strong performance in protein backbone generation tasks, offering unprecedented capabilities for de novo protein design. However, despite their generation quality, these models are constrained by slow sampling, often requiring hundreds of iterative steps. This computational bottleneck limits their practical utility in large-scale protein discovery, where thousands to millions of candidate structures are needed. To address this challenge, we explore the techniques of score distillation, which has shown great success in reducing the number of sampling steps in the vision domain while maintaining high generation quality. However, a straightforward adaptation of these methods results in unacceptably low designability. We introduce Score Protein identity Distillation (SPID), which resolves this incompatibility by combining few-step generation with inference-time noise scaling. SPID adapts the Score identity Distillation (SiD) framework to both diffusion- and flow-based models without requiring access to pretraining data. Applied to the Proteina flow-matching model, our 16-step generator achieves 94.4% designability, matching the 400-step teacher, while delivering more than a 20-fold reduction in effective backbone-generation time and maintaining comparable diversity and novelty. SPID generalizes across unconditional generation, fold-class conditional generation, and motif scaffolding, and extends to equivariant diffusion architectures, achieving significant reduction in generation time with comparable generation quality to the teacher in all tasks. The resulting reduction in inference cost could facilitate large-scale in silico protein design, thereby advancing diffusion-based models toward real-world protein engineering applications. The PyTorch implementation is available at https://github.com/LY-Xie/SiD_Protein
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。