GenUnfold

Rapidly Predict Protein Mechanical Unfolding Trajectory via a Physics-Guided Diffusion Model

GenUnfold is the codebase for our ICML 2026 paper:

GenUnfold: Rapidly Predict Protein Mechanical Unfolding Trajectory via a Physics-Guided Diffusion Model
Yiyuan Zhang, Cailong Hua, Vinitendra Singh, Joseph M. Muretta, James M. Ervasti, and Murti V. Salapaka
Proceedings of the 43rd International Conference on Machine Learning (ICML 2026), Seoul, South Korea, PMLR 306.

GenUnfold is a physics-guided diffusion framework for predicting full protein mechanical unfolding force--extension (F--x) trajectories from protein sequence and native structure. It is designed to accelerate the characterization of protein mechanostability by generating statistically meaningful unfolding trajectory ensembles in seconds rather than relying solely on labor-intensive single-molecule force spectroscopy (SMFS) experiments or computationally expensive steered molecular dynamics (SMD) simulations.

Figure 1. Main framework of GenUnfold for physics-guided protein mechanical unfolding trajectory generation.

Highlights

• Full trajectory generation: predicts complete mechanical unfolding F--x curves rather than only scalar properties such as maximum unfolding force.
• Evolution--topology co-conditioning: combines global protein identity from ESM-2 sequence embeddings with local, mechanics-aware structural representations.
• Physics-biased attention: injects structure-derived topology priors into the attention logits, allowing the model to learn residue--residue dependencies and force-propagation pathways.
• Scalable ensemble sampling: generates repeated stochastic trajectories for downstream estimation of unfolding force and unfolding energy distributions.
• Benchmark for protein unfolding generation: establishes a systematic benchmark on the Biomolecule Stretching Database (BSDB) with transformer and protein-language-model baselines.

Why GenUnfold?

Protein mechanical unfolding is central to many biological processes, including load-bearing and force-transmitting systems such as muscle and cytoskeletal proteins. SMFS directly measures force--extension trajectories, but obtaining enough curves for reliable statistics can be slow and costly. SMD simulations provide atomistic or coarse-grained alternatives, but repeated trajectory sampling remains expensive.

GenUnfold addresses this bottleneck by learning the conditional distribution of unfolding trajectories:

[ p_\theta(F \mid S, X, R_f, G), ]

where F is the force--extension trajectory, S is the amino-acid sequence, X is the native structure, R_f contains residue-level physicochemical and geometric features, and G contains structure-derived physical topology priors.

Method Overview

GenUnfold uses a diffusion transformer backbone with two complementary conditioning streams:

1. Coevolutionary context
   ESM-2 residue embeddings are mean-pooled into a global protein representation and combined with the diffusion timestep embedding. This global vector modulates the diffusion transformer through adaptive layer normalization.

2. Mechanical representation
   A multi-head physics-biased attention encoder processes residue-level features together with physical topology priors from the native structure. The resulting local mechanical representation is injected into the diffusion backbone through cross-attention.

The physical topology priors include:

• Distance map (D): pairwise residue distances.
• Stiffness map (K): local contact stiffness approximated from residue distances.
• Effective resistance map (R): global mechanical coupling derived from the graph Laplacian pseudoinverse.

A lightweight CNN fuses these priors into learnable per-head attention biases, and a timestep-dependent gate dynamically modulates their influence during denoising.

Figure 2. Physics-biased attention module for learning mechanically informed residue representations from structure-derived topology priors.

Paper Results

On the BSDB benchmark, GenUnfold achieved state-of-the-art performance among the evaluated diffusion baselines. The result comparison below summarizes the FID, Force-JSD, and Energy-JSD of the evaluated models. Lower values are better.

Figure 3. Benchmark comparison of GenUnfold and baseline diffusion models using FID, Force-JSD, and Energy-JSD.

The model also showed improved out-of-distribution generalization on unseen CATH superfamilies and generated force and energy distributions that closely matched the corresponding simulated ensembles across structural classes.

Dataset

The main benchmark is built from the Biomolecule Stretching Database (BSDB) and cross-referenced with PDB structures. The curated dataset used in the paper contains:

• 16,495 unique proteins
• 32,990 force--extension curves
• A strict 80/10/10 protein-level train/validation/test split
• Precomputed conditioning features:
  • ESM-2 sequence embeddings
  • Residue-level physicochemical and structural features
  • Distance, stiffness, and effective-resistance topology maps

Repository Layout

.
├── README.md
├── requirements.txt
├── data                      # Optional local dataset directory
├── scripts
│   ├── run.py                # Training/inference entrypoint; edit local paths before use
│   ├── inference.ipynb       # Notebook example for inference
│   ├── results.py            # Plotting and evaluation utilities used for paper figures
│   ├── data                  # Sample arrays, labels, and metadata
│   ├── 1emb                  # Example PDB-derived features and outputs
│   └── sample_results        # Example output figures
├── src
│   ├── data_processing       # Dataset loading, preprocessing, and feature extraction
│   ├── models                # Diffusion architectures, protein encoders, and schedulers
│   ├── training              # Trainers, losses, and training helpers
│   ├── evaluation            # Metrics and visualization utilities
│   ├── analysis              # Curve and property analysis utilities
│   ├── utils.py              # Pipeline, training, and inference glue code
│   └── __init__.py
└── trained_models
    └── best_model.pt         # Example checkpoint; config.yaml is expected alongside it

Installation

git clone https://github.com/SalapakaLab-SIMBioSys/Gen_Unfold.git
cd Gen_Unfold

python -m venv venv
source venv/bin/activate        # Linux/macOS
# venv\Scripts\activate.bat    # Windows CMD
# venv\Scripts\Activate.ps1    # Windows PowerShell

pip install -r requirements.txt

Additional dependencies may be required depending on the protein language model or structure-feature pipeline used in your local setup.

Usage

1. Training

Edit the local paths and configuration in scripts/run.py, then run:

python scripts/run.py

The training entrypoint calls train_pipeline(...) from src/utils.py. A YAML configuration file and local dataset paths should be provided before launching training.

2. Inference / Curve Generation

The inference pipeline uses curve_prediction(...) from src/utils.py. It expects:

• A pretrained model directory containing best_model.pt and config.yaml.
• A PDB/mmCIF file or PDB-derived structure path.
• The chain identifier.
• A feature directory for precomputed or newly generated structure features.

Example:

from Gen_Unfold.src import curve_prediction

curve_prediction(
    pretrained_model_path="/path/to/checkpoint_dir",
    pdb_id_or_path="/path/to/1emb.cif",
    chain="A",
    feature_path="/path/to/feature_dir",
    save_path="/path/to/predictions.npy",
    num_samples=1024,
    device="cuda",
)

The output is an ensemble of generated force--extension curves that can be used to estimate downstream mechanical statistics such as unfolding-force and unfolding-energy distributions.

3. Plotting and Evaluation

Use scripts/results.py to reproduce representative comparison plots and evaluation summaries:

python scripts/results.py

Example plots are available in scripts/sample_results/.

Citation

If you use this repository or build on GenUnfold, please cite:

@inproceedings{zhang2026genunfold,
  title     = {GenUnfold: Rapidly Predict Protein Mechanical Unfolding Trajectory via a Physics-Guided Diffusion Model},
  author    = {Zhang, Yiyuan and Hua, Cailong and Singh, Vinitendra and Muretta, Joseph M. and Ervasti, James M. and Salapaka, Murti V.},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  series    = {Proceedings of Machine Learning Research},
  volume    = {306},
  address   = {Seoul, South Korea},
  year      = {2026}
}

Notes

• This repository is research-oriented. Some scripts contain local path placeholders and should be edited before use.
• trained_models/best_model.pt is an example checkpoint. The inference pipeline expects a config.yaml in the same model directory.
• Sample data for plotting and example outputs are stored under scripts/data/ and scripts/sample_results/.

License

This project is released under the MIT License.