Frame Averaging for Invariant and Equivariant Network Design

1 INTRODUCTION . Many tasks in machine learning ( ML ) require learning functions that are invariant or equivariant with respect to symmetric transformations of the data . For example , graph classification is invariant to a permutation of its nodes , while node prediction tasks are equivariant to node permutations . Consequently , it is important to design expressive neural network architectures that are by construction invariant or equivariant for scalable and efficient learning . This recipe has proven to be successful for many ML tasks including image classification and segmentation ( LeCun et al. , 1998 ; Long et al. , 2015 ) , set and point-cloud learning ( Zaheer et al. , 2017 ; Qi et al. , 2017a ) , and graph learning ( Kipf & Welling , 2016 ; Gilmer et al. , 2017 ; Battaglia et al. , 2018 ) . Nevertheless , for some important instances of symmetries , the design of invariant and/or equivariant networks is either illusive ( Thomas et al. , 2018 ; Dym & Maron , 2020 ) , computationally expensive or lacking in expressivity ( Xu et al. , 2018a ; Morris et al. , 2019 ; Maron et al. , 2019 ; Murphy et al. , 2019 ) . In this paper , we propose a new general-purpose framework , called Frame Averaging ( FA ) , that can systematically facilitate expressive invariant and equivariant networks with respect to a broad class of groups . At the heart of our framework , we build on a basic fact that arbitrary functions φ ∶ V → R , Φ ∶ V →W , where V , W are some vector spaces , can be made invariant or equivariant by symmetrization , that is averaging over the group ( Yarotsky , 2021 ; Murphy et al. , 2018 ) , i.e. , ψ ( X ) = 1 ∣G∣ ∑g∈G φ ( g−1 ⋅X ) or Ψ ( X ) = 1 ∣G∣ ∑g∈G g ⋅Φ ( g−1 ⋅X ) . ( 1 ) where G = { g } denotes the group , ψ ∶ V → R is invariant and Ψ ∶ V → W is equivariant with respect to G. Furthermore , since invariant and equivariant functions are fixed under group averaging , i.e. , ψ = φ for invariant φ and Ψ = Φ for equivariant Φ , the above scheme often leads to universal ( i.e. , maximally expressive ) models ( Yarotsky , 2021 ) . However , the challenge with equation 1 is that when the cardinality of G is large ( e.g. , combinatorial groups such as permutations ) or infinite ( e.g. , continuous groups such as rotations ) , then exact averaging is intractable . In such cases , we are forced to approximate the sum via heuristics or Monte Carlo ( MC ) , thereby sacrificing the exact invariance/equivariance property for computational efficiency , e.g. , Murphy et al . ( 2018 ; 2019 ) define heuristic averaging strategies for approximate permutation invariance in GNNs ; similarly , Hu et al . ( 2021 ) and Shuaibi et al . ( 2021 ) use MC averaging for approximate rotation equivariance in GNNs . The key observation of the current paper is that the group average in equation 1 can be replaced with an average over a carefully selected subsetF ( X ) ⊂ G while retaining both exact invariance/equivariance and expressive power . Therefore , if F can be chosen so that the cardinality ∣F ( X ) ∣ is mostly small , averaging over F ( X ) results in both expressive and efficient invariant/equivariant model . We call the set-valued function F ∶ V → 2G , a frame , and show that it can successfully replace full group averaging if it satisfies a set equivariance property . We name this framework Frame Averaging ( FA ) and it serves as the basis for the design of invariant/equivariant networks in this paper . We instantiate the FA framework by considering different choices of symmetry groups G , their actions on data spaces V , W ( manifested by choices of group representations ) , and the backbone architectures ( or part thereof ) φ , Φ we want to make invariant/equivariant toG . We consider : ( i ) MultiLayer Perceptrons ( MLP ) , and Graph Neural Networks ( GNNs ) with node identification ( Murphy et al. , 2019 ; Loukas , 2020 ) adapted to permutation invariant Graph Neural Networks ( GNNs ) ; ( ii ) Message-Passing GNN ( Gilmer et al. , 2017 ) adapted to be invariant/equivariant to Euclidean motions , E ( d ) ; ( iii ) Set network , DeepSets and PointNet ( Zaheer et al. , 2017 ; Qi et al. , 2017a ) adapted to be equivariant or locally equivariant to E ( d ) ; ( iv ) Point cloud network , DGCNN ( Wang et al. , 2018 ) , adapted to be equivariant to E ( d ) . Theoretically , we prove that the FA framework maintains the expressive power of its original backbone architecture which leads to some interesting corollaries : First , ( i ) results in invariant universal graph learning models ; ( ii ) is an E ( d ) invariant/equivariant GNN that maintain the power of message passing ( Xu et al. , 2018a ; Morris et al. , 2019 ) ; and ( iii ) , ( iv ) furnish a universal permutation and E ( d ) invariant/equivariant models . We note that both the construction and the proofs are arguably considerably simpler than the existing alternative constructions and proofs for this type of symmetry ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ; Dym & Maron , 2020 ) . We experimented with FA on different tasks involving symmetries including : point-cloud normal estimation , beyond 2-Weisfeiler-Lehman graph separation , and n-body dynamics predictions , reaching state of the art performance in all . 2 FRAME AVERAGING . In this section we introduce the FA approach and theory in generic formulation , while in the next section we show the power of this general approach by instantiating it to ( seemingly different ) problems of interest . 2.1 FRAME AVERAGING FOR FUNCTION SYMMETRIZATION . Let φ ∶ V → R and Φ ∶ V →W be some arbitrary functions , where V , W are normed linear spaces with norms ∥⋅∥V , ∥⋅∥W , respectively . For example , φ , Φ can be thought of as neural networks . We consider a group G = { g } that describes some symmetry we want to incorporate into φ , Φ . The way the symmetries g ∈ G are applied to vectors in V , W is described by the group ’ s representations ρ1 ∶ G → GL ( V ) , and ρ2 ∶ G → GL ( W ) , where GL ( V ) is the space of invertible linear maps V → V ( automorphisms ) . A representation ρi preserves the group structure by satisfying ρi ( gh ) = ρi ( g ) ρi ( h ) for all g , h ∈ G ( see e.g. , Fulton & Harris ( 2013 ) ) . As customary , we will sometimes refer to the linear spaces V , W as representations . Our goal is to make φ into an invariant function , namely satisfy φ ( ρ1 ( g ) X ) = φ ( X ) , for all g ∈ G and X ∈ V ; and Φ into an equivariant function , namely Φ ( ρ1 ( g ) X ) = ρ2 ( g ) Φ ( X ) , for all g ∈ G and X ∈ V . We will do that by averaging over group elements , but instead of averaging over the entire group every time ( as in equation 1 ) we will average on a subset of the group elements called a frame . Definition 1 . A frame is defined as a set valued function F ∶ V → 2G ∖ ∅ . 1 . A frame is G-equivariant if F ( ρ1 ( g ) X ) = gF ( X ) , ∀X ∈ V , g ∈ G , ( 2 ) where as usual , gF ( X ) = { gh ∣ h ∈ F ( X ) } , and the equality in equation 2 should be understood as equality of sets . 2 . A frame is bounded over a domainK ⊂ V if there exists a constant c > 0 so that ∥ρ2 ( g ) ∥op ≤ c , for all g ∈ F ( X ) and all X ∈K , where ∥⋅∥op denotes the induced operator norm over W . Figure 1 provides an illustration . How are equivariant frames useful ? Consider a scenario where an equivariant frame is easy to compute , and furthermore its cardinality , ∣F ( X ) ∣ , is not too large . Then averaging over the frame , denoted ⟨⋅⟩F and defined by ⟨φ⟩F ( X ) = 1 ∣F ( X ) ∣ ∑g∈F ( X ) φ ( ρ1 ( g ) −1X ) ( 3 ) ⟨Φ⟩F ( X ) = 1 ∣F ( X ) ∣ ∑g∈F ( X ) ρ2 ( g ) Φ ( ρ1 ( g ) −1X ) ( 4 ) provides the required function symmetrization . In Appendix A.1 we prove : Theorem 1 ( Frame Averaging ) . Let F be a G equivariant frame , and φ ∶ V → R , Φ ∶ V →W some functions . Then , ⟨φ⟩F is G invariant , while ⟨Φ⟩F is G equivariant . Several comments are in order : First , the invariant case ( equation 3 ) is a particular case of the equivariant case ( equation 4 ) under the choice of W = R and the trivial representation ρ2 ( g ) ≡ 1 . Second , in this paper we only consider X and frame choices F for which F ( X ) are finite sets . Nevertheless , treating the infinite case is an important future research direction . Third , a trivial choice of an equivariant frame is F ( X ) ≡ G , that is , taking the frame to be the entire group for all X ∈ V ( for infinite but compact G the sum in the FA in this case can be replaced with Harr integral ) . This choice can be readily checked to be equivariant , and turns the FA equations 3 , 4 into standard group averaging operators , equation 1 . The problem with this choice , however , is that it often results in an intractable or challenging computation , e.g. , when the group is large or infinite . In contrast , as we show below , in some useful cases one can compute a manageable size frame and can use it to build invariant or equivariant operators in a principled way . Let us provide a simple example for Frame Averaging : consider V = Rn , W = R , and G = R with addition as the group action . We choose the group actions1 in this case to be ρ1 ( a ) x = x + a1 , and ρ2 ( a ) b = b + a , where a , b ∈ R , x ∈ Rn , and 1 ∈ Rn is the vector of all ones . We can define the frame in this case using the averaging operator F ( x ) = { 1 n 1Tx } ⊂ G = R. Note that in this case the frame contains only one element from the group , in other cases finding such a small frame is hard or even impossible . One can check that this frame is equivariant per Definition 1 . The FA of φ ∶ Rn → R would be ⟨φ⟩F ( x ) = φ ( x− 1 n ( 1Tx ) 1 ) in the invariant case , and ⟨φ⟩F ( x ) = φ ( x − 1 n ( 1Tx ) 1 ) + 1 n 1Tx in the equivariant case . Incorporating G as a second symmetry . An important use case of frame averaging is with the backbones φ , Φ already invariant/equivariant w.r.t . some symmetry group H and our goal is to make it invariant/equivariant to H ×G . For example , say we want to add G = E ( 3 ) equivariance to permutation invariant set or graph functions , i.e. , H = Sn . We will provide sufficient conditions for the FA to provide this desired invariance/equivariance . First , let us assumeH is acting on V andW by the representations τ1 ∶H → GL ( V ) and τ2 ∶H → GL ( W ) , respectively . Assume φ is H invariant and Φ is H equivariant . We say that representations ρ1 and τ1 commute if ρ1 ( g ) τ1 ( h ) X = τ1 ( h ) ρ1 ( g ) X for all g ∈ G , h ∈H , andX ∈ V . If ρ1 and τ1 commute then the map γ1 ∶H ×G→ GL ( V ) defined by γ1 ( h , g ) = τ1 ( h ) ρ1 ( g ) is a representation of the group H ×G . Second , we would need that the frame F ( X ) is invariant to H , that is F ( τ1 ( h ) X ) = F ( X ) . We show a generalization of Theorem 1 : Theorem 2 ( Frame Average second symmetry ) . Assume F is H-invariant and G-equivariant . Then , 1 . If φ ∶ V → R is H invariant and ρ1 , τ1 commute then ⟨φ⟩F is G ×H invariant . 2 . If Φ ∶ V →W isH equivariant and ρi , τi , i = 1,2 , commmute then ⟨Φ⟩F isG×H equivariant . 1Note that since these are affine maps they are technically not representations but have an equivalent representation using homogeneous coordinates . Therefore , FA is also valid with affine actions as used here . Right actions . Above we used left actions for the definition of equivariance . There are other flavors of equivariance , e.g. , if one of the actions is right . For example , if g multiplies F ( X ) from the right , then equivariance will take the form : F ( ρ1 ( g ) X ) = F ( X ) g−1 , ∀X ∈ V , g ∈ G ( 5 ) and accordingly ⟨φ⟩F ( X ) = 1 ∣F ( X ) ∣ ∑g∈F ( X ) φ ( ρ1 ( g ) X ) , ⟨Φ⟩F ( X ) = 1 ∣F ( X ) ∣ ∑g∈F ( X ) ρ2 ( g ) −1Φ ( ρ1 ( g ) X ) ( 6 ) are G invariant and equivariant , respectively . Efficient calculation of invariant frame averaging . There could be instances of the FA framework ( indeed we discuss such a case later ) where ∣F ( X ) ∣ is still too large to evaluate equations 3,4 . In the invariant case , there is a more efficient form of FA , that can potentially be applied . To show it , let us start by defining the subgroup of symmetries of X , i.e. , its stabilizer . The stabilizer of an elementX ∈ V is a subgroup ofG defined byGX = { g ∈ G ∣ ρ1 ( g ) X =X } . GX naturally induces an equivalence relation ∼ on F ( X ) , with g ∼ h ⇐⇒ hg−1 ∈ GX . The equivalence classes ( orbits ) are [ g ] = { h ∈ F ( X ) ∣g ∼ h } = GXg ⊂ F ( X ) , for g ∈ F ( X ) , and the quotient set is denoted F ( X ) /GX . Theorem 3 . Equivariant frame F ( X ) is a disjoint union of equal size orbits , [ g ] ∈ F ( X ) /GX . The proof is in A.3 . The first immediate consequence of Theorem 3 is that the cardinality of F ( X ) is at-least that of the stabilizer ( intuitively , the inner-symmetries ) of X , namely ∣F ( X ) ∣ ≥ ∣GX ∣ . Therefore , there could be cases , such as when X describes a symmetric graph , where ∣F ( X ) ∣ could be too large to average over . A remedy comes from the following observation : for every h ∈ [ g ] , we have that h = rg , r ∈ GX , and φ ( ρ1 ( h ) −1X ) = φ ( ρ1 ( g ) −1ρ1 ( r ) −1X ) = φ ( ρ1 ( g ) −1X ) , since also r−1 ∈ GX . Therefore the summands in equations 3 are constant over orbits , and we get ⟨φ⟩F ( X ) = 1 mF ∑ [ g ] ∈F ( X ) /GX φ ( ρ1 ( g ) −1X ) , ( 7 ) where mF = ∣F ( X ) /GX ∣ = ∣F ( X ) ∣ / ∣GX ∣ . This representation of invariant FA requires only mF = ∣F ( X ) ∣/∣GX ∣ evaluations , compared to ∣F ( X ) ∣ in the original FA in equation 3 . Approximation of invariant frame averaging . Unfortunately , enumerating F ( X ) /GX could be challenging in some cases . Nevertheless , equation 7 is still very useful : it turns out we can easily draw a random element from F ( X ) /GX with uniform probability . This is an immediate application of the equal orbit size in Theorem 3 : Corollary 1 . Let F ( X ) be an equivariant frame , and g ∈ F ( X ) be a uniform random sample . Then [ g ] ∈ F ( X ) /GX is also uniform . Therefore , an efficient approximation strategy is averaging over uniform samples , gi ∈ F ( X ) , i ∈ [ k ] , ⟪φ⟫F ( X ) = 1 k k ∑ i=1 φ ( ρ1 ( gi ) −1X ) . ( 8 ) This approximation is especially useful , compared to the full-blown FA , when mF = ∣F ( X ) ∣/∣GX ∣ is small , i.e. , when ∣GX ∣ is large , or X has many symmetries . Intuitively , the smaller mF the better the approximation in equation 8 . A partial explanation to this phenomenon is given in Appendix A.4 , while an empirical validation is provided in Section 5.2 .

Summary: The paper proposes to make any neural network equivariant by symmetrizing over a subset of the group, rather than over whole group. If the subset selection F(X), depending on input X, is equivariant (gFX=FgX), then the symmetrization is equivariant. The authors furthermore prove: 1)	When interested in invariant prediction, the subset can be chosen in the quotient G/G_X, where G_X is the stabilizer subgroup of X. 2)	When symmetrizing with a random subsample of F(X), the probability of a particular subsample that deviates from symmetrizing with all of F(X) by less than some epsilon, is bounded below. 3)	When using the symmetrization of a universal model, the resulting model class is universal in the class of equivariant functions.