Repository navigation
Add an RFD4-Proteina LoRA training example - #15155
kierandidi wants to merge 12 commits into
Conversation
Trains PEFT adapters on the denoiser with the released recipe, saved in the format RFD4ProteinaPipeline.load_lora_weights reads.
d23d344 to
024d284
Compare
|
Hi @kierandidi, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Once the PR links an issue (or gets the |
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Looking good. Some minor feedback.
In general, we should try to follow the patterns already establised in the existing training scripts for LoRA, when it comes to loading and saving adapters.
Could you also show us the results from a successful run?
Save the adapter the way the current LoRA examples do: the PEFT state dict goes out unconverted, and `_collate_lora_metadata` records rank, alpha and target modules alongside it. Resolve `--resume_from_checkpoint` against `--output_dir`, so naming a checkpoint works as well as passing its path.
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Thanks! Some further comments.
Can we also see some validation examples where LoRA training provided an advantage?
Train on 40 of the fifty published structures and score the remaining ten with the adapter toggled on and off over identical noise draws. Held-out flow-matching loss falls 11%, on every held-out structure.
|
@kierandidi This is cool!
cc @sayakpaul as well |
Every run now writes a model card next to the adapter, naming the base checkpoint and the structures trained on, with a snippet that generates from the result. `--push_to_hub`, `--hub_model_id` and `--hub_token` upload the adapter and that card at the end of training, as the other LoRA examples do.
Fixes #15162
Adds a LoRA training example for RFD4-Proteina, with an Accelerate script, README and requirements. It trains the denoiser on annotated CIFs while keeping the autoencoder frozen, and saves adapters through the standard Diffusers LoRA helpers.
The recipe uses rank 32 / alpha 64 and flow-matching losses for backbone coordinates and local latents. It leaves out the native trainer's auxiliary sequence and insertion losses. CIF preparation is cached, and unusable structures are filtered before distributed sharding.
Cyclic peptide check
Trained on 255 CPSea complexes for 1,000 steps on two H200s, starting from weights matched to the updated native release. We tested the PD-1 cyclic peptide inputs with the same 16 seeds per chemistry at steps 0, 500 and 1,000. PD-1 was excluded from fine-tuning.
Median contacts (heavy-atom pairs within 5 Å) rose from 187.5→229.5, 150.5→262 and 160.5→220 for head-to-tail, disulfide and isopeptide designs. Contacts per peptide atom also increased. Shading shows 95% seed-bootstrap intervals; loss uses EMA α=0.02 over faint raw values.
This is one target and one training run. Head-to-tail closure stayed at 16/16, but disulfide fell from 12/16 to 6/16 and isopeptide from 5/16 to 1/16. More contacts alone don't establish better binding or overall design quality.