E(n) Equivariant Normalizing Flows

This paper introduces a generative model equivariant to Euclidean symmetries : E ( n ) Equivariant Normalizing Flows ( E-NFs ) . To construct E-NFs , we take the discriminative E ( n ) graph neural networks and integrate them as a differential equation to obtain an invertible equivariant function : a continuous-time normalizing flow . We demonstrate that E-NFs considerably outperform baselines and existing methods from the literature on particle systems such as DW4 and LJ13 , and on molecules from QM9 in terms of log-likelihood . To the best of our knowledge , this is the first flow that jointly generates molecule features and positions in 3D . 1 Introduction . Leveraging the structure of the data has long been a core design principle for building neural networks . Convolutional layers for images famously do so by being translation equivariant and therefore incorporating the symmetries of a pixel grid . Analogously , for discriminatory machine learning tasks on 3D coordinate data , taking into account the symmetries of data has significantly improved performance ( Thomas et al. , 2018 ; Anderson et al. , 2019 ; Finzi et al. , 2020 ; Fuchs et al. , 2020 ; Klicpera et al. , 2020 ) . One might say that equivariance has been proven an effective tool to build inductive bias into the network about the concept of 3D coordinates . However , for generative tasks , e.g. , sampling new molecular structures , the development of efficient yet powerful rotation equivariant approaches—while having made great progress—is still in its infancy . A recent method called E ( n ) Equivariant Graph Neural Networks ( EGNNs ) ( Satorras et al. , 2021 ) is both computationally cheap and effective in regression and classification tasks for molecular data , while being equivariant to Euclidean symmetries . However , this model is only able to discriminate features on nodes , and can not generate new molecular structures . In this paper , we introduce E ( n ) Equivariant Normalizing Flows ( E-NFs ) : A generative model for E ( n ) Equivariant data such as molecules in 3D . To construct E-NFs we parametrize a continuous-time flow , where the first-order derivative is modelled by an EGNN . We adapt EGNNs so that they are stable when utilized as a derivative . In addition , we use recent advances in the dequantization literature to lift the discrete features of nodes to a continuous space . We show that our proposed flow model significantly outperforms its non-equivariant variants and previous equivariant generative methods ( Köhler et al. , 2020 ) . Additionally , we apply our method to molecule generation and we show that our method is able to generate realistic molecules when trained on the QM9 dataset . ⇤Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 Background . Normalizing Flows A learnable invertible transformation x = g✓ ( z ) and a simple base distribution pZ ( such as a normal distribution ) yield a complex distribution pX . Let f✓ = g 1✓ , then the likelihood of a datapoint x under pX can be exactly computed using the change of variables formula : pX ( x ) = pZ ( z ) |det Jf ( x ) | , where z = f ( x ) , ( 1 ) where Jf is the Jacobian of f✓ . A particular type of normalizing flows are continuous-time normalizing flows ( Chen et al. , 2017 , 2018b ; Grathwohl et al. , 2018 ) . These flows use a conceptual time direction to specify an invertible transformation as a differential equation . The first order differential is predicted by a neural network , referred to as the dynamics . The continuous-time change of variables formula is given by : log pX ( x ) = log pZ ( z ) + Z 1 0 Tr J ( x ( t ) ) dt , where z = x+ Z 1 0 ( x ( t ) ) dt , ( 2 ) where x ( 0 ) = x and x ( 1 ) = z . In practice , Chen et al . ( 2018b ) ; Grathwohl et al . ( 2018 ) estimate the trace using Hutchinson ’ s trace estimator , and the integrals can be straightforwardly computed using the torchdiffeq package written by Chen et al . ( 2018b ) . It is often desired to regularize the dynamics for faster training and more stable solutions ( Finlay et al. , 2020 ) . Continuous-time normalizing flows are desirable because the constraints that need to be enforced on are relatively mild : only needs to be high order differentiable and Lipschitz continuous , with a possibly large Lipschitz constant . Equivariance in Normalizing Flows Note that in this context , the desirable property for distributions is often invariance whereas for transformations it is equivariance . Concretely , when a function g✓ is equivariant and a base distribution pZ is invariant , then the distribution pX given by x = g✓ ( z ) where z ⇠ pZ is also invariant ( Köhler et al. , 2020 ) . As a result , one can design expressive invariant distributions using an invariant base distribution pZ and an equivariant function g✓ . Furthermore , when g✓ is restricted to be bijective and f✓ = g 1✓ , then equivariance of g✓ implies equivariance of f✓ . In addition , the likelihood pX can be directly computed using the change of variables formula . 2.1 Equivariance . Equivariance of a function f under a group G is defined as Tg ( f ( x ) ) = f ( Sg ( x ) ) for all g 2 G , where Sg , Tg are transformations related to the group element g , where in our case Sg = Tg will be the same . In other words , f being equivariant means that first applying the transformation Sg on x and then f yields the same result as first applying f on x and then transforming using Tg . In this paper we focus on symmetries of the n-dimensional Euclidean group , referred to as E ( n ) . An important property of this group is that its transformations preserve Euclidean distances . The transformations can be described as rotations , reflections and translations . Specifically , given a set of points x = ( x1 , . . . , xM ) 2 RM⇥n embedded in an n-dimensional space , an orthogonal matrix R 2 Rn⇥n and a translation vector t 2 Rn , we can say f is rotation ( and reflection ) equivariant if , for z = f ( x ) we have Rz = f ( Rx ) , where Rx is the shorthand2 for ( Rx1 , . . . , RxM ) . Similarly f 2To be precise , in matrix multiplication notation ( Rx1 , . . . , RxM ) = xRT , for simplicity we use the notation Rx = ( Rx1 , . . . , RxM ) and see R as an operation instead of a direct multiplication on the entire x. is translation equivariant if z+ t = f ( x+ t ) , where x+ t is the shorthand for ( x1 + t , . . . , xM + t ) . In this work we consider data defined on vertices of a graph V = ( x , h ) , which in addition to the position coordinates x , also includes node features h 2 RM⇥nf ( for example temperature or atom type ) . Features h have the property that they are invariant to E ( n ) transformations , while they do affect x . In other words , rotations and translations of x do not modify h. In summary , given a function f : x , h 7 ! zx , zh we require equivariance of f with respect to the Euclidean group E ( n ) so that for all orthogonal matrices R and translations t : Rzx + t , zh = f ( Rx+ t , h ) ( 3 ) E ( n ) Equivariant Graph Neural Networks ( EGNN ) ( Satorras et al. , 2021 ) consider a graph G = ( V , E ) with nodes vi 2 V and edges eij . Each node vi is associated with a position vector xi and node features hi as the ones defined in previous paragraphs . Then , an E ( n ) Equivariant Graph Convolutional Layer ( EGCL ) takes as input the set of node embeddings hl = { hl 0 , . . . , hl M 1 } , coordiante embeddings xl = { xl 0 , . . . , xl M 1 } at layer l and edge information E = ( eij ) and outputs a transformation on hl+1 and xl+1 . Concisely : hl+1 , xl+1 = EGCL [ hl , xl , E ] . This layer satisfies the equivariant constraint defined in Equation 3 . The equations that define this layer are the following : mij = e ⇣ hl i , hl j , xl i xl j 2 ⌘ and mi = X j 6=i eijmij , ( 4 ) xl+1 i = xl i + X j 6=i xl i xl j x ( mij ) and hl+1i = h hl i , mi . ( 5 ) A neural network hL , xL = EGNN [ h0 , x0 ] composed of L layers will define the dynamics of our ODE flow . In our experiments we are not provided with an adjacency matrix , but only node information . Therefore , we use the edge inference module eij = inf ( mij ) introduced in the EGNN paper that outputs a soft estimation of the edges where inf : Rnf ! [ 0 , 1 ] 1 resembles a linear layer followed by a sigmoid function . This behaves as an attention mechanism over neighbor nodes . 3 Related Work . Group equivariant neural networks ( Cohen and Welling , 2016 , 2017 ; Dieleman et al. , 2016 ) have demonstrated their effectiveness in a wide variety of tasks . A growing body of literature is finding neural networks that are equivariant to transformations in Euclidean space ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ; Horie et al. , 2020 ; Finzi et al. , 2020 ; Hutchinson et al. , 2020 ; Satorras et al. , 2021 ) . These methods have proven their efficacy in discriminative tasks and modelling dynamical sytems . On the other hand , graph neural networks ( Bruna et al. , 2013 ; Kipf and Welling , 2016 ) can be seen as networks equivariant to permutations . Often , the methods that study Euclidean equivariance operate on point clouds embedded in Euclidean space , and so they also incorporate permutation equivariance . Normalizing Flows ( Rippel and Adams , 2013 ; Rezende and Mohamed , 2015 ; Dinh et al. , 2015 ) are an attractive class of generative models since they admit exact likelihood computation and can be designed for fast inference and sampling . Notably Chen et al . ( 2018a , b ) ; Grathwohl et al . ( 2018 ) introduced continuous-time normalizing flows , a flow that is parametrized via a first-order derivative over time . This class of flows is useful because of the mild constraints on the parametrization function compared to other flow approaches ( Dinh et al. , 2017 ; Kingma and Dhariwal , 2018 ) . There are several specialized methods for molecule generation : Gebauer et al . ( 2019 ) generate 3D molecules iteratively via an autoregressive approach , but discretize positions and use additional focus tokens which makes them incomparable in log-likelihood . Gómez-Bombarelli et al . ( 2016 ) ; You et al . ( 2018 ) ; Liao et al . ( 2019 ) generate discrete graph structures instead of a coordinate space and Xu et al . ( 2021 ) generate molecule positions only . Some flows in literature Noé et al . ( 2018 ) ; Li et al . ( 2019 ) ; Köhler et al . ( 2020 ) model positional data , but not in combination with discrete properties . Recently various forms of normalizing flows for equivariance have been proposed : Köhler et al . ( 2019 , 2020 ) propose flows for positional data with Euclidean symmetries , Rezende et al . ( 2019 ) introduce a Hamiltonian flow which can be designed to be equivariant using an invariant Hamiltonian function . Liu et al . ( 2019 ) ; Biloš and Günnemann ( 2020 ) model distributions over graphs , with flows equivariant to permutations . Boyda et al . ( 2020 ) introduce equivariant flows for SU ( n ) symmetries . For adversarial networks , Dey et al . ( 2021 ) introduced a generative network equivariant to 90 degree image rotations and reflections . Our work differs from these approaches in that we design a generalpurpose normalizing flow for sets of nodes that contain both positional and invariant features while remaining equivariant to E ( n ) transformations . This combination of positional and non-positional information results in a more expressive message passing scheme that defines the dynamics of our ODE . A by-product of our general approach is that it also allows us to jointly model features and structure of molecules in 3D space without any additional domain knowledge .

Summary: The paper introduces a new type of equivariant normalizing flow by using the CNF framework with a modified equivariant graph NN as a dynamics function. The equivariance discussed is wrt Euclidean symmetries, which are important for physical systems. They show that their flow beats prior equivariant baselines on synthetic data. Moreover, the model is capable of transforming spacial, categorical and ordinal features simultaneously. This allows the sampling of molecular configurations with positions, atom types, and charges.