AlphaFold Distillation for Protein Design

Igor Melnyk, Aurelie Lozano, Payel Das, Vijil Chenthamarakshan

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Inverse Protein Folding Design, Protein Design, Model Distillation, AlphaFold, Protein Folding

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We distilled AlphaFold to generate structural confidence scores (pTM, pLDDT) for any protein sequence and applied it to inverse folding design

Abstract:

Inverse protein folding, the process of designing sequences that fold into a specific 3D structure, is crucial in bio-engineering and drug discovery. Traditional methods rely on experimentally resolved structures, but these cover only a small fraction of protein sequences. Forward folding models like AlphaFold offer a potential solution by accurately predicting structures from sequences. However, these models are too slow for integration into the optimization loop of inverse folding models during training.
To address this, we propose using knowledge distillation on folding model confidence metrics, such as pTM or pLDDT scores, to create faster and end-to-end differentiable distilled model. This model can then be used as a structure consistency regularizer in training the inverse folding model. Our technique is versatile and can be applied to other design tasks, such as sequence-based protein infilling.
Experimental results show that our method outperforms non-regularized baselines, yielding up to 3\% improvement in sequence recovery and up to 45\% improvement in protein diversity while maintaining structural consistency in generated sequences. Anonymized code for this work is available at https://anonymous.4open.science/r/AFDistill-28C3

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 8347