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
Venue: ICLR 2026 Poster (2026)
Keywords: Protein design
Self-supervised Learning
3D Geometry
Rigidity
Conformational ensembles
Flow matching
SE(3)-equivariance
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a rigidity-aware self-supervised learning framework for protein design and conformational ensemble generation.

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. The code is available at: https://github.com/ZhanghanNi/RigidSSL.git.

Published as a conference paper at ICLR 2026
RIGIDITY-AWAR GEOMETRIC PRETRAINING FOR PROTEIN DESIGN AND CONFORMATIONAL ENSEMBLES
Zhanghan Ni Yanjing Li
Zeju Qiu
Bernhard Schölkopf"
Hongyu Guo
Weiyang Liu 3.6
Shengchao Liu
1 University of Illinois Urbana- Chamapa-ha
University of Washington
3 MPI for Intelligent Systems, Tübingen
*National Research Council of Canada
6 University of Ottawa
'The Chinese University of Hong Kong
*Equal contribution
ABSTRACT
Generative models have recently advanced 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 solution; (2) Current pretraining meth- ods 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 (Rigidity-Awan Self- Sufervieer Learning), a geomet- ric pretraining framework that front-loads geometry learning prior to generative finetuning. Phase I (RigidSSII-Perturh) learns geometric priors from 432K struc- tures 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 unconditional generation. Furthermore, RigidSSI -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. The code is available this repository.
INTRODUCTION
Proteins are large and complex biomolecules whose three-dimensional structures determine their diverse biological functions, enabling precise molecular activities such as substrate specificity, binding affinity, and catalytic activity (LaPelusa & Kaushik, 2025). The ability to design proteins with desired properties has the potential to revolutionize multiple disciplines. Examples include novel therapeutics and vaccines in medicine and sustainable biomaterials in materials science (Listov al., 2024; Miserez et al., 2023). In recent years, rapid advances in deep generative modeling have opened new directions for protein science (Huguet et al., 2024; Yim et al., 2023). More specifically, generative modeling has been widely used for protein-related biology tasks, including de novo protein design (Bose et al., 2023; Watson et al., 2023), motif scaffolding (Wang et al., 2022a), and ensemble generation (Jing et al., 2024). showing the potential of geometric learning to model structure and dynamics at scale.
Challenge Learning geometry and generation simultaneously brings in challenges in model generalization. Despite significant progress in protein generative modeling, existing end-to-end frameworks often require the model to jointly learn the fundamental geometry of proteins and the complex mechanism of structure generation within single objective (Huguet et al., 2024; Yim et al., 2023). This tight coupling can lead to inefficient optimization and limit generalization to novel or design tasks. To alleviate this issue, existing approaches typically follow three directions. One line of work repurposes large structure prediction models to inherit their geometric priors (Watson et al., 2023), but this strategy incurs substantial computational overhead and limits

---

Published as a conference paper at ICLR 2026
(a)
(b)
translational perturb
90
91
rotational perturb
go
91
frame
View Construction Perturb
View Construction MD
(c)
2. Rigidity-leeee Interpolation
1.Canonicalization
Camonicalization
gt
A.Bi-directional Pretraining
9o
91
Figure 1: Overview of RigidSSL. (a) View construction in SigidSSI-PPertuii translational noise in IR3 and rotational noise in SO(3) applied to generate perturbations in the rigid body motion group SE(3). (b) View construction in RigidSSL-MD: perturbed states obtained by sampling conformational frames from MD trajectories. (c) Rigidity-based pretraining RigidSSL: proteins are canonicalized into 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.
architectural flexibility. Another line of work designs sophisticated architectures to compensate for the lack of priors (Wang et al., 2024), resulting complex and specialized models achieve competitive performance. Motivated by the success of pretraining (Jiao et al., 2022; Liu et al., 2021; 2022a; 2023b; Zaidi et al., 2022), a more scalable alternative is to front-load geometric understanding prior to downstream generation by first pretraining models to acquire general understanding of protein geometry. Such pretrained representations can then serve as a foundation for diverse generative tasks. However, the effectiveness of this paradigm eritically depends how protein geometry is represented, which determines what structural priors can be learned and transferred.
Pretraining requires global and efficient representations. Realizing effective pretraining paradigm places stringent requirements on the underlying representations, which must support scalable computation while faithfully capturing global geometry. However, most geometric pretraining methods rely on local, non-rigid atomic or fragment-level representations (Chen et al., 2023; Guo al., 2022; Hermosilla & Ropinski, 2022: Zhang et al., 2022) that primarily capture short-range geometric patterns. While this locality is often sufficient for property prediction, it can under-represent global folding geometry, limiting structural abstraction and thus the transferability of learned representations to generative design tasks (Bouatta et al., 2021; Skolnick et al., 2021). Achieving representations that are both global and efficient nontrivial, in part because all-atom modeling is computationally expensive, which limits scalability and learning efficiency on large geometric datasets.
Challenge 2:
Challenge 3: Pretraining requires diverse and multi-scale data. In addition to the choice of protein backbone representation, the availability and quality of the dataset present another critical bottleneck for effective cructurrebassd pretraining. Large-scale structure databases such the AlphaFold Protein Structure Database (AFDB) (Varadi et al., 2022) and Protein Data Bank (PDB) (Berman et al., 2000) together provide hundreds of millions of protein structures encoding rich geometric regularities
2

---

Published as a conference paper at ICLR 2026
of natural folds, which is an invaluable resource for geometric pretraining. However, these datasets are dominated by static structures and largely overlook the intrinsic conformational flexibility of proteins. As result, models pretrained solely on such static snapshots primarily learn the geometry of static states, failing to capture near-native fluctuations or transitions among metastable conformations. This limitation highlights a fundamental challenge for aructtur-based pretraining: the requirement for diverse, multi-scale data, including conformational ensembles derived from molecular dynamics (MD) simulations, that span both rigid geometric motif's and dynamic variability to support generative modeling with improved conformational diversity.
Our contributions. To address these challenges, we propose RigidSSL (Rigidity-Aware Self- Supervised Learning). a two-stage geometric pretraining framework for protein structure generation. RigidSSI. explicitly follows the backbone modeling via residue-level translations and rotations in SE(3), which substantially reduces degrees of freedom and allows the model to learn geometric priors under physical constraints. Building on such representation, RigidSSL employs a two-phase pretraining strategy that sequentially integrates multi-scale structural information from static and dynamic data. Phase (RigidSSSL-Perturb) pretrains on 432K static AFDB structures with simulated perturbations in SE(3) to emulate broad but coarse conformational variation (Figure (a)); Phase II (RigidSSL-MDD pretrains on 1.3K MD trajectories to provide physically realistic transitions, refining the representation toward true dynamical flexibility (Figure 1(b)). In both phases, operat- ing under a canonicalized reference system, RigidSSI. adopts a bi-directional, rigidity- flow matching objective to jointly optimize translational and rotational dynamics, serving as a surrogate for maximizing mutual information between paired conformations (Figure ((c)).
To verify the effectiveness of RigidSSL, we consider two types of downstream tasks: (1) protein design, including protein unconditional generation and motif scaffolding, and (2) conformational ensemble generation. In unconditional generation, RigidSSL variants improve designability by up to 43% while enhancing novelty and diversity. Notably, RigidSSL- effectively generalizes to long protein chains of 700-800 residues, yielding the best stereochemical quality as reflected by the lowest Clashscore and MolProbity score. In motif scaffolding, RigidSSI--Perturb improves the average success rate by 5.8%. In ensemble generation, RigidSSL variants learn to represent a wider and more physically realistic conformational landscape of G protein receptors (GPCRs), collectively achieving the best performance on 7 out of 9 metrics. Overall, RigidSSL substantially improves physical plausibility and generative diversity in downstream design tasks.
Related Work. We briefly review the most related works here and include a more detailed discussion in Appendix A. Geometric representation learning has been extensively explored for both small molecules (Coors et al., 2018; Gasteiger et al., 2020; Schütt et al., 2018) and proteins (Fan et al., 2022; Jing et al., 2020; Wang al., 2022b), with methods categorized into SE(3)-invariant and SE(3)-equivariant models (Liu et al., 2023a). Existing pretraining methods for small molecules optimize mutual information between or within modalities (Jiao et al., 2022; Liu al., 2021; 2022a; 2023b; Zaidi et al., 2022). Protein-specife pretraining has evolved through diffusion- based methods (Guo et al., 2022), contrastive learning on substructures via GearNet (Zhang et al., 2022) and ProteinContrast (Hermosilla & Ropinski, 2022), and MSA-based strategies like MSAGPT (Chen et al., 2024). Finally, joint protein-ligand pretraining approaches, including CoSP (Gao et al., 2022) and MBP (Yan et al., 2023), model binding interactions directly.
2 PRELIMINARIES
Protein Structure Representation. A protein composed of a sequence of residues, each containing three backbone atoms in the order N, C.. and C. Accordingly, the protein backbone with L residues can be represented as L x [N,C., C). Although there rotational freedom around the N-C, (Ø) and Ca-C' (1)) bonds, following the common simplification used in AlphaFold2 (Jumper et al., 2021), we treat each residue as a rigid body. In this formulation, bond lengths and bond angles within the backbone are fixed to idcalized values, and only torsional rotations (ò, bh w) are allowed to vary. Because each residue is modeled as a rigid body, its configuration in 3D space can be described entirely by its position and orientation. Thus, a protein chain can be modeled as sequence of rigid residues, where each residue is parameterized by a translation and a rotation. Formally, we represent a protein structure as a sequence of rigid transformations graw = [Traw, 1- = [(Fraw, Traw. i) H=r where the subscript "raw" indicates the raw and uncanonicalized coordinates from the database,
3

---

Published as a conference paper at ICLR 2026
and each residue i is described by a translation vector & E R and rotation matrix r, SO(3). specifying its position and orientation in 3D space. Later, we will adopt a camonicalization process to align each protein structure to its invariant pose, Le., g = [T: = {(t,, r:)) We define z; := & as the 3D coordinates of the C. atom at residue i, so {a; ii gives the Ca trace of the backbone. In what follows, we will use I and H; interchangeably according to different contexts. Furthermore, each residue is also associated with an amino acid identity A; E (1....,21 ,21), standard amino acids plus one unidentificd type. More details can be found in Appendix C.I.
representing the 20
Structure Encoder. Invariant Point Attention (IPA), first introduced in AlphaFold2 (Jumper et al., 2021), I), a geometric attention mechanism designed to encode protein backbone representations. is provably invariant under rigid Euclidean transformations due to the fact that the inner product of two rigidity bases stays constant under SE(3) transformations. We adopt IPA as the base model both pretraining and downstream tasks, More details can be found in Appendix E.
Flow Matching. Flow Matching (FM) has been an expressive generative model (Albergo Vanden- Eijnden, 2022; Lipman et al., 2022; Liu et al., 2022b). It directly minimizes the discrepancy between a learnable vector field (x) and a target velocity field we (2) that transports samples along probability path from simple prior distribution at 0 to target distribution at / 1. The training objective is
(0) 2e- ==p((() [1)x(x)) ue (x))
(1)
where (z(z) is the model output and u, (z) is the ideal (but generally unknown) velocity field that induces the time-dependent distribution pe (x). However, this objective is often intractable because computing (I)) requires knowledge of the exact marginal distributions and their dynamics.
To address this issue, Conditional Flow Matching (CFM) (Lipman al., 2022) introduces conditioning on known target sample x1 define a tractable conditional probability path from To to ø1. The objective then becomes:
LCvM = E0-M(0,11 (1 (xll))) (mpel))))) [Ilv:(*) - ue (x; (x;xx1))) x1))
(2)
where (x11) ) denotes the marginal distribution of targets and pe | 21) defines the interpolated distribution between 2o and i at time f. In this setting, u, (x x1) is known velocity field (e.g.. from linear or spherical interpolation), making the training objective tractable. This formulation enables an efficient training of conditional generative models by leveraging path-based supervision.
METHOD: RIGIDSSL
RigidSSL learns transferable geometric representations of proteins by integrating large-scale static structures with dynamic conformational data. While databases such as the PDB and AFDB pro- vide high-quality but static snapshots, native proteins continually fluctuate around equilibrium states (Miller & Phillips, 2021; Nam & Wolf Wotz-Watz, 2023). To capture both equilibrium geometry and near-native dynamics, RigidSSI. adopts a two-phase rigidity-based pretraining strategy (Figure 1), comprising three components: (1) frame canonicalization; (2) view construction; and (3) an objec- tive function based on rigidity-guided flow matching. Two phases pretraining contain: Phase I (RigidSSL-Perturb) distills large-scale geometric regularities from 432K AFDB structures via rigid body perturbations SE(3). Phase II (RigidSSL-MD) refines these representations using 1.3K MD trajectories that capture physically realistic motions. Across both phases, proteins are canonicalized into inertial reference frames, and the model is optimized with a bi-directional, rigidity-awarr flow matching objective that jointly optimizes translational and rotational dynamics to maximize mutual information between conformations.
3.1 REFERENCE FRAME CANONICALIZZATION
Protein structures in databases exist under arbitrary coordinate systems. To establish a canonical reference system, align each protein structure to its reference inertial frame, convention widely adopted in the community (Guo et al., 2025; Li et al., 2025). Recall from Section 2 that during mod- eling, treat each protein structure 9 sequence of L rigid residues, where each residue is repre- sented with a translation and rotation transformation. i.e., Iraw {Traw, ) //= Recall that the subscript "raw" indicates the raw and uncanonicalized coordinates from the database. Based this, the global alignment proceeds in two steps.
{(Fraw.s Fraw :) )
4

---

Published as a conference paper at ICLR 2026
Translational Alignment. We align the protein to its center of mass, F = defines the origin of the inertial frame. Each residue's 's coordinate is shifted as 24 = Braw, T. and the translation vector becomes E = trww, - T.
E/=1 Fraw, which
Rotational Alignment. We align the protein to its principal axes, defined by the axes of the inertial frame. More concretely, we compute the inertia tensor Σ I|= - xxar]) where I3 the 3 x 3 identity matrix. We then perform the eigendecomposition I = VAVT, where the columns of V are the eigenvectors forming the principal axes of inertia. Additional geometric constraints are applied to obtain deterministic V E (3(3) (see Appendix C.2). Then, the aligned rotation matrix is computed as r, = Traw,""
The canonicalized protein structure is represented as g= - (f Frame alignment serves canonicalization step before Perrttb and RigidSSL-MD, as will be introduced next. By expressing all protein structures in a consistent reference frame, we ensure that the rotational and translational interpolation paths in SH(3) also reside in consistent reference system, improving the generalizability of RigidSSL.
3.2
VIEW CONSTRUCTIONS IN A TWO-PHASE PRETRAINING FRAMEWORK
3.2. PIIASE I: RIGIDSSL-PEETURE
RigidSSL-Perturb adopts the massive AFDB (Varadi et al., 2022) with 432K unique protein structures. each training iteration, RigidSSL-Perturb applies simulated perturbation each AFDB structure g" to generate second view g¹, as shown in Figure (()). Concretely, for gd parameterized as collection of aligned rigid bodies (Ti)
((T,, RigidSSI independently
constructs perturbed view for each rigid body T; (h,rr) obtain g². This residue-wise perturbation handles translation t; and rotation separately.
Translation Perturbation. 7.0 is perturbed by adding Gaussian noise in the Euclidean space:
7. r. + 2 1(0,
Rotation Perturbation. To generate physically plausible structural variations in protein backbones, we employ the isotropic Gaussian distribution on the special orthogonal group SO(3) (IGsors)) This choice is guided by three key considerations: (I) protein dynamics fundamentally arise from thermal Brownian motion, which I95o(3) naturally models in the rotational domain (Yanagida et al., 2007), (2) rotations form a -uelldeide manifold, requiring manifold-aware sampling to maintain geometric validity, and (3) since proteins explore their entire conformational landscape through continuous rotational changes, we need smooth, continuous sampling across all rotational states without singularities. Methodologically. IGso(3) is parametrized by mean rotation SO(3) and a concentration parameter E R,. The density function of this distribution is expressed as:
p(ri,,))
Z(c) exp (=)eee((+)))
(4)
where Z()) is a normalization constant. To sample a rotation r' ~ IG (1,, e²), we follow the axis- angle parameterization and sampling method of Leach et al. (2022) (see Appendices D.2 and D.3). Then, the rotational rigid perturbation is obtained via:
=
IGsoro
Note that perturbations are applied canonicalized proteins to ensure consistent geometric effects, as detailed in Appendix D.4.
Summary. Building upon Equations (3) and (5), we define perturbation function Perturb SE(3) - SE(33 that applies both translational and rotational noise to individual rigid body frames:
T Perturb., (T: = (T. 2. .)). where = ~ Io), r II
(6)
The perturbed view g is then constructed the set of transformed frames --
5

---

Published as a conference paper at ICLR 2026
3.2.2 PHASE II: RIGIDSSL-MD
As for phase II, RigidSSL-MD aims learn about more realistic protein structural dynamics. is trained on the ATLAS dataset (Vander Meersche et al., 2024), which contains 1.3K MD trajectories. MD simulations sample the energy landscape of proteins by numerically integrating Newton's equations of motion, producing trajectories that reflect structural fluctuations governed by physical force fields. To capture such fluctuations, we form two views g and y' from time-separated snapshots within the same trajectory.
In trajectory, the raw snapshot at time s can be denoted as Fraw = [(Era, ia Faw.:))es We extract pairs (Fraw,s; Fraw,s+8) with fixed interval 8. The choice of ē controls the scale of conformational variation: smaller values capture only minor thermal vibrations, whereas larger values may reflect substantial conformational rearrangements. In this work, we set 5 2 yield views that represent conformational fluctuations. Canonicalizing the raw snapshots gives Fs, pstt, yielding g" = = {(T', ((Grrr)) ri) and n² == pst6 = {(,, -.45 ))
3.3 RIGID FLOW MATCHING FOR MULTI-VIEW PRETRAINING: RIGIDSSL
In Sections 3 and 3.2.2, we introduce two sequential phases, each constructing set of paired views ((g".g") }. Our goal is to learn a generalizable latent space by maximizing the mutual information (MI) between the constructed views, i.e., L MI(g". .g),,, encourages the model to capture the underlying geometric and structural patterns that are preserved across different views.
To address this, we follow Liu et al. (2021) and adopt surrogate objective based on conditional likelihoods:
(7) where 9º r;)),, and g¹ {(E] r. ") denote two rigid body views of protein, with translations and rotations r. For later calculation, convert the rotation matrix into quaternion q for interpolation between views. Thus, the objective becomes
[ log p(g" g + log" l
logp((?)).. .q: (7, 9. )) + log .q. ( ()).
(8)
To optimize Equation (8), adopt the Flow Matching framework (Albergo & Vanden- Fijjdenn, 2022; Lipman et al., 2022; Liu et al., 2022b). Unlike generic flows, our model respects the inductive bias that each residue behaves as a rigid body during interpolation. The objective is to learn a velocity field that drives the system from gd to g² (and vice versa), via an intermediate state at time T E [0, 1]. We decompose this into translation and rotation components, as described below.
LERP for Translation in R33 We interpolate the translations between residues using linear interpo- lation (LERP):
T' LERP(E ",',,)))))))) = ri' (1-
where 7 E [o, 1] is the interpolation parameter.
SLERP for Rotation in SO(3). For rotations, we interpolate quaternions using spherical linear interpolation (SLERP):
SLERP("'.''.)) sin((1 - sin((1-)0))
+ sin(ro)g'
(10)
sin(e)
where 7 E [o. 1] is the interpolation parameter and Φ is the angle between qa and q'
Recall that we are using IPA as our backbone model, ug. We want up to learn the true flow both translation and rotation (quaternion) through flow matching, i.e., [ug,R,, (4,so = v .T T).
Thus, the objective function for one direction g" - g² is:
Po = Lgs + Lsocs) = |[" 1-'_"" 0 Ug
- dr
SLERP
-) ue.soca (11)
Final Objective. We apply the same formulation in the reverse direction, yielding the final loss:
C - C= -0"
=g" + -a"
(12)
The final objective employed sequentially over the two phases described in Sections 3.2.1 and 3.2.2, which we refer to as RigidSSL--Pertus and RigidSSL-MDD, respectively.

---

Published as a conference paper at ICLR 2026
Table 1: Comparison of Designability (fraction with scRMSD < 2.0 A). Novelty (max. TM-score 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.
Model
Pretraining
Designability
Novelty
Diversity
Fraction (1)
avg. max TM ())
pairwise TM (L)
MaxCluster (f)
FrfmeDiff None
0.775 + 0.060
0.555 + 0.079
0.565 + 0.016
0,033
Fiff
BeoSSL-EEB-NCE
0.725 + 0.048
0.615 + 0.054
0.597 + 0.010
0.033
FiffeDiff
LIfoSS-IIffONEE
0.650 + 0.058
0.613 = 0.068
0.568 = 0.013
0.033
FiffeDiff
GeoSSL-RR
0.700 1 0.065
0.579 1 0.073
0.604 1 0.015
0.089
FiffeDiff
RigidSSL- -Perturr (ours)
0.875 + 0.050
0.494 + 0.066
0.534 + 0.011
0.033
FifmeDiff
RigidSSI--MD (ours)
0.700 + 0.062
0.657 + 0.084
0.471 0.009
0.156
FoldFlow-2
None
0.329 = 0.013
0.810 = 0.004
1620
0.183
FoldFlow-2
BOOSSL-EBM-NE
0.424 1 0.014
0.790 1 0.009
0.626 0.002
0.225
FoldFlow-2
GeoSSI-IIffNNC..
0.333 + 0.012
0.786 + 0.005
0.631 0.001
0.052
FoldFlow-2
GeoSSI-RR
0.344 + 0.012
0.787 + 0.005
0.601 0.003
0.137
FoldFlow-2
RigidSSL-Pertun (ours)
0.758 + 0.016
0.770 = 0.003
0.650 0.001
0.252
FoldFlow-2
RigidSSL-MD (ours)
0.584 1 0.018
0.782 0.004
0.613 0.002
0.318
4 EXPERIMENTS AND RESULTS
We pretrain the IPA module under the RigidSSL framework, following Appendix F, We evaluate on diverse protein structure generation tasks against four baselines: a randomly initialized model and three coordinate-awen geometric sulf-supervised learning methods that maximize mutual information. Specifically, GeoSSL-InfoNCE and GeoSSL-EBM-NCE employ contrastive objectives that align positive pairs while repelling negatives (Liu et al., 2021; Oord et al., 2019), whereas GeoSSL-RR adopts generative objective, leveraging intra perta-proeiin supervision by reconstructing one view from its counterpart in the representation space (Liu et al., 2021). A comparison of pretraining and downstream training configurations can be found in Appendix F.1.
4.1 DOWNSTREAM: UNCONDITIONAL PROTEIN STRUCTURE GENERATION
We evaluate two representative generative models for protein backbones, FrameDiff and FoldFlow-2, both built on IPA and rigid body representations. FrameDiff (Yim al., 2023) generates novel backbones by defining a diffusion process directly on the manifold of residue-wise SE(3) frames. Built on IPA to update embeddings. it learns an -equivvauva score function that drives the reverse diffusion dynamics. FoldFlow-2 (Huguct et al., 2024) is a sequence-augmenter matching model. Its architecture consists of two IPA blocks, with the first encoding protein structures into latent representations and the second decoding multimodal inputs into S SE(3) vector fields used for backbone generation. We pretrain the IPA module FrameDiff and FoldFlow-2 with RigidSSL, and subsequently finetune them for monomer backbone generation following their original training objectives (see Appendix F). We then assess the generated samples terms of designability, diversity, novelty, Fréchet Protein Structure Distance (FPSD). Fold Jensen-Shannon Divergence (FJSD), and Fold Score (IS). See Appendix F.4.2 for metric details.
SE(3)-
As shown in Table 1, compared with pretraining or with GeoSSL-EBM-NCE GeOSSL-InfoNCE. and GeoSSL-RR, FrfmeDiff pretrained with RigidSSL-PPrturr achieves substantially higher perfor- mance in designability, novelty, and diversity. Specifically, relative to the unpretrained FifmeDiff, RigidSSL-Perturb improves mean designability by 10%, mean novelty by 6.1%, and mean diversity by 3.1 3.1% (as measured by pairwise TM). However, while RigidSSL-MD further improves mean diver- sity by 9.4% relative to the unpretrained model, it leads to decreased performance in designability and novelty. We provide further discussion of this phenomenon in Section 5. Similarly, for FoldFlow-2. RigidSSL-Perturb improves mean designability by 42.9%, mean novelty by 4%, and mean diversity by 6.9% (as measured by MaxCluster) when compared with the unpretrained model. We hypothesize that the increased diversity across FrfmeDiff and FoldFlow-2 by RigidSSL-MD arises from the ATLAS dataset, as offers an exhaustive and non-redundant coverage of the PDB conformational space. Given the gap between computational oracle and real-world wet lab validation, generating structurally varied set of candidates is especially valuable as it increases the probability of identifying functional hits while avoiding the resource- resouree-intensiee failure mode of testing redundant, non-viable designs. Furthermore, we examine the distribution of secondary structures, as shown in Figure 2. FoldFlow-2 pretrained with RigidSSL-MDD exhibits the greatest secondary structure diversity, fol-
7

---

Published as a conference paper at ICLR 2026
(a) (Pretraining)
(B) PDB (Training)
Unpretrained
(d) GeoSSL-EDM-NCE
2X0
600
230
pchellc
o-bell
sihelic
(e) GeoSSL-IntONEE
(1) GeoSSL-RR
(g) RigidsSL-Perrtrb
(h) Rigid5SL-MD
3.50
62
230
13
6.
02 0.6
04
03 44
Figure 2: Distribution of secondary structure elements (a-helices, f-sheets, and coils) in protein structure database (a-b) and in designable proteins (scRMSD < 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.
lowed by RigidSSLL-Pertub.. While the unpretrained model and the BOooSS--EBM-NCE-pretaaiii model generate predominantly a-helix structures, RigidSSL yields substantially broader spectrum that includes coils and mixed cx-helix/3-sheet compositions.
Next, we quantify alignment between the cenerated-structuue distribution and the reference distribu- tions at both the feature and fold-class levels using the probabilistic metrics of Geffner et al. (2025) (Table 3). FPSD, measured in GearNet (Zhang et al., 2022) feature space, reveals that RigidSSI-MMD produces structures whose learned representations most closely match natural protein distributions, achieving the best scores for FrameDiff and second-best for FoldFlow-2 when compared against both PDB and AFDB. This suggests that incorporating structural dynamics encourages the model to generate samples that occupy similar regions of the protein structure manifold as natural proteins. Interestingly, the topology level, while RigidSSL-MD achieves the best fJSD across both archi- tectures (4.99 for FrameDiff, 5.07 for FoldFlow-2 against PDB), its IS scores (5.35 and 5.54) are lower than those of GeoSSL-InfoNCE (6.86 and 6.01). This indicates that RigidSSL-MD generates structurally diverse proteins whose collective fold distribution closely matches the reference, but individual samples may explore a broader range of conformations within each fold class rather than converging to canonical representatives, leading to softer classifier predictions and consequently lower fold scores.
Furthermore, we conducted case study on the generation of long protein chains to explore the full potential of RigidSSL unconditional generation. this setting, RigidSSIL-Perturb enables the model to generate ultra-long proteins of 700-800 residues that remain stereochemically accurate and self-consistent, achieving the best MolProbity score and Clashscore among all pretraining methods. This demonstrates that RigidSSL effectively captures global structural patterns that enable stable and physically realistic long-chain generation. Details are provided Appendix F.5.
4.2
CASE STUDY: ZERO-SHOT MOTIF SCAFFOLDING
Motif scaffolding refers designing protein structure that supports a functional motif with fixed geometry. Following Wang et al. (2022a), we formulate this as an "inpainting" task: the model must generate viable scaffold backbone given only the motif coordinates, without prespecifying the global topology.
We adapted the benchmark suite from Watson al. (2023), which includes 22 targets, in a zero-shot setting, i.c., without spse-ipcc training. Detailed evaluation metrics and results provided in Appendix F.6. RigidSSI- -Perturb achieved the highest average success rate of
8

---

Published as a conference paper at ICLR 2026
13.19%. demonstrating significant improvement over the unpretrained model (9.35%) Notably, RigidSSL-Peertub exhibited superior robustness on difficult targets that demand longer or more extended scaffolds. For example, the 5TRV long target, RigidSSL-Perturts Perttrb achieved a success rate of 51 %, outperforming the next best method, GeoSSLL-IffNCEE, by 21%. While RigidSSI achieved a lower average performance (10.08%) than the perturbation variant, it demonstratee distinct inductive biases, yielding best performance on specific targets such as 2KL8 and i17ra_ge.
4.3 CASE STUDY: GPCR CONFORMATIONAL ENSEMBLE GENERATION
AlphaFlow (Jing et al., 2024) introduces novel approach to ensemble generation by re-engineering AlphaFold2 into Now-matching--asse generative model. Instead of predicting single structure, AlphaFlow learns a continuous flow field that transforms noisy protein conformations into realistic structures, allowing it to sample from the equilibrium ensemble distribution. In particular, AlphaFlow incorporates MD data by finetuning trajectories, where each frame represents near-equilibrium structure along the protein's dynamic landscape. This enables the network to learn how proteins fluctuate across their conformational landscape.
ensembles of protein conformations sampled from MD
To investigate how RigidSSL's geometric pretraining affect this process, we selected GPCRs as a representative target. They are particularly difficult to model because their activation involves multiple metastable states connected by slow, collective rearrangements of transmembrane helices, driven by subtle allosteric couplings. This makes their conformational landscape highly anisotropic and multimodal, posing a stringent test for generative models of protein dynamics (Aranda-García et al., 2025). Training and evaluation details can be found in Appendix F.7.1.
We evaluate the generated ensembles across three tiers of increasing biophysical complexity: predict- ing flexibility, distributional accuracy, and ensemble observables. As shown in Table 6, RigidSSL fundamentally reorganizes the unpretrained AlphaFlow's 's learned conformational manifold for GPCRs. For flexibility prediction, RigidSSL-Pertuu to the aggregate diversity of the GPCR MD ensembles. It generates conformations with a pairwise RMSD (2.20) and all-atom RMSF (1.08) that best approximate ground truth targets (1.55 and 1.0, respectively), outperforming the unpretrained model and pretraining baselines. In terms of distri- butional accuracy, although GeoSSI. variants achieve marginally lower overall root mean W2 and MD PCA W2 distances, RigidSSL- -Perturb minimizes the Joint PCA Wy-distance (17.53). Crucially, both RigidSSL variants increase the geometric alignment of generated motion modes with MD trajectories, jointly achieving the highest success rate (5.97%) for > 0.5 principal component cosine similarity. Finally, assess complex ensemble observables. This represents the most stringent evaluation category, requiring accurate modeling of sidechain fluctuations, transient structural associations, and higher-order thermodynamic properties. By explicitly incorporating MD dynamics, RigidSSL-MD outperforms all baselines in this regime. It achieves the highest Jaccard similarity for weak contacts (0.43) and cryptically exposed residues (0.71), alongside the highest Spearman correlation (p 0.03) for the solvent exposure mutual information matrix. Ultimately, RigidSSL-Pertun primarily enforces collective mode coherence as a structural regularizer, while RigidSSL-MDD successfully captures higher-order biophysical statistics to produce realistic contact and exposure profiles.
reduces spurious flexibility and yields the closest match
DISCUSSION
RigidSSL-Peerurb Improves Geometric Quality of Learned Structures. In the self-supervised pretraining step in Phase I, we employ data augmentation by adding noise to the original geometry to construct a perturbed geometry as the second view. This choice is motivated by three considerations. (1) Gaussian perturbation can be interpreted as a form of "masking" applied to atom positions in 3D Euclidean space. Analogous to how self-supervised learning masks image patches in computer vision or tokens in NLP, noise injection in geometry partially obscures atomic coordinates and encourages the model to infer consistent global features across corrupted inputs. (2) At small noise scales, this augmentation mimics systematic errors that arise during data collection prediction. Protein structures from crystallography, cryo-EM, or computational databases such as AFDB often contain small uncertainties in atom placement. Exposing the model to controlled perturbations helps become robust to such imperfections. (3) Gaussian perturbations can also approximate natural conformational fluctuations around stable equilibrium. Proteins dynamic molecules, and their

---

Published as a conference paper at ICLR 2026
functional states often involve small deviations from the ground-state structure. Training with noisy perturbations introduces the model to this variability, allowing it to capture flexible yet physically plausible structural patterns.
As shown in Table 1, Phase I pretraining improves designability, indicating that the model's outputs shift toward more reliable distribution, but this comes the cost of reduced diversity compared to Phase II. By maximizing mutual information between noisy views, the model learns representations that emphasize stable, fold-defining features while downweighting minor structural deviations. Down- stream generators built on such representations therefore tend to produce more reliable and foldable backbones. However, this objective can also suppress conformational variability by collapsing nearby structures into single representation, thereby limiting the diversity of generated folds.
RigidSSL-MD Promotes Structural Diversity and Biophysical Fidelity. In the second phase, RigidSSL-MMD incorporates MD trajectories to build augmented views. While this improves structural diversity, reduces designability compared to perturbation- based pretraining (Table 1). The trade-off arises because MD pretraining is not purely sule--uper in the sense that training views are drawn from force-field-eee sice-field-generated trajectories, inheriting biases that may misalign with de novo design goals and even cause negative transfer, seen in prior studies (Liu et al., 2023a). In contrast, the same physically grounded bias improves higher-order biophysical properties (Table 6). Another explanation for the lower designability is that most structure prediction models are trained on stable, static conformations. While MD trajectories capture conformational dynamics to enrich structural flexibility, they also bias the model toward generating metastable conformations rather than the absolute ground states expected by folding and inverse-folding pipelines, thereby hindering re-folding performance. Overall, RigidSSL-MD emphasizes diversity and physical fidelity, whereas RigidSSL- Perturb favors geometric quality described by the designability metric. These two paradigms may therefore serve as complementary strategies, depending on the downstream design objectives.
ACKNOWLEDGEMENT
We are grateful to Travis Wheeler and Ken Youens-Clark for maintaining MDRepo, and to the anonymous ICLR reviewers for helpful feedback. We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Zeju Qiu.
USE OF LARGE LANGUAGE MODELS
In this work, we used large language models (e.g., ChatGPT) to refine our writing, assist in formulating &Tix equations, and identify relevant literature.
ETHICS STATEMENT
This research fully adheres to the ICLR Code of Ethics. The study does not involve human subjects or the use of personal sensitive data. All datasets and code utilized and released conform to their respective licenses and terms of use. The contributions in this work are foundational and not raise issues related to fairness, privacy, security, potential misuse. We confirm that all ethical considerations have been thoroughly addressed.
REPRODUCIBILITY STATEMENT
We are committed to making our work easy to reproduce. All necessary details for replicating our main experimental results, including data access, experimental setup, model configurations. evaluation metrics, and checkpoints, publicly available this repository. Users follow documentation and scripts to accurately reproduce the results, ensuring transparency and scientific rigor.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Michael Samuel Albergo and Eric Vanden-Fijnden. Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations, September 2022.
David Aranda-García, Tomasz Maciej Stepniewski, Mariona Torrens- Adrian García-Recio, Marta Lopez-Balastegui, Brian Medel-Lacruz, Adrián Morales Pastor, Alejandro Peralta-García, Miguel Dieguez-Eceolazz David Sotillo-Nüñe' Large scale investigation of GPCR molecular dynamics data uncovers allosteric sites and lateral gateways. Nature Commuiccations, 16(1):0020, 2025. 9
Helen M Berman, John Westbrook, Zukang Feng. Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic Acids Research, 28(1): 235-242, 2000.
Avishek Joey Bose, Tara Akhound Sadeeh. Guillaume Huguet, Kilian Fatras, Jarrid Rector- Bcooor-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Sto3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:23 2023. 1
Nazim Bouatta, Peter Sorger, and Mohammed AlQuraishi. Protein structure prediction by Alphafold2: Are attention and symmetries all you need? Acta Crystallographica. Section D. Structural Biology, 77(Pt 8):982 991, August 2021. ISSN 2059-7983. doi: 10. 10.100//2009799220075111 2
Bo Chen, Zhilei Bei, Xingyi Cheng, Pan Li, Jie Tang, and Le Song. MSAGPT: Neural prompting protein structure prediction via msa generative pre-training. Advances in Neural Information Processing Systems, 07777504437344 2024. 3, 17
Can Chen, Jingho Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self-supervised learning. Bioinformatics, 39(4):btad1 189, 2023. 2
Hua Cheng, R. Dustin Schaeffer, Yuxing Liao, Lisa N. Kinch, Jimin Pei, Shuoyong Shi, Bong-Hyun Kim, and Nick V. Grishin. ECOD: An evolutionary classification of protein domains. PLOS Computational Biology, 10(12))e December 2014. ISSN 1553-7358. doi: 10.1371/journal 1003.60 URL www.hhtt:::/e /journaas ploss. org/ wtt.scccmpoooo/ariiieieeeeeeeeee 10.1371/ 19.13311jouzzne journall pcbi htt.:cib...09000 1003926. Publisher: Public Library of Science. 22
Benjamin Coors, Alexandru Paul Condurache, and Andreas Geiger. Spherenet: Learning spherical representations for detection and classification in omnidirectional images. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 318-533, 2018. 3,
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai. Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel, et Robust deep learning based protein sequence design using ProteinMPNN, Science, 378(6615 (7866615) 5):49-56, 2022. 23, 25
Ian W. Davis, Andrew Leaver-Fay, Vincent B. Chen, Jeremy N. Block, Gary J. Kapral, Xueyi Wang, Laura W. Murray, W. Bryan Arendall, Jack Snoeyink, Jane Richardson, and David C. Richardson. MolProbity: All-atom contacts and structure validation for proteins and nucleic acids. Nucleic Acids Research, 35W Server (ssue):W375-11 July 2007. ISSN 1362-4962. doi:
10.1093/nargkm2211 22,
Hehe Fan, Zhangyang Wang, Yi Yang, and Mohan Kankanhalli. Continuous-discret convolution for geometry-sequeee -equence modeling in proteins. International Conference on Learning Representations, 2022. 3, 17
Zhangyang Gao, Cheng Tan, Lirong Wu, and Stan Z Li. CoSP: Co-supervised pretraining of pocket and ligand. arXiv preprint arXiv:2206. 12241, 2022. 3, 17
Johannes Gastciger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03/23, 2020. 3, 17
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach. Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, et al. Proteina: Scaling flow-based protein structure generative models. arXiv preprint arXiv:2503.00710 2025. 8, 23

---

Published as a conference paper at ICLR 2026
Hongyu Guo, Yoshua Bengio, and Shengchao Liu. AssembleFlow: Rigid flow matching with inertial frames for molecular assembly. In International Conference on Learning Representations, 2025. URL httptt///ppperriieeeeee net wttp://eee////nmeeeeeeeeeee jaw.uuniddddkcceeaaaaaaaa 4
Yuzhi Guo, Jiaxiang Wu, Hehuan Ma, and Junzhou Huang. Self-supervised pre-tranning for pro- tein embeddings using tertiary structures. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36(6), pp. 6801-6809, 2022. 2,
M. J. Harvey, G. Giupponi, and G. De Fabritüi. ACEMD: Accelcrating bio-molecular dynamics in the microsecond time-scale, February 2009. URL httt:://araxii org. www/aass/000 0827. arXiv:0902.0827 [physics]. 26
Herbert, A. and Stembberg M. MaxCluster: A tool for protein structure comparison and clustering, 2008. URL https:ttp/..... /ww. sbg. bio. ic. ac www..k//aaaaaeec htt///aacllsstt 23
Pedro Hermosilla and Timo Ropinski. Contrastive representation learning for 3D protein structures. arXiv preprint arXiv:2205. 15675, 2022. 2,
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neural Information Processing Systems, 33:6840-6851, 2020. 21
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laune Laufer, Pablo Lemos, Riashat Islam, Cheng Chong-Hao Liu, Jarrid Rector-Brooks, Tara Akhound- -aaoonn,SSddgh,, Michael Bronstein, al. esquence-augmenner SE(3)-flow matching for conditional protein backbone generation. Advances in Neural Information Processing Systemx, 2024. 1, 7, 22, 24, 27
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. 3D equivariant molecular graph
pretraining. arXiv preprint arXiv:2207.0S8: 08824,, 2022. 2, 3, 17
Bowen Jing, Stephan Eismann. Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411, 2020. 3,
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. AlphaFold meets flow matching for gener- ating protein ensembles, September 2024. URL http: arXiv:2402.04845 [q-bio]. 9,
04845
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583-58) 2021. 3, 4, 18, 20
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International
Conference on Learning Representations, 2015. 22
Kerlen T. Korbeld, Vsevolod Viliuga, and Maximilian Fürst. Limitations of the refolding pipeline for de novo protein design, December 2025. URL https:// http://wwwbbbi.xxo/// www.xiii.iiv.................... org/ content/ 10. 64898/2025. 22. 09.09...... 693122v1. 25
Andrew LaPelusa and Ravi Kaushik. Physiology, proteins. In StatPearls. StatPearls Publishing, Trea- sure Island (FL), 2025. URL
ncbi nlm, nih, gww.goo///oooooooo
Adam Leach, Sebastian M. Schmon, Matteo T. Degiacomi, and Chris G. Willcocks. Denoising diffusion probabilistic models on SO(3) for rotational alignment. In Proceedings ICLR 2022 Workshop on Geometrical and Topological Representation Learning, March 2022.
Haorui Li, Weitao Du, Yuqiang Li, Hongyu Guo, and Shengchao Liu. InertialAR: Autoregressive 3D molecule generation with inertial frames. arXiv preprint arXiv:2510.274977 2025. 4
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam FarandZarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 17936637)11133-11333 March 2023. doi: 10. 11111///cienccc.de25774 URL https: //www/wwww///////////// science. www.eeeneeeeeeeeeeeeeeeeeeeeeeee wttr///oo/// doi/10. 1126/ /126/scienee science.. ade2574. Publisher: American Association for the Advancement of Science. 23, 25
12

---

Published as a conference paper at ICLR 2026
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, arXiv preprint arXiv:2210.02747, 2022. 4, 6
Dina Listov, Casper A. Goverde, Bruno E. Correia, and Sarel Jacob Fleishman. Opportunities and challenges in design and optimization of protein function. Nature Reviews Molecular Cell Biology, 25(8):639-66 653, August 2024. ISSN 1471-0080. doi: 00.1038//4180-0222
000000022-0001---- URL https:/ / www. naturet.e.... cow att///rticccc a1180-1115-000 024-0000--2--------0- Publishing Group.
Publisher: Nature
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre- training molecular graph representation with 3D geometry. arXiv preprint arXiv:2110.07728, 2021. 2, 3, 6, 17
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)-invariant denoising distance matching. arXiv preprint arXiv:2206. 13602, 2022a. 2, 3, 17
Shengehao Liu, Weitao Du, Yanjing Li, Zhuoxinran Li, Zhiling Zheng, Chenru Duan, Zhiming Ma, Omar Yaghi, Anima Anandkumar, Christian Borgs, Jennifer Chayes, Hongyu Guo, and Jian Tang. Symmetry-informee geometric representation for molecules, proteins, and crystalline materials.
arXiv preprint arXiv:23366.93377 2023a. 3, 10, 17
Shengehao Liu, Weitao Du, Zhi-Ming Ma, Hongyu Guo, and Jian Tang. A group symmetric stochastic differential equation model for molecule multi-modal pretraining. In International Conference on Machine Learning, pp. 21497-21526. PMLR, 2023b. 3
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow, arXiv preprint 2rXiv:2209. 22000003,0,,,, 2022b. 4, 6
Mitchell D. Miller and George N. Phillips. Moving beyond static snapshots: Protein dynamics and the Protein Data Bank. The Journal of Biological Chemistry, 296:100749, May 2021. ISSN 0021-9258, doi: 10.1016/j.jb..2222
URL https: /ttc//mccmmmm/////// www.ncciin... nlm. giv.
Ali Miserez, Jing Yu, and Pezhman Mohammadi. Protein-based biological materials: Molecular design and artificial production. Chemical Reviews, 1236))20499-2111, March 2023. ISSN 0009-2665. doi: 10. chemrev. 2c0 0 621. Publisher: American Chemical Society.
URL https: : fooi.org/100 1ttp//.....1//cc..../
Kwangho Nam and Magnus Wolf-Watz. Protein dynamics: The future is bright and complicated!
Structural Dynamics, 10(1)x01430,, February 2023. ISSN 2329-7778. doi: (0.1066/4.0000177 URL https: /tcc///cc//////// ncbi. n nlm, htt.://i.....o.ttt..ceee///////e AMM.:t..2.211///M9992222
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, January 2019. URL [cs].
abs// 807. 03748. arXiv: arXiv:1807.03744 1807.03748
Ismael Rodríguez-Espigares, Mariona Forrens-Fontanals, Johanna K. S. Tiemann, David Aranda- García, Juan Manuel Ramirez-Anguita, Tomasz Maciej Stepniewski, Nathalie Worp, Alejandro Varela-Rial, Adrián Morales-Pastor, Brian Medel-Lacruz, Gáspár Pándy-Szekeres,, Eduardo Mayol, Toni Giorgino, Jens Carlsson, Xavier Deupi, Slawomir Filipek, Marta Filizola, José Carlos Gómez- Tamayo, Angel Gonzalez, Hugo Gutiérrez-de Terán, Mireia Jiménez-Rosés Willem Jespers, Jon Kapla, George Khelashvili, Peter Kolb, Dorota Latek, Maria Morti-Solano, Pierre Matricon, Minos- Timotheos Matsoukas, Przemyslaw Miszta, Mireia Olivella, Laura Perez-Benito Davide Provasi, Santiago Ríos, Iván R. Torrecillas, Jessica Sallander, Agnieszka Szyler, Silvana Vasile, Harel Weinstein, UIrich Zachariae, Peter W. Hildebrand, Gianni De Fabritiis, Ferran Sanz, David E. Gloriam, Arnau Cordomi, Ramon Guizà-González, and Jana Selent. GPCRmd uncovers dynamics of the 3D-GPCRome. Nature Methods, 17(8):777-787 August 2020. ISSN 1548- 7105. doi: (00.:////44112200000844-4444 URL https: www. http://wwww..t..e/e//t////: nawuraateee.... com/ articles/ 841592 020 8 y.. Publisher: Nature Publishing Group. 26
Amitava Roy, Ethan Ward, Illyoung Choi, Michele Cosi, Tony Edgin, Travis S Hughes, Md Shafayet Islam, Asif M Khan, Aakash Kolekar, Mariah Rayl, Isaac Robinson, Paul Sarando, Edwin Skid- more, Tyson I. Swetnam, Mariah Wall, Zhuoyun Xu, Michelle I. Yung, Nirav Merchant, and
13

---

Published as a conference paper at ICLR 2026
Travis J Wheeler. MDRepo open data warehouse for community molecular dynam- ics simulations of proteins. Nucleic Acids Research, 530D1)).477-D486, January 2025. ISSN 1362- 4962. doi: 10.1093/nar/gkaell 109. URL https: http://.........../o// gaae1199. 22, 26
Kristof T Schütt, Huziel Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018. 3,
Jeffrey Skolnick, Mu Gao, Hongyi Zhou, and Suresh Singh. AlphaFold 2: Why it works and its implications for understanding the relationships of protein sequence, structure, and function. Jour- nal of Chemical Information and Modeling, 61(10):4827 4831,, October 2021. ISSN 1549-9596. doi: 10.102 w/t.//////....m.//////.//o 1c01114. URL https: //mc://///.///////////////// pmc. ncbi nlm. nih www..oo/artic......oo
PMC8592092/..
Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, and Pietro Liò. 3D infomax improves GNNs for molecular property prediction, June 2022. URL http: /aaxivivviv..////aa//aaa///aaavv org/ abs////////////// 2110.01000001111000211021011120 04126. 11.04:2110.04126 [cs]. 17
Michel Van Kempen, Stephanic S Kim, Charlotte Tumescheit, Milot Mirdita, Jeongjae Lee, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Fast and accurate protein structure search with Foldseek. Nature Biotechnology. 42(2):243-24) 246, 2024. 23
Yann Vander Meersche, Gabriel Cretin, Aria Gheeraert, Jean-Christoppe Gelly, and Tatiana Ga- lochkina. ATLAS: Protein flexibility description from atomistic molecular dynamics simulations. Nucleic Acids Research, 52(D1))D384-D333 2024. 6, 22
Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, Augustin Zidek, Tim Green, Kathryn Tunyasuvunakool, Stig Petersen, John Jumper, Ellen Clancy, Richard Green, Ankur Vora, Mira Lutfi, Michael Figurnov, Andrew Cowie, Nicole Hobbs, Pushmeet Kohli, Gerard Kleywegt, Ewan Birney, Demis Hassabis, and Sameer Velankar. AlphaFold Protein Structure Database: Massively expanding the structural coverage of protein-sequence space with high- accuracy models. Nucleic Acids Research, 50(DI):D439 D444, January 2022. ISSN 0305-1048, doi: 10. 10.109//aar/gaab1000 URL hwwph:ttpt/.om /oi org 100 1093 nar /gkall 066. 2,5
Chentong Wang, Yannan Qu, Zhangzhi Peng, Yukai Wang, Hongli Zhu, Dachuan Chen, and Longxing Cao. Proteus: Exploring protein structure generation for enhanced designability and efficiency. bioRxiv, 2024. 2
Jue Wang. Sidney Lisanza, David Juergens, Doug Tischer, Joseph L. Watson, Karla M. Castro, Robert Ragotte, Amijai Saragovi, Lukas F. Milles, Minkyung Back, Ivan Anishchenko, Wei Yang. Derrick R. Hicks, Marc Expòsit, Thomas Schlichthaerle, Jung-Ho Chun, Justas Dauparas, Nathaniel Bennett, Basile I. M. Wicky, Andrew Muenks, Frank DiMaio, Bruno Correia, Sergey Ovchinnikov, and David Baker. Scaffolding protein functional sites using deep learning. Science (New York, N.Y.). 37766604)3887-33999 July 2022a. ISSN 0036-8075. doi: 10.1 10.1126//cience.abb33 URL https:tp///.mmc pmc. ncbi nlm. nih gov/ articiestes... /PMC 6219922221 694/ 1, 8
Limei Wang, Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning hierarchical protein representations via complete 3D graph networks. arXiv preprint arXiv:2207.12600 2022b. 3, 17
Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew Borst, Robert Ragotte, Lukas Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Back, and David Baker. De novo design of protein structure and function with REdiffusion. Nature, 620r7976):108--1
August 2023, ISSN 1476-4687. doi: 10. 00000///4115602--00000 https: / www. nature. com/ www.aiaticlleseeo.e a 18866666 02 5-5-8. Publishing Group. 1, 8, 24
URL
Publisher: Nature
14

---

Published as a conference paper at ICLR 2026
Jiaxian Yan, Zhaofeng Ye, Ziyi Yang. Chengqiang Lu, Shengyu Zhang, Qi Liu, and Jiezhong Qiu. Multi-task bioassay pre-training for potteinlligand binding affinity prediction. arXiv preprint NAXYv:23066048880 2023. 3, 17
Toshio Yanagida, Masahiro Ueda, Tsutomu Murata, Seiji Esaki, and Yoshiharu Ishii. Brownian motion, fluctuation and life. Bio Systems, 88(3):228-242, April 2007. ISSN 0303-2647, doi: 10.1016/j. bos.:yoostssss.20.........
Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. arXiv preprint urXiv:2302.02277, 2023. 1, 7, 22
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzzlez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.00133 00133, 2022. 2, 3, 17
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. arXiv preprint 223:::2003.0612,, 2022. 2, 3, 8, 17, 23
15

---

Published as a conference paper at ICLR 2026
Appendix
Table of Contents
A Related Works
17
B Group Symmetry and Equivariance
17
C Protein Backhone Parameterization and Canonicalization
18
C.I Local Frame Construction Canonicalization Details
18
D Rotation Perturbation in RigidSSL-Perturb
18
D.1 Parametrizations of SO(3)
18
D.2 Isotropic Gaussian Distribution on SO(3): I9so(3)
19
D.3 Rotation Sampling from IGso(3)
20
D.4 Rotation Perturbation Effects in R
20
E Base Protein Encoder
20
F
Experiment Details
22
F11
Training Overview
22
F2
Pretraining: Dataset
................................
22
F3
Pretraining: Optimization
22
F4
Downstream: Unconditional Generation
22
E5
Case Study: Long Chain Generation
24
F66
Case Study: Zero-Shot Motif Scaffolding
24
F.7
Case Study: GPCR Conformational Ensemble Generation
26
Ablation Studies
27
16

---

Published as a conference paper at ICLR 2026
A
RELATED WORKS
Geometric Representation for Small Molecules and Proteins. wide range of representation learning models with invariance equivariance properties have been proposed for both small molecules (Coors et al., 2018; Gasteiger et al., 2020; Schütt et al., 2018) and proteins (Fan et al., 2022; Jing al., 2020; Wang al., 2022b), targeting tasks such property prediction and molecular generation, As a comprehensive benchmark, Geom3D (Liu et al., 2023a) offers a systematic roadmap for geometric representations of small molecules and proteins. It categorizes mainstream approaches into three major families: SE(3) - models, SE(()-equivariant models with spherical frame bases, and Squiivaqiiaaiian models with vector frame bases.
Geometric Pretraining for Small Molecules. The GraphMVP (Liu et al., 2021) proposes one con- trastive objective (EBM-NCE) and one generative objective (variational representation reconstruction, VRR) to optimize the mutual information between the topological and conformational modalities. 3D InfoMax (Stärk et al., 2022) is a special case of GraphMVP, where only the contrastive loss is considered. GeoSSL (Liu et al., 2022a) proposes maximizing the mutual information between noised conformations using an SE(3)-invariant denoising score matching, and parallel work (Zaidi et al., 2022) is special case of GeoSSL. using only one denoising layer. 3D-EMGP (Jiao et al., 2022) is parallel work, yet it is E(3)-qquivariant, which needlessly satisfies the reficction-quiivariatt constraint in molecular conformation distribution.
Geometric Pretraining for Proteins. Several parallel efforts have been made toward geometric pretraining for proteins. Guo et al. (2022) focuses on maximizing mutual information between diffusion trajectory representations of structurally correlated conformers. GearNet (Zhang et al., 2022) introduces contrastive self-supervised framework trained on AFDB v1, aiming maximize the similarity between representations of substructures within the same protein through sparse edge message passing. Similarly, ProteinContrasl (Hermosilla & Ropinski, 2022) proposes a contrastive learning approach trained on 476K PDB structures to learn effective 3D protein representations. More recently, MSAGPT (Chen et al., 2024) presents a novel pretraining strategy based on Multiple Sequence Alignments (MSAs), where the model learns to generate MSAs by capturing their statistical distribution, thereby improving structure prediction for proteins with few or no known homologs.
Binding Pretraining on Stable Positions. In addition to pretraining on proteins and small molecules separately, several studies have explored joint pretraining strategies that model protein-ligane interac- tions directly. CoSP (Gao et al., 2022) introduces a co-supervised framework that simultaneously learns 3D representations of protein pockets and ligands using a gated geometric message passing network, incorporating contrastive loss to align real pocket-ligand pairs and employing chemical similarity-enhanced negative sampling strategy to improve performance. MBP (Yan et al., 2023) addresses the scarcity of high-quality 3D structural data by constructing pretraining dataset called ChEMBL-Dock, including over 300K protein-ligand pairs and approximately 2.8 million docked 3D structures; MBP employs multi-task pretraining approach, treating different affinity measurement types (e.g.. IC50, Ki, Kd) as separate tasks and incorporating pairwise ranking within the same bioassay to mitigate label noise, thereby learning robust and transferable structural knowledge for binding affinity prediction.
B GROUP SYMMETRY AND EQUIVARIANCE
In this article, a 3D molecular graph represented by 3D point cloud. The corresponding symmetry group is SE(3), which consists of translations and rotations. Recall that we define the notion of cquivariance functions in R in the main text through group actions. Formally, the group SE(3) is said to act on R if there is a mapping á SE(3) x R Rd satisfying the following two conditions:
1. if e SE(3) the identity element, then
dle, r) = for Vr R
2. if 92 E SE(3), then
@(g1, (g2, r)) o(9192,7) for Vr R
Then, there is a natural SE(3) action on vectors r in R by translating r and rotating r for multiple times. For g SE(3) and r E R3, we denote this action by gr. Once the notion of group action is
17

---

Published as a conference paper at ICLR 2026
defined, say function f R R that transforms R is equivariant if: flgr) = gf(r), for r
On the other hand, f F R is invariant, if is independent of the group actions:
f(gr) frr,, for rrER.
C PROTEIN BACKBONE PARAMETERIZATION AND CANONICALIZATION
C.I LOCAL FRAME CONSTRUCTION
Following AlphaFold2 (Jumper al., 2021), we construct local coordinate frames from atomic coordinates in protein structures. Each residue frame is defined as rigid transformation (f,r), where E' E R3 is the translation vector and r E SO(3) is the rotation matrix. Frames derived from the backbone atoms N, C.. and C, with coordinates Ii, I2, and I3, respectively. The translation vector f's set to the C. position 22, while the rotation matrix r is computed using Stram-Schmin orthogonalization in Algorithm 1. As a result, each residue in a protein { re) is thus associated with orthonormal frame centered at its backbone atom.
Algorithm Frame construction via Gram-Schmice orthogonalization
Input: Three atomic positions 21, 42, 23 R
2: Oufput: Translation vector 8' e IR3 and rotation matrix r E 1x3
3: č +
4: 02 31 22
5: é1 + /
6: 22 + 02
02
El
7: ē + T. 72
8: è é 62
9: r + E 62 03
10: [+ 02
return (T,r)
C.2 CANONICALIZATION DETAILS
The eigendecomposition in Section 3.1 requires two additional steps to yield deterministic, valid rotation matrix V SO(3).
Axis ordering. We sort the eigenvectors by their corresponding eigenvalues in ascending order, establishing a deterministic ordering the principal axes.
Right-handedness.
Eigenvectors are only guaranteed to form an orthogonal basis (V 0())), and may include reflection when dei(V) 1. We enforce V E SO(3) by flipping one axis via V + V diag(1, 1, det(detV)))) V)), which negates the third eigenvector when det(V) -1.
D ROTATION PERTURBATION IN RIGIDSSL-PERTURE
Recall that structure is modeled sequence of L. rigid residues, with each residue represented by translation and rotation transformation: {fe, H= In RigidSSL-Perturb, we independently perturb each residue's 's translation f; and rotation r; This section details the mathematical foundation, implementation, and implications of rotation perturbation.
D. PARAMETRIZATIONS OF SO(3)
The rotation group SO(3) admits several equivalent representations. We review its group structure, Lie algebra, anis-angle parameterization, and the exponential and logarithm maps, which together provide the foundations for defining and sampling rotational perturbations in Pirtub..
18

---

Published as a conference paper at ICLR 2026
D.1.1 ROTATION GROUP AND MANIFOLD STRUCTURE
The special orthogonal group SO(3) {R Raxs R² = det(R) 1} forms 3-dimensional compact Lie group manifold representing all possible rotations in 3D space. As manifold, SO(3) is diffeomorphic to real projective space RP.,, obtained by identifying antipodal points on the 3-sphere S
D.1.2 LIE ALGEBRA AND TANGENT SPACE
The Lie algebra s0(3) is the tangent space to SO(3) the identity matrix J. It consists of all 3 skew-symmetric matrices: 50(3) = {X E R333 X -X). Elements of so(3) can interpreted
infinitesimal rotations angular velocities. A remarkable isomorphism exists between s3 and
Ww3
W2
R3 via the "hat map": A IR - so(3), wn
433
W1
The inverse "vee map"
Wwg
0
denoted V so(3) + R3
D.1.3 AXIS-ANGLE REPRESENTATION OF ROTATION
Axis-angle rotation representations exist in two related forms:
Axis-angle representation. A geometric pair (n, (f) where S² is unit vector representing the rotation axis and 0 E [0, =] is the rotation angle. This representation lives in the product space 8 x [0, π
with identifications: (n, 0) ~ (-n, 0) and (n, .) (-n, T)-
A Axis-angle vector (rotation vector). An algebraic vector w = An E R where the direction specifies the rotation axis and the magnitude |س specifies the rotation angle. This representation belongs to ball of radius in R with antipodal points on the boundary identified. The axis-angle vector is precisely the exponential coordinates in so(3) under the isomorphism provided by the hat map. That is, for rotation vector w, the corresponding element in 50(3) is w
1.1.4 EXPONENTIAL MAP
The exponential map exp s0(3) - SO(3) connects the Lic algebra to the Lie group: exp(X) = Ee X For anis-angle vector ên, Rodrigues' formula provides closed-for form expression for the corresponding rotation matrix: exp(w>) = I + sin ~ + 1-cos0
6a )2 When this
reduces to the identity matrix J. The exponential map is surjective but not injective globally and wraps around when | π.
D.1.5 LOGARITHM MAP
The logarithm map log : SO(3) + s3() the (local) inverse of the exponential map. For rotation matrix R E SO(3), where trace(/t) + 11 log(R) 2iin 0 (R- π²²) where cos traes( R)
2
The logarithm map yields an axis-angle vector representation when composed with the vee map: (log(R) R)
D.2 ISOTROPIC GAUSSIAN DISTRIBUTION ON SO(3): IG
With axis-angle representation of rotations, IG consists of two components:
Axis Distribution. The axis distribution is uniform over the unit sphere Sa, with PDF: p(n) - where n E S2 is unit vector representing the axis of rotation.
Angle Distribution. The angle distribution for 0 E [0, n has the probability density function:
1 cocos
cos
(2l + 1)e 1141) sin + ())
p())
(13)
sin (
l=0
19

---

Published as a conference paper at ICLR 2026
where is the concentration parameter of IGso(3). For small values of e where numerical convergence of the infinite sum becomes difficult, an approximation can be used:
8 - c x ((0 2ne of + (0 2 22 $
p(0) (1 cOs(O)) VIT f
(14)
2sin (*)
π
As oc, p(0) approaches the uniform distribution on SO(3): 0
Junifoem (0)
(15)
π
D.3 ROTATION SAMPLING FROM Igso(3)
A sample from (g (j, e²) can be decomposed as r = pre where ~ IGsor3) (²) Therefore, we first sample from an identity centerd distribution and then rotate by μ. Specifically, the sampling process is as follows:
Sample the rotation axis. Generate a random unit vector E S2 uniformly by sampling a point (z, V, 2) from the standard normal distribution (0, Ia) and normalizing to unit length.
Sample the rotation angle. Sample 0 E [0, x| from the density p(0) described above using inverse transform sampling. Specifically, numerically compute the CDF F(0) = fu p(t) dt, sample Uniform(0, 1]. and compute 0 (u) through numerical interpolation.
Construct the exponential coordinates. Form the axis-angle vector w On E R This vector rep- resents elements in the Lie algebra so( via the hat map: wn E so(3), detailed in Appendix D.1.
Apply the exponential map. Convert w rotation matrix using Rodrigues' formula:
Te = exp(w>> = I sin B
1 cos
(w~)².
(16)
0²
Apply the mean rotation. Compute the final rotation as:
pre.
(17)
D.4 ROTATION PERTURBATION EFFECTS IN R3
When applying rotational perturbations sampled from IGso(3), the geometric effect in R3 varies with initial rotation despite consistent geodesic distances on SO(3). For rotations R Rg SO(3) and point p E 33, with perturbations of equal magnitude:
R R exp(w:)) R Rf exp(w>)),
(18)
where Way Wa E 3 are axis-angle vectors with ||all| = |w/l| = The resulting displacements
A, R2, Rop R2O (exp(w<)) I) p:
(19)
A = R.p RRp = R (cxp(wx)) I)
(20)
generally have | # |Ab], as displacement magnitude depends on the initial rotation and point position.
E BASE PROTEIN ENCODER
We adopt the IPA module from AlphaFold2 (Jumper et al., 2021) as the base protein encoder. Specifically, the edge embedding initialized using a distogram, which discretizes the continuous pairwise distances into one-hot vector representation. For any two residues i and j, the Euclidean distance d, = ||x; x, ll2 is categorized into one of Noins discrete bins. The 1-thh component of the resulting embedding vector is determined by the indicator function I(-):
EdgeEEmeed(d:; (dz)a T(4 dig < ux). for t 1 1,..., Nuins:
(21)
20

---

Published as a conference paper at ICLR 2026
where (42, u1a) are the lower and upper bounds of the k-th distance bin. In our model, set Noins = 22, with bins linearly spaced between minimum of x 10-5 Á and maximum of 20.0 Á. The final bin extends to infinity to capture all larger distances.
Next, like normal attention mechanisms, generate the scalar and point- boint-based query, key, and value features for each node using bias-free linear transformations:
9: k. ve
h
W1 Si.
(22)
q
hp hp
Vi hp
W2 8):
(23)
b W3 Zajo
(24)
where W1 E (lod)xd, W2 E x) W3 E R and s; E R is the input node embedding. Here, denotes the number of attention heads, and p is the number of reference points used per head in the 3D attention mechanism. Each point-based feature (e.g., "") is reshaped into [h. 3), representing a set of 3D vectors per head that enable spatial reasoning in the IPA module.
key component then, is calculating the attention score:
k
TT o x2/011))
a softmax,
WL
b/
Σ
T, o q
(25)
d
where 7' o (r y + z) represents the transformation.
As shown in Equation (25), IPA introduces an additional bias term besides the bias bij from the pair- wise feature, Ep T, -. T, k/00 which explicitly models the spatial relationship between nodes after applying the local transformations T; and T₁. When the transformed point T; q, and T K differ significantly, it indicates large spatial discrepancy between node i and node j. and the resulting attention score is correspondingly reduced.
We then calculate the attention score and update the node representation:
6
Zijs
(26)
0
h
dij V"
h
(27)
0,
hp
T, o
Σ a (T vz7e) (I)Ov?T))
(28)
6, Linear (concata.p (o) of th [(a/1)))
(29)
In 1nop (E, a (I, o v>r))
))
the model first applies the global transformation T3 to each
neighboring point vector hp weighs them by the attention scores at,, and then aggregates the results in the global frame. The result is then transformed back to the local coordinate frame of node i using T.-. This ensures that the update is the same in the local frame though global transform is applied.
To make it compatible with our flow matching setting, after each IPA block, we embed the current time step using sinusoidal positional embedding (Ho et al., 2020) denoted as F; and add that to s; complete the full update:
S. + (.(().
(30)
Finally, we map the updated node representation back to the translation and quaternion:
Fi Linear(s:),
(31)
where x; and y represent the translation and rotation (quaternion) respectively.
21

---

Published as a conference paper at ICLR 2026
F EXPERIMENT DETAILS
F.I TRAINING OVERVIEW
Table 2 summarizes the training setups for method and all baselines, including data sources, dataset sizes, protein length ranges, training objectives, and computational resources. Additional implementation and training details are provided in the subsequent sections.
Table Overview of pretraining and downstream training configurations.
Mallodd
Stage
Data Souree
Dats Size
Length
Obyective
Compute Time
Pretmining
AFDB LaiPonKK)
452,194
60 512
Contrative (IafoNCE)
(4xA1000)
Geo$S1 JERM-NCE
Protnining
AFDB (CmiProKR))
452,194
60-512
Contrnsive (FBM-NCT))
Pretrainig
AFDB (UniProtKD)
60-512
(Representation Recoestruction)
(4xA100)
RigidSSL (Oun)
Peetraining
ATDB (UniProiKD)
452,194
Generative (RigdSSLL)
2.75 days
RigidSSL-MI (Ows)
Pretraning (Pluse
ATLAS
u.390
Geserative (RigidSSL)
1388 days (1
PrameDitf
Finesning
PDB
20.312
60-512
SE,3) Diftusion
days (4 (4x:11000)
FoldFlou-2
Finctening
PDB AFD8
20.312
60- 384
SP(3) Flo Matching
4 4days AIx)A
F.2 PRETRAINING: DATASET
RigidSSL-Perturb. For the first phase of our pretraining strategy, we use version of AFDB. Specifically, we utilize AFDB's coverage for the UniProtkB/Swiss Prot section of the UniProt KnowledgeBase (UniProtKB), which comprises 542.378 proteins. We filter sequences by length, retaining only those between 60 and 512 residues, yielding 432,194 proteins for training.
To determine the optimal translation and rotation noise levels for constructing perturbed views in RigidSSI-PPrttufi we conducted ablation study (see Appendix G) and selected a translation noise scale of α = 0.03 and a rotation noise scale of 0.5.
RigidSSL.-MD. For the second pretraining phase, we utilize the ATLAS dataset, obtained from MDRepo (Roy et al., 2025; Vander Meersche et al., 2024). This dataset contains 100 ns all-atom MD simulation triplicates for 1,390 protein chains, which were selected provide an exhaustive and non-redundant sampling of the PDB's ' conformational space based on the Evolutionary Classification of Domains (ECOD) (Cheng et al., 2014). All simulations were performed with GROMACS using the CHARMM36m force field; each protein was solvated in a triclinic box with TIP3P water and neutralized with Na+ /CI- Nat/CI- ions a 150 m.M concentration. We filter sequences by length, retaining only those between 60 and 512 residues. To generate training samples for RigidSSI-MDD. we extracted 80 pairs of adjacent conformations separated by a 2 ns interval from each trajectory to reflect the conformational fluctuations of the protein structures.
F.3 PRETRAINING: OPTIMIZATION
For training RigidSSL-Perturt and RigidSSL-MD, we use Adam optimizer (Kingma & Ba, 2015) with learning rate of 0.0001. RigidSSL-Perturb was trained for 2.75 days one NVIDIA H100 GPU with batch size 1. RigidSSL--MD was trained for 1.88 days on one NVIDIA H100 GPU with batch size of 1.
F,4 DOWNSTREAM: UNCONDITIONAL GENERATION
F.4.1 TRAINING AND INFERENCE FOR UNCONDITIONAL GENERATION
We trained Frame and FoldFlow-2 using the same dataset preprocessing, model architecture, and optimization hyperparameters in their respective original works (Huguet et al., 2024; Yim et al., 2023), but warm-started the IPA modules with pretrained weights. FrameDiff was trained for one week on four NVIDIA H100 GPUs, whereas FoldFlow-2 was trained for four days on two NVIDIA A100 80 GB GPUs.
For the results in Table we generate structures by sampling backbones of lengths 100-300 residues with FrameDiff and 100-600 residues with FoldFlow-2, in increments of 50, using the inference hyperparameters specified in the original works (Huguet et al., 2024; Yim et al., 2023). For the long-chain study in Appendix F.5, follow the same inference procedure but extend the target lengths to 700 and 800 residues. We then select the structure with the highest self consisteesn and assess its stereochemical quality using MolProbity (Davis et al., 2007).
22

---

Published as a conference paper at ICLR 2026
F.4.2 EVALUATION METRICS FOR UNCONDITIONAL GENERATION
Designability is self-consistency metric that examines if there exist amino acid sequences that can fold into the generated structure, Specifically, we employ ProteinMPNN (Dauparas et al., 2022) design sequences (i.e. reverse fold) for FrameDiff/Foldlov- structures, which are then folded with ESMFold (Lin et al., 2023) to compare with the original structure. We quantify self-consisteney through scRMSD and employ scRMSD < 2 A as the criterion for being designable.
Novelty assesses whether the model can generalize beyond the training and produce novel backbones dissimilar from those in PDB. We use FoldSeek (Van Kempen et al., 2024) to search for similar structures and report the highest TM-scores (measure of similarity between two protein structures) of samples to any chain in PDB.
Diversity measures structural differences among the generated samples. We quantify diversity in two ways: (1) the length-averaged mean pairwise TM-score among designable samples, and (2) the number of distinct structural clusters obtained using MaxCluster (Herbert, A. & Sterbberg, M., 2008), which hierarchically clusters backbones using average linkage with seguence-independen structural alignment at T'-sscore threshold of 0.7. We report Max diversity as the proportion of unique clusters, i.e., (number of clusters) / (number of samples).
However, these standard protein-structue generation benchmarks only emphasize sample-level criteria but do directly assess whether generator matches target distribution of structures. Following Geffner et al. (2025), we additionally report in Table 3 three distribution-level metrics derived from fold classifier: the Fréchet Protein Structure Distance, the fold Jensen-Shannon divergence, and the fold score.
Fold classifier and representations. Let denote a protein backbone structure. We use GearNet- based fold classifier Po(- v) (Zhang et al., 2022) that outputs probability distribution over fold classes. The classifier also induces a non-linear feature extractor (prr) (we use the last-layer features of the classifier), which we treat as embedding of protein backbones.
Let Kgen and Ne denote sets (or empirical distributions) of generated and reference structures, respectively. Define the empirical feature moments:
Agen ExwxXgm [p(x)], Egen E ((p(x)) Pgen) (p(() - Mgen) 7T,
(32)
Prae Ex.xx [((x)], Exef Ex. ((p(x) Mrot) (p(()) Mref)
We also define the marginal predicted label distributions:
Pgen (.) = Bez-apm [po(: (: [)),
(33)
Piet (-) Ex-x. pol. (pol x)].
Fréchet Protein Structure Distance FPSD compares generated and reference backbones in the non-linear feature space p(z) by approximating each feature distribution Gaussian and measuring the Fréchet distance between these Gaussians:
FPSD(X (egen Xret)
||Mgen
Mreell2 2 + tr Egen / + Eret 2'gee Dret) 122.
(34)
Lower FPSD indicates closer alignment between generated and reference distributions in feature space.
Fold Jensen-Shannon divergence fISD) measures discrepancy in the categorical fold-label space by computing the Sasenn-Shanoon divergence between marginal predicted label distributions:
EJSD(Agen) Kret) == 10 x Dus(Pgen Pret).
(35)
The Jensen-Shannon divergence is
Das(P = (PIII) + (Q)) M = (P Q)
(36)
We multiply by 10 for reporting convenience. Lower fJSD indicates closer agreement of generated vs. reference fold distributions. When fold labels are hierarchical (c.g., C/A/T), fJSD can be computed at each level; we report the average across levels.
23

---

Published as a conference paper at ICLR 2026
Fold score The fold score summarizes both per-sample confidence and across sample diversity under the fold classifier. Let Pzen (-) = Nx [Pp(: x)] be the marginal label distribution over generated
samples. Then
(geen)) == exp Eg [DKL (ps(- |
(>) Pgea(-))])
(37)
A higher IS indicates that individual samples are assigned sharp fold distributions while the aggregate covers broad range of fold classes.
Table 3: Comparison of FPSD, fS, and IJSD. Lower is better for FPSD/UJD; higher is better for IS. Best mean values are in bold and second-bee are underlined.
Model
Pretraining
FPSD vs.
IS
DS
PDB (1)
AFDB (1)
(C/A/T) (1)
PDB (C/A/T) ())
AFDB (C/A/T) (
Frlm
None
908.743
840.682
1.052/1.22744:499
3,174/3 902/5,
2.209/3.119/4 807
Fifm
EBM-NCE
827,919
270.448
L.188/2.43/6.0 025
3.017/3 593/5.079
2.068/2.80744.444
FifmeDiff
GeoSSL-InfoNCE
858.332
806.346
1.002/2.0006.88
3.360/3. 817/88877/.1199
1.369/2.96544411
Fiff
GeoSSL-RR
892.657
846, 222
1 1.002/1851/6... 56.5559
3,364/3.474/5 037
2.373/2.452/4.. 21//2//55//433555 357
FrameDiff
RigidSSI. Perturb (ours)
901.424
843.407
1.003/1.35/44100
3.359/3 605/5.417
2.370/2.6 222022632/4.44 687
Fifm
RigidSSL-Mi (ours)
776.297
701.705
1,01/11 834/5 345
3,374/3
552/4.987
2.383/2.560 560/4.151
FoldFlow-2
None
1273.363
1198.292
1,09871
464/4.855
2.8 1.813/3.9555 91573.5555.611111 61
1.90373, 184/3.2)
FoldFlow-2
Geo8SI.EEBN -EBM-NCH
1250.866
1201.031
1.04/11
565/55665
3,099/4 www.//.111/.... 834
2.131/3,
F
GeoSSL-InfoNCC
1182.036
1127.343
1.075/1. 78866.005
2.940/3.4925 5666
1.996/2.022
202/2/.022/222220 207
FoldFlow-2
GeoSSL-RR
2532.639
2485. 163
16681/2.4978 497/4.115
1.219/3 240
0.885/3.279/5.. /5.222 227
FoldF1oo-2
RigidSSI. Perturb (ours)
959,982
875.549
5/1 647/4 404
3,288/3
921/5 .94
2,30//31 7/4.622
FoldFlow-2
RigidSSL-MD (ours)
1113,829
1070.094
L L198/2
2499/ .35
2.623/318555.073
1.711/2.22448810
F.5 CASE STUDY: LONG CHAIN GENERATION
To explore the potential of RigidSSL, we conducted a case study on the generation of long protein (details provided in F.4.1). Whereas the original FoldFlow-2 analysis was restricted to sequences of 100-300 residues, extend the evaluation to ultra-long sequences of 700 and 800 residues under different pretraining methods. We visualize the best structure in terms of scRMSD among 50 generated samples in Figure 3. Notably, only FoldFlow-2 pretrained with RigidSSL-Perturb is able to produce self-consistent and thus designable structures.
To further understand the stereochemistry of the generated structures,
MolProbity (Davis et al., 2007) to compute: (1) Clashscore, which is the number of serious steric overlaps (>0.4 Å) per 1000 atoms, and (2) the MolProbity score, which is a composite metric integrating Clashscore, rotamer, and Ramachandran evaluations. As shown in Table 4, FoldFlow-2 pretrained with RigidSSL-- Perturb achieves the best clashscore and MolProbity score, indicating that it effectively captures the bicphysical constraints underlying long, complex protein structures despite being trained only on shorter sequences (60-384 residues). This demonstrates that RigidSSL-Perturb substantially improves the model's ability maintain stereochemical accuracy at long sequence lengths, suggesting robust learning of global structural patterns that generalize beyond the training regime.
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.
Menie
Length
Unpretrained 132.53
EBM-NNE
GooSSL-RR
RigidSSL-Periu GeGNSE-I
RigiaSLLKKU
Clashscore (1)
700
45.31
96.38
35,02
21.36
37.16
800
16
34.98
39.35
MolProbity score (1)
700
2.64
2.22
2,51
221
1.82
2.05
800
2.38
2.1
2.34
1.91
2.08
F.6 CASE STUDY: ZERO-SHOT MOTIF SCAFFOLDING
We evaluate motif scaffolding by adapting the RFdiffusion protocol (Watson et al., 2023). For each target motif, we generate 100 backbone scaffolds using FoldFlow-2 models pretrained under different frameworks and finetuned for unconditional backbone generation. Unlike the approach in Huguet et al. (2024), we did not employ additional finetuning via pseudo-label training. Motif residues are held fixed during generation by setting their corresponding entries in the fixed mask to 1.0, while scaffold residues are sampled from the learned flow. We use 50 inference steps with noise scale 0.1 and minimum time of miin 0.01. The motif structure is incorporated conditioning signal by
24

---

Published as a conference paper at ICLR 2026
Unpretraline
GeOSSL-EBM-NCE
GeoSSL-InfONCE
GeoSSL-RR
Riglds5
Rigid55L-MD
17A1A
scRMSD:
SCNMSD: 26.49
SCRMED: 17.92
sameso: 20.26
SOWSD:: 20.24
servE:: 152A
soNMSO: 12.19
Length
Figure 3: FoldFlow-2 generated structures (orange) compared against ProteinMPNN ESMFold refolded structures (grey). Columns denote pretraining methods, and rows denote sequence lengths 700 and 800.
centering coordinates at the motif center mass. each timestep, the model predicts vector fields for all residues, but only scaffold coordinates are updated while motif coordinates remain unchanged, thereby sampling scaffolds from distribution conditioned on the fixed motif geometry.
For each generated backbone scaffold, perform inverse folding with ProteinMPNN (Dauparas et al., 2022) to design amino acid sequences. We sample K== sequences per backbone at sampling temperature 0.1. Motif residue identities are fixed during sequence design, and ProteinMPNN designs only the scaffold region. We assess designability by refolding the TroteinMPNN-desiiee with ESMFoll (Lin et al., 2023) and computing structural similarity metrics. A scaffold is designated as designable if at least one of the eight refolded structures satisfies the following success criteria:
sequences
global Co RMSD < 4.0 Å, motif Ca RMSD < 3.0 Å, and scaffold Ca RMSD < 4.0 Å. We emphasize that these scRMSD thresholds relaxed to reasonable extent relative the original benchmark to provide greater resolution among methods in the challenging zero-shot setting (Korbeld al., 2025). Accordingly, results intended for internal comparison across pretraining frameworks rather than for direct comparison to prior benchmarks in the literature. RMSDs are computed under optimal superposition for each residue subset. We report designability as the fraction of the 100 generated scaffolds per target that satisfy these criteria; results are shown in Table 5.
Table 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 underlined.
Target
Unpretrained
BBMSS-EEM--CE
GeoSSL-InfoNCe
GeoSSL-RR
RigidSSL-
RigidSSL-MD
15
IPRW
II
11
IQIG
IYCR
10
2KL8
15
5
23
14
17
34
3IX
4JHW
STUS
16
12
36
22
STPN
SIRV Joog
20
11
30
17
51
25
STRV med
17
20
sbort
25
11
5WN9
10
6E6R Jong
10
20
15
10
6B6R short
3
6EX7 long
6EX7_mg med
6EXZ short
2
7
7MRX_12 128
0
3
7MRX 60
7MRX_8 85
24
20
rsv. site4
Average
9.346
7,077
13.192
4,077
15.192
10,077
25

---

Published as a conference paper at ICLR 2026
F.7 CASE STUDY: GPCR CONFORMATIONAL ENSEMBLE GENERATION
F.7.1 TRAINING AND INFERENCE FOR GPCR ENSEMBLE GENERATION
For the case study on GPCR ensemble generation, we utilize the GPCRmd dataset hosted on MDRepo et al., 2020; Roy et al., 2025), which provides all-atom MD simulations across all structurally characterized GPCR classes. To focus on intrinsic receptor flexibility, we specifically use the apo form of each receptor, which was generated by removing the ligand from its binding pocket prior to simulation. Each protein was embedded in lipid bilayer, solvated with water, and neutralized with ions. All systems were simulated for 500 ns in three replicas using the ACEMD software (Harvey al., 2009). To generate the training and evaluation samples for RigidSSI., processed these trajectories into final dataset comprising 523 training, 68 validation, and 67 testing MD trajectories, each subsampled 100 frames. We trained AlphaFlow from scratch using the same optimization hyperparameters as the original work, warm-starting the IPA module with pretrained weights. Sampling follows the procedure described in the original work.
F.7.2 EVALUATION METRICS FOR GPCR ENSEMBLE GENERATION
To evaluate the generated GPCR ensembles against the ground-truth ensembles, we adopt the metrics used in Jing et al. (2024); we provide brief definitions below and refer the reader to the original work for full details.
Let X = (=()), ( (M) denote a predicted ensemble of M structures and y = (
y (K))
the ground-true truth MD ensemble of K structures, where each (mm), y E N are the 3D coordi- nates of a protein with N atoms. All structures RMSDD to common reference structure prior to analysis.
Pairwise RMSD. Mean pairwise Ca-RMSD within ensemble: M
PairRMSD(2))
Σies RMSD(2 z(0)..
2
All-atom RMSE. For atom n with ensemble mean T.. RMSF, V 1M-1,, 1|
(m) F. 12
Root Mean Wa-Distance (RMWD). Fit 3D Gaussians N( En, to each atom's positional distribution. The squared 2- 2-Wasserstein distance hetween Gaussians is (3(, = 1
12ll² + Tr(2; + Σg 22,, 1//2) Then RMWD(X, = VN- Σ., W... - 12 where T = N Σ. I 22 (translation) and V2 = - Σ., Tr > 12) (variance).
This decomposes as RMWD = 2
PCA Wa-Distance. Project flattened C, coordinates onto the top 2 principal components and compute W2 (in Á RMSD units) between the projected distributions. MD PCA: PCs computed from Y alone. Joint PCA: PCs computed from the equally weighted pooling of X' and y
% PC-sim > 0.5. Let of" uyx,07 of be the top principal components of each ensemble. PC-sim (e " [[e)"]]). We report the percentage of targets with PC-sim > 0.5.
Weak Contacts. A C. pair (i.j) is in contact if ||æ; - ==,ll < 8 Å. Weak contacts are those in contact in the crystal structure but dissociating in > 10% of ensemble members: W(X) = {(i,j) o
1, (x) < 0.9). Agreement is measured by Jaccard similarity J(A, B) = |An B //AAAUB
Exposed Residue. Sidechain SASA is computed via Shrake Shrake-Ruple Rupley (probe radius 2.8 Å). Cryptically exposed residues are buried in the crystal (SASA ref < 2.0 A²) but exposed (SASA, > 2.0 A²) in 10% of ensemble members. Agreement is measured by Jaccard similarity.
Exposed MI Matrix. For binary exposure indicators en 1[SASA (m) > 2.0 A compute pairwise mutual information Ml; = En, (ty(a, of X' and Y is measured by Spearman correlation p.
Piy(:, b) log P.y(a,)) p,( (b) Agreement between MI matrices
26

---

Published as a conference paper at ICLR 2026
Table 6: Evaluation of GPCR MD ensembles generated with pretrained AlphaFlow variants. Best results shown bold.
Metric
Unpretrained
GeoSSL-EEDM LBM-NCE GrSSSL-EalONCE
GeoSSL
GeuSSL-RR
RighSSL SeuuSSL-Peturrbooo (ouns)
RighSSL MD (ours)
Predicting Becibility
Paiewise RMSD (=1
2.31
2.45
asom RMSF (=1.0)
1.30
1.19
1.22
1.17
1,08
1.28
Roxt Wi-list
18.36
18.46
17.79
1783
17.90
18,10
Distributional
Lo Vaxinnee coatals
1.45
1.50
1.42
1.52
1.44
1.60
accurncy
MD PCA Wi-dist I
3.15
2,97
3.16
3,14
3.10
3.22
Wi-dist
17.55
17.67
% PC-sim 0.5 /
0.00
1.49
4,48
4.48
5.07
5.97
Welk
0.36
0.36
0,43
Ensemble observalles
Exposed residue
0.67
0.58
0.58
0.42
Exposed MI matrix
0.01
0.00
0.00
0.02
0000
0.03
G ABLATION STUDIES
We present ablation studies of RigidSSI. to isolate the contributions of our pretraining framework from the effects of increased data volume, and evaluate specific design choices in the view construction process.
First, to disentangle the benefits of our pretraining method from the advantage of simply training on a vastly larger dataset, compare FoldFlow-2 baseline trained from scratch combined AFDB and PDB dataset against our standard pipeline, which pretrains on AFDB via RigidSSL-Perturt and then finetunes on PDB. As shown in Table 7, RigidSSL achieves higher designability and diversity, along with comparable novelty, while requiring 100,000 fewer training steps. This demonstrates that the observed performance gains stem fundamentally from RigidSSI. rather than just an increase in data scale.
Table 7: Comparison of Designability (fraction of proteins with scRMSD < 2.0 Å), Novelty (max. TM-score to PDB), and Diversity (avg. pairwise TM-score and MaxCluster diversity) of structures generated by FoldFlow-2 trained PDB AFDB from scratch versus models first pretrained 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 Huguet et al. (2024).
Designability
Novelty
Diversity
Model
Downstream Dataset
Pretraining
Steps
Fraction (?)
TM (
pairwise TM
MaxCluster (f)
FoldFlow
PDB + AFDB
Nome
500k
0.738 = 0.014
0.764
0X003
0.657
0.001
0.250
FoldHlow-2
PDB
RigadSSL-Peruub
400k
0.758 0.016
0.770
0.003
0.650
0.001
0.252
Second, we investigate the impact of translational and rotational noise scales in RigidSSL-Perturb (Section 3.2.1). As shown in Figure 4, increasing either noise scale leads to more steric clashes and reduced bond validity. The resulting FoldFlow-2 models trained with these noise settings are reported in Table 8.
Table Comparison of Designability (fraction of proteins with scRMSD 2.0 Å). Novelty (max. TM-score PDB), Diversity (avg. pairwise TM-score) structures generated by FoldFlow-2 pretrained with RigidSSL- Perturb under different noise scales. All metrics include standard errors.
Designability
Novelty
Diversity
Translation Noise
Rotation Noise
Length
Fraction (f)
avg. max TM (L)
pairwise TM (L)
0.01
0.5
100-600
0.336 + 0.012
0.768 0.010
0.635 + 0.002
0.03
0.5
100-600
0.758 + 0.016
0.770 + 0.003
0.650 + 0.001
0.05
0.5
100-600
0.589 + 0.014
0.769 0.005
0.654 + 0.002
0.5
0.5
100-600
0.382 + 0.015
0.74S 0.005
0.649 1 0.001
0.5
0.75
100-600
0.660 + 0.014
0.763 + 0.004
0.644 1 0.001
0.5
1.0
100-600
0.352 1 0.025
0.772 1 0.008
0.614 0.003
1.0
0.75
100-600
0.460 1 0.016
0.773 0.005
0.663 1 0.001
2.0
0.75
100-600
0.347 1 0.014
0.797 0.006
0.624 1 0.002
27

---

Published as a conference paper at ICLR 2026
Figure Impact translation and rotation noise scale on protein structure validity.
28


Reviews:
[{"review_id": "zbZ0aUvtcc", "reviewer": "Reviewer_LsQ9", "summary": "This work proposes a pre-training procedure for diffusion and flow-matching based protein design methods. The pre-training method consists of two phases. The first phase, RigidSSL-Perturb, adds Gaussian noise to the translation and rotation component of the residual frames of a protein's backbone. The second phase, RigidSSL-MD, uses pairs of structures from the ATLAS dataset, which contains molecular dynamics simulations of proteins.", "questions": "Is high designability of RigidSSL-Perturb potentially only achieved by a high amount of alpha helices?\n\nThe augmentation of RigidSSL-Perturb, which adds Gaussian noise to the coordinates and rotation, could destroy physical plausibility of the proteins. How is the noise level controlled to prevent this?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Incorporating MD simulations into a pre-training step of protein design methods is an interesting and novel contribution.", "weaknesses": "The experiments demonstrate that RigidSSL-Perturb outperforms the baselines for designability and novelty, while RigidSSL-MD outperforms the baselines in diversity. However, RigidSSL-MD is not the methods of choice with respect to designability and novelty. These results limit the appilcability of the approach, as there seems no practical advantage of RigidSSL-MD over RigidSSL-Perturb, which is merely a simple data-augmentation of the data with Gaussian noise.\n\nThe examples of generated structures in figure 3 for RigidSSL seem to solely consist of alpha helices, a common bias in generative models for proteins (compare e.g. with Wagner et al. 2024). Please also report alpha helix and beta strand content for each method.\n\nReferences:\nWagner, S., Seute, L., Viliuga, V., Wolf, N., Gräter, F., & Stühmer, J. (2024). Generating highly designable proteins with geometric algebra flow matching. Advances in Neural Information Processing Systems, 37, 77987-78026.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H24jmIThg7", "reviewer": "Reviewer_yn5Y", "summary": "This paper proposes a pretraining method for protein generation by viewing each residue structure as a rigid body and using flow matching. Specifically, it first defines an inertial frame as a reference frame for each protein and aligns protein backbone structure with the inertial frame; then, it perturbs each protein's translation and rotation by randomly sampling transforms from an Euclidean space and a special orthogonal group to form two views; next, it samples trajectory segments from a molecular dynamics dataset; finally, it uses flow matching to generate each view from the other views. The experiments are conducted on a protein generation task and the results outperform the state-of-the-art methods.", "questions": "1. In Section 3.2.1, $g^0$ and $g^1$ represent the original state of one protein and the perturbed state by adding random noise; in Section 3.2.2, $g^0$ and $g^1$  represent different conformations sampled at different time steps along a dynamic trajectory. This inconsistency is confusing.\n\n2. In Section 3.2.1, the other view is constructed by adding random noises sampled from Euclidean space and the SO(3) group. How to make sure the perturbed view is biologically valid?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. It utilizes the structural information available in large-scale protein datasets to pretrain a protein generation model in an unsupervised manner.\n\n2. It achieves superior protein generation performance to the compared approaches.", "weaknesses": "1. This paper is more engineering-oriented. Though it achieves superior performance across several models on a protein generation benchmark, it seems to contain little new algorithms or architectures. Reference frame definition and flow matching are widely used across many areas, including, but not limited to, machine learning, computer vision, and computational biology.\n\n2. The construction of two different conformation views is a little bit new, but the motivation for such a construction is unclear. We can directly sample two time steps in a dynamic trajectory and generate protein structures at one time step to the other.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ghmPIK6dBC", "reviewer": "Reviewer_8zoL", "summary": "This paper proposes a rigidity-based geometric pretraining method for protein backbone generation. For perturbation-based view construction on the large-scale AFDB dataset and molecular dynamics trajectories of the ATLAS dataset, they canonicalize proteins into a reference frame and use bi-directional flow matching to pretrain the IPA module, which is shared by FrameDiff, FoldFlow-2.", "questions": "**Q1. I'm curious whether the parameters of these equivariant models based on IPA is hard to scale? For rigidity-based geometric pretraining, thanks to dataset like AFDB, we can scale up the data size during the pretrain phase easily. But can we scale the equivariant models based on IPA module like Proteina[1], SimpleFold[2]?**\n\n\n[1].Proteina: Scaling Flow-based Protein Structure Generative Models\n\n[2].SimpleFold: Folding Proteins is Simpler than You Think", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The SE(3) rigidity pretraining for protein backbone generation is reasonable. Perturbations on rigidity align with the protein's natural conformational fluctuations, which can be interpreted as a masking-like paradigm. The MD snapshots used for pretraining are novel and interesting.", "weaknesses": "**W1. It is hard to determine whether the performance gains stem from the introduction of new data (e.g., AFDB, ATLAS) or the proposed rigidity-based geometric pretraining method. (My main concern)**\n\nThe impact of RigidSSL-MD on diversity has been analyzed in lines 408–411 and Section 5, I think the improvements of diversity are attributed to the new data in the ATLAS dataset. \n\nFor RigidSSL-Perturb, both FrameDiff and FoldFlow2 achieved improvements in designability and novelty. However, \n\n(1). FrameDiff was trained only on a small PDB dataset (\\~20K), while RigidSSL-Perturb incorporates new data from the AFDB (\\~432k). \n\n(2). The training length range for FoldFlow2 is 60-384, while RigidSSL-Perturb ranges is 60-512. For evaluations, the generated length is 100-600. This is unfair. \n\nI suggest authors:\n\n(1). Show FoldFlow-2 + RigidSSL-Perturb results for 100-300 instead of 100-600.\n\n(2). Add a comparison: train baselines Framediff and FoldFlow-2 on the same AFDB and ATLAS datasets using their original training objectives. In this case, comparison on 100-600 is OK.\n\n(3). Use a table in the main body of the paper to present data sources, numbers, protein length ranges, training objectives, training time, and other relevant information of various methods (including training, pretraining, and fine-tuning methods). This will make the paper's contributions clearer.\n\n**Note**: In my opinion, if my understanding is correct, previous CV/NLP tasks usually have scarce data labels, so we adopt self-supervised learning to utilize large amounts of unlabeled data. However, the unconditional protein backbone generation task here is different. We can directly utilize data from AFDB for training the original SE(3) flow matching. FoldFlow-2 [1] also validated that directly using new AFDB data (+PDB, \\~160K filtered data in total) for training yields performance improvements. This is why I have the ‘W1’ concern. However, I still believe that the rigidity-based geometric pretraining could be beneficial. Directly learning noise to protein means simultaneously learning the fundamental rules of protein geometry and the complex task of de novo protein design, which may be challenging. RigidSSL decouples these challenges by first pretraining on large-scale data to learn a robust geometric representation that serves as a good initialization. This motivation is reasonable for me, but the authors may need to clarify the specific benefits of RigidSSL according to suggestions (1)-(3).\n\n[1].Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation\n\n**W2. Referring to RigidSSL-Perturb as phase 1 and RigidSSL-MD as phase 2 is confusing. Particularly, they are two independent pre-training methods in the experiment part. What are the results of their combined use? Can you provide further experimental results and analysis?**\n\n**W3. Can you show some results on the conditional generation experiments, e.g.,  motif scaffolding?**\n\n\nI will raise my score if my concerns are resolved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HpchqOHfaN", "reviewer": "Reviewer_1qeX", "summary": "This paper introduces a pretraining method for an unconditional protein backbone structure generation model. Experimental results demonstrate that this approach enhances the performance of FrameDiff and FoldFlow2 in generating long proteins and high-quality proteins that are not purely helical.", "questions": "1. The batch size during the pretraining stage is set to 1, and no analysis related to batch size is provided. Most existing generation models demonstrate that a larger batch size often leads to better training performance. Why a larger batch size was not used in this work?\n2. Can this approach naturally extend to the transformer decoder framework like AlphaFold 3?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The problem addressed by this method is highly important. A major issue with unconditional protein generation models is that the generated designable proteins are predominantly helical. This method alleviates this problem to some extent through pretraining.\n\n2. The pretraining strategy used in this method is relatively general and trains fast.\n\n3. The experimental section follows established conventions and provides strong support for the claims made in the paper.", "weaknesses": "1. Some descriptions are somewhat confusing. For example, it is unclear whether Phase 1 and Phase 2 represent two different strategies or if they are sequentially related.\n\n2. The case study lacks illustrations or examples of designable proteins with a low proportion of helices.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)