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: They propose a model for approximating invariant functions for Lie groups, the Lie algebra convolutional network (L-conv). While convolutional networks for ordinary groups require some kind of discretization, their model does not require it. Also, Lie algebra generators in our L-conv can be learned to automatically discover symmetries.