CONDITIONAL INVARIANCES FOR CONFORMER IN-VARIANT PROTEIN REPRESENTATIONS

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 treat 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 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 efficiency.

1. 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 sciences, group equivariant graph neural networks and transformers have specifically found applications to small molecules, as well as large molecules (e.g. proteins) with tremendous success (Klicpera et al., 2020; Anderson et al., 2019; Fuchs et al., 2020; Hutchinson et al., 2021; Satorras et al., 2021; Batzner et al., 2021) . 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., 2021) . Traditionally, proteins have been modeled using standard 3D CNNs (Karimi et al., 2019; Pagès et al., 2019) , graph neural networks (GNNs) (Kipf & Welling, 2016; Hamilton et al., 2017), and transformers (Vaswani et al., 2017) 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 to the downstream task. And the transformations they are invariant to do not depend on the input. For instance, invariance to the Euclidean group restricts the protein representation to 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 a specific set of protein conformations. For example, 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 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?



. More recently, several works Jing et al. (2020; 2021); Jumper et al. (2021); Hermosilla et al. (2021) have enriched the above models with neural networks that are equivariant (invariant) to transformations from the Euclidean and rotation groups.

