Conditional Invariances for Conformer Invariant Protein Representations
Authors: Balasubramaniam Srinivasan
Vassilis N. Ioannidis
Soji Adeshina
Mayank Kakodkar
George Karypis
Bruno Ribeiro
Venue: Submitted to ICLR 2023 (2023)
Keywords: Invariances
Conditional Invariances
input dependent invariances
proteins
protein representation learning
conformer invariant representations
graph neural networks
group invariant neural networks

Abstract: Representation learning for proteins is an emerging area in geometric deep learning. Recent works have factored in both the relational (atomic bonds) and the geometric aspects (atomic positions) of the task, notably bringing together graph neural networks (GNNs) with neural networks for point clouds. The equivariances and invariances to geometric transformations (group actions such as rotations and translations) so far treats large molecules as rigid structures. However, in many important settings, proteins can co-exist as an ensemble of multiple stable conformations. The conformations of a protein, however, cannot be described as input-independent transformations of the protein: Two proteins may require different sets of transformations in order to describe their set of viable conformations. To address this limitation, we introduce the concept of conditional transformations (CT). CT can capture protein structure, while respecting the restrictions posed by constraints on dihedral (torsion) angles and steric repulsions between atoms. We then introduce a Markov chain Monte Carlo framework to learn representations that are invariant to these conditional transformations. Our results show that endowing existing baseline models with these conditional transformations helps improve their performance without sacrificing computational cost.

Under review as a conference paper at ICLR 2023
CONDITIONAL INVARIANCES FOR CONFORMER IN- VARIANT PROTEIN REPRESENTATIONS
Anonymous authors
Paper under double-bling review
ABSTRACT
Representation learning for proteins is an emerging area in geometric deep learn- ing. Recent works have factored in both the relational (atomic bonds) and the geometric aspects (atomic positions) of the task, notably bringing together graph neural networks (GNNs) with neural networks for point clouds. The equivariances and invariances to geometric transformations (group actions such as rotations and translations) so far treat large molecules as rigid structures. However, in many important settings, proteins can co-exist an ensemble of multiple stable con- formations. The conformations of a protein, however, cannot be described as inpu-indeependent transformations of the protein: Two proteins may require differ- ent sets of transformations in order to describe their set of viable conformations. To address this limitation, introduce the concept of conditional transformations (CT). CT can capture protein structure, while respecting the constraints on dihedral (torsion) angles and steric repulsions between atoms. We then introduce Markov chain Monte Carlo framework to learn representations that are invariant to these conditional transformations. Our results show that endowing existing baseline models with these conditional transformations helps improve their performance without sacrificing computational efficiency.
INTRODUCTION
The literature on geometric deep learning has achieved much success with neural networks that explicitly model equivariances (or invariances) to group transformations (Cohen & Welling, 2016
Maron et al. 2018 Kondor & Trivedi 2018 Finzi et al. 2020). Among applications to physical sci-
ences, group equivariant graph neural networks and transformers have specifically found applications
to small molecules, as well large molecules (e.g. proteins) with tremendous success (Kliepera
et al.
2020 Anderson et al. 2019 Fuchs et al., 2020 Hutchinson et al., 2021: Satorras et al.) 2021 Batzner et al.
2021 D. Specifically, machine learning for proteins (and 3D macromolecular structures
in general) is a rapidly growing application area in geometric deep learning. (Bronstein et al. 2021
Gerken et al. 202T). Traditionally, proteins have been modeled using standard 3D CNNs (Karimi
ct al. 2019: Pagès al. 2019), graph neural networks (GNNs) (Kipl & Welling 2016 Hamilton et al. 2017). ), and transformers Vassaan et al. 2017). More recently, several works Jing et al. (2020)
2021); Jumper al. (202I); Hermosilla et al. (202T) have enriched the above models with neural networks that are equivariant (invariant) to transformations from the Euclidean and rotation groups. While equivariance (invariance) to the transformations in these groups are necessary properties for the model, unfortunately, they are limited to only capture rigid transformations of the input object.
However, these models may not yet account for all invariances of the input pertinent the downstream lask. And the transformations they are invariant to do not depend on the input. For instance, invariance the Euclidean group restricts the protein representation act as if the protein were a rigid structure, regardless of the protein under consideration. However, treating proteins as rigid structures may not be optimal for many downstream tasks. And different proteins may have different types of conformations (protein 3D structures with flexible side chains) (Harder et al.) 2010, Gainza et al. 2012 Miao Cao. 2016) The existing rigid body assumption in protein representations may hurt these methods in datasets & downstream tasks that require protein representations to be invariant to specific set of protein conformations. For example, for most proteins, regardless of their side chain conformation (as long viable) under consideration their protein fold class/ other scalar properties remain the same, their mutation (in)stabilityy remains unaltered, protein ligand binding affinity (apart from changes at the ligand binding site) remain the same, etc.
In light of this limitation of current methods, a question naturally arises: Is it possible to learn conformer-invariant protein representations?

---

Under review as a conference paper at ICLR 2023
Side Chein
Q
Racesona Alone
Single anina add in protett
Figure I: Magnified image of side chain of a single generic amino acid (here. with 6 atoms in the side chain) in protein molecule. A protein molecule typically contains tens hundreds of amino acids
Our Approach: We propose representation method where the set of symmetries in the task are input-dependent, which we denote conditional invariances. For our specific application, we model every protein as a directed forest, where every amino acid forms directed tree. We leverage the constructed directed forest of the protein to sample viable protein conformations (where viability is checked with Protein structure validation tools such Molprobity (Davis et al. 2007 Chen et al.) 20100), which is additionally coupled with a Markov chain Monte Carlo (MCMC) framework to create data augmentations to train the neural network to learn conformer invariant representations.
Our contributions can be summarized as follows:
We provide guiding principles for defining conditional invariant representations as inductive biases, which serves as generalization of group invariant neural networks.
We then provide a principled strategy to sample conformations for any given protein from the support of its protein conformer distribution (which consists of all its viable protein conformations) which captures their true flexibility. Viable conformations respect domain specific constraints such as those on dihedral angles, steric repulsions, among others. This is particularly useful the scenario when the set of viable transformations is different across different proteins.
Further, we develop an MCMC-based learning framework which guides the sampling of protein conformations such that the (asymptotic empirical) average representation obtained by all viable conformations of protein are identical.
Finally, we perform experimental evaluation of our proposal, where endowing baseline models with proposed strategy (via an implicit data augmentation scheme) shows noticeable improvements on multiple different classification and regression tasks on proteins.
CONDITIONAL INVARIANCES FOR PROTEINS
The symmetry of an object is the set of transformations that leaves the object invariant. The notion of symmetries, expressed through the action group on functions defined on some domain, has served as a fundamental concept of geometric deep learning. In this section, we start by defining the concept of conditional symmetries that are both computationally efficient and useful for representing protein conformations. Some group theoretic preliminaries are presented in the Appendix (A.2
conditional symmetries and then proceed specifically tailor these
Definition 2.1 (Conditionally wymmeriicinvariann functions). A function f $2 I, is said to be conditionally symmetric-invant invariant if
f(tz x) f(x) .tz E Sz, Vx E Ω
where Sp is a set of transformations unique to element x and by Ω - Ω.
It is easy see that conditionally invariant functions are a generalization of group invariant functions where Sz = G Vr E Ω where G' is a group. The above definition is motivated by the fact that representations for proteins may not necessarily be limited be invariant just group actions, but to a more general set of protein specific transformations. We detail this motivation next.
protein molecule is composed of amino acids (say n amino acids), where each atom the protein belongs to one amino acid e {1...., nt.. Excluding the hydrogen atoms, every amino acid contains four atoms known as the backbone atoms (see Figure T. and other atoms which belong to the side chain. Traditionally, protein structures are solved by X-ray crystallography or cryo-EM and the obtained structure is normally considered a unique 3D conformation of the molecule. Molecules available via the Protein Data Bank (Berman et al. 2000) generally include only unique sets of coordinates to demonstrate the 3D structures, which lead to the unique answer as "gold standard' in structure prediction, such as protein structure prediction competitions CASP (Kryshtafovych et al. 2019). Under these assumptions protein side-chain conformations are usually assumed to be clustered into rotamers, which rigid conformations represented by discrete side-chain dihedral angles.
2

---

Under review as a conference paper at ICLR 2023
However, works over the past decade (Harder et al. 2010: Gainza et al. 2012 Miao & Cao 2016) have observed that - while the backbone atoms of all n amino acids in a protein molecule together (i.e. 4n atoms) form rigid structure, its side chains can exhibit flexible (continuous and discrete) conformations beyond clustered rotamers. The underlying goal of our work is to capture this inherent flexibility of proteins and ensure different conformers of the same protein get identical represen- tations (More details- - Appendix 10). The above problem of capturing protein conformations is further compounded by the fact that protein conformations are often times unique the protein i.e. conformations exhibited by protein are input (protein) specific. The desiderata, therefore is a model which outputs conformation invariant representations when the conformations (symmetries) exhibited by protein varies across proteins when access to only single protein conformer is available.
We denote a protein as a tuple p = (V. xs, X,) (data from pdb files) where is the set of nodes (every atom is a node) in the protein, X, E Rmxd with the atoms, Xp E Rmmx3 (where m denotes the number of atoms in the protein) are the positional coordinates associated with the atoms in the protein. Without loss of generality, we number the nodes in V = (1...., m following the same ordering of the rows in X.. X,p
1 is the scalar atom feature matrix associated
Definition 2.2 (Rigid Backbone Protein Conformations). For an m--tom protein p with atom posi- tional coordinates X, E R.xx3 (recall that by definition p contains information about Xp). where Cp C IR denotes the set of viable conformations of p, which keeps the positional coordinates associated with the backbone fixed. We use Tp C mmx333 (where 3x3 matrix is associated with every single atom) to denote the set of non-isometric transformations of p, which yield the set Cp i.e. Vcp Cp: 36p Tp, s.t. for atom with index i E {1,...,m) the new position of atom i is X, Xpli)tp cp16), Xpli]
R Ix33 tyli] E R333
Tp forms set of transformations which acts on the protein atomic coordinates via matrix multiplica- tion. Two elements of T (with corresponding matrices of individual atoms and then aggregating for the protein as a whole) may may not combine via matrix multiplication (as the binary operation) to result in another element of Tp i.e. matrix multiplication is not necessarily closed in T.. A concrete example of the non group structure of protein conformations is provided in Appendix A.3
Unfortunately, sampling transformations from Tp to obtain viable protein conformations is an unattainable task, since this would entail first sampling a transformation from Rmx3x33 and then verifying if the transformation is viable (and the vast majority of transformations are not viable). We will address this issue using an MCMC method- which is then combined with MCGD training of the neural network (Sun et al. 2018) to learn conformation invariant protein representations where different viable conformations are obtained via MCMC are used in every batch. Markov Chain Gradient Descent (MCGD) is variant of Stochastic Gradient Descent (SGD), where the samples are drawn from the trajectory of a Markov Chain. We note that, while there have been prior Monte Carlo and MCMC based techniques for sampling protein conformations (Boomsma et al., 2013 Irbäck & Mohanty| 2006 Vitalis & Pappul (009). the overarching goals of the existing MCMC methods are significantly different compared to ours. The existing methods define distribution over the conformations (based on the properties of the bonds, etc.) and then sample highly probable conformations from this distribution. Our method, is much simpler and only requires that the chain being used is ergodic and is invariant to the actual form of the distribution. Therefore, we devise a simple chain whose transitions are easier to sample and don't 't require us to compute conditionals of an energy based model. However, we also note that the existing Markov chains can be used as drop-in replacements to sample conformations as part of our framework. Before specifying our MCMC procedure, we specify the invariance requirements for learning protein representations.
symmetric-invariant function for proteins (per Definition 2.1 1, should be invariant to (i) transfor- mations which change its atomic coordinates to any of its rigid backbone conformations in Cp (ii) transformations which change the atomic coordinates of all its atoms via rigid body transformations group actions such rotations/ translations (iii) combination of the two above which transforms it into one of its viable conformations and is then further acted upon by rigid body transformations.
3 OBTAINING VIABLE CONFORMATIONS
Before we describe our MCMC procedure to sample atomic coordinates from the set of rigid backbone conformations. need way sample from the local neighborhood of conformation.
3

---

Under review as a conference paper at ICLR 2023
g 10 a
gl + .n
&
g 52 of
dlv g/
3
Figure 2: Directed tree corresponding to set with points which exhibits conditional invariances. Our proposed model (not limited to proteins), for example, allows node 2 and its descendants to transform its coordinates about node 1 (its parent) upon actions from group Clon - 91 E G1. In practice, for protein molecules we use G S3(3) and cf.) = (0(3)vi * 1. We note that in protein molecules not all transformations about node would be allowed due to steric repulsions between atoms well potential overlaps of atoms.
3.1 A SIMPLE PROTEIN CONFORMATION SAMPLING STRATEGY
Akin to Irbäck & Mohanty 2006), to efficiently sample candidate conformation from Cp. will follow a multi-step approach: First, we (i) construct a directed tree for every amino acid in the protein capture the inherent flexibility in protein structures. Then, we (ii) leverage directed forest of the protein (described next) to transform its atomic coordinates and check for viability.
Part Directed forest construction. Using the input protein molecule, we construct a directed forest using the following three step procedure:
1. Each amino acid in the protein gets its own directed free. The atoms in the amino acid's backbone form the base (root) nodes (i.e., there can be multiple base nodes in every amino acid). This is illustrated as node in Figure 2]s tree. The set of base nodes of all the amino acids in the protein are jointly referred as the base nodes of the protein (or equivalently of the directed forest).
2. The atoms adjacent to the amino acid's backbone via covalent bond become their immediate children in the tree. For example, the node SC1 forms the child of aC in the Figure I
3. The other covalent bonds of the children establish the grandchildren of the root, and so on, until all the atoms in the amino acid are part of the free. For example. SC2 and SC3 become the children of the node SC1 in the directed tree constructed for Figure
Figure is an illustration of some directed tree constructed with the above procedure, where node 1 is the root node, with nodes 2, 4 as its children and so It is important to note that, the above procedure ensures that in the constructed free, each (non-root) node has a single parent and may have arbitrarily many children. Further, cycles/rings which present in the molecule are broken on the basis of bond length, i.e., larger bonds are chosen before smaller bonds. Ties between same-length bonds are broken arbitrarily. While multiple directed forests can be created for given protein since bonds of equal length are broken arbitrarily, in our implementation, given protein, the tree is fixed because the directed forest for every protein is created exactly once (in the preprocessing step) where the ties are broken deterministically based on their ordering in the "pdb file".
Next, we shall look at how the atomic coordinates are transformed using the directed tree.
Part 2. Conditional transformations of atomic coordinates. Let X, be the available input atomic coordinates of the protein. To obtain a new candidate protein conformation (say X, E Rmx3), we sample uniformly, one of the directed trees from the forest. Next, we transform the atomic coordinates of the subset of atoms in this directed tree. This set is denoted by A C V. That is, in the candidate conformation, X
Xp(i), Vi E V A
equivalently [p] I, Vi V A.
Starting with the root node, we use breadth first search to traverse the directed tree and update all the descendants of the node currently being considered. For the backbone atoms, we leave Xp[i| X11 1), i.c., atomic coordinates are unchanged for backbone atoms. Next, use
(or equivalently tp pointed sets to describe approximate transformations on a given protein through a directed tree.
Definition 3.1 (Pointed sets). A pointed set is an ordered pair (X, where X is a set and xo X called the basepoint.
For instance, let M,-- = [X, : j J, where N₂ is the set containing and all its descendants in its directed tree. Then, (M, X,(i) is pointed set.
4

---

Under review as a conference paper at ICLR 2023
Let the parent of node i be denoted by o and let Gl X,(o)) (in practice, SO(3)) be the group from which actions E G( are sampled uniformly that can transform the positional coordinates of node and its descendants about its parents in the directed tree, where go xp(1) [o] is the transformed atomic coordinates of the parent of node i (we use BFS so a top down approach). In the case of the
root node (for example in Figure the actions are just drawn from Groot (or G1 in Figure 2
The associated transformations of the atomic coordinates of the atoms in N; can be given by the mapping h; (Mr, Xp [?]) + ([IIN) x3 9o. Xplo] + 9s (Xp [ Xp p pll))) node j E N, are updated as:
where the coordinates of
X + ge Xp[o]+ gi (Xp Xplo
(I)
If the node j A i, its coordinates, i... X;bi] can be further updated as we traverse the tree. When G ((Mo-Xx(']) S3))) Vi E V. every node in the directed tree can be rotated about its parents as (Appendix A.4) -which is the side chain of the amino acid shown in Figure |
shown in Figure S
Algorithm 1 Sampling a viable conformation of Cp
1: Input: Conformation cp of protein
2: Obtain the directed forest corresponding to the protein with n directed trees where n is the number of amino acids in the protein.
3: repeat
4: Sample y (NNIF((1,,,) ,..,,,,,..,......
n }) and obtain the corresponding directed tree 5: Traverse through directed tree via BFS and update atomic coordinates via Equation (I) to obtain 2, a candidate conformation of cy- where group actions are always sampled uniformly from SÓ(3).
6: Check if c is valid protein conformation via protein structure validation tools
7: until c, is valid conformation
8: Return: E. a valid conformation of Cp
However, each candidate X obtained via the above transformation process may not belong to Cp- To that end protein structural validation tools may be used to check the validity of these candidates. Molprobity (Davis et al.] 2007 Chen et al., 2010 Williams et al. 2018) is such a toihh outputs scores corresponding different metrics (such number of dinedral (torsion) angles outside the allowed threshold, number of atom-atom clashes arising due to steric repulsions, etc) which can be compared with the originally provided conformation X,, Note, that while Molprobity is not perfect and we may end up developing models more invariant than just to viable conformations, alternatively (more precise) tools including those which allow for the evaluation of molecular force- fields, can be used to check validity. Additionally, the goal of this work is not to use Molprobity/ other protein structure validation tools as a part of generative models and more invariant model does not necessarily affect performance on unseen but valid protein test data.
Our algorithm to sample viable conformations summarized in Algorithm T1 The repeat- tooo proposes conformation in the neighborhood of the current conformation Cp which is only accepted when the proposed conformation &, is a valid conformation. As such it is an acceptance-rejection sampling algorithm. We will see later that the actual sampling probability distribution is not required to be known by our procedure. Our MCMC procedure defined next only requires Algorithm T to follow certain conditions which we argue it follows.
3.2 SAMPLING CONFORMERS VIA MCMC
We now describe the MCMC procedure that, starting at any conformer configuration cp E Cp we will, in steady state, sample conformations according a unique stationary distribution which depends on the protein/ where Cp is the set of all viable conformations of as defined in Definition 2.2
To that end we first define the transition kernel of the Markov chain. Algorithm D samples a conformation de in the neighborhood of a given conformation Cp using a directed forest. The neighborhood A (cp)) of the conformation Gp is loosely defined as the set of all possible conformations dp which may result from running Algorithm I with cp as input. We denote the probability distribution induced by Algorithm| as the transition kernel k(c/,c,)) which has support N (cp)..
Definition 3.2 (Conformer Sampling Markov Chain (CSMC) p,). We define the conformer sampling
Markov chain C, as time-homogenous Markov chain over the state space Cp with transition kernel
5

---

Under review as a conference paper at ICLR 2023
defined above, where Cp is the set of valid conformations associated with given protein p as defined in Definition
The consistency of our learning procedure does not depend on the precise definition of K. However, our procedure relies on the fact that any conformer 2 E Cp can be sampled by k in finite number of steps, starting from conformer Cp- We use this fact to show that the chain &p converges to unique stationary distribution regardless of the initial conformer. (All proof's in Appendix [A.5).
Proposition 3.3. Given the CSMC o from Definition 3.2 whose transitions governed by which implicitly defined by Algorithm [ for any pair of conformers cp. G E Cp. there exists Tp < co, independent of Cp. such that P (Cp, (p) > O, where P is the Tp step transition probability.
Next, we show the existence of a unique steady state distribution of our Markov chain.
Proposition 3.4. The CSMC pp defined in Definition 3.2 uniformly ergodic if Proposition 3.3 is satisfied. Specifically there exists a unique steady state distribution Fp such that for all Cp E Op [| Φ, (-)l) I CR," where C < 0o and < are constants that depend on pp. 2, the n step Iransition probability and || is the & norm.
Given that now have the ability to draw samples from Markov chain which achieves a unique stationary distribution 7p on Cp. we shall leverage this next, in our learning framework to learn conformer invariant representations of proteins.
4 LEARNING FRAMEWORK
We shall employ a learning strategy. where the viable conformations obtained via the Markoy chain to learn function which outputs conformer invariant representations.
Let D ((x), yj) yز}} be the input data (where we have single conformer for every protein in the dataset). We shall consider single data point (æz, yJ) henceforth, where x, (V, Xp) is the input protein. We consider a supervised learning setting where yj is the associated target. We consider both classification and regression tasks.
Let Cj = {cj.} } be the set of viable conformations of protein Tj. We only consider viable confor- mations which defer only in the atomic coordinates matrix Xp. For given protein zy. we shall use Ty to denote protein z, but with X, of the original protein modified by Cj. and use S, = {æ,, to denote the set of all viable 2j. i.e. the state space.
Let f Ω + Rd. d network model such as 3D CNN, GNN, Transformer, LSTM, etc.) which takes a protein (in the form XV, Xp)) as input and outputs protein representations. The function f need output conformer invariant representations of the protein Zj. Then, simple way to obtain conformer invariant representations of protein ž (apart from using trivial functions such as a constant function or function independent of X,) is computing its expected value over all its viable conformations. We shall denote to denote a function which outputs conformer invariant representations of protein.
be any function with learnable parameters 05, (for e.g. any neural
necessarily
Flzz; 0f) Ex=a, (((;;
where (-)-) is the steady state distribution of the Markov chain associated with the protein zz- As such, F(2y; 0f) = F(af; ff) for any viable protein conformation z`J E Sy. Subsequently, to learn an optimal f (which denote by f*), wish minimize the loss, defined follows:
L(D;05,00)
Σ i Lly p ((((((;)))))
1=1
1 lim Σ i(y pl h (
(3)
where convex loss function e.g. cross entropy loss, p is some differentiable function with parameters 8º (in practice, an MLP) and lim k & f(r); can be employed as an asymp- totically unbiased and consistent estimate of where x² is the 5th sample from the Markov chain corresponding to the protein 3y

---

Under review as a conference paper at ICLR 2023
Since Equation 3 is computationally intractable, employ a surrogate for the loss given by:
N
J(D:01,0P) lim
L(u, P(L
(4)
N
Observe that in Equation 4.. the expectation over the conformations is now outside the and p func- tions while still remaining conformation invariant (however, the optimal parameters corresponding minimizing are different from minimizing T.). Following (Murphy et al., (209b]a). a). we note that, when p the identity function, Equation 44) serves as an upper bound for Equation (3) via Jensen's inequality. However, learning representations by averaging over multiple conformations (in training) still computationally expensive (back-prop and high variance issues). Next, show how MCGD (Sun et al. 2018) can be leveraged to make this tractable (more details in Appendix A.10).
Making the case for simplicity, we next show convergence properties using the full batch setting. Note that is procedure and proofs are equally applicable in the mini-batch setting with minor modifications.
Full-batch training:
We consider the data as single batch B {(11, 11),,... .... (z,, Uj).,.. (IN, yN)} and for a data point (xz, Uj), denote the corresponding steady state distributions of corresponding Markov chains using γy and the state space as Sy.
We denote all learnable parameters of p o f as 0 (0f,0") and consider 9 E θ R", m > 0 such that the function f may be non-convex for some values of O.
We consider sequence of step sizes (7/x T which satisfy:
Σ Tk to0, Σnn
72 < +00
(5)
k
and follow gradient scheme where parameters are updated as:
O2 = 0h-1 7 Zn, k 0
(6)
N V To 1 1):00 and 2 is the k - 1 sample from the
where Z
Markov chain corresponding to the data point (æz, Hj) and employ the Markov chain gradient descent procedure (Sun et al. 2018) where the loss in epoch k, k >> 0 of training is obtained by
j N Σ
and minimizes the surrogate loss given Equation (4 when is sufficiently large so that the Markov chain has converged its unique steady state distribution.
Next, we list the set of assumptions we make to ensure convergence of our conformer invariant MCGD procedure.
Assumption 4.1. We make the following assumptions:
1. For any O and x Sj, the function is differentiable Vj 2. supece../es, V e po f(x))]]) too i.c. the gradients are bounded.
3. Vuly E Sy. V0,, 0g E +, VVVo, po ((x)) Vo, po f(a))1l < Tl|0; - 02l1 for some 2
i.e., the gradients are I, Lipschtzz.
4. E. Pay
Ve po f( V E Ezz~a, 1essaal00 [p ((=))) [(,))]
Next, formally state the convergence of our optimization procedure to optimal parameters 0* which yield conformer invariant protein representations.
Proposition 4.2. Let the step sizes satisfy ()) and the function parameters f be updated (6) and Assumption 4.1 hold, then the MCGD optimization enjoys properties of almost sure convergence to the optimal 0.
The use of Markov chain gradient descent procedure to optimize the conformer invariant learning procedure optimizes the objective // in Equation (4), and thus has the following implication on how outputs should be calculated at inference time:
7

---

Under review as a conference paper at ICLR 2023
Table I: GNN(GCN), GVP-GNN, E(N) GNN Bascline vs Conformation Invariant Strategies for multiple different tasks on proteins from the ATOM3D dataset. Corresponding the metrie, 1 indicates that higher better, while 4 indicates that lower is better. Bold values indicate best results for given row. The values for GNN were obtained from Townshend et al. (220)) and for the GVP-GNN from Jing et al. (02T). Gray colored cells indicates that the augmented model outperforms the baseline model.
MCMC Augmented GNN (Ours)
MCMC Augmented GVP-GNN (Ours)
MCMC Augmented E(N) GNN (Ours)
Baseline (GNN[GCN))
Bascline (GVP GNN)
Baseline (E(N) GNN)
Task
Metric
PSR Global R.
0.755 + 0.004
0.761 + 0.004
0.845 + 0.004
0.852 0.006
0.827 + 0.004
0.852 0.004
LEP
AUROC1
0,740 + 0.010
0,672 + 0.012
0,628 + 0.055
0.704 # 0.039
0,677 + 0.014
0.714 + 0.005
LBA
RMSE 1
1,570 +: 0.025
1,519 + 0.022
1.594 + 0.073
1.455 = 0.007
1.392 : 0.001
1.384 + 0.011
MSP
A
0.621 +: 0.009
0.662 + 0.008
0.680 + 0.015
0.857 = 0.049
0.652 + 0.006
0.843 + 0.037
Inference time:
We estimate f using an empirical average of evaluated over conformations (in practice, average final layer [before softmax] representations) visited by the Markov chain:
(z)
1 k
:0)
where o.)) is the ith state visited by the Markov chain started T.j. Since the Markov chain has unique steady state distribution, J& (xyy 0*) is both asymptotically unbiased and consistent. That is lim (1) e. [(z;
k=000
5 RESULTS
We evaluate our proposed augmentation procedure multiple different tasks from the ATOM3D dataset (Townshend et al. 2020). Results are provided in Table [
TASKS ON ATOM3D DATASETS
ATOM3D (Twnshennd et al. 2020 is a unified collection of datasets concerning the 3D structure of proteins. These datasets are specifically designed provide a benchmark for ML methods which operate on 3D molecular structure, and represent a variety of important structural, functional, and engineering tasks. As part of our evaluation, we perform the (a) Protein Structure Ranking (PSR) (b) Ligand Efficacy Prediction (LEP) (c) Protein Mutation Stability Prediction (MSP) (d) Ligand Binding Affinity (LBA). We describe each the tasks detail in Appendix[A.] For all the datasets and tasks we report the same metrics as proposed by the authors in Townshend et al. (2020).
MODEL, BASELINES AND DISCUSSION
We endow the vector gated GVP-GNI model (Jing et al. 2021 GNN using GCN Kipp & Welling 2016) layers and the E(n) GNN /Satorass et al. (2021) with conditional transformations from our proposed MCMC method. It is important to note that when the positions of the atoms altered the protein graph which are inputs to the base encoders are changed appropriately in every epoch. As an ablation study, we also provide a strategy where all the transformations are created from "gold standard" X, rather than via the MCMC method, i.c., the MCMC is restarted at every epoch during training the Appendix.
Looking at Table [] we note that models augmented with our proposed conformer invariance frame- work, in general, outperforms the baseline models in multiple tasks for which conformer invariance is requirement. In the LEP task, the proposed addition outperforms the baseline models for the E(n) GNN and the GVP-GNN, but not for the GNN(GCN). The LEP result for the GNN(GCN) an oddity here, where GNN which doesn't it explicitly incorporate any atomic positional information outperforms all other models with information about atomic coordinates.
Additional results and ablation studies are presented in Appendix A.9
COMPUTATIONAL COMPLEXITY & SCALABILITY:
The directed forest for every protein is computed only once preprocessing step (a fixed Rz 5-10 minutes per dataset). At every epoch, for a given protein, we select only one among all its constituent amino acids & perform sequence of 3 x 3 matrix multiplications to obtain new conformer. If k denotes the number atoms in the side chain of an amino acid, and height a directed tree is in the order of log(k), the computational complexity obtain a conformation is O log k) (where k zs 15
8

---

Under review as a conference paper at ICLR 2023
in avg in our datasets), & this can be run in parallel for all proteins in mini-batch Molprobity (run in parallel again, currently runs extemally on a web server via a command line call) adds 2s delay per minibatch (we could reduce to milliseconds if we integrated Molprobity into our code). The rejection rate from Molprobity is also very small (less than reject in average across 100 molprobity calls). As an example, training (100 epochs) GVP-GNN- MSP requires Ra 45 min, and LBA requires = 34 min. In both cases our method requires an additional 4 min overhead. A complete training time table n " "Reidi EDJ WORGGP for all models & datasets can be found in Appendix A.11
Here, we present high level summary of works related to ours. A more detailed comparison to the works mentioned below and others) are presented in Appendix A.6
Group Equivariant Neural Networks:
Group equivariant and invariant neural networks (Cohen
& Welling. 2016: Lenssen et al. 2018 Kondor & Trivedi 2018 Finzi et al. 2020 Hutchinson
et 2022 |uuhs al.] [020] 2021 Dehmamy |l.]|2000 help capture discrete and continuous
2021
symmetries of elements (e.g. images) by introducing group theoretic constraints as inductive biases
in the neural network. While group equivariant neural network capture global symmetries, local symmetries of manifold spaces can be captured via gauge equivariant networks (Cohen et al. 2019]
De Haan et al. 2020). (Lenssen et al.) 2018: Gerken et al.] 202T; Bronstein et al., 2021) provide complete review of the theoretical aspects and applications of group equivariant neural networks.
Graph Neural Networks: Graph Neural Networks (Kipf & Welling, 2016 Hamilton et al. 2017 Battaglia et al. 2018 Xu et al.) 2018) have gained renewed focus over the past few years and have found applications in recommender systems, biology, chemistry, and many other real world problems which can be formulated as graphs and currently serve the state of the art in majority of node and graph classification/ regression tasks. Graph neural networks work on the principles of permutation equivariance/ invariances (also groups) and have exploited a message passing framework to learn powerful and expressive representations of nodes/ graphs.
These networks combine continuous symmetries (lie groups) with permutation equivariances and has found applications with resounding success on small molecules (Anderson al., [219 Kllicpera et al. 2020; [Satorras et al.] [021 Batzner et al., 2021) which exhibit rigid body characteristics. Employing (Farina & Slade 2021). would make the neural network excessively invariant and allow the protein to be more flexible (allows unviable conformations) than it truly is. More recently, they have also been applied to learning representations of proteins, which is discussed below.
Group Equivariant Graph Neural Networks:
Monte Carlo and MCMC Methods for Sampling Protein Conformations: There has been lot of prior work- (Boomsma al., 2013 Olsson et al.) [2013 Antonov et al., [016 Irbäck Mohanty 2006, Vitalis& & Pappul 2009) which sample protein conformations by either defining/inheriting a distribution over the conformations (majorly based on the properties of the bonds, etc.) These
methods can seamlessly be plugged into our framework as long as the MCMC is ergodic.
Protein representation has gained a lot of attention especially with the tremendous successes of Alphafold and Alphafold2. A variety of neural network architectures including 3D CNNs, LSTM's and Transformers (treating the protein as a sequence) as well as graph neural networks have been employed to exploit the rigid body symmetries
Neural Networks for Representation Learning of Proteins:
of proteins (Karimi et al. 2019 Pagès et al. 2019 Ingraham et al. 2019 Strokach et al. 2020
Baldassarre al. 2021 Hermosilla al. 202 Jing et al. 2020: 2021)
Generative Models:
Conformation generation models (Mansimov et al. 2019 Simm et al. 2020
Ganea et al. 2021: Xuu et al. 22221ba Shi et al.) [021] Luo al. [2021] have also gained attention recently with the goal to predict 3d structure of molecules given their 2d structure our objective in this work is very different, but be used to improve predictions of the aforementioned models.
CONCLUSIONS
This work addresses the limitations of current protein representation learning methods which are unable to learn conformer invariant representations - and hence unable to capture the inherent flexibility present in protein side chains pertinent to many downstream tasks. To address these, we introduced conditional transformations to capture protein structure, while respecting the restrictions posed by constraints on dihedral (torsion) angles and steric repulsions between atoms. Subsequently, we introduced a Markov chain Monte Carlo based framework to learn representations that are invariant to these conditional transformations.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. arXiv preprint arXiv: 990600015, 2019.
Lubomir Antonov, Simon Olsson, Wouter Boomsma, and Thomas Hamelryck. Bayesian inference of protein ensembles from saxs data. Physical Chemistry Chemical Physics, 18(8)5583,-5888, 2016.
Federico Baldassarre, David Menéndez Hurtado, Arne Hlofsson, and Hossein Azizpour. Graphqa: protein model quality assessment using graph convolutional networks. Bioinformatics, 37(3): 360-366, 2021.
Peter Battaglia, Jessica B Hamrick, Victor Bapst. Alvaro Sanchez-Gonzalez Vinicius Zambaldi. Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 1806.00060000 2018.
Simon Batzner, Tess E Smidt, Lixin Sun, Jonathan Mailoa. Mordechai Kornbluth, Nicola Molinari. and Boris Kozinsky. Se (3vequrvari graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXiv:2101.03164, 2021.
Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic acids research, 28(1):
235-242, 2000.
Wouter Boomsma, Jes Frellsen, Tim Harder, Sandro Bottaro, Kristoffer E. Johansson, Pengfei Tian, Kasper Stovgaard, Christian Andrectta, Simon Olsson, Jan Valentin, et al. Phaistos: A framework for markov chain monte carlo simulation and inference of protein structure. Journal of computational chemistry, 34(19):1677-1705 2013.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2 13478, 2021.
Vincent B Chen, W Bryan Arendall, Jeffrey J Headd, Daniel Keedy, Robert M Immormino, Gary J Kapral, Laura W Murray, Jane Richardson, and David C Richardson. Molprobity: all-atom structure validation for macromolecular crystallography. Acta Crystallographica Section D: Biological Crystallography, 66(1):12-21, 2010.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International
conernaeeecnlnofeeeeneccc conference
machine learning, pp. 2990-2999. PMLR, 2016.
Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral cnn. In International Conference on Machine Learning, pp. 1321- 1330. PMLR, 2019.
Ian W Davis, Andrew Leaver-Fay, Vincent B Jeremy N Block, Gary Kapral, Xueyi Wang, Laura W Murray, W Bryan Arendall III, Jack Snoeyink, Jane Richardson, et al. Molprobity: all-atom contacts and structure validation for proteins and nucleic acids. Nucleic acids research, 35(suppl_2):W311 W383, 2007.
Pim De Haan, Maurice Weiler, Taco Cohen, and Max Welling. Gauge equivariant mesh cnns: Anisotropic convolutions geometric graphs. arXiv preprint arXiv:2003.05422 2020.
Nima Dehmamy, Yanchen Liu, Robin Walters, and Rose Yu. Lie algebra convolutional neural networks with automatic symmetry extraction. 2020.
Francesco Farina and Emma Slade. Symmetry-drivven graph neural networks. arXiv preprint arXiv:2105.14058, 2021.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. International Conference Machine Learning, pp. 3165-3176. PMLR, 2020.
10

---

Under review as a conference paper at ICLR 2023
Richard A Friesner, Jay L Banks, Robert B Murphy, Thomas A Halgren, Jasna Klicic, Daniel T Mainz, Matthew P Repasky, Eric H Knoll, Mee Shelley, Jason K Perry, et al. Glide: new approach for rapid, accurate docking and scoring. 1. method and assessment of docking accuracy. Journal of 2004.
medicinal chemistry,
Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. Se (3)ttransfrr::::: 3d roto-translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Fabian B Fuchs, Edward Wagstaff, Justas Dauparas, and Ingmar Posner. Iterative (3)atrssffrrmers. arXiv preprint arXiv:2102 13419, 2021.
VL Fushchich and RZ. Zhdanov. Conditional symmetry and reduction of partial differential equations. Ukrainian Mathematical Journal, 44(7)8875-886, 1992.
Pablo Gainza, Kyle E Roberts, and Bruce Donald. Protein design using continuous rotamers. PLoS computational biology, 8 s(1):e1002335 2012.
Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances Neural Information Processing Systems, 34, 2021.
Jan E Gerken, Jimmy Aronsson, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, Christoffer Petersson, and Daniel Persson. Geometrie deep learning and equivariant neural networks, arXiv preprint arXiv:2105.13900 2021.
William Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 1025-1035, 2017.
Tim Harder, Wouter Boomsma, Martin Paluszewski, Jes Frellsen, Kristoffer E Johansson, and Thomas Hamelryck. Beyond rotamers: a generative, probabilistic model of side chains in proteins. BMC bioinformatics, 11(1):1-13, 2010.
Pedro Hermosilla, Marco Schäfer, Matej Lang, Gloria Fackelmann, Pere-Pau Vázquez, Barbora Kozlikova, Michael Krone, Tobias Ritschel, and Timo Ropinski. Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. In International Conference on Learning Representations, 2021.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning graphs. Advances in neural information processing systems, 33:22118-22133, 2020.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: equivariant self-attention for lie groups. International Conference Machine Learning, pp. 4533-4543, PMLR, 2021.
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. 2019.
Anders Irbäck and Sandipan Mohanty. Profasi: a monte carlo simulation package for protein folding and aggregation. Journal of computational chemistry, (1113))51-1155555 2006.
Justina Jankauskaite, Brian Manénez--Grría, Justas Dapkünas, Juan Fernández- Recio, and Iain H Moal. Skempi 2.0: updated benchmark of changes in protein protein binding energy, kinetics and thermodynamics upon mutation. Bioinformatics, 35(3)462-446, 2019.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323-2332. PMLR, 2018.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint 1, 2020.
11

---

Under review as a conference paper at ICLR 2023
Bowen Jing. Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint (XXi:22006008433 2021.
Anthony Joseph. The theory of conditional invariance. i. Proceedings of the Royal Society of London. Series Mathematical and Physical Sciences, 305(1482):405-444 427, 1968.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool. Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, pp- 2021.
Olav Kallenberg. Foundations of modern probability. Springer Science Business Media, 2006.
Mostafa Karimi, Di Wu, Zhangyang Wang, and Yang Shen. Deepaffinity: interpretable deep learning of compound parpoonudd--oteiin affinity through unified recurrent and convolutional neural networks. Bioinformatics, 35(18):3133) ()3383)3333333 2019.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: 02900,, 2016.
Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03 2020.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning, pp. 2747-2755. PMLR, 2018
Andriy Kryshtafovych, Torsten Schwede, Maya Topf, Krzysztof Fidelis, and John Moult. Critical assessment of methods of protein structure prediction (casp)-round round xiii. Proteins: Structure, Function, and Bioinformatics, 87(12):1011-15 2019.
Jan Eric Lenssen, Matthias Fey, and Pascal Libuschewski. Group equivariant capsule networks. arXiv preprint arXiv: 180.000066.0 2018
Zhihai Liu, Yan Li, Li Han, Jie Li, Jie Liu, Zhixiong Zhao, Wei Nie, Yuchen Liu, and Renxiao Wang. Pdb-wide collection of binding data: current status of the pdbbind database. Bioinformatics, 31)):: 405-412, 2015.
Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular conformation via dynamic graph score matching. Advances in Neural Information Processing Systems, 34, 2021.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using a deep generative graph neural network. Scientific reports, 911: 2019.
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09900 2018
Jonathan Masci, Davide Boscaini, Michael Bronstein, and Pierre Vandergheynst. Geodesic con- volutional neural networks riemannian manifolds. In Proceedings of the IEEE international conference on computer vision workshops, pp. 37-45, 2015.
Sean P Meyn and Richard I. Tweedie. Markov chains and stochastic stability. Springer Science & Business Media, 2012.
Zhichao Miao and Yang Cao. Quantifying side-chain conformational variations in protein structure. Scientific reports, 6(1):1-10, 2016.
Christopher Morris, Nils M Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. arXiv preprint arXiv:2007.08663. 2020.
Ryan Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Relational pooling for graph representations. In International Conference on Machine Learning, pp. 4673. PMLR, 2019a.
12

---

Under review as a conference paper at ICLR 2023
Ryan L. Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Janossy pooling: Learning deep permutation permutation-invariant invariant functions for variable-size inputs. In International Confer- ence on Learning Representations, 2019b. URL /ttpttttss////eeeeeeeee..eeee
BJluy22cFF
Simon Olsson, Jes Frellsen, Wouter Boomsma, Kanti Mardia, and Thomas Hamelryck. Inference of structure ensembles of flexible biomolecules from sparse, averaged data. PloS one, 8(11): e79439, 2013.
Peter Olver and Evgenii M Vorob ev. Nonclassical and conditional symmetries. CRC handbook of Lie group analysis of differential equations, 3:291-328, 1996.
Guillaume Pagès, Benoit Charmettant, and Sergei Grudinin. Protein model quality assessment using 3d oriented convolutional neural networks. Bioinformatics, 35(18):3313-3315 2019.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. (n) equivariant graph neural networks. arXiv preprint arXiv:2102.09844, 2021.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, pp. 9558-9568. PMLR, 2021.
Gregor NC: Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hernández-LLobto. Symmetry- aware actor-critic for 3d molecular design. arXiv preprint 0rXiv:2011.127477 2020.
Alexey Strokach, David Becerra, Carles Corbi- Verge, Albert Perez-Riba, and Philip M Kim. Fast and flexible protein design using deep graph neural networks. Cell Systems, 11 11(4):402-411 2020.
Tao Sun, Yuejiao Sun, and Wotao Yin. On markov chain gradient descent. Advances in neural information processing systems, 31, 2018.
Jonathan Taylor, Allan D Jepson, and Kiriakos N Kutulakos. Non-rigid structure from locally-rigid motion. In 2010 JEEE Computer Society Conference on Computer Vision and Pattern Recognition, pp. 2761-2768 IEEE, 2010.
Raphael JL Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Brandon Anderson, Stephan Eismann, Risi Kondor, et al. Atom3d: Tasks molecules in three dimensions. arXiv preprint arXiv:2012.04000 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Andreas Vitalis and Rohit V Pappu. Methods for monte carlo simulations of biomacromolccules. Annual reports in computational chemistry, 3649-76,220 2009.
Renxiao Wang, Xueliang Fang, Yipin Lu, and Shaomeng Wang. The pdbbind database: Collection of binding affinities for protein- ligand complexes with known dine-dimmnsional structures. Journal of medicinal chemistry, 47 47122):977729880 2004.
Christopher Williams, Jeffrey J Headd, Nigel W Moriarty, Michael G Prisant, Lizbeth L Videau, Lindsay N Deis, Vishal Verma, Daniel A Keedy, Bradley Hintze, Vincent B Chen, et al. Molpro- bity: More and better reference data for improved all-atom structure validation. Protein Science, 27(1):293-315 2018.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint 18Xiv:1810.00826. 2018
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2 /0240, 2021a.
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel programming. In International Conference on Machine Learning, pp. 11537-11547. PMLR, 2021b.
13

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 COMMONLY USED ACRONYMS IN THE PAPER
1. G -group
2. S₂ Set of transformations unique element
3. & Transformation action which transforms the element
4. Number of atoms in the protein under consideration
5. n Number of amino acids in the protein under consideration
6. X, Matrix of Scalar features associated with atoms in the protein p
7. Xp Matrix of atomic coordinates associated with atoms in the protein
8. Cp Set of conformations of protein
9. SO(3) Group of all rotations 3D space
A.2 GROUP THEORY PRELIMINARIES
Definition A.I (Group). A group is G equipped with binary operation :: GG G G obeying the following axioms:
for all 91. E G. 91
(closure).
for all 91, 92, Da E G. g (g2 ga) = (g) g2) 93 (associativity).
there is a unique G such that e g
for all G (identity).
for all G there exists g E G such that g g.---= --
g (inverse).
Definition A.2 (Group invariant functions). Let G be group acting on vector space V. We say that a function f: V + G-invariant if flg z) = f()) Vx E V,g E G.
Definition A.3 ((Left) Group Action). For group G with identity element e, and is set, (left) group action of G' on X is function G x X + X that satisfies the following two conditions:
1. Identity: (e,, x) = x, Vx E
2. Compatibility: agg, ah, x)) = a(gh,
We will use a short hand of g x for (g, when the action being considered is clear from context.
A.3 EXAMPLE DEMONSTRATING NON GROUP STRUCTURE OF A SET OF PROTEIN CONFORMATIONS
Consider X to be protein with atoms and m amino acids and let the set of viable conformations of X as X2,.X2.... X2,.. X,,... ".....Xp.- Let X, be the conformation available in our dataset.
In each step of the transition, only few atoms <<< n) (atoms in single amino acid of the entire protein- where the protein made up of multiple hundreds of amino acids (m) in general) are subjected to an action from the SO(3) group here. The positions of all other atoms (outside that amino acid) in the side chain remain unaltered. So, the xxx3x3 matrix most of the 3x3 entries are the identity matrix of 3 dimension.
Let Ti R38333 where i E {1,2,..., .P} be the transformation which yields conformation X, from Xi.
Now consider T? (subscript of since we start conformation X1) which takes X1 to X₂ and 2 which takes X, to X3 where T? and 73, do not act the same amino acid in the protein. Now, however, consider the case where performing T? and T (or the other way around) sequentially would result in where atoms in different amino acids would overlap (or come too
14

---

Under review as a conference paper at ICLR 2023
close each other causing steric repulsion) therefore resulting in non viable conformation Le. a composition of T and T acting on X1. would not be present in the set of all allowed transformations Ti [T) ................................ Ti.... TT). Therefore the set I is not closed and doesn't 't form group.
Also for two different conformations X1 and X2. their allowed transformation sets Tí, T2 will not identical (in the above example 7 72). It also easy to construct case where U T; is not a group and all actions from this group acting on any given X; will not necessarily result in viable/ valid conformation.
A.4 FLEXIBILITY ALLOWED BY OUR PROPOSED MODEL
Our proposed model allows every node in the directed be rotated about its parents. For example, for the side chain shown in Figure T (Main Paper), the allowed flexibility is Figure
R 6
SCB
SC5
R5
SC3
SC2
R
R4
SC1
R2
R1
&C
Figure 3: Maximum flexibility allowed by our candidate sampling process when the group associated with every node in the tree is SO(3) (the special orthogonal group in 3 dimensions). While not candidate is likely to be accepted such candidate generation process provides the flexibility for every node to be rotated about its immediate parent while preserving hond lengths.
A.5 PROOFS OF PROPOSITIONS
First, we restate and prove Proposition
Proposition A.4. Given the CSMC & from Definition 3.2 whose transitions are governed by k which is implicitly defined by Algorithm] 1 as described above. For any pair of conformers Cp. c E Cp. there exists Tp co, independent of Cp. such that pTy (cp. (b) 0, where PT is the Tp step transition probability.
Proof. Proof by construction. We prove the proposition by showing that one can construct a path (c = Cp....., , c)) such that of E Cp and x(c ((c(tt))((())) > O, for all 0 < i < t, and t < Tp The trivial case where cp = & is proved since every group contains the identity element - sampling the identity element for every node the directed tree yields the same conformation. Since we consider only non backbone transforming conformations, for the non trivial case, maximum of m 4n atoms can differ in positions between any two conformations where m is the number of atoms in the protein and n is the the number of amino acids in the protein n. Both n are finite and we are dealing with continuous conformers (and continuous group actions about every node groups are closed under their associated binary action, and SO(3) is path connected). So we can traverse between conformers (until we reach the desired conformer) sequentially in finite number of steps, by using the constructed directed forest selecting amino acid (which doesn't 't violate the viability). fixing the positions of all other amino acids in the protein and rotating the side chain atoms in single conformer to the final desired state. While this process may result in some side chains
15

---

Under review as a conference paper at ICLR 2023
being visited multiple times (due to viability constraints), considering continuous conformers and the SO(3) group (which is path connected) ensures we will never reach state of deadlock. The second condition is satisfied because every group action has inverse and only use transformations from SO(3) for every node in the directed trees.
Next, we restate and and prove Proposition 3.4
Proposition A.5. The CSMC p defined in Definition 3.2 is uniformly ergodic Proposition 3.3 is satisfied. Specifically there exists a unique steady stale distribution Fp such that for all cp E Cp. are constants that depend on Dp. P the
I| by (Cp:-) - Ty(-)] CRR, where C < oo and n step transition probability and || is the & norm.
Proof. By Proposition 3.3 Φ satisfies Doeblin's condition as defined in page 396 of Meyn & Tweedie (2012) which states that for Cp: Sp.C. Cp. P'r 4. (Cp. (y)) for some > 0 The uniform
ergodicity then holds due to Theorems 16.2.3 and 16.2.1 from Meyn & Tweedie (2012).
Next, we restate and prove Proposition 4.2 We also restate the required assumptions for the proposition.
Assumption A.6. We make the following assumptions:
1. For any E θ and r Sj, the function is differentiable vjj
2. su (lVo a [()))] < +oo i.e. the gradients are bounded.
3. Vu Sj. 10,, 02 O, VV, po (x Vo, po f(8))l L110 - 02ll for some i.e., the gradients are L- L-piscchit.
4. E Po fx
Bejwa, f(
Proposition A.7. Let the step sizes satisfy 5 and the function parameters be updated as O and Assumption 4 hold, then the MCGD optimization enjoys properties of almost convergence to the optimal 8.
Proof. Given that each protein has an associated time homogeneous Markov Chain with a unique steady state, independent of other proteins, the set of proteins in a batch also form Markov chain with unique steady state. We then leverage Corollary (Page 12) of Sun et al. (2018) along with Proposition 4.2 to ensure almost sure convergence to the optimal 0.
A.6 EXTENDED RELATED WORK
Here, we elaborate on the related works section in Section 6
Group Equivariant and Invariant Neural Networks: Group equivariant and invariant neu- ral networks (Cohen Welling 2016 Lenssen et al., 2018 Kondor & Trivedi, 2018 Finzi al.) 2020 Hutchinson et al., [021] Fuchs et al. 2020 2021 Dehmamy et al. (0020) help capture discrete and continuous groups symmetries of elements (e.g. images, point clouds). In this work, we learn representations which are invariant to symmetries which are not just groups, but to input dependent sets transformations. To the best of our knowledge, our work is the first to consider conditional (input dependent) invariances.
Prior works on learning invariant models have leveraged Monte Carlo procedures (Finzi et al. 2020:Murphy al. 2019b) to learn be invariant to transformations of the input. Alternatively, our work constructs a Markov Chain with a unique steady state and leverages MCGD (Sun et al. 2018) to make it computationally tractable. Secondly, the aformentioned approaches are limited to being invariant to transformations from any specified Lie group with a surjective exponential map/ permutation group, while our work is not limited groups. Thirdly, our theory ensures that every example/ object in the dataset can have a different set of input dependent, conditional transformations
We note that this is simplified version of the actual statement which is defined the r-algebra over Cp denoted by (C,)). Our proof holds when de E (C,))
16

---

Under review as a conference paper at ICLR 2023
and in fact can be seen generalization of the above works. In fact, the ablation study that perform (Results provided in Appendix [A.9. Table 4) uses a Monte Carlo estimator and our MCMC procedure yields better performance than the Monte Carlo estimator.
While group equívariant neural network capture global symmetries, local symmetries of manifold spaces can be captured via gauge equivariant networks (Cohen et al. 2019 De Haan al., 2020)) Gauge symmetries require the manifold to be smooth - which is not the case for proteins. Moreover, different proteins have different sets of viable conformations, which would not be able to captured by standard gauge equivariant neural networks.
(Lenssen et al.) 2018 Gerken et al., 2021] Bronstein et al.) (02T) provide a complete review of the theoretical aspects and a wide variety of applications of group equivariant neural networks. A more comprehensive theoretical analysis of input dependent conditionally invariant neural networks is planned for future work.
Graph Neural Networks: Graph Neural Networks (GNNs) (Kipf & Welling 2016 Hamilton et al.
2017 Battaglia et al. 2018: Xu et al. (018)) have gained renewed focus over the past few years and have found applications in recommender systems, biology, chemistry, and many other real world problems which can be formulated as graphs and currently serve as the state of the art in majority of node and graph classification/ regression tasks. Graph neural networks work on the principles of permutation equivariance/ invariances (Murphy et al. 2019a) (also groups) and have exploited a message passing framework to learn powerful and expressive representations of nodes/ graphs
Both molecular graphs (both small molecules and macromolecules) as well as graphs based on intra molecular distances have been used with GNNs, to achieve state of the art for many molecular datasets and tasks (Hu et al.) [2020] Morris et al. 2020). Here, we leverage three graph based neural networks as baselines for our model.
Group Equivariant Graph Neural Networks: Group Equivariant GNNs combine contin- uous symmetries (lie groups such as SE(3), E(3), SO(3)) with permutation equivariances and has found applications with resounding success on small and large molecules (Anderson et al. 2019 Klicpera al., 2020 Satorras et al. 202T Batzner et al. (02T))
However, the methods are only able to capture rigid body characteristics of molecules and while capturing the above lic group symmetries is also able to capture input depedent transformations. Employing (Farina & Slade (02T). would make the neural network excessively invariant and allow the protein be more flexible (allows unviable conformations) than truly is. More recently, they have also been applied to learning representations of proteins, which is discussed below.
Monte Carlo and MCMC Methods for Sampling Protein Conformations: There has been a lot of prior work - (Boomsma et al. 2013) Olsson et al. 2013; Antonov et al. 2016 Irbäck & Mohanty 2006; Vitalis & Pappul 2009) which sample protein conformations by internal coordinate transformations. However, The goals of the existing MCMC methods are significantly different compared to ours. The existing methods define/inherit distribution over the conformations (majorly based on the properties of the bonds, etc.) and then aim to sample highly probable conformations from this distribution. Our method, is much simpler and only requires that the chain being used is ergodic and is invariant to the actual form of the distribution. We note that the existing Markov chains can seamlessly be used as drop-in replacements sample conformations as part of our framework as long as the MCMC is ergodic. We consider studying the impact of different MCMC methods (which sample from different distributions) and their influence on the performance in different tasks as important future work.
Neural Networks for Representation Learning of Proteins: Protein representation has gained a lot of attention especially with the tremendous successes of Alphafold and Alphafold2. A variety of neural network architectures including 3D CNNs, LSTM's and Transformers (treating the protein as a sequence) as well as graph neural networks have been employed to exploit the rigid body
symmetries of proteins (Karimi al. 2019 Pagès al. 2019 Ingraham et al. 2019 (Srooaach et al.)
2020 Baldassarre et al. 2021 1] Hermosilla et al. 2021 Jing et al. 2020 (021).
In this work, while we use GNNs and Group Invariant GNNs part of the model, we note that
can equally replace them with CNNs, LSTMs, Transformers and other models used for proteins without any change in the underlying theory.
Tree Construction Method for Molecules: Jin et al. 2018 in their work JTVAE, propose
17

---

Under review as a conference paper at ICLR 2023
procedure construct tree for molecules. We note that JTVAE is more general purpose approach for constructing trees from molecules and may not be the best suited for proteins where the backbone atoms are largely observed (experimentally) to be rigid.
Protein conformation generation models (Mansimov et al.] 2019 Simm et al. 2020; Ganea et al. 2021 Xu et al. 202 a Shi et al. [022]] Luo et al. (021)) have also recently gained attend where the goal of the model is predict 3d structure of molecules given input 2d structure our objective in this work is completely different, but can be used to improve predictions of the aforementioned models.
Generative Models:
While our model is explicitly not a generative model, our framework can be leveraged towards generative modeling with the help of tools such as noise outsourcing (Chapter 6) (Kallenberg 2006
and we see this as important future work.
Non-Rigid Body Dynamies: Non-Rigid Body Dynamics of objects has long been studied both by physicists and in the fields of computer vision to understand and capture the geometric deformations of objects (Taylor et al., (2000) Masci et al., 2015). To the best of our knowledge, there exists no prior work in deep learning which captures the non rigidity of protein molecules (which cannot be modeled O* manifolds). As important future work, we would like to study the impact of leveraging input dependent conditional invariances for modeling other geometric objects (which are C* manifolds) as well as images and robotics (e.g. the symmetries for humanoid different from that of a tractor).
Unrelated Work with similar names: Non classical and conditional symmetries of solutions to ODE's, PDE's have been discussed in the past these works while they share a similar title, have very little in common as we are not dealing with jet spaces or manifolds (Joseph) 1968; Fushchich & Zhdanov 1992; Olver & Vorob' ev 1996).
A.7 DETAILS ABOUT DATASETS AND TASKS
In this section, we describe briefly each of the datasets (and their associated tasks). Information about the splits and license information is provided in Table
This task utilizes data from the structural models submitted to the Critical Assessment of Structure Prediction competition (CASP Kryshtafovych al.) (2019) a blind protein structure prediction competition) to rank protein structures from the experimentally determined structure of the protein. The problem is formulated as a regression task, where we predict the global distance test of each structural model from the experimentally determined structure. As prescribed by the dataset authors, the dataset is split by competition years.
PSR:
MSP: The goal of this task is to identify mutations that stabilize a protein's interactions which forms an important step towards the design of new proteins. This task is significant as probing mutations experimentally techniques
labor-intensive Atom3D (Townshend al.) 2020 derives this dataset by collecting single-point mutations from the SKEMPI database (Jankauskaité et al. 2019) and model each mutation into the structure to produce mutated structure. The learning problem is then formulated as a binary classification task where the goal is to predict whether the stability of the complex increases as a result of the mutation. We employ the same splits as suggested by the dataset authors wherein the protein complexes are split such that no protein in the test dataset has more than 30% sequence identity with any protein in the training dataset.
This task deals with the problem of predicting the strength (affinity) of a candidate drug molecule's 's interaction with a target protein. The dataset is constructed using the PDBBind database (Wang et al.) 2004 Liu et al.) 2015). curated database containing protein-ligand complexes from the PDB and their corresponding binding strengths (affinities). The task is formulated as a regression task with the goal to predict py log (K), where K is the binding affinity in Molar units. The splits are created such that no protein in the test dataset has more than 30% sequence identity with any protein in the training dataset.
LBA:
The shape of protein impacts whether a protein is in an on or off state which plays an important role in predicting the shape protein will favor during drug design. This dataset is obtained by curating proteins from several families with both "active" and "inactive" state structures, and model 527 small molecules with known activating or inactivating function using the program Glide (Friesner
LEP:
18

---

Under review as a conference paper at ICLR 2023
al. 2004). The task is formulated binary classification task where the goal is to predict whether a molecule bound to the structures will be an activator the protein's 's function or not. We use the same split recommended by the ATOM3D authors.
Table 2: Summary of the datasets
Task
Train
Val
Test
Original Source
License
MSP
247
SKEMPI (Jankuukkiil el
2019
Creative Commons CC-BY
Creative Commons ComCommercial-ooDere
LBA
3563
448
452
PDBBlind Wang et al.
2004
(CC-BY-NC-A)
LEP
304
110
104
PDB Berman al.
2000
Creative Commons CC-BY
PSR
25400
2800
16099
CASP Kryshtatovych al.
2019
Creative Commons CC-BY
A.8
EXPERIMENTAL SETUP
The code for the baseline models (GVP-GNN (Jing et al. 2021), B(N) GNN (Satorras et al. 2021 and GNN(GCN) (Townshend et al. 2020 Kipf & Welling 2016)) were used as provided by the authors (licenses as dictated by the code authors). Our conformer invariance implementation is in Py PyTorch using Python 3.8. We also leverage networkx to create the directed forests. For all three models we tune the hyperparameters learning rate (E {0.1, 0.01, 0.001, 0.0001 }) and mini batch size (E (4, 8, 16, 32, 64)). For the E(n) GNN model since there have been previous models for the aforementioned protein tasks we also tune the number of GNN layers (E {4, 5, 6, 7)) as a hyper parameter. The experiments were all performed on Tesla V100 GPU's. For more details refer the code provided.
A.9 ADDITIONAL RESULTS
In Table 3 we present the results, including additional datasets and tasks, than presented in the main paper. In the LEP task, the proposed addition outperforms the baseline models for the E(n) GNN and the GVP-GNN, but not for the GNN(GCN), The LEP result for the GNN(GCN) is an oddity here, where a GNN which doesn't 't incorporate any rigid body transformations outperforms all other models.
In Table present ablation study, where provide strategy where all the transformations created from "gold standard" X, rather than via the MCMC method, i.c., the MCMC is restarted at every epoch during training. From the table, we note that the MCMC method tends outperform the non MCMC method, which can be attributed to the guarantees provides to the learning framework.
Table 3: GNN(GCN), GVP-GNN, E(N) GNN Baseline vs Conformation Invariant Strategies for multiple different tasks on proteins from the ATOM3D dataset. Corresponding to the metric, indicates that higher better, while J indicates that lower better. Bold values indicate best results for given row. The values for GNN were obtained from Townshend et al. (2020) and for the GVP- GNN from Jing et al. (2021). Gray colored cells indicates that the augmented model outperforms the baseline model.
MCMC Augmented GNN (Ours)
MCMC Augmented GVP-GNN (Ours)
MCMC Augmented E(N) GNN (Ours)
Baseline (GNN[GCNT)
Baseline (GVP GNN)
Baseline (E(N) GNN)
Metric
PSR
Global R,
0.755 + 0.004
0.761 0.004
0.845 0.004
0.852 0.006
0.827 0.004
0.852 0.004
LEP
AUROC1
0.740 +: 0.010
0.672 0.012
0.628 0.055
0.704 0.039
0.677 0.014
0.714 0.005
LBA
RMSE J
1.570 +: 0.025
1.519 0.022
1.594 + 0.073
1.435 0.007
1.392 # 0.001
1.384 0.011
MSP
AUROC1
0.621 +: 0.009
0.662 0.008
0,680 + 0.015
0.857 0.049
0.652
0.006
0.843 + 0.037
A.
CASE AGAINST USING AVERAGE REPRESENTATIONS IN TRAINING AND FOR REQUIRING SINGLE INVARIANT REPRESENTATION
The case for the requirement of single conformer invariant representations - can seen from the fact that different protein conformations, say X1, X₂ of the same protein X, may be seen during
19

---

Under review as a conference paper at ICLR 2023
Table GVP-GNN, GNN (GCN) Baseline Ablation vs Conformation Invariant Strategies for four different tasks proteins from the ATOM3D ( (Townshend et al. 2020) dataset. Corresponding to the metric, indicates that higher better, while 4 indicates that lower is better. Bold values indicate best results for a given row. Gray colored cells indicates that the augmented model outperforms the baseline model.
Baseline
Non MCMC
MCMC Augmented GVP-GNN
Baseline (GNN)
Non MCMC
MCMC
Angmented GVP-GNN
Augmented
Augmented
GNN)
Jing et al. 2021
Townshend ot
2020
(Ablation)
(Ours)
(Ablation)
(Ours)
PSR
Global R, 1
0.845
0.004
0.806 0.011
0.852 0.006
0.755
0.004
0.766
0.001
0.761
0.004
AUROCT
0.628
0.055
0.740
0.012
LBA
RMSE
1,594 + 0.073
1.635
0.007
1.435 0.007
1.570
0.025
1,520
0.022
1,519 0.022
MSP
AUROC1
0.680 = 0.015
0.799 0.016
0.857
0.049
0.621 = 0.009
0.610 0.021
0.662 0.008
train and test phases. Without conformer invariant representations this may lead different representations and therefore different predictions for the same protein (bad).
One may argue, that an approximate conformer invariant neural network can be learned by averaging representations of multiple Monte Carlo conformations during training. However, this is computationally expensive this would need to back-propagate and update parameters for multiple conformations for every protein in every epoch bad as this leads to an exponential overhead. On the other hand, procedure with MCGD, in every epoch uses only conformation from Markov Chain and still ensures convergence to optimal parameters. during inference still do average representations over multiple conformations (from Markov chain) to output conformer invariant representations. which again due to MCGD training procedure and the unique steady state of our Markov Chain, ensures confomer invariant representations are achieved which yield better performance than current state of the art on multiple datasets and tasks. As stated in the main paper, our framework can work well with other Markov chains as well. We chose the proposed chain purely because it is simpler than the existing methods (in that it doesn't 't require a distribution over conformations to be assumed) and as such is easier to sample from.
is important to note that,
A.11 TRAINING TIME
In T we present the training time (for 100 epochs) for each of the baseline models as well as well for models with our proposed conformer invariance framework addition. From the table, we note that, on an average the increase in training time over 100 epochs 3-4 minutes which is negligible in comparison to the training time (without the proposed framework addition).
Table 5: Training Time (in minutes for 100 epochs rounded to closest integer) GNN(GCN). GVP-GNN, E(N) GNN Baseline vs Conformation Invariant Strategies for multiple different tasks on proteins from the ATOM3D dataset. Lower is better.
MCMC Augmented
MCMC Augmented GVP-GNN (Ours)
MCMC Augmented E(N) GNN (Ours)
Baseline (GNN[GCND)
Baseline (GVP GNN)
Baseline (E(N) GNN)
Task
GNN (Ours)
PSR
184
190
1112
1118
485
492
LEP
6
8
8
9
10
11
LBA
26
28
34
38
59
62
MSP
58
61
45
49
125
128
20


Reviews:
[{"review_id": "XgZ9_7s9ul", "reviewer": "Reviewer_z3cU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors proposed a \"conformer invariant representation learning\" for protein property prediction in this submission. The basic idea is to sample viable conformation based on the given input protein and conformation. MCMC for sampling conformer distribution is integrated to representation learning via the existing Monte-Carlo Gradient Descent (MCGD) algorithm. The authors have also reported empirical results for four protein property prediction tasks on the ATOM3D data set. ", "summary_of_the_review": "The authors developed a learning framework that may help protein representation learning considering possible conformation transformations for given input proteins. This reviewer has concerns on both reported theoretical and empirical results in the current version as detailed above. ", "strength_and_weaknesses": "The authors developed a learning framework that may help protein representation learning considering possible conformation transformations for given input proteins. The procedure of generating and sampling viable conformations is integrated by MCGD to protein representation learning that may help achieve better protein property prediction tasks. \n\nThis reviewer has some concerns: \n\n1. The theoretical analyses for the ergodicity and the convergence to the steady-state conformer distributions seems to be problematic. Based on the description of the procedure, the sampled conformers will be validated by \"structure validation tools such as Molprobity\". Introducing such procedures may violate the desired ergodicity critical in propositions 3.2 and 3.3. Based on the presentation, this reviewer was not sure whether the proposed directed forest construction procedure will always guarantee that all the viable conformers can indeed be included. If not, the procedure does not really have the guarantee to be \"conformer invariant\" exactly. \n\n2. The theoretical analysis appears to be dependent on Definition 2.2  of \"rigid backbone protein conformations\" while in the main text and supplement, the authors emphasized that the proposed work is for non-rigid transformations. The authors may clearly state the limitations of the proposed procedure. \n\n3. The presented empirical results are limited with some tasks that the proposed procedure under-performing the baseline without considering \"conformer invariance\". Also, the proposed procedure appears to share some similarity with augmentation tricks. The authors may want to move the results in the appendix with augmentations into the performance comparison experiments of the main text to check whether the proposed procedure is meaningful or significant. Also, if the authors emphasized that considering more flexible conformers is critical, the performance comparison with other models considering rigid transformations may need to be checked. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "TrrJ2CzVzKp", "reviewer": "Reviewer_K2LR", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper introduces an MCMC-based framework that introduces conditional transformations that respect certain 3d structural properties of proteins with the motivation to learn conformation invariant representation. ", "summary_of_the_review": "Although the proposed MCMC based conditional sampling strategy was well thought out, some key concepts were not well explained in the main text. Experiments suggest the proposed sampling strategy improves over several baseline from the literature when these baselines were trained with conformation invariant samples generated with the proposed sampling strategy, but experimental evidence is not very compelling about whether this improvement is due to the networks learning better representations or due to change in some other aspects of the learning process involving MCGD.   ", "strength_and_weaknesses": "Strengths:\n\n+ MCMC-based strategy for learning representations that are invariant to conditional transformations\n\nWeaknesses:\n\n- Most technical details are moved to Appendix to save space and the extra space is used to explain details of these details moved to Appendix, which makes it difficult to follow certain sections of the paper.\n\n- Experimental results are included with not much insight. It was not clear whether the improvement over the baseline techniques for four tasks were due to proposed conditional sampling or other aspect of the learning changed by MCGD algorithm. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "O3PmnvDsDN", "reviewer": "Reviewer_W5rQ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The manuscript discusses the need to consider a broader class of invariances in protein modelling than the rotation and translational invariances typically considered. The paper introduces the concept of conditional invariance, defined by transformations that modify the sidechain degrees of freedom of a protein, while leaving the backbone degrees of freedom unaltered. The authors propose using a Markov chain Monte Carlo procedure to generate alternative sidechain conformations during training such that the internal representation in a supervised learning task becomes invariant to the sidechain conformation of the input structure.", "summary_of_the_review": "While the paper is technically sound, the motivatation for the proposed method is not convincing. Furthermore, the results section does not fully convince me that the complexity of the method (i.e. the MCMC sampling during training) is necessary to achieve the reported gains.", "strength_and_weaknesses": "The paper does a good job of formally describing the problem and the details of the proposed algorithm. The manuscript seems technically sound.  However, I struggle with the motivation for the method. The authors introduce the conditional invariances to sidechain conformations as a fundamental property on par with rotational and translation invariances of the global structure of the molecule, but this equivalence seems shaky. The standard desire for invariance to rotation and translation of a molecular structure has a clear motivation: we wish any result to be independent on the arbitrarily chosen coordinate system of our 3D coordinate system. Or in other words, the protein structure that we care about has 6 degrees of freedom fewer than when we parameterize it using Cartesian coordinates of all its atoms. From a physical perspective, if we are considering intrinsic properties of a protein, for instance its stability, the position and rotation of a protein is fundamentally irrelevant. All physical interactions between atoms in the molecule are independent of the global orientation and translation.\n\nIn contrast, the invariance towards sidechain conformations is less clear cut. The authors state that \"for most proteins, regardless of their side chain conformation (as long as viable) under consideration – their protein fold class/ other scalar properties remain the same, their mutation (in)stability remains unaltered, protein ligand binding afﬁnity (apart from changes at the ligand binding site) remain the same, etc.\". While it is true that the fold class stays the same because it is typically defined only based on the backbone conformation of a protein, the remaining properties (stability, binding affinity), are certainly dependent on side-chain conformations. The keyword here seems to be \"as long as viable\", which implicitly establishes the equivalence of different structures. But what is deemed \"viable\" will depend on the physical quantity we wish to predict, and the resolution at which we are modelling the system. The desired level of \"invariance\" will thus depend on the modelling task. For some tasks, the exact sidechain orientations will not be important - for other tasks, they will be critical.\n\nRather than consider the set of \"equivalent\" sidechain conformations, the standard approach in statistical physics would be to consider the *distribution* over sidechain structures. Any property of interest (e.g. prediction of a downstream task) could then be considered an expectation under this distribution. For instance, if we observe the physical system at some temperature in a fixed volume, the Boltzmann distribution will gives us such a distribution. In the case of this manuscript, the authors are interested in keeping the backbone degrees of freedom fixed - and would thus consider this conditional distribution of sidechain given backbone. The above would be valid for a molecular forcefield parameterizing the energy of a system - but one could presumably make a similar argument for MolProbity: that it induces a probability distribution over valid sidechain conformations, over which you then calculate the expectation. As far as I can see, this would lead to a similar procedure as the one the authors are proposing, but where you use the validity scores of molprobility as an energy rather than considering all sidechains conformations as equivalent. This approach would be easier to defend, because it does not rely on exact equivalences between structural states.\n\nAlternatively, the method could be sold purely as a principled way to do data augmentation, for increased performance in a downstream task. \n\nI would suggest that the authors rewrite \"1. introduction\" and parts of the \"2 Conditional Invariances for Proteins\" so that the rotational and translation equivariances of the rigid body case are not presented as equivalent as the conditional invariances that they discuss here. As sketched out above, it would probably be most fruitful if the authors get rid of the concept of \"conditional invariance\" alltogether, because these invariances are task specific and only approximate - and that they instead present the results in terms of expectations over distributions over sidechain distributions. If the authors stick with \"conditional invariance\", they should make the limitations of this concept clear - and clarify that the underlying goal is to do meaningful data augmentation in protein systems, and perhaps let that be the driving motivation in the paper.\nIn short, my main problem with the paper is that the conditional invariances are elevated to something fundamental about the physical system - when in reality it is merely a consequence of the fact that certain details are irrelevant for a particular downstream task.  \n\nFinally, for completeness, small fluctuations of the *backbone* structure will have similar properties as the sidechain \"invariances\" - that they will in many cases be \"equivalent\" with respect to a downstream prediction task.  It would make sense that the authors mention that they make the simplifying choice of considering only the sidechain conformers in this paper, ignoring this other source of fluctuation.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Fas-W1uFxKE", "reviewer": "Reviewer_zqSH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Summary of the work\nThe paper studies protein representation learning from the 3D structure. Recent works have shown that although the back-bond of protein 3D structure is fixed, the side chains associated with each back-bone atom are flexible and may correspond to different conformations.\n\nThe authors proposed a method to augment the current 3D structures with randomly sampled conformations from the gold 3D structure. The samples are used as additional data for both training and inference in downstream learning tasks where the representation of the proteins is approximated by the average of the representation of the conformations.\n\nThey show that this simple data augmentation approach improves the performance of supervised downstream tasks that only relies on 3D CNN. 3D GCN GNN to transform the gold conformations in the Atomic3D benchmark dataset.\n\nEven though methods for sampling confirmation using MCMC exist in the literature, a novelty of the work relies on the new sampling approach that constructs for each back-bond atom a directed tree by breaking the loops in the side chains associated with the back-bond atom. Sampling was done by the sampling point set of the nodes in the trees starting from the root nodes down to the leaves, where the later point set is conditionally dependent on the parent's conformations. An MCMC approach is used to sample conformations of the side chains from these point sets.\n\n\n", "summary_of_the_review": "\nI see this work as interesting and may have a good impact on the field. But some points need to improve especially stronger support for the proposed conformation sampling approach. I proposed additional experiments that the authors should consider to strengthen the support evidence, I will be happy to change my score once these additional experiments are added to the papers and the results are still significant concerning the new baselines.  see this work interesting and may have a good impact in the field. But there are points that need to improve especially a stronger support for the proposed conformation sampling approach. I proposed additional experiments that the authors should consider to strengthen the support evidence, I will be happy to change my score once these additional experiments are added to the papers and the results are still significant with respect to the new baselines. ", "strength_and_weaknesses": "Strength\nThis is an interesting and novel idea on conformation sampling.\n\nGood experimental results compared to the baseline that does not leverage augmented data.\n\nWeaknesses\n\n\nI see a lot of strong assumptions without thorough validation and supported evidence both theoretically and empirically\n\nStep 4 in Algorithm 1: trees were sampled independently. Is there any evidence that the conformation of side chains of different amino acids are independent of each other?\n\nRepresenting an amino acid as a tree requires breaking circles (if exists) at a random point, different breaking points result in different tree structures. The current conformation sampling approach assumes conditional independence between a node and all of its grand-parent given the parent nodes. This assumption seems a strong assumption that may lead to ill-approximation of the true conformation distribution. Is there any way to validate that the distribution of the sampled data fits the right distribution of conformation?\n\n\nSince evidence about the approximation of the true conformation distribution is not provided, to demonstrate that the proposed conformation sampling approach based on the atom forest is a more effective data augmentation approach than other simple data augmentation approaches, the following simple  data augmentation should be considered as a baseline to compare to:\n+ For each back-bond atom, keep the back-bond atom fixed, and randomly transform all the side-chain nodes associated with that back-bond atom using a random rotation matrix.\n\nThis baseline will validate that the data augmentation process proposed in the paper is an effective data augmentation approach compared to this simple data augmentation baseline.\n\n\nA lot of relevant existing works in protein representation using amino acid sequences are ignored. When protein 3D structure data is scarce because of expensive acquisition, a lot of work on protein representation learning relies on linear amino acid sequences. Although this representation of protein does not tell much about the 3D conformations of protein explicitly, the availability of a large database of an amino acid sequence is very relevant for self-supervised learning to be trained on. The discussion of these works in the related work and comparison to these works should be conducted carefully. I would suggest the authors compare to at least the following baseline representation trained on large amino acid sequence databases, https://github.com/facebookresearch/esm:\n+ ESM-1b\n+ ESM-MSA-1b  \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "AxDXjwo4HP", "reviewer": "Reviewer_bwoG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper aims to learn a conditional invariant representation for protein structure by introducing a data augmentation approach. The approach augments the protein 3D conformations by applying conditional invariant transformations to sub-structures in protein with a novel Markov chain Monte Carlo (MCMC) sampling. The paper demonstrates improved performance over a GNN baseline that inputs the initial 3D structure of protein without data argumentation.", "summary_of_the_review": "In summary, the authors introduce a novel conditional invariant approach to argument the protein conformations and improve the performance in property prediction tasks. Despite the nice theory, the performance improvement seems marginal and there are some realistic assumptions in the direction forest construction approach. I feel this work requires further improvement and more baselines to demonstrate its usefulness in realistic tasks. ", "strength_and_weaknesses": "Strength:\n- The directed forest construction and conditional transformation of sub-structures provide a principled approach to sample conformations given a protein structure.\n- The authors formally show that the MCGN optimization ensures the convergence of conformer invariant learning procedure.\n\nWeakness:\n- The approach assumes that the protein molecule can be decomposed into a directed forest structure. This assumption is often not true for realistic proteins. The authors use some arbitrary rules to deal with cycles/rings: “cycles/rings which are present in the molecule are broken on the basis of bond length, i.e., larger bonds are chosen before smaller bonds. Ties between same-length bonds are broken arbitrarily.”The resulting conformations are likely unrealistic due to these arbitrary rules.\n- Consequently, the authors use a validity checker (Molprobity) to filter invalid conformations. It means that the sampled structures largely depend on the validity checker. It is unclear if the conditional invariant sampling approach provides benefits compared with a prior MCMC sampling approach. It will also be beneficial if the authors can report the acceptance rate of the MCMC sampling.\n- The performance improvement over a GNN baseline is marginal (Table 1). It is unclear if the proposed approach can outperform a much simpler data argumentation baseline, e.g., random Gaussian noise + validity filtering.\n- Have the authors considered leveraging substructures similar to the JTVAE paper by Jin et al. (ICML 2018) for the construction of the directed forest? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject