Rigidity-Aware Geometric Pretraining for Protein Design and Conformational Ensembles

Zhanghan Ni; Yanjing Li; Zeju Qiu; Bernhard Schölkopf; Hongyu Guo; Weiyang Liu; Shengchao Liu

cs.LGMar 2, 2026

Rigidity-Aware Geometric Pretraining for Protein Design and Conformational Ensembles

Authors: Zhanghan Ni, Yanjing Li, Zeju Qiu, Bernhard Schölkopf, Hongyu Guo, Weiyang Liu, Shengchao Liu

Organizations: University of Illinois Urbana-Champaign · University of Washington · MPI for Intelligent Systems, Tübingen · National Research Council of Canada · University of Ottawa · The Chinese University of Hong Kong

Abstract

Generative models have recently advanced $\textit{de 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 $\textbf{RigidSSL}$ ( $\textit{Rigidity-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.

Figures & tables

Figure 1: Overview of RigidSSL. (a) View construction in RigidSSL-Perturb: translational noise in $\mathbb{R}^{3}$ and rotational noise in $\operatorname{SO}(3)$ are applied to generate perturbations in the rigid body motion group $\mathrm{SE}(3)$ . (b) View construction in RigidSSL-MD: perturbed states are obtained by sampling conformational frames from MD trajectories. (c) Rigidity-based pretraining in RigidSSL: proteins are canonicalized into a reference frame, intermediate states are constructed via interpolation of translations and rotations for each rigid residue frame, and bi-directional flow matching is applied for pretraining. Details can be found in Section 3 .

Model	Pretraining	Designability	Novelty	Diversity
		Fraction ( $\uparrow$ )	avg. max TM ( $\downarrow$ )	pairwise TM ( $\downarrow$ )	MaxCluster ( $\uparrow$ )
FrameDiff	None	$\underline{0.775\pm 0.060}$	$\underline{0.555\pm 0.079}$	$0.565\pm 0.016$	0.033
FrameDiff	GeoSSL-EBM-NCE	$0.725\pm 0.048$	$0.615\pm 0.054$	$0.597\pm 0.010$	0.033
FrameDiff	GeoSSL-InfoNCE	$0.650\pm 0.058$	$0.613\pm 0.068$	$0.568\pm 0.013$	0.033
FrameDiff	GeoSSL-RR	$0.700\pm 0.065$	$0.579\pm 0.073$	$0.604\pm 0.015$	$\underline{0.089}$
FrameDiff	RigidSSL-Perturb (ours)	$\mathbf{0.875\pm 0.050}$	$\mathbf{0.494\pm 0.066}$	$\underline{0.534\pm 0.011}$	0.033

Table 1: Comparison of Designability (fraction with scRMSD

\leq 2.0

Å), Novelty (max. TM-score to PDB), and Diversity (avg. pairwise TM-score and MaxCluster diversity). Reported with standard errors. Best mean values are in bold and second-best are underlined .

Figure 2: Distribution of secondary structure elements ( $\alpha$ -helices, $\beta$ -sheets, and coils) in protein structure database (a-b) and in designable proteins (scRMSD $\leq$ 2.0 Å) generated by FoldFlow-2 under different pretraining methods (c-h). Plots of the structure database are color-coded by sequence length, whereas those of the generated structures are color-coded by scRMSD.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Stage	Data Source	Data Size	Length	Objective	Compute & Time
GeoSSL-InfoNCE	Pretraining	AFDB (UniProtKB)	432,194	60 $-$ 512	Contrastive (InfoNCE)	4 days (4 $\times$ A100)
GeoSSL-EBM-NCE	Pretraining	AFDB (UniProtKB)	432,194	60 $-$ 512	Contrastive (EBM-NCE)	4 days (4 $\times$ A100)
GeoSSL-RR	Pretraining	AFDB (UniProtKB)	432,194	60 $-$ 512	Generative (Representation Reconstruction)	3 days (4 $\times$ A100)
RigidSSL-Perturb (Ours)	Pretraining (Phase I)	AFDB (UniProtKB)	432,194	60 $-$ 512	Generative (RigidSSL)	2.75 days (1 $\times$ H100)
RigidSSL-MD (Ours)	Pretraining (Phase II)	ATLAS	1,390 traj.	60 $-$ 512	Generative (RigidSSL)	1.88 days (1 $\times$ H100)
FrameDiff	Finetuning	PDB	20,312	60 $-$ 512	SE(3) Diffusion	7 days (4 $\times$ H100)

Appendix

Table 2: Overview of pretraining and downstream training configurations.

Model	Pretraining	FPSD vs.		fS	fJSD vs.
		PDB ( $\downarrow$ )	AFDB ( $\downarrow$ )	(C/A/T) ( $\uparrow$ )	PDB (C/A/T) ( $\downarrow$ )	AFDB (C/A/T) ( $\downarrow$ )
FrameDiff	None	908.743	840.682	1.052 /1.827/4.249	3.174 /3.902/5.454	2.209 /3.119/4.807
FrameDiff	GeoSSL-EBM-NCE	827.919	770.448	1.108 / 2.243 /6.025	3.017 /3.593/5.079	2.068 /2.807/4.408
FrameDiff	GeoSSL-InfoNCE	858.332	806.346	1.002/ 2.060 / 6.862	3.360/3.817/5.195	2.369/2.965/4.419
FrameDiff	GeoSSL-RR	892.657	846.222	1.002/1.851/ 6.559	3.364/ 3.474 / 5.037	2.373/ 2.452 / 4.357
FrameDiff	RigidSSL-Perturb (ours)	901.424	843.407	1.003/1.735/4.105	3.359/3.605/5.417	2.370/2.632/4.687

Appendix

Table 3: Comparison of FPSD, fS, and fJSD. Lower is better for FPSD/fJSD; higher is better for fS. Best mean values are in bold and second-best are underlined .

Metric	Length	Unpretrained	GeoSSL-EBM-NCE	GeoSSL-RR	GeoSSL-InfoNCE	RigidSSL-Perturb	RigidSSL-MD
Clashscore ( $\downarrow$ )	700	152.53	45.11	96.38	55.02	21.36	37.16
Clashscore ( $\downarrow$ )	800	172.36	61.70	123.16	74.98	26.42	39.35
MolProbity score ( $\downarrow$ )	700	2.64	2.22	2.51	2.21	1.82	2.05
MolProbity score ( $\downarrow$ )	800	2.69	2.38	2.61	2.34	1.91	2.08

Appendix

Table 4: Comparison of Clashscore and MolProbity scores for 700- and 800-residue proteins generated by FoldFlow-2 under different pretraining methods. The best values are shown in bold and the second-best are underlined .

Figure 3: FoldFlow-2 generated structures (orange) compared against ProteinMPNN $\rightarrow$ ESMFold refolded structures (grey). Columns denote pretraining methods, and rows denote sequence lengths of 700 and 800.

Target	Unpretrained	GeoSSL-EBM-NCE	GeoSSL-InfoNCE	GeoSSL-RR	RigidSSL-Perturb	RigidSSL-MD
1BCF	11	2	29	0	47	15
1PRW	5	1	11	1	11	7
1QJG	11	9	17	7	8	10
1YCR	9	10	8	5	19	12
2KL8	15	5	23	14	17	34
3IXT	25	22	23	8	25	19

Appendix

Table 5: Per-target designability comparison across different geometric pretraining methods. Each entry reports the number of successful designs out of 100 trials for the specified target. Best values are shown in bold and second-best are underlined .

	Metric	Unpretrained	GeoSSL-EBM-NCE	GeoSSL-InfoNCE	GeoSSL-RR	RigidSSL-Perturb (ours)	RigidSSL-MD (ours)
Predicting flexibility	Pairwise RMSD (=1.55)	2.37	2.32	2.31	2.48	2.20	2.66
Predicting flexibility	All-atom RMSF (=1.0)	1.30	1.19	1.22	1.17	1.08	1.28
Distributional accuracy	Root mean $\mathcal{W}_{2}$ -dist. $\downarrow$	18.36	18.46	17.79	17.83	17.90	18.10
	$\hookrightarrow$ Translation contrib. $\downarrow$	18.26	18.31	17.75	17.76	17.83	18.05
	$\hookrightarrow$ Variance contrib. $\downarrow$	1.45	1.50	1.42	1.52	1.44	1.60
	MD PCA $\mathcal{W}_{2}$ -dist. $\downarrow$	3.15	2.97	3.16	3.14	3.10	3.22

Appendix

Table 6: Evaluation of GPCR MD ensembles generated with pretrained AlphaFlow variants. Best results are shown in bold .

				Designability	Novelty	Diversity
Model	Downstream Dataset	Pretraining	Steps	Fraction ( $\uparrow$ )	avg. max TM ( $\downarrow$ )	pairwise TM ( $\downarrow$ )	MaxCluster ( $\uparrow$ )
FoldFlow-2	PDB + AFDB	None	500k	$0.738\pm 0.014$	$0.764\pm 0.003$	$0.657\pm 0.001$	$0.250$
FoldFlow-2	PDB	RigidSSL-Perturb	400k	$0.758\pm 0.016$	$0.770\pm 0.003$	$0.650\pm 0.001$	$0.252$

Appendix

Table 7: Comparison of Designability (fraction of proteins with scRMSD

\leq

2.0 Å), Novelty (max. TM-score to PDB), and Diversity (avg. pairwise TM-score and MaxCluster diversity) of structures generated by FoldFlow-2 trained on PDB + AFDB from scratch versus models first pretrained on AFDB and finetuned on PDB. Reported with standard errors. Results for the former are evaluated using the official model checkpoint provided by the authors of ? .

			Designability	Novelty	Diversity
Translation Noise	Rotation Noise	Length	Fraction ( $\uparrow$ )	avg. max TM ( $\downarrow$ )	pairwise TM ( $\downarrow$ )
0.01	0.5	100–600	$0.336\pm 0.012$	$0.768\pm 0.010$	$0.635\pm 0.002$
0.03	0.5	100–600	$0.758\pm 0.016$	$0.770\pm 0.003$	$0.650\pm 0.001$
0.05	0.5	100–600	$0.589\pm 0.014$	$0.769\pm 0.005$	$0.654\pm 0.002$
0.5	0.5	100–600	$0.382\pm 0.015$	$0.748\pm 0.005$	$0.649\pm 0.001$
0.5	0.75	100–600	$0.660\pm 0.014$	$0.763\pm 0.004$	$0.644\pm 0.001$

Appendix

Table 8: Comparison of Designability (fraction of proteins with scRMSD

\leq

2.0 Å), Novelty (max. TM-score to PDB), and Diversity (avg. pairwise TM-score) of structures generated by FoldFlow-2 pretrained with RigidSSL-Perturb under different noise scales. All metrics include standard errors.

Figure 4: Impact of translation and rotation noise scale on protein structure validity.

Explore similar work

Proteo-R1: Reasoning Foundation Models for De Novo Protein Design
May 1, 2026Fang Wu, Weihao Xuan, Heli Qi +26Generative Modeling 3D Molecular Generation
Stanford University
ProteinZero: Self-Improving Protein Generation via Online Reinforcement Learning
Jun 9, 2025Ziwen Wang, Jiajun Fan, Ruihan Guo +3Protein Language Models Molecular Generation
Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, Urbana, IL, USA · DOE Center for Advanced Bioenergy and Bioproducts Innovation, University of Illinois Urbana-Champaign, Urbana, IL, USA
Designing Protein Structures from Natural Language
Oct 9, 2026Joseph Lee, Yanjing Li, Bo Qiang +7Protein Design Generative Modeling
University of Washington · National University of Singapore · Allen Institute for AI (Ai2) +1

May 1, 2026J/K move · Enter open · S save

Proteo-R1: Reasoning Foundation Models for De Novo Protein Design

Fang Wu, Weihao Xuan, Heli Qi, Hanqun Cao+25

Stanford University

Generative Modeling 3D Molecular Generation Molecular Generation Protein Design Large Language Model-Guided Scientific Discovery Conditional Generative Modeling

Open paper AskPDF arXiv

Deep learning in de novo protein design has achieved atomic-level fidelity. However, existing models remain largely non-deliberative: they directly synthesize molecular geometries without explicitly reasoning about which residues or interactions are functionally essential. As a result, design decisions are entangled with continuous sampling dynamics, limiting interpretability, controllability, and systematic reuse of biochemical knowledge. We introduce Proteo-R1, a reasoning-guided protein design framework that explicitly decouples molecular understanding from geometric generation. Proteo-R1 adopts a dual-expert architecture in which a multimodal large language model (MLLM) serves as an understanding expert, analyzing protein sequences, structures, and textual context to identify key functional residues that govern binding and specificity. These residue-level decisions are then passed as hard constraints to a separate diffusion-based generation expert, which performs conditional co-design while respecting the fixed interaction anchors. This factorization mirrors how human experts approach molecular engineering: first, reasoning about critical interactions, then optimizing geometry subject to those constraints. By operationalizing reasoning as explicit residue-level commitments rather than latent textual guidance, Proteo-R1 achieves stable, interpretable, and modular integration of LLM reasoning with state-of-the-art geometric generative models. Code, data, and demos are available at https://smiles724.github.io/r1/.

Abstract

Explore similar work

Proteo-R1: Reasoning Foundation Models for De Novo Protein Design

ProteinZero: Self-Improving Protein Generation via Online Reinforcement Learning

Designing Protein Structures from Natural Language