Lie Algebra Convolutional Neural Networks with Automatic Symmetry Extraction

Many machine learning ( ML ) tasks involve data from unfamiliar domains , which may or may not have hidden symmetries . While much of the work on equivariant neural networks focuses on equivariant architectures , the ability of the architecture to discover symmetries in a given dataset is less studied . Convolutional Neural Networks ( CNN ) ( LeCun et al. , 1989 ; 1998 ) incorporate translation symmetry into the architecture . Recently , more general ways to construct equivariant architectures have been introduced ( Cohen & Welling , 2016a ; b ; Cohen et al. , 2018 ; Kondor & Trivedi , 2018 ) . Encoding equivariance into an ML architecture can reduce data requirements and improve generalization , while significantly reducing the number of model parameters via parameter sharing ( Cohen et al. , 2019 ; Cohen & Welling , 2016b ; Ravanbakhsh et al. , 2017 ; Ravanbakhsh , 2020 ) . As a result , many other symmetries such as discrete rotations in 2D ( Veeling et al. , 2018 ; Marcos et al. , 2017 ) and 3D ( Cohen et al. , 2018 ; Cohen & Welling , 2016a ) as well as permutations ( Zaheer et al. , 2017 ) have been incorporated into the architecture of neural networks . Many existing works on equivariant architectures use finite groups such as permutations in Hartford et al . ( 2018 ) and Ravanbakhsh et al . ( 2017 ) or discrete subgroups of continuous groups , such as 90 degree rotations in ( Cohen et al. , 2018 ) or dihedral groups DN in Weiler & Cesa ( 2019 ) . Ravanbakhsh ( 2020 ) also proved a universal approximation theorem for single hidden layer equivariant neural networks for Abelian and finite groups . General principles for constructing group convolutional layers were introduced in Cohen & Welling ( 2016b ) , Kondor & Trivedi ( 2018 ) , and Cohen et al . ( 2019 ) , including for continuous groups . A challenge for implementation is having to integrate over the group manifold . This has been remedied either by generalizing Fast Fourier Transforms ( Cohen et al. , 2018 ) , or using irreducible representations ( irreps ) ( Weiler et al. , 2018a ) either directly as spherical harmonics as in Worrall et al . ( 2017 ) or using more general Clebsch-Gordon coefficients ( Kondor et al. , 2018 ) . Other approaches include discretizing the group as in Weiler et al . ( 2018a ; b ) ; Cohen & Welling ( 2016a ) , or solving constraints for equivariant irreps as in Weiler & Cesa ( 2019 ) , or approximating the integral by sampling ( Finzi et al. , 2020 ) . The limitations in all of the approaches above are that : 1 ) they rely on knowing the symmetry group a priori , and 2 ) require encoding the whole group into the architecture . For a continuous group , it is not possible to encode all elements and we have to resort to discretization or a truncated sum over irreps . Our work attempts to resolve the issues with continuous groups by using the Lie algebra ( the linearization of the group near its identity ) instead of the group itself . Unlike the Lie group which is infinite , the Lie algebra usually has a finite basis ( notable exception being Kac-Moody Lie algebras for 2D Conformal Field Theories ( Belavin et al. , 1984 ) in physics ) . Additionally , we show that the Lie algebra basis can be learned during training , or through a separate optimization process . Hence , our architecture , which generalizes a group convolutional layer , is potentially capable of learning symmetries in data without imposing inductive biases . Learning symmetries in data was tackled in restricted settings of mostly commutative Lie groups as in Cohen & Welling ( 2014 ) and 2D rotations and translations in Rao & Ruderman ( 1999 ) and Sohl-Dickstein et al . ( 2010 ) or permutations ( Anselmi et al. , 2019 ) . However , the symmetries learned by the architecture are not necessarily familiar spatial symmetries . As we show in the case of linear regression , the symmetries may correspond to transformations preserving the statistics of the data . Specifically , we show a general relation between the symmetries of linear regression and a deformed orthogonal group preserving the covariance matrix . Such symmetries of the probability distribution and ways to incorporate them into the architecture were also discussed in Bloem-Reddy & Teh ( 2019 ) . The work that is closest in spirit and setup to ours is Zhou et al . ( 2020 ) which uses meta-learning to automatically learn symmetries . Although the weight-sharing scheme of Zhou et al . ( 2020 ) and their encoding of the symmetry generators is different , their construction does bear some resemblance to ours and we will discuss this after introducing our architecture . Contributions Our main contributions can be summarized as follows • We propose a group equivariant architecture using the Lie algebra , introducing the Lie algebra convolutional layer ( L-conv ) . • In L-conv the Lie algebra generators can be trained to discover symmetries , and it outper- forms CNN on domains with hidden symmetries , such rotated and scrambled images . • Group convolutional layers on connected Lie groups can be approximated by multi-layer L-conv , and Fully-connected , CNN and graph convolutional networks are special cases of L-conv . • In linear regression , we show analytical relations between symmetries in and orthogonal groups preserving covariance of data . 1 EQUIVARIANCE IN SUPERVISED LEARNING . Consider the functional mapping yi = f ( xi ) of inputs X = ( x1 , . . . , xn ) to outputs Y = ( y1 , . . . , yn ) . We assume each input x ∈ Rd×m where Rd are the “ space ” dimensions and Rm the “ channels ” , and y ∈ Rc ( or Zc2 for categorical variables ) . We assume a group G acts only on the space factor ( Rd , shared among channels ) of x through a d-dimensional representation Td : G→ GLd ( R ) mapping each g to an invertible d× d matrix . The map Td must be continuous and satisfy Td ( u ) Td ( v ) = Td ( uv ) for all u , v ∈ G ( Knapp , 2013 , IV.1 ) . Similarly , let G act on y via a c-dimensional representation Tc . To simplify notation , we will denote the representations simply as uα ≡ Tα ( u ) . A function f solving yi = f ( xi ) is said to be equivariant under the action of a group G by representations Tc , Td if ucy = ucf ( x ) = f ( udx ) ∀u ∈ G ⇔ f ( x ) = ucf ( u−1d x ) . ( 1 ) Lie Groups and Lie Algebras The full group of invertible d × d matrices over R is the general linear group , denoted as GLd ( R ) . It follows that every real d- dimensional group representation Td ( G ) ⊂ GLd ( R ) . If T is a “ faithful representation ” ( i.e . T ( u ) 6= T ( v ) if u 6= v ) , thenG ⊂ GLd ( R ) . In our problem , we only know the group G through its linear action on X or on output hl of layer l in a neural network . Therefore we may assume the representation Tdl acting on h l with largest dl is faithful . GLd ( R ) is a Lie group and all of its continuous subgroups G are also Lie groups . Next , we will first briefly review the Lie algebra of a Lie group and then use it to introduce our equivariant architecture . Notation Unless stated or obvious , a in Aa is an index , not an exponent . We write matrix products as A · B ≡ ∑ aA aBa Recall that x ∈ Rd×m . For a linear transformation A : Rd1 → Rd2 acting on the spatial index or the channel index , we will use one upper and one lower index as in Aµνhµ = [ A · h ] ν . We will use ( a , b , c ) for channels , and ( µ , ν , ρ ) for spatial , and ( i , j , k ) for Lie algebra basis indices . We will occasionally keep explicit summation ∑ i for clarity . For any Lie group G ⊂ GLd ( R ) , group elements infinitesimally close to the identity element can be written as u = I + iLi . We can find a basis Li which span the “ Lie algebra ” g of the group G , meaning they are closed under commutations [ Li , Lj ] = LiLj − LjLi = ∑ k fij kLk ( 2 ) with fijk called the structure constants of the Lie algebra . The Li are called the infinitesimal generators of the Lie group . They define vector fields spanning the tangent space of the manifold of G near its identity element I . For elements u ∈ GI in the connected component GI ⊂ G containing the identity , an exponential map exp : g→ GI can be defined such that u = exp [ t · L ] . For matrix groups , if G is connected and compact , the matrix exponential defined through a Taylor expansion 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 ( including GI ) every element u can be written as a product u = ∏ α exp [ tα · L ] using the matrix exponential ( Hall , 2015 ) . We will use this fact to modify existing results about group equivariant architectures to introduce L-conv . 2 GROUP EQUIVARIANT ARCHITECTURE . Consider the case where the solution f to y = f ( x ) is implemented as a feedforward neural network . Denote the linear output of layer l by hl = Fl ( hl−1 ) ∈ Rdl×ml , with h0 = xi . Kondor & Trivedi ( 2018 ) ; Cohen & Welling ( 2016a ) showed that a feedforward neural network is equivariant under the action of a group G if and only if each of its layers implement a group convolution ( G-conv ) given by hl = ( fl ∗ gl ) ( hl−1 ) = ∫ G dµ ( u ) fl ( u −1 l−1hl−1 ) gl ( u ) ( 3 ) where dµ ( u ) is the Haar measure on the group manifold , ul is an appropriate dl-dimensional representation of G , and gl : G → Rml×ml+1 is a set of convolution kernels , where ml and ml+1 are the number of input and output channels ( filters ) . Here fl are point-wise activation functions . In the formalism of Kondor & Trivedi ( 2018 ) , we may consider hl to be a map from the homogeneous space Xl to Rml where Xl is a space such that the space of maps LR [ Xl ] is parameterized by Rdl . To use equation 3 in practice , we need to make the integral over the group manifold tractable . Although , Finzi et al . ( 2020 ) , which approximates the integral overG by sampling , approach the idea of working in the Lie algebra , they use a logarithm map to define the Lie algebra . We take a different approach and formulate the architecture in terms of the Lie algebra basis . We also allow the architecture to learn basis Li for the symmetries . For continuous symmetry groups we will now show that using the Lie algebra can approximate G-conv without having to integrate over the full group manifold , thus alleviating the need to discretize or sample the group . The key point is the discreteness and finiteness of the Lie algebra basis for many Lie groups . Proposition 1 . Let G be a Lie group and Td : G → Rdl a dl-dimensional representation of G. If a convolution kernel gl : G → Rml−1×ml has support only on an infinitesimal η neighborhood of identity , a G-conv layer of equation 3 can be written in terms of the Lie algebra . Proof : Linearization over η yields ul ≈ I+ ·L in equation 3 , with Li ∈ Rdl×dl . The inverse is then u−1l ≈ I − · L+O ( 2 ) . Since gl has support only in an η neighborhood of identity , fixing a basis Li , we can reparametrize gl ( I + ·L ) = g̃l ( ) as a function over the Lie algebra g̃l : g→ Rml×ml+1 . Dropping the layer index l for brevity , we have ( f ∗ g ) ( h ) ≈ ∫ G g̃ ( ) f ( ( I − · L ) h ) dnL ≈ ( W 0I −W · L ) h · ∇hf ( h ) ( 4 ) W 0 ≡ ∫ G g̃ ( ) dnL , W i ≡ ∫ G ig̃ ( ) dnL . ( 5 ) where the number of W i is nL , the number of Li . When f ( h ) is a homogeneous function like ReLU or linear activation , h · ∇hf ( h ) = af ( h ) . In general , since we are designing the layers , we can choose h · ∇hf ( h ) = σ ( h ) to be a desired activation function . Note that in G-conv , the convolutional kernels g ( u ) are the trainable weights . The difficulty with learning them is that one has to parametrize and integrate the convolution over the group G. The major advantage of working in the Lie algebra instead of G is that in equation 4 the integrals of g ( ) separate from f ( h ) , leaving us with a set of integrated weights W . This means that instead of the architecture needing to implement the integral ∫ G g ( u ) f ( uh ) dµ , we only need to learn a finite number of convolutional weights W , like the filter weights in CNN . Additionally , we can learn Li . Indeed , having learnable Li to learn hidden symmetries is one of main features we are proposing , as we show in our experiments . L-conv Layer The construction in equation 4 is at the core of our results . h ∈ Rdl−1×ml−1 has components haµ , where µ ∈ { 1 , · · · , dl−1 } is the spatial index and a ∈ { 1 , · · ·ml−1 } the channel index . We define the Lie algebra convolutional layer ( L-conv ) as L-conv : FL ( h ) = σ ( nL∑ i=0 Li · h ·W i + b ) FL ( h ) a µ = σ ∑ ν , c , i [ Li ] ν µ h c ν [ W i ] a c + ba ( 6 ) where σ is the activation function , L0 = I and W i ∈ Rml×ml−1 are the convolutional weights , and the second line shows the components ( all symbols are indices , not exponents ) . Note that , similar to discussion in Weiler et al . ( 2018a ) , an arbitrary nonlinear activation σ may not keep the architecture equivariant under G , but linear activation remains equivariant . Also , note that in equation 3 , fl is the activation σl−1 of the previous layer . In defining L-conv we are including the activation σ after the convolution , which is why σ appears outside in equation 6 instead of as σ ( h ) . Since in connected Lie groups larger group elements can be constructed as u = ∏ α exp [ tα ·L ] ( Hall , 2015 ) from elements near identity , it follows that any G-conv layer can be constructed from multiple L-conv layers on these groups . Proposition 2 . Any G-conv layer on a connected Lie group can be approximated by a multi-layer L-conv to arbitrary accuracy . Proof : To show this we need to take two steps . First , we discretize the support of g ( u ) in equation 3 to infinitesimally small sets G = ⋃ kGk , each having support around an element uk ∈ Gk . This allows us to write ( f ∗ g ) ( h ) ≈ ∑ k ∆µkg ( uk ) f ( u −1 k h ) , where ∆µk = ∫ Gk dµ . Next , we write each uk as a product . On connected groups for any uk ∈ G , we can write uk = ∏l α=1 vα where vα = exp [ tα · L ] ≈ ( I + tα · L ) are elements close to identity . Then , f ( uh ) can be written as a composition l G-conv layers , each with weights gα ( u ) = δ ( u − vα ) ( defined using a faithful representation of G ) f ( u−1k h ) = f ( ∏ α vαh ) = f ( Fl ◦ · · · ◦ F1 ( h ) ) , Fα ( h ) = ∫ G dµ ( u ) fα ( uαh ) gα ( u ) ( 7 ) Since uα are near identity , each Fα can be converted to an L-conv layer with linear activation and weights W = tα and W 0 = 1 . Thus , the G-conv layer can be approximated as sum and composition of L-conv layers . Each Gk can be considered one output channel for the multi-layer L-conv and g ( uk ) are weights of a final aggregation layer . So far we have shown that L-conv layers could replace G-conv for connected groups . The advantage of using L-conv layers is that for many potential symmetry groups , such as orthogonal and unitary groups , the number of generators Li is finite and often a small number of them may be sufficient . For instance , consider SO ( 2 ) rotations on p× p images . The flattened image has d = p2 dimensions and the number of GLd ( R ) generators is p4 . However the induced representation of SO ( 2 ) on Rd has a single generator . In fact , in many domains we expect symmetries hidden in the data to be much lower in dimensionality than the data itself . Additionally , in the other extreme , where there is no restriction on the symmetry , meaning G = GLd ( R ) , we observe that L-conv becomes a fully-connected layer , as shown next . Proposition 3 . A fully-connected neural network layer can be written as an L-conv layer using GLd ( R ) generators , followed by a sum pooling and nonlinear activation . Proof : The generators of GLd ( R ) are one-hot E ∈ Rd×d matrices Li = E ( α , β ) which are non-zero only at index i = ( α , β ) 1 with elements written using Kronecker deltas GLd ( R ) generators : Lνi , µ = [ E ( α , β ) ] νµ = δµαδνβ ( 8 ) Now , consider the weight matrix w ∈ Rm×d and bias b ∈ Rm of a fully connected layer acting on h ∈ Rd as F ( h ) = σ ( w · h+ b ) . The matrix element can be written as wνb = ∑ µ ∑ α , β wαb 1β [ E ( α , β ) ] ν µ = ∑ µ ∑ α , β W b,1 ( α , β ) [ E ( α , β ) ] ν µ = ∑ µ W b,1 · [ L ] νµ ( 9 ) meaning an L-conv layer with weights W b,1 ( α , β ) = w α b 1β ( 1 input channel , and 1 being a vector of ones ) followed by pooling over µ is the same as a fully connected layer with weights w. Thus , interestingly , more restricted groups , rather than large , meaning groups with fewer Li , lead to more parameter sharing . Indeed , as we show below , the most restricted L-conv are graph convolutional networks . It is also easy to see that familiar CNN can be encoded as L-conv . Note that , generally shift operators used in CNN are thought of as group elements for a discrete permutation group . However , the same shift operators Li can be used to create continuous fractional shifts as ( I + αLi ) / ( α+ 1 ) . Since shifts commute with each other , they form a basis for the Lie algebra of an Abelian subgroup of GLd ( R ) . Proposition 4 . CNN is a special case of L-conv where the generators are shift operators . Proof : In 1D CNN , for a kernel size k we have k shift operators Li given by Liνµ = δµ , ν−i . Plugging this into equation 6 and doing the sum over ν , we recover the convolution formula F ( h ) aµ = σ ∑ i , c hcµ−iW i , a c + b a ( 10 ) In higher dimensional CNN , i covers the relevant flattened indices for the required shifts . For instance , a p× q image , when flattened becomes a vector of length pq . To cover a ( 2 , 2 ) convolutional kernel , we need shifts to be by 0 , 1 , q , q + 1 pixels . Thus , we need four Li given by Li ν µ = δµ , ν−si , si = [ 0 , 1 , q , q + 1 ] i ( 11 ) This means that the number of generators Li is the size of the convolutional kernel in CNN . Generalization to higher dimensions is straightforward . In CNN , covering larger patches is equivalent to g ( u ) covering a larger part of the group . But large convolutional kernels are not generally used . Instead , we achieve larger receptive fields by adding more CNN layers , similar to Proposition 2 . 1We may also label them by single index like i = α+ βd , but two indices is more convenient . Connection with Graph Convolutional Networks Finally , we note that the L-conv layer equation 6 has the structure of a Graph Convolutional Networks ( GCN ) Kipf & Welling ( 2016 ) with multiple propagation rules Li derived from graph adjacency matrices and µ , ν indexing graph vertices . Thus , an L-conv with a single L is a GCN . In fact , there is a deeper connection here , as we show now . Proposition 5 . A GCN with propagation rule A ∈ Rd×d is equivalent to an L-conv , with the group being 1D flows generated by A− I . Proof : We can define a linear flow similar to a diffusion equation h ( t+ δt ) = Ahν ( t ) ⇒ dh ( t ) dt = λ ( A− I ) h ( t ) ( 12 ) where λ = δt/dt sets the time scale . Thus L = λ ( A− I ) is the generator a 1D group of flows with elements u = exp [ tL ] , a subgroup of the group of diffeomorphisms on Rd , with a single generator L. Thus , a GCN with propagation rule A is an L-conv using L0 = I and L1 = L with the same convolutional weights for L0 , L1 . Hence , the most restricted L-conv based on 1D flow groups with a single generator are GCN . These flow groups include Hamiltonian flows and other linear dynamical systems . CNN can also be interpreted as a GCN with multiple propagation rules , where each shift operator is a subgraph of the grid network .

Summary: The paper meticulously builds a theoretical framework for Lie-algebra convolutional layers and then goes on to show how CNNs, GCNs and FC layers are a special case of L-convs. The paper also demonstrates how the underlying generators can be learnt from data and provide convincing supporting experimental results. The proposed L-conv layers also use much fewer parameters compared to previous works.