This model, presented in the paper Distilled Protein Backbone Generation, introduces an approach for protein backbone generation using score distillation techniques. It significantly reduces sampling time while maintaining high generation quality, offering a more efficient method for de novo protein design.

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, while achieving notable performance in generation quality, these models are limited by their generating speed, often requiring hundreds of iterative steps in the reverse-diffusion process. 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. Through extensive study, we have identified how to appropriately adapt Score identity Distillation (SiD), a state-of-the-art score distillation strategy, to train few-step protein backbone generators which significantly reduce sampling time, while maintaining comparable performance to their pretrained teacher model. In particular, multistep generation combined with inference time noise modulation is key to the success. We demonstrate that our distilled few-step generators achieve more than a 20-fold improvement in sampling speed, while achieving similar levels of designability, diversity, and novelty as the Proteina teacher model. This reduction in inference cost enables large-scale in silico protein design, thereby bringing diffusion-based models closer to real-world protein engineering applications.

Below are some protein structures sampled from our 16-step distilled generator.

Code: GitHub Repository

To generate protein structures unconditionally with specified lengths (number of residues), run

Replace 0 with the desired gpu device to use. If not specified, gpu0 will be used.

By default, the 16-step generator is used for generation. If a different model is desired, please make sure to change the model_path, nstep, and other relevant parameters in run_generate_distilled_proteina.sh as needed. Unless otherwise specified, the generated structures will be saved in protein_out/proteina_(un)cond_distilled_{nstep}step_noise{noise_scale}_output based on the parameters nstep and noise_scale as well as the conditional flag.

For fold-class-conditional generation, simply uncomment the --conditional flag in run_generate_distilled_proteina.sh. The lengths and CATH codes will be sampled from the empirical joint distribution as specified in training/proteina/configs/experiment_config/inference_cond_sampling_specific_codes.yaml.

The few-step generators distilled from Proteina are provided here.

If you find this open source release useful, please cite in your paper:

This repo is heavily built upon EDM, SiD, and Proteina. We thank the authors for their excellent work and for making their code publicly available.