RFDiffusion Training Implementation

The original RFdiffusion release provides inference code but not its training pipeline. This repository implements the published training protocol and extends it for custom protein datasets.

Tested: A TCR–pMHC training implementation.

Experimental: generalized templates for other protein systems.

Environment-specific: cluster, storage, and authentication configuration.

• Implemented training protocol following the RFDiffusion supplement
• Loss functions custom backbone frame loss and L2D distance prediction
• PyTorch Lightning wrapper for clean training loops
• Weights & Biases integration for experiment tracking
• DeepSpeed support for 16-bit training (vs original 32-bit) - 2x memory reduction
• Azure Blob Storage integration for large-scale data handling

Used this protocol to fine-tune RFDiffusion for TCR-pMHC complex generation. Shown here is a random example from a hold-out set.

Before Training:

After Training:

(Note: the peptide is in red)

Losses: RFDiffusion employs two primary loss functions: dFrame and L2D. The dFrame loss was implemented following the methodology described in the supplementary materials. For the L2D loss, inconsistencies were found between the binning strategy outlined in the supplement and the actual code implementation (specifically in the PARAMS variable). To resolve this discrepancy, I opted to follow the implementation from the original RosettaFold2 repository, though this choice may introduce potential inaccuracies given the conflicting documentation.

Generalized Templates: I removed TCR/peptide/MHC specific parts and sketched out generalized code for dataset.py, train_rfdiffusion.py, test_diffusion.py that I have not tested. You'll need to adapt them for your protein system.

Compute Enviornment This implementation will not build out of the box and requires adaptation to your specific computing environment.

Require customization for your specific use case

• dataset.py - Generalized protein dataset template

• workflows/train_rfdiffusion.py - Main training pipeline template

• test_diffusion.py - Runner script

• custom - includes the custom losses and utils

1. Customize the templates: Update dataset.py and workflows/train_rfdiffusion.py to suit your custom problem

   • E.g. Define your protein regions and diffusion strategy. Modify paths and project settings in training_script.py

2. Configure for your cluster:

   • Configure DockerFile + pyproject.toml + pixi enviornment
   • Update storage paths and authentication
   • Set-up WandB if desired.

3. Launch training:

   python workflows/train_rfdiffusion.py

• PyTorch Lightning
• Weights & Biases
• DeepSpeed
• RFDiffusion (original package)

Adapted from RFDiffusion by Hussein Al-Asadi @ Adaptive Biotechnologies, with special thanks to Andrew FigPope @ Adaptive Biotechnologies for writing the PodMapper for GPU training, helping with DeepSeed, and setting up the Pixi Enviornment.