upload ckpts

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,27 @@
 ---
 license: apache-2.0
 ---
+
+<div align=center>
+
+## ⚡️ReQFlow: Rectified Quaternion Flow for Efficient and High-Quality Protein Backbone Generation
+[![arXiv](https://img.shields.io/badge/arXiv-2502.14637-b31b1b?style=flat&logo=arxiv)](https://arxiv.org/abs/2502.14637)
+[![Github](https://img.shields.io/badge/Github-Code-blue?style=flat&logo=github)](https://github.com/AngxiaoYue/ReQFlow)
+[![Hugging Face](https://img.shields.io/badge/Hugging%20Face-Model-orange?style=flat&logo=huggingface)](https://huggingface.co/AngxiaoYue/ReQFlow/tree/main)
+
+</div>
+
+Our **ReQFlow** achieves **state-of-the-art (SOTA)** performance 🎯 in protein backbone generation while requiring significantly fewer sampling steps and substantially reducing inference time. For example, it is **37×** faster than RFDiffusion and **62×** faster than Genie2 when generating a backbone of length 300, demonstrating both its effectiveness and efficiency 🚀🚀🔥.
+<p align="center" style="margin-bottom: -2.8px;">
+  <img src="assets/scheme_reqflow.png" width="64.4%">
+  <img src="assets/log_speed_10_v2.png" width="30%">
+</p>
+<p align="center" style="margin-top: 0px;">
+  <img src="assets/prot_vis.png" width="95%">
+</p>
+Ready to dive in? We’ve got checkpoints and step-by-step guides to get you inferencing in no time—playtime starts now! ✨
+
+## 🚀🚀🔥
+🔗arxiv link: [https://arxiv.org/abs/2502.14637](https://arxiv.org/abs/2502.14637)
+
+🌳github link: [https://github.com/AngxiaoYue/ReQFlow](https://github.com/AngxiaoYue/ReQFlow)

qflow_pdb/config.yaml

@@ -0,0 +1,149 @@
+data:
+  task: hallucination
+  dataset: pdb
+  loader:
+    num_workers: 4
+    prefetch_factor: 10
+  sampler:
+    max_batch_size: 128
+    max_num_res_squared: 1000000
+interpolant:
+  min_t: 0.01
+  twisting:
+    use: false
+  rots:
+    corrupt: true
+    sample_schedule: exp
+    exp_rate: 10
+  trans:
+    corrupt: true
+    batch_ot: true
+    sample_schedule: linear
+    sample_temp: 1.0
+    vpsde_bmin: 0.1
+    vpsde_bmax: 20.0
+    potential: null
+    potential_t_scaling: false
+    rog:
+      weight: 10.0
+      cutoff: 5.0
+  sampling:
+    num_timesteps: 100
+    do_sde: false
+  self_condition: ${model.edge_features.self_condition}
+experiment:
+  debug: false
+  seed: 123
+  num_devices: 4
+  warm_start: null
+  warm_start_cfg_override: true
+  resume_cfg_path: null
+  resume_ckpt_path: null
+  resume: true
+  training:
+    mask_plddt: true
+    bb_atom_scale: 0.1
+    trans_scale: 0.1
+    translation_loss_weight: 2.0
+    t_normalize_clip: 0.9
+    rotation_loss_weights: 1.0
+    aux_loss_weight: 1.0
+    aux_loss_use_bb_loss: true
+    aux_loss_use_pair_loss: true
+    aux_loss_t_pass: 0.5
+  wandb:
+    name: null
+    project: null
+  optimizer:
+    lr: 0.0001
+  trainer:
+    overfit_batches: 0
+    min_epochs: 1
+    max_epochs: 1000
+    accelerator: gpu
+    log_every_n_steps: 1
+    deterministic: false
+    strategy: ddp
+    check_val_every_n_epoch: 2
+    accumulate_grad_batches: 2
+  checkpointer:
+    save_last: false
+    save_top_k: -1
+    dirpath: null
+  inference_dir: null
+shared:
+  seed: 123
+  max_cache_size: 100000
+  samples_per_eval_length: 5
+  num_eval_lengths: 8
+  max_eval_length: 256
+  min_motif_percent: 0.05
+  max_motif_percent: 0.5
+scope_dataset:
+  seed: ${shared.seed}
+  csv_path: ./metadata/scope_metadata.csv
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  add_plddt_mask: false
+  max_eval_length: ${shared.max_eval_length}
+  inpainting_percent: 1.0
+  samples_per_eval_length: ${shared.samples_per_eval_length}
+  num_eval_lengths: ${shared.num_eval_lengths}
+  filter:
+    max_num_res: 128
+    min_num_res: 60
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+pdb_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  cluster_path: ./metadata/pdb.clusters
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  inpainting_percent: 1.0
+  add_plddt_mask: false
+  max_eval_length: ${shared.max_eval_length}
+  samples_per_eval_length: ${shared.samples_per_eval_length}
+  num_eval_lengths: ${shared.num_eval_lengths}
+  filter:
+    max_num_res: 512
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+model:
+  node_embed_size: 256
+  edge_embed_size: 128
+  symmetric: false
+  node_features:
+    c_s: ${model.node_embed_size}
+    c_pos_emb: 128
+    c_timestep_emb: 128
+    max_num_res: 2000
+    timestep_int: 1000
+    embed_chain: false
+  edge_features:
+    single_bias_transition_n: 2
+    c_s: ${model.node_embed_size}
+    c_p: ${model.edge_embed_size}
+    relpos_k: 64
+    feat_dim: 64
+    num_bins: 22
+    self_condition: true
+    embed_chain: false
+    embed_diffuse_mask: true
+  ipa:
+    c_s: ${model.node_embed_size}
+    c_z: ${model.edge_embed_size}
+    c_hidden: 128
+    no_heads: 8
+    no_qk_points: 8
+    no_v_points: 12
+    seq_tfmr_num_heads: 4
+    seq_tfmr_num_layers: 2
+    num_blocks: 6

qflow_pdb/qflow_pdb.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db617f19a92e929a06b46a9194f04110cdb5f86e6ca09d28141833fb18fb380b
+size 201852896

qflow_scope/config.yaml

@@ -0,0 +1,156 @@
+data:
+  task: hallucination
+  dataset: scope
+  loader:
+    num_workers: 4
+    prefetch_factor: 10
+  sampler:
+    max_batch_size: 128
+    max_num_res_squared: 500000
+interpolant:
+  min_t: 0.01
+  twisting:
+    use: false
+  rots:
+    corrupt: true
+    sample_schedule: exp
+    exp_rate: 10
+  trans:
+    corrupt: true
+    batch_ot: true
+    sample_schedule: linear
+    sample_temp: 1.0
+    vpsde_bmin: 0.1
+    vpsde_bmax: 20.0
+    potential: null
+    potential_t_scaling: false
+    rog:
+      weight: 10.0
+      cutoff: 5.0
+  sampling:
+    num_timesteps: 100
+    do_sde: false
+  self_condition: ${model.edge_features.self_condition}
+experiment:
+  debug: false
+  seed: 123
+  num_devices: 4
+  warm_start: null
+  warm_start_cfg_override: true
+  resume_cfg_path: null
+  resume_ckpt_path: null
+  resume: true
+  training:
+    mask_plddt: true
+    bb_atom_scale: 0.1
+    trans_scale: 0.1
+    translation_loss_weight: 2.0
+    t_normalize_clip: 0.9
+    rotation_loss_weights: 1.0
+    aux_loss_weight: 1.0
+    aux_loss_use_bb_loss: true
+    aux_loss_use_pair_loss: true
+    aux_loss_t_pass: 0.25
+  wandb:
+    name: null
+    project: null
+  optimizer:
+    lr: 0.0001
+  trainer:
+    overfit_batches: 0
+    min_epochs: 1
+    max_epochs: 1000
+    accelerator: gpu
+    log_every_n_steps: 1
+    deterministic: false
+    strategy: ddp
+    check_val_every_n_epoch: 2
+    accumulate_grad_batches: 2
+  checkpointer:
+    save_last: false
+    save_top_k: -1
+    dirpath: null
+  inference_dir: null
+  batch_ot:
+    enabled: true
+    cost: kabsch
+    noise_per_sample: 1
+    permute: false
+shared:
+  seed: 123
+  max_cache_size: 100000
+  samples_per_eval_length: 5
+  num_eval_lengths: 8
+  max_eval_length: 256
+  min_motif_percent: 0.05
+  max_motif_percent: 0.5
+scope_dataset:
+  seed: ${shared.seed}
+  csv_path: ./metadata/scope_metadata.csv
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  add_plddt_mask: false
+  max_eval_length: ${shared.max_eval_length}
+  inpainting_percent: 1.0
+  samples_per_eval_length: ${shared.samples_per_eval_length}
+  num_eval_lengths: ${shared.num_eval_lengths}
+  filter:
+    max_num_res: 128
+    min_num_res: 60
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+pdb_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  cluster_path: ./metadata/pdb.clusters
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  inpainting_percent: 1.0
+  add_plddt_mask: false
+  max_eval_length: ${shared.max_eval_length}
+  samples_per_eval_length: ${shared.samples_per_eval_length}
+  num_eval_lengths: ${shared.num_eval_lengths}
+  filter:
+    max_num_res: 512
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+model:
+  node_embed_size: 256
+  edge_embed_size: 128
+  symmetric: false
+  node_features:
+    c_s: ${model.node_embed_size}
+    c_pos_emb: 128
+    c_timestep_emb: 128
+    max_num_res: 2000
+    timestep_int: 1000
+    embed_chain: false
+  edge_features:
+    single_bias_transition_n: 2
+    c_s: ${model.node_embed_size}
+    c_p: ${model.edge_embed_size}
+    relpos_k: 64
+    feat_dim: 64
+    num_bins: 22
+    self_condition: true
+    embed_chain: false
+    embed_diffuse_mask: true
+    use_rbf: true
+    num_rbf: 32
+  ipa:
+    c_s: ${model.node_embed_size}
+    c_z: ${model.edge_embed_size}
+    c_hidden: 128
+    no_heads: 8
+    no_qk_points: 8
+    no_v_points: 12
+    seq_tfmr_num_heads: 4
+    seq_tfmr_num_layers: 2
+    num_blocks: 6

qflow_scope/qflow_scope.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54f117c7c08fc060948b0f72b65dcc2d4859f38d143177e7c9b5f37be1fd0372
+size 201846688

reqflow_pdb_rectify/config.yaml

@@ -0,0 +1,159 @@
+data:
+  task: hallucination
+  dataset: pdb
+  rectify: true
+  distill: false
+  loader:
+    num_workers: 32
+    prefetch_factor: 10
+  sampler:
+    max_batch_size: 128
+    max_num_res_squared: 1000000
+interpolant:
+  min_t: 0.01
+  twisting:
+    use: false
+  rots:
+    corrupt: true
+    sample_schedule: exp
+    exp_rate: 10
+  trans:
+    corrupt: true
+    batch_ot: true
+    sample_schedule: linear
+    sample_temp: 1.0
+    vpsde_bmin: 0.1
+    vpsde_bmax: 20.0
+    potential: null
+    potential_t_scaling: false
+    rog:
+      weight: 10.0
+      cutoff: 5.0
+  sampling:
+    num_timesteps: 100
+    do_sde: false
+  self_condition: ${model.edge_features.self_condition}
+experiment:
+  is_training: true
+  debug: false
+  seed: 123
+  num_devices: 4
+  warm_start: null
+  warm_start_cfg_override: true
+  resume_cfg_path: null
+  resume_ckpt_path: null
+  resume: false
+  training:
+    mask_plddt: true
+    bb_atom_scale: 0.1
+    trans_scale: 0.1
+    translation_loss_weight: 2.0
+    t_normalize_clip: 0.9
+    rotation_loss_weights: 1.0
+    aux_loss_weight: 1.0
+    aux_loss_use_bb_loss: true
+    aux_loss_use_pair_loss: true
+    aux_loss_t_pass: 0.5
+  wandb:
+    name: null
+    project: null
+  optimizer:
+    lr: 0.0001
+  trainer:
+    overfit_batches: 0
+    min_epochs: 1
+    max_epochs: 2000
+    accelerator: gpu
+    log_every_n_steps: 1
+    deterministic: false
+    strategy: ddp
+    check_val_every_n_epoch: 1
+    accumulate_grad_batches: 2
+  checkpointer:
+    dirpath: null
+    save_last: true
+    save_top_k: -1
+  inference_dir: null
+shared:
+  seed: 123
+  max_cache_size: 100000
+  min_motif_percent: 0.05
+  max_motif_percent: 0.5
+scope_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  rectify_csv_path: null
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  add_plddt_mask: false
+  min_eval_length: 60
+  max_eval_length: 128
+  inpainting_percent: 1.0
+  samples_per_eval_length: 10
+  num_eval_lengths: 10
+  filter:
+    max_num_res: 128
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+pdb_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  rectify_csv_path: null
+  cluster_path: ./metadata/pdb.clusters
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  inpainting_percent: 1.0
+  add_plddt_mask: false
+  min_eval_length: 100
+  max_eval_length: 300
+  samples_per_eval_length: 10
+  num_eval_lengths: 10
+  filter:
+    max_num_res: 512
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+model:
+  node_embed_size: 256
+  edge_embed_size: 128
+  symmetric: false
+  node_features:
+    c_s: ${model.node_embed_size}
+    c_pos_emb: 128
+    c_timestep_emb: 128
+    max_num_res: 2000
+    timestep_int: 1000
+    embed_chain: false
+  edge_features:
+    single_bias_transition_n: 2
+    c_s: ${model.node_embed_size}
+    c_p: ${model.edge_embed_size}
+    relpos_k: 64
+    feat_dim: 64
+    num_bins: 22
+    self_condition: true
+    embed_chain: false
+    embed_diffuse_mask: true
+  ipa:
+    c_s: ${model.node_embed_size}
+    c_z: ${model.edge_embed_size}
+    c_hidden: 128
+    no_heads: 8
+    no_qk_points: 8
+    no_v_points: 12
+    seq_tfmr_num_heads: 4
+    seq_tfmr_num_layers: 2
+    num_blocks: 6

reqflow_pdb_rectify/reqflow_pdb_rectify.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4044dd79283228cde045ee903659430289e14b78cbe031865b5f94949ff30a0e
+size 201851744

reqflow_scope_rectify/config.yaml

@@ -0,0 +1,160 @@
+data:
+  task: hallucination
+  dataset: scope
+  rectify: true
+  distill: false
+  loader:
+    num_workers: 32
+    prefetch_factor: 10
+  sampler:
+    max_batch_size: 128
+    max_num_res_squared: 500000
+interpolant:
+  min_t: 0.01
+  twisting:
+    use: false
+  rots:
+    corrupt: true
+    sample_schedule: exp
+    exp_rate: 10
+  trans:
+    corrupt: true
+    batch_ot: true
+    sample_schedule: linear
+    sample_temp: 1.0
+    vpsde_bmin: 0.1
+    vpsde_bmax: 20.0
+    potential: null
+    potential_t_scaling: false
+    rog:
+      weight: 10.0
+      cutoff: 5.0
+  sampling:
+    num_timesteps: 100
+    do_sde: false
+  self_condition: ${model.edge_features.self_condition}
+experiment:
+  is_training: true
+  debug: false
+  seed: 123
+  num_devices: 4
+  warm_start: null
+  warm_start_cfg_override: true
+  resume_cfg_path: null
+  resume_ckpt_path: null
+  resume: false
+  training:
+    mask_plddt: true
+    bb_atom_scale: 0.1
+    trans_scale: 0.1
+    translation_loss_weight: 2.0
+    t_normalize_clip: 0.9
+    rotation_loss_weights: 1.0
+    aux_loss_weight: 1.0
+    aux_loss_use_bb_loss: true
+    aux_loss_use_pair_loss: true
+    aux_loss_t_pass: 0.25
+  wandb:
+    name: null
+    project: null
+  optimizer:
+    lr: 0.0001
+  trainer:
+    overfit_batches: 0
+    min_epochs: 1
+    max_epochs: 1000
+    accelerator: gpu
+    log_every_n_steps: 1
+    deterministic: false
+    strategy: ddp
+    check_val_every_n_epoch: 1
+    accumulate_grad_batches: 2
+  checkpointer:
+    dirpath: null
+    save_last: true
+    save_top_k: -1
+  inference_dir: null
+shared:
+  seed: 123
+  max_cache_size: 100000
+  min_motif_percent: 0.05
+  max_motif_percent: 0.5
+scope_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  rectify_csv_path: null
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  add_plddt_mask: false
+  min_eval_length: 60
+  max_eval_length: 128
+  inpainting_percent: 1.0
+  samples_per_eval_length: 10
+  num_eval_lengths: 10
+  filter:
+    max_num_res: 128
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+pdb_dataset:
+  seed: ${shared.seed}
+  csv_path: null
+  cluster_path: ./metadata/pdb.clusters
+  max_cache_size: ${shared.max_cache_size}
+  cache_num_res: 0
+  inpainting_percent: 1.0
+  add_plddt_mask: false
+  min_eval_length: 100
+  max_eval_length: 300
+  samples_per_eval_length: 10
+  num_eval_lengths: 10
+  filter:
+    max_num_res: 512
+    min_num_res: 60
+    max_coil_percent: 0.5
+    rog_quantile: 0.96
+    oligomeric:
+    - monomeric
+    num_chains:
+    - 1
+  min_motif_percent: ${shared.min_motif_percent}
+  max_motif_percent: ${shared.max_motif_percent}
+model:
+  node_embed_size: 256
+  edge_embed_size: 128
+  symmetric: false
+  node_features:
+    c_s: ${model.node_embed_size}
+    c_pos_emb: 128
+    c_timestep_emb: 128
+    max_num_res: 2000
+    timestep_int: 1000
+    embed_chain: false
+  edge_features:
+    single_bias_transition_n: 2
+    c_s: ${model.node_embed_size}
+    c_p: ${model.edge_embed_size}
+    relpos_k: 64
+    feat_dim: 64
+    num_bins: 22
+    self_condition: true
+    embed_chain: false
+    embed_diffuse_mask: true
+    use_rbf: true
+    num_rbf: 32
+  ipa:
+    c_s: ${model.node_embed_size}
+    c_z: ${model.edge_embed_size}
+    c_hidden: 128
+    no_heads: 8
+    no_qk_points: 8
+    no_v_points: 12
+    seq_tfmr_num_heads: 4
+    seq_tfmr_num_layers: 2
+    num_blocks: 6

reqflow_scope_rectify/reqflow_scope_rectify.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4e22235e074d9f039cfc2853b2307b61d46a19489ea1656ee33bc970108d298
+size 201851872