# Rigidity-Aware Geometric Pretraining for Protein Design and Conformational Ensembles

arXiv: https://arxiv.org/abs/2603.02406 · HTML: https://mlsanity.com/paper/2603.02406/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 58 references that are in the map.
- First posted on arXiv 2 Mar 2026.

## Builds on (most cited first, 26 of 26)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Learning from Protein Structure with Geometric Vector Perceptrons](https://mlsanity.com/paper/2009.01411/) (2020)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [SE(3) diffusion model with application to protein backbone generation](https://mlsanity.com/paper/2302.02277/) (2023)
- [Protein Representation Learning by Geometric Structure Pretraining](https://mlsanity.com/paper/2203.06125/) (2022)
- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [3D Infomax improves GNNs for Molecular Property Prediction](https://mlsanity.com/paper/2110.04126/) (2021)
- [Pre-training via Denoising for Molecular Property Prediction](https://mlsanity.com/paper/2206.00133/) (2022)
- [AlphaFold Meets Flow Matching for Generating Protein Ensembles](https://mlsanity.com/paper/2402.04845/) (2024)
- [Molecular Geometry Pretraining with SE(3)-Invariant Denoising Distance Matching](https://mlsanity.com/paper/2206.13602/) (2022)
- [Proteina: Scaling Flow-based Protein Structure Generative Models](https://mlsanity.com/paper/2503.00710/) (2025)
- [Learning Hierarchical Protein Representations via Complete 3D Graph Networks](https://mlsanity.com/paper/2207.12600/) (2022)
- [Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation](https://mlsanity.com/paper/2405.20313/) (2024)
- [Structure-aware Protein Self-supervised Learning](https://mlsanity.com/paper/2204.04213/) (2022)
- [A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining](https://mlsanity.com/paper/2305.18407/) (2023)
- [Contrastive Representation Learning for 3D Protein Structures](https://mlsanity.com/paper/2205.15675/) (2022)
- [Symmetry-Informed Geometric Representation for Molecules, Proteins, and Crystalline Materials](https://mlsanity.com/paper/2306.09375/) (2023)
- [CoSP: Co-supervised pretraining of pocket and ligand](https://mlsanity.com/paper/2206.12241/) (2022)
- [Multi-task Bioassay Pre-training for Protein-ligand Binding Affinity Prediction](https://mlsanity.com/paper/2306.04886/) (2023)
- [MSAGPT: Neural Prompting Protein Structure Prediction via MSA Generative Pre-Training](https://mlsanity.com/paper/2406.05347/) (2024)
- [InertialAR: Autoregressive 3D Molecule Generation with Inertial Frames](https://mlsanity.com/paper/2510.27497/) (2025)

## Abstract (arXiv)

Generative models have recently advanced textitde novo protein design by learning the statistical regularities of natural structures. However, current approaches face three key limitations: (1) Existing methods cannot jointly learn protein geometry and design tasks, where pretraining can be a solution; (2) Current pretraining methods mostly rely on local, non-rigid atomic representations for property prediction downstream tasks, limiting global geometric understanding for protein generation tasks; and (3) Existing approaches have yet to effectively model the rich dynamic and conformational information of protein structures. To overcome these issues, we introduce RigidSSL (textitRigidity-Aware Self-Supervised Learning), a geometric pretraining framework that front-loads geometry learning prior to generative finetuning. Phase I (RigidSSL-Perturb) learns geometric priors from 432K structures from the AlphaFold Protein Structure Database with simulated perturbations. Phase II (RigidSSL-MD) refines these representations on 1.3K molecular dynamics trajectories to capture physically realistic transitions. Underpinning both phases is a bi-directional, rigidity-aware flow matching objective that jointly optimizes translational and rotational dynamics to maximize mutual information between conformations. Empirically, RigidSSL variants improve designability by up to 43% while enhancing novelty and diversity in unconditional generation. Furthermore, RigidSSL-Perturb improves the success rate by 5.8% in zero-shot motif scaffolding and RigidSSL-MD captures more biophysically realistic conformational ensembles in G protein-coupled receptor modeling.