Automatic Symmetry Discovery with Lie Algebra Convolutional Network

1 Introduction . Incorporating symmetries into a deep learning architecture can reduce sample complexity , improve generalization , while significantly decreasing the number of model parameters ( Cohen et al. , 2019b ; Cohen & Welling , 2016b ; Ravanbakhsh et al. , 2017 ; Ravanbakhsh , 2020 ; Wang et al. , 2020 ) . For instance , Convolutional Neural Networks ( CNN ) ( LeCun et al. , 1989 , 1998 ) implement translation symmetry through weight sharing . General principles for constructing symmetry-aware group equivariant neural networks were introduced in Cohen & Welling ( 2016b ) , Kondor & Trivedi ( 2018 ) , and Cohen et al . ( 2019b ) . However , most work on equivariant networks requires knowing the symmetry group a priori . A different equivariant model needs to be re-designed for each symmetry group . In practice , we may not have a good inductive bias and such knowledge of the symmetries may not be available . Constructing and selecting the equivariant network with the appropriate symmetry group becomes quite tedious . Furthermore , many existing works are limited to finite groups such as permutations Hartford et al . ( 2018 ) ; Ravanbakhsh et al . ( 2017 ) ; Zaheer et al . ( 2017 ) , 90 degree rotations Cohen et al . ( 2018 ) or dihedral groups DN and E ( 2 ) Weiler & Cesa ( 2019 ) . 1Code : github.com/nimadehmamy/L-conv-code 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . For a continuous group , existing approaches either discretize the group Weiler et al . ( 2018a , b ) ; Cohen & Welling ( 2016a ) , or use a truncated sum over irreducible representations ( irreps ) Weiler & Cesa ( 2019 ) ; Weiler et al . ( 2018a ) via spherical harmonics in Worrall et al . ( 2017 ) or more general Clebsch-Gordon coefficients Kondor et al . ( 2018 ) ; Bogatskiy et al . ( 2020 ) . These approaches are prone to approximation error . Recently , Finzi et al . ( 2020 ) propose to approximates the integral over the Lie group by Monte Carlo sampling . This approach requires implementing the matrix exponential and obtaining a local neighborhood for each point . Both parametrizing Lie groups for sampling and finding irreps are computationally expensive . Finzi et al . ( 2021 ) provide a general algorithm for constructing equivariant multi-layer perceptrons ( MLP ) , but require explicit knowledge of the group to encode its irreps , and solving a set of constraints . We provide a novel framework for designing equivariant neural networks . We leverage the fact that Lie groups can be constructed from a set of infinitesimal generators , called Lie algebras . A Lie algebra has a finite basis , assuming the group is finite-dimensional . Working with the Lie algebra basis allows us to encode an infinite group without discretizing or summing over irreps . Additionally , all Lie algebras have the same general structure and hence can be implemented the same way . We propose Lie Algebra Convolutional Network ( L-conv ) , a novel architecture that can automatically discover symmetries from data . Our main contributions can be summarized as follows : • We propose the Lie algebra convolutional network ( L-conv ) , a building block for constructing group equivariant neural networks . • We prove that multi-layer L-conv can approximate group convolutional layers , including CNNs , and find graph convolutional networks to be a special case of L-conv . • We can learn the Lie algebra basis in L-conv , enabling automatic symmetry discovery . • L-conv also reveals interesting connections between physics and learning : equivariant loss generalizes important Lagrangians in field theory ; robustness and equivariance can be expressed as Euler-Lagrange equations and Noether currents . Learning symmetries from data has been studied in limited settings for commutative Lie groups as in Cohen & Welling ( 2014 ) , 2D rotations and translations in Rao & Ruderman ( 1999 ) , Sohl-Dickstein et al . ( 2010 ) or permutations ( Anselmi et al. , 2019 ) . In the non-commutative case , GeoManCEr ( Pfau et al. , 2020 ) uses data points related by small transformations to learn non-abelian Lie groups , but it does not introduce an equivariant layer architecture . ( Zhou et al. , 2020 ) propose a general method for symmetry discovery . Yet , their weight-sharing scheme and the symmetry generators are very different from ours . Our approach use much fewer parameters and has a direct interpretation using Lie algebras ( SI B.3 ) . Benton et al . ( 2020 ) propose Augerino to learn a distribution over data augmentations . It also involves Lie algebras , but is restricted to a subgroup of 2D affine transformations and requires matrix logarithm and sampling ( SI B.3 ) . In contrast , our approach is simpler and more general . Our approach uses composition of small transformations to achieve large transformations . In this sense bears some resemblance to symnets ( Gens & Domingos , 2014 ) , but the rest of the construction is different . 2 Background . We review the core concepts L-conv builds upon : equivariance , group convolution and Lie algebras . Notations . Unless explicitly stated , a in Aa is an index , not an exponent . We use the Einstein summation AaBab = ∑ aA aBab = [ AB ] b , where a repeated upper and lower index are summed . Equivariance . Let S be a topological space on which a Lie group G ( continuous group ) acts from the left , meaning for all x ∈ S and g ∈ G , gx ∈ S. We refer to S as the base space . Let F , the “ feature space ” , be the vector space F = Rm . Each data point is a feature map f : S → F . The action of G on the input of f induces an action on feature maps . For “ scalar ” features , for u ∈ G , the transformed features u · f are given by u · f ( x ) = f ( u−1x ) . ( 1 ) Denote the space of all functions from S to F by FS , so that f ∈ FS . Let F be a mapping to a new feature space F ′ = Rm′ , meaning F : FS → F ′S . We say F is equivariant under G if G acts on F ′ and for u ∈ G , we have u · ( F ( f ) ) = F ( u · f ) . ( 2 ) Group Convolution . Kondor & Trivedi ( 2018 ) showed that F is a linear equivariant map if and only if it performs a group convolution ( G-conv ) . To define G-conv , we first lift x to elements in G ( Kondor & Trivedi , 2018 ) . Specifically , we pick an origin x0 ∈ S and replace each point x = gx0 by g. We will often drop x0 for brevity and write f ( g ) ≡ f ( gx0 ) . Let κ : G→ Rm ′ ⊗ Rm be a linear transformation from F to F ′ . G-conv is defined as [ κ ? f ] ( g ) = ∫ G κ ( g−1v ) f ( v ) dv = ∫ G κ ( v ) f ( gv ) dv , ( 3 ) We denote the Haar measure on G as dv ≡ dµ ( v ) for brevity . Equivariance of G-conv . G-conv in equation 3 is equivariant ( Kondor & Trivedi , 2018 ) . By definition , for w ∈ G we have [ κ ? w · f ] ( g ) = ∫ G κ ( v ) w · f ( gv ) dv = ∫ G κ ( v ) f ( w−1gv ) dv = [ κ ? f ] ( w−1g ) = w · [ κ ? f ] ( g ) ( 4 ) Existing works on equivariance networks implement ∫ G by discretizing the group or summing over irreps . We take a different approach and use the infinitesimal generators of the group . While a Lie group G is infinite , usually it can be generated using a small number of infinitesimal generator , comprising its “ Lie algebra ” . We use the Lie algebra to introduce a building block to approximate G-conv . Figure 1 visualizes a Lie group , Lie algebra and the concept we discuss below . Lie algebra . Let G be a Lie group , which includes common continuous groups . Group elements u ∈ G infinitesimally close to the identity element I can be written as u ≈ I + iLi ( note Einstein summation ) , where Li ∈ g with the Lie algebra g = TIG is the tangent space of G at the identity element . The Lie algebra has the property that it is closed under a Lie bracket [ · , · ] : g× g→ g [ Li , Lj ] = cij kLk , ( 5 ) which is skew-symmetric and satisfies the Jacobi identity . Here the coefficients cijk ∈ R or C are called the structure constants of the Lie algebra . For matrix representations of g , [ Li , Lj ] = LiLj − LjLi is the commutator . The Li are called the infinitesimal generators of the Lie group . Exponential map . If the manifold of G is connected 2 , an exponential map exp : g → G can be defined such that g = exp [ tiLi ] ∈ G. For matrix groups , if G is connected and compact , the matrix 2When G has multiple connected components , these results hold for the component containing I , and generalize easily for mutli-component groups such as Zk ⊗G ( Finzi et al. , 2021 ) . exponential is such a map and it is surjective . For most other groups ( except GLd ( C ) and nilpotent groups ) it is not surjective . Nevertheless , for any connected group every g ∈ G can be written as a product g = ∏ a exp [ t i aLi ] ( Hall , 2015 ) . Making t i a infinitesimal steps dt i ( s ) tangent to a path γ from I to g on G yields the surjective path-ordered exponential in physics , denoted as g = P exp [ ∫ γ dtiLi ] ( SI A , and see Time-ordering in Weinberg ( 1995 , p143 ) ) . Pushforward . Li ∈ TIG can be pushed forward to L̂i ( g ) = gLig−1 ∈ TgG to form a basis for TgG , satisfying the same Lie algebra [ L̂i ( g ) , L̂j ( g ) ] = cijkL̂k ( g ) . The manifold of G together with the set of all TgG attached to each g forms the tangent bundle TG , a type of fiber bundle ( Lee et al. , 2009 ) . L̂i is a vector field on TG . The lift maps L̂i to an equivalent vector field on TS , which we will also denote by L̂i . Figure 1 illustrates the flow of these vector fields on TG and TS . 3 Lie Algebra Convolutional Network . We can use the Lie algebra basis Li ∈ g to construct the Lie group G with the exponential map . Similarly , we show that Lie algebras can also serve as building blocks to construct G-conv layers . We propose the Lie algebra convolutional network ( L-conv ) . The key idea is to approximate the kernel κ ( u ) using localized kernels which can be constructed using the Lie algebra ( Fig . 2 ) . This is possible because the exponential map is a generalization of a Taylor expansion . We show that a G-conv whose kernel is concentrated near the identity can be expanded in the Lie algebra . Let δη ( u ) ∈ R denote a normalized localized kernel , meaning ∫ G δη ( g ) dg = 1 , and with support on a small neighborhood of size η centered around the identity I ( i.e. , δη ( I + iLi ) → 0 if ‖ ‖2 > η2 ) . We pick δη ( v ) ∼ θ ( η2 −‖ ‖2 ) , for v = I + iLi ∈ TIG and δη ( v ) = 0 for all other v /∈ TIG ( θ ( · ) being the Heaviside step function ) . Let κ0 : G→ Rm ′ ⊗ Rm be given by [ κ0 ] b a ( u ) = [ W 0 ] c a δη ( u ( I − [ i ] b c Li ) ) ( 6 ) where W 0 ∈ Rm′ ⊗ Rh and i ∈ Rh ⊗ Rm are constants , and we choose ∣∣ [ i ] ab ∣∣ < η . Note that ( I + iLi ) ( I − jLj ) = I + [ − ] iLi +O ( η2 ) . Therefore , ∫ id δη ( ( I + iLi ) ( I − jLj ) ) = i ( 7 ) The localized kernels κ0 can be used to approximate G-conv . Linear expansion of G-conv with localized kernel . We can expand a G-conv whose kernel is κ0 ( u ) = W 0δη ( u ) in the Lie algebra of G to linear order . With v = I + iLi , we have ( see SI A ) Q [ f ] ( g ) = [ κ0 ? f ] ( g ) = ∫ G dvκ0 ( v ) f ( gv ) = ∫ ‖ ‖ < η dv κ0 ( v ) f ( gv ) = W 0 ∫ d δη ( v ) [ f ( g ) + igLi · d dg f ( g ) +O ( 2 ) ] = W 0 [ I + igLi · d dg ] f ( g ) +O ( η2 ) ( 8 ) with W 0 ∈ Rm′ ⊗ Rh and i ∈ Rh ⊗ Rm , as before . Here d is the integration measure on the Lie algebra g = TIG induced by the Haar measure dv on G. Interpreting the derivatives . In a matrix representation of G , we have gLi · dfdg = [ gLi ] β α df dgβα = Tr [ [ gLi ] T df dg ] . This can be written in terms of partial derivatives ∂αf ( x ) = ∂f/∂xα as follows . Using xρ = gρσx σ 0 , we have df ( gx0 ) dgαβ = xβ0∂αf ( x ) , and so L̂if ( x ) ≡ gLi · df dg = [ gLi ] α βx β 0∂αf ( x ) = [ gLix0 ] · ∇f ( x ) ( 9 ) Hence , for each Li , the pushforward gLig−1 generates a flow on S through the vector field L̂i ≡ gLi · d/dg = [ gLig−1x ] α∂α ( Fig . 1 ) . Lie algebra convolutional ( L-conv ) layer . Equation 8 states that for a kernel localized near the identity , the effect of the kernel can be summarized in W 0 and iL̂i . Note that we do not need to perform the integral over G explicitly anymore . Instead of working with a kernel κ0 , we only need to specify W 0 and i . Hence , in general , we define the Lie algebra convolution ( L-conv ) as Q [ f ] ( x ) = W 0 [ I + iL̂i ] f ( x ) = W 0 [ I + i [ gLix0 ] α∂α ] f ( x ) ( 10 ) Being an expansion of G-conv , L-conv inherits the equivariance of G-conv , as we show next . Proposition 1 ( Equivariance of L-conv ) . With assumptions above , L-conv is equivariant under G. Proof : First , note that the components of L̂i transform as [ L̂i ( vx ) ] α = [ vgLix0 ] α = vαβ L̂i ( x ) β , while the partial transforms as ∂/∂ [ vx ] α = [ v−1 ] γα∂γ . As a result in L̂i = [ gLix0 ] α∂α all factors of v cancel , meaning for v ∈ G , L̂i ( vx ) = L̂i ( x ) . This is because of the fact that L̂i ∈ TS is a vector field ( i.e . 1-tensor ) and , thus , invariant under change of basis . Plugging into equation 10 , for w ∈ G w ·Q [ f ] ( x ) = Q [ f ] ( w−1x ) = W 0 [ I + iL̂i ( w −1x ) ] f ( w−1x ) = W 0 [ I + iL̂i ( g ) ] f ( w−1x ) = W 0 [ I + iL̂i ( g ) ] w · f ( x ) = Q [ w · f ] ( x ) ( 11 ) which proves L-conv is equivariant . Examples . Using equation 9 we can calculate L-conv for specific groups ( details in SI A.2 ) . For translationsG = Tn = ( Rn , + ) , we find the generators become simple partial derivatives L̂i = ∂i ( SI A.2.2 ) , yielding f ( x ) + α∂αf ( x ) . For 2D rotations ( SI A.2.1 ) the generator L̂ ≡ ( x∂y − y∂x ) = ∂θ , which is the angular momentum operator about the z-axis in quantum mechanics and field theories . For rotations with scaling , G = SO ( 2 ) ×R+ , we have two Li , one L̂θ = ∂θ from so ( 2 ) and a scaling with Lr = I , yielding L̂r = x∂x + y∂y = r∂r . Next , we discuss the form of L-conv on discrete data .

Summary: The paper proposes a method to build equivariant neural networks by defining the layers in terms of the generators of the Lie group to which the NN is to be equivariant to. At the core lies the idea that any function (or conv kernel) with local support can be expanded in the Lie algebra basis via a Taylor approximation-like approach. Then, by observing that any function can be approximated with a basis of (shifted) local functions, it can be concluded that the presented approach in theory can approximate any group convolution. In addition, the authors the investigate the option to let go of the user-imposed constraint of being equivariant to a specific group, and instead let the neural network learn the generators (of the unknown group) in addition to the learnable weights that parametrize the layers. This is an open problem and the authors make promising steps in this direction.  Finally, the authors draw connections to physics, but I must admit, this part was too much for me and I did not get it, but possibly could have by investing more time in the paper.