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: **Summary and Contributions**:  The paper introduces a framework called Frame Averaging (FA) that can adapt existing backbone architectures to become invariant/equivariant to new symmetry types. It achieves this by averaging over an input-dependent frame which outputs a subset of groups. Frame averaging is often much more efficient to compute than averaging over the entire group, while at the same time, guarantees exact invariance/equivariance.  On the technical side, the paper also proves that FA-based models have the same expressive power as the original backbone architectures.  On the empirical side, the paper provides new classes of models using FA such as universal Euclidean motion invariant point cloud networks / Message Passing GNNs, and demonstrates their practical effectiveness on several tasks.