Skip to content

Latest commit

 

History

4 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

PAR: Protein Autoregressive Modeling via Multi-Scale Structure Generation

🏆 ICML 2026 Oral Presentation

teaser

Abstract. We present protein autoregressive modeling (PAR), the first multi-scale autore- gressive framework for protein backbone generation via coarse-to-fine next-scale prediction. Using the hierarchical nature of proteins, PAR generates structures that mimic sculpting a statue, forming a coarse topology and refining structural details over scales. To achieve this, PAR consists of three key components: (i) multi- scale downsampling operations that represent protein structures across multiple scales during training; (ii) an autoregressive transformer that encodes multi-scale information and produces conditional embeddings to guide structure generation; (iii) a flow-based backbone decoder that generates backbone atoms conditioned on these embeddings. Moreover, autoregressive models suffer from exposure bias, caused by the training and the generation procedure mismatch, and substantially de- grades structure generation quality. We effectively alleviate this issue by adopting noisy context learning and scheduled sampling, enabling robust backbone gener- ation. Notably, PAR exhibits strong zero-shot generalization, supporting flexible human-prompted conditional generation and motif scaffolding without requiring fine-tuning. On the unconditional generation benchmark, PAR effectively learns protein distributions and produces backbones of high design quality, and exhibits favorable scaling behavior. Together, these properties establish PAR as a promising framework for protein structure generation.

Setup 🚀

You can use the pre-configured proteina environment, or build the environment using the following script:

# 1. make sure conda is available before running the following

# 2. change DATA_PATH and CKPT_HOME to your own path in prepare_env.sh
DATA_PATH=/your/own/data
CKPT_HOME=/your/own/ckpt/home

# 3. create conda env and activate it
source prepare_env.sh && conda activate PAR

# 4. install MPNN and ESMFold for inference
source script_utils/download_pmpnn_weghts.sh  # Prepare PMPNN weights for designability metrics

For training and evaluation, additional meta files are required. Please refer to the proteina setup for more details.

Download Model Checkpoints 🤗

PAR model checkpoints are available on HuggingFace:

Download the checkpoints from ByteDance-Seed/PAR and set pretrain_ckpt_path to the downloaded checkpoint path for finetuning or inference.

Download D_FS And PDB Data 📦

For the D_FS dataset, you need to download the d_fs_index.txt following proteina setup. If you properly download the data files as in proteina setup, you will find the d_FS_index.txt under $DATA_PATH/D_FS/, which contains the indices of the D_FS dataset as created by Proteina team.

Check the data config configs/datasets_config/afdb/d_FS.yaml and make sure overwrite: False to avoid overwriting your processed data.

See example_download_data.ipynb for more details.

Statistics: The D_FS dataset used for pretraining contains 0.6M structures, as processed by Proteina. This dataset ensures both high quality (pLDDT > 80) and structural diversity, with sequence lengths ranging from 32 to 256 residues. We follow Proteina and split it by 98:19:1 for training, validation and testing. The PDB dataset used for finetuning contains 21K designable (sc-RMSD <= 2) single chains after filtering. For detailed criterion, please refer to the Proteina paper.

Unconditional Pretraining 🧠

We use hydra to compose the training config. Taking configs/experiment_config/training_ca_par_60m.yaml as the example, we introduce the config structure as follows:

# training_ca_par_60m.yaml

...
# the following lines will load the model and data config
defaults:
  - /model: ca_par_60m  # caflow or frameflow
  - _self_
  - /datasets_config: afdb/d_FS

# by default, we use self conditioning, scheduled sampling, and noisy context learning for training
training:
  self_cond: true
  fold_cond: false
  schedule_sampling: true
  noisy_context_learning: true

The model config is located at configs/model/ca_par_60m.yaml

# ca_par_60m.yaml

...
# the following lines will load ca_af3_60M_notri_par as the transformer encoder and caflow as the flow decoder, which follow the same architecture as Proteina
# by default, we do not use aux loss and pair representation
# browse the configs/model/nn and configs/atom_diff/nn folders for more options
defaults:
  - _self_
  - nn: ca_af3_60M_notri_par
  - /atom_diff: caflow_60m

# the default scale configuration
protein_decomposer:
  structure_dim: 3  # Ca-only
  num_local_centroids: [64, 128, 256]

You can simply reproduce our 60m pretraining using the following lines:

source prepare_env.sh && conda activate PAR
python proteinfoundation/train_par.py

This will train a 60m model on D_FS dataset for 200k steps with learning rate 0.0001.

We also provide a detailed training script run/run_training_par.sh for you to customize training.

# Enviroments handled inside
bash run/run_training_par.sh

PDB Finetuning 🔧

Change the following line in configs/experiment_config/finetune_ca_par_400m.yaml to specify the pretrained checkpoint.

pretrain_ckpt_path: /your/pretrained/model.ckpt

Then run the following lines to reproduce our PDB finetuning experiments:

source prepare_env.sh && conda activate PAR
python proteinfoundation/train_par.py experiment_config=finetune_ca_par_400m

By default, we only use the designable PDB structures for finetuning, as configured in configs/datasets_config/pdb/pdb_train.yaml:

filter_designable: True

This requires a pre-computed designable set of PDB codes, specified in proteinfoundation/datasets/pdb_data.py.

designable_path = 'configs/datasets_config/pdb/pdb_train_designable.pkl'

You can also custimize the filtering criterion by modifying the parameters in configs/datasets_config/pdb/pdb_train.yaml.

You can also modify parameters in run/run_training_par.sh for finetuning. Since the PDB dataset is much smaller than D_FS, we recommend using a smaller learning rate and fewer steps for finetuning.

opt.lr=0.00001
opt.max_steps=10000
opt.val_check_interval=1000
log.checkpoint_every_n_steps=1000

Sampling & Evaluation 📊

Use run/run_inference_par.sh to sample and evaluate the model. This script will execute the following steps:

  1. Sample 500 structures in multiple gpus for computing designability.
  2. Run the MPNN-ESMFold pipeline for folding the sampled structures in multiple gpus.
  3. Sample 5000 structures in multiple gpus for computing FPSD and fS as proposed in Proteina.
  4. Aggregate the folding results and compute sc-RMSD and designability.
  5. Collect the sampled structures and compute FPSD and fS.

We provide example outputs under example_output. After uncompressing rmsd.tar.gz and fid.tar.gz, you will find files in the following structure:

example_output/
  - inference_fid_par/samples_fid
    - *.pdb
  - inference_rmsd_par/
    - **/*.pdb
  - results_inference_fid_par.csv
  - results_inference_rmsd_par.csv

To evaluate designability of your own samples, you can use fold_inference_data.py to run the MPNN-ESMFold pipeline for folding the sampled structures.

# args: 
# 1: number_of_gpus
# 2: gpu_id
# 3: inference_path
python fold_inference_data.py 1 0 your_inference_path

This code will split the .pdb samples under your_inference_path into number_of_gpus parts, and run folding with one part on cuda:gpu_id.

The diversity and secondary structure are computed only for the designable structures. Run example_compute_metrics.ipynb to obtain the metrics.

Generate With Prompt ✨

We provide a notebook example_prompt_generation.ipynb with example structures shown in the paper. This notebook will read a structure in PDB format, downsample it to 16 points, and use a 60m 5-scale model to generate the complete structure. You can also feed your 16 point prompt to the model directly. Before running, uncompress example_prompt/prompt.tar.gz, which contains the prompt structures used in the paper (fig. 3).

barrel/
  - 1_all_scales.pdb
  - 1_ref.pdb
  - 1_rmsd_xxx.pdb

The *_rmsd_xxx.pdb files are the generated structures. The *_ref.pdb files are the reference structures. The *_all_scales.pdb files store the generated structures of all scales.

0-Shot Motif Scaffold 🧩

We provide a notebook example_motif_scaffold.ipynb to demonstrate how does the model work. In this example, we read the test set and extract the last 20 residues as the motif segment, and generate the rest of the structure.

We also provide a script run/motif_scaffold_par.sh to execute the motif benchmark as in RFDiffusion (data located in ./motif_data), with compute_motif_success.py to compute the success rate. For new motif tasks, you can put the pdb file under ./motif_data and add corresponding contig string and min/max lengths in MOTIF_CONTIGS and MOTIF_LENGTH in proteinfoundation/motif_inference_par.py.

Citation

If you use this codebase in your research, please cite the following paper:

    @article{qu2026par,
        title={Protein Autoregressive Modeling via Multiscale Structure Generation},
        author={Yanru Qu and Cheng-Yen Hsieh and Zaixiang Zheng and Ge Liu and Quanquan Gu},
        journal={arXiv preprint arXiv:2602.04883},
        year={2026},
        url={https://arxiv.org/abs/2602.04883}, 
    }

About

(ICML 2026 Oral) Protein Autoregressive Modeling via Multiscale Structure Generation, Official Implementation from ByteDance Seed

Topics

Resources

Stars

18 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages