DistilledProteina

Distilled Protein Backbone Generation

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.

Abstract

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

Checkpoints

Every distilled generator is released in two formats:

import pickle, torch

# .pkl -- the generator is the 'ema' entry, ready to call
with open("16-step-checkpoint.pkl", "rb") as f:
    G = pickle.load(f)["ema"]

# .pt -- construct the wrapper first, from the config of the teacher it was distilled from
from training.networks import ProteinaWrapper
G = ProteinaWrapper("proteina/configs/experiment_config/", "inference_ucond_200m_notri")
G.load_state_dict(torch.load("16-step-checkpoint.pt"))

Run from the repository root so that training.* and torch_utils.* resolve; the .pkl snapshots reference those modules when unpickled. generate_distilled_proteina.py accepts either format and dispatches on the file extension.

Unconditional generators (SiD)

Distilled from the M^no-tri_FS Proteina teacher (inference_ucond_200m_notri). The step count in the name is the number of sampling steps the generator was trained for; pass the matching --nstep at inference.

Diff-Instruct generators

The same distillation recipe with the Diff-Instruct objective in place of SiD, released as a step-count scan so the two objectives can be compared at matched K. All are the kimg-36 snapshot of their run.

Provenance note. diffinstruct-10-step-checkpoint comes from an earlier training run than the other six, which were produced together in a single later scan under one recipe. The recipe is the same, but it is a different run, so treat the K=10 point as nominally rather than strictly matched to its neighbours.

Motif scaffolding

Distilled from the motif-conditioned Proteina teacher (proteina_v1.7_DFS_60M_notri_motif_scaffolding, config inference_motif_sid). Smaller than the unconditional generators because the teacher is the 60M model.

Long-chain generation

Distilled from the long-chain Proteina teacher (proteina_v1.6_DFS_200M_notri_long_chain_generation, config inference_long_len) and evaluated at 300, 400 and 500 residues. Designability over 100 samples per length:

γ = 0.45 is the better setting at every length and the margin widens with length, but no individual length is significant at n=100 (pooled over 300 samples per arm: 82.0% → 87.7%, +5.7 pp, p = 0.053). Cluster diversity, novelty and PDB-FPSD move in the same direction at all three lengths.

Genie backbone

Distilled from a Genie teacher rather than Proteina, to show the method is not specific to one backbone. Sample it with generate_genie_distilled_sweep.py, not generate_distilled_proteina.py.

Naming note. This file was previously circulated as genie-16step.pkl, which was a misnomer — the generator was trained with 10 sampling steps. It is named for its true step count here.

Usage

Unconditional inference

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

bash run_generate_distilled_proteina.sh 0

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.

Motif scaffolding

Give the sampler a motif PDB and a contig string describing the scaffold layout:

python generate_distilled_proteina_motif.py \
    --model_path motif-16step-checkpoint.pkl \
    --out_dir outputs/motif_scaffolds \
    --contig_string "10-20/A14/10-20" \
    --motif_pdb_path design25/1qjg.pdb \
    --motif_task_name 1QJG \
    --nsamples 20 --nstep 16 --noise_scale 0.4

eval_rfdiffusion_benchmark.py runs the same sampler across the standard 25-case motif-scaffolding benchmark and reports motif RMSD alongside scaffold scRMSD.

Long-chain generation

Long chains need the long-chain teacher config; the default config will not load this checkpoint correctly:

python generate_distilled_proteina.py \
    --model_path long-chain-32step-checkpoint.pkl \
    --out_dir outputs/long_chain \
    --lengths 300,400,500 --num_batch 10 --batch_size 10 \
    --nstep 32 --noise_scale 0.45 \
    --config_name inference_long_len

Genie backbone

python generate_genie_distilled_sweep.py \
    --ckpt genie-10step-checkpoint.pkl \
    --out_dir outputs/genie_samples \
    --gamma 0.6 --nstep 10 \
    --lengths 50 100 --batch_size 10 --niter 5

See README_GENIE.md in the code repository for the Genie training and sampling path.

Evaluation

All of the above produce structures only. Score them with run_protein_evaluation.py so that designability, diversity, novelty and FPSD are computed through the same pipeline for every model.

Citation

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

@misc{xie2025distilledproteinbackbonegeneration,
      title={Distilled Protein Backbone Generation},
      author={Liyang Xie and Haoran Zhang and Zhendong Wang and Wesley Tansey and Mingyuan Zhou},
      year={2025},
      eprint={2510.03095},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2510.03095},
}

Acknowledgement

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.