Efficient Equivariant Network

1 Introduction . In the past few years , convolutional neural networks ( CNNs ) have been widely used and achieved superior results on multiple vision tasks , such as image classification [ 31 , 55 , 51 , 22 ] , semantic segmentation [ 3 ] , and object detection [ 44 ] . A compelling explanation of the good performance of CNNs is that their built-in parameter sharing scheme brings in translation equivariance : shifting an image and then feeding it through a CNN layer is the same as feeding the original image and then shifting the resulted feature maps . In other words , the translation symmetry is preserved by each layer . Motivated by this , Cohen and Welling [ 9 ] proposed Group Equivariant CNNs ( G-CNNs ) , showing how convolutional networks can be generalized to exploit larger groups of symmetries . Following G-CNNs , researchers have designed new neural networks that are equivariant to other transformations like rotations [ 9 , 61 , 24 , 49 ] and scales [ 65 , 53 ] . However , G-CNNs still have two main drawbacks : 1 ) In the implementation , G-CNNs would introduce extra dimensions to encode new transformations , such as rotations and scales , thus have a very high computational cost . 2 ) Although G-CNNs achieve group equivariance by sharing kernels , like vanilla CNNs , they lack the ability to adapt kernels to diverse feature patterns with respect to different spatial positions , namely , the spatial-agnostic problem [ 68 , 39 , 70 , 71 , 54 , 67 , 36 ] . ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Some previous works focus on solving these two problems . Cheng et al . [ 4 ] proposed to decompose the convolutional filters over joint steerable bases to reduce model size . However , it is essentially G-CNNs which still have the inherent spatial-agnostic problem . To incorporate dynamic filters , one solution is introducing attention mechanism into each convolution layer in G-CNNs without disturbing inherent equivariance [ 48 , 45 ] . The cost is that they introduce extra parameters and increase the complexity of space and time . Another solution is to replace group convolution layers with standalone self-attention layers by designing a specific position embedding to ensure equivariance [ 47 , 26 ] . However , the self-attention mechanism suffers from quadratic memory and time complexity , because it has to compute the attention score at each pair of inputs . Actually , Cohen et al . [ 7 ] , Kondor et al . [ 29 ] and Bekkers [ 1 ] revealed that an equivariant linear layer is essentially a convolution-like operation . Inspired by this , we further discover that a general feature-extraction layer , either linear or non-linear , being equivariant is equivalent to that the feature aggregation mechanism between each pair of inputs only depends on the relative positions of these two inputs . Based on this observation , we propose a generalized framework of previous equivariant models , which includes G-CNNs and equivariant attention networks as special cases . Under this generalized framework , we design a new equivariant layer to conquer the aforementioned difficulties . Firstly , to avoid quadratic computational complexity , the feature aggregation operator is explicitly decomposed into a kernel generator and an encoder which takes one single feature as the input . Since our kernels are calculated based on input features , they are essentially dynamic rather than being spatial-agnostic . In addition , we decouple the feature aggregation mechanism across spatial and extra geometric dimensions to reduce the inter-channel redundancy in convolution filters [ 4 ] and further accelerate computation . Extensive experiments show that our method can process data very efficiently and perform significantly better than previous works using lower computational cost . As our method is parameter Efficient , computational Efficient , data Efficient and Equivariant , we name our new layer as E4-layer . We summarize our main contributions as follows : • We propose a generalized framework of previous equivariant models , which includes GCNNs and attention-based equivariant models as special cases . • Under the generalized framework , we explicitly decompose the feature aggregation operator into a kernel generator and an encoder , and further decouple the spatial and extra geometric dimensions to reduce computation . • Extensive experiments verify that our method is also data efficient and performs competitively with lower computational cost . 2 Related Work . Vanilla CNNs [ 34 ] are naturally translation equivariant . More symmetries are considered to be exploited into the network for different tasks , such as rotations over plane [ 9 , 66 , 35 , 12 , 61 , 59 , 49 , 37 , 52 , 4 , 41 , 2 , 10 , 58 , 21 , 24 ] , rotations over 3D space [ 62 , 16 , 57 , 64 , 14 , 60 , 15 , 50 , 28 , 6 ] , scaling [ 65 , 40 , 53 , 46 ] , symmetries on manifold [ 8 , 11 ] , and other general symmetry groups [ 17 , 56 , 18 ] . These works accomplish equivariance by constraining the linear mappings in layers , followed by pointwise non-linearities to enhance their expressive power . In general , researchers [ 29 , 7 , 1 ] pointed out that an equivariant linear mapping can always be written as a convolution-like integral , i.e. , G-CNNs in practice . However , their theory is still limited to linear cases . As works [ 68 , 39 , 70 , 71 , 54 , 67 , 36 ] point out the spatial-agnostic problem of CNNs and attention mechanisms [ 25 , 63 , 43 , 13 , 20 ] achieve impressive results on various vision tasks , researchers start to consider non-linear equivariant mapping . Romero et al . [ 48 , 45 ] directly reweighted the convolution kernels with attention weights generated by features and obtained non-linear equivariant models . However , compared with G-CNNs , these methods introduce extra parameters and operations , resulting in an even heavier computational burden . Also , some works [ 47 , 26 , 23 ] proposed group equivariant self-attention [ 43 , 13 ] . Fuchs et al . [ 19 ] incorporated self-attention into 3D equivariant networks and proposed SE ( 3 ) -Transformers . However , since their filters are essentially calculated based on a pair of inputs , the computational complexity is quadratic . In this work , we further extend the linear equivariant theory to a more general situation , including non-linear cases . Under the framework , we design a new equivariant layer to solve both the spatial- agnostic problem in convolution-based equivariant models and heavy computation cost problem in most equivariant models . 3 A Unified Framework of Previous Group Equivariant Models . In this section , we first briefly review two representative group equivariant models : the linear model G-CNNs [ 9 ] , and the non-linear model equivariant self-attention [ 47 , 26 ] . Then , we propose a general framework of previous equivariant models based on the inner relationship among these specific models . 3.1 Equivariance . Equivariance indicates that the outputs of a mapping transform in a predictable way with the transformation of the inputs . Formally , a group equivariant map Ψ satisfies that ∀u ∈ G , Ψ [ Tu [ f ] ] = T ′u [ Ψ [ f ] ] , ( 1 ) where G is a transformation group , f is an input feature map , and Tu and T ′u are group actions , indicating how the transformation u acts on the input and output features , respectively . Besides , since we hope that two transformations u , v ∈ G acting on the feature maps successively is equivalent to the composition of transformations uv ∈ G acting on the feature maps directly , we require that TuTv = Tuv , where uv is the group product of u and v. The same is the case with T ′u . Now we examine the specific form of the transformation group G. In this work , we focus on the analysis of 2D images defined on R2 . Consequently , we are most interested in the groups of the form G = R2 o A , resulting from the semi-product ( o ) between the translation group R2 and a group A acts on R2 , e.g. , rotations , scalings and mirrorings . This family of groups is referred to as affine groups and their group product rule is : uv = ( xu , au ) ( xv , av ) = ( xu + auxv , auav ) , ( 2 ) where u = ( xu , au ) and v = ( xv , av ) , in which xu , xv ∈ R2 and au , av ∈ A . For ease of implementation , following [ 9 ] , we take A as the cyclic group C4 or the dihedral group D4 , then G becomes p4 or p4m . As for the group action , we employ the most common regular group action in this work , i.e. , Tu [ f ] ( v ) = f ( u−1v ) . ( 3 ) Here , we only care about the group action over the feature maps defined on G , because we always use a lifting operation to lift the input images defined on R2 to the feature maps on G , where the equivariance can be preserved properly , as will be shown in Section 3.2 . 3.2 G-CNNs . Let f ( l ) : X → RCl and W : G → RCl+1×Cl be the input feature and the convolutional filter in the l-th layer , respectively , where Cl denotes the channel number of the l-th layer . X is taken as R2 for the first layer , and taken as G for the following layers . Then for any g ∈ G , the group convolution [ 29 , 7 , 1 ] of f ( l ) and W on G at g is given by f ( l+1 ) ( g ) = Ψ [ f ( l ) ] ( g ) = ∫ X W ( g−1g̃ ) f ( l ) ( g̃ ) dµ ( g̃ ) , ( 4 ) where µ ( · ) is the Haar measure . When X is discrete , Eqn . ( 4 ) can be rewritten as f ( l+1 ) ( g ) = ∑ g̃∈X W ( g−1g̃ ) f ( l ) ( g̃ ) . ( 5 ) G-CNNs essentially generalize the translation equivariance of conventional convolution to a more general group G. In fact , the first layer maps the 2D images to a function defined on G , while the following layers map one feature map on G to another . As a result , the computational complexity of the first layer and the following layers are of the order O ( k2|A| ) and O ( k2|A|2 ) , respectively , where k is the kernel size in the spatial space . As a result , G-CNNs have a much larger computational cost when A is large , especially for the intermediate layers . In this work , we employ the first layer of G-CNNs as a lifting operation , and focus on reducing the computation of the latter layers . 3.3 Equivariant Attention Networks . Group Equivariant Self-Attention ( G-SA ) [ 47 , 26 ] is a representative method of equivariant attention networks , whose form can be simplified as follows : f ( l+1 ) ( g ) = ∑ g̃∈G Softmaxg̃ [ h T Q ( f ( l ) ( g ) ) ( hK ( f ( l ) ( g̃ ) ) + Pg−1g̃ ) ] hV ( f ( l ) ( g̃ ) ) , ( 6 ) where hV : RCl → RCl+1 , and hQ , hK : RCl → Rd are the embedding functions of values , querys and keys , respectively , which are neural networks in the most general case . d is the dimension of the low dimensional embeddings , and Pg−1g̃ ∈ Rd encodes the relative positions of the query f ( l ) ( g ) and the key f ( l ) ( g̃ ) .

Summary: This paper introduces a functional abstraction for group-convolutional networks employing first- and second-order features. Based on this functional abstraction, a new functional form for second-order group-convolutional networks is devised, which can be made more parameter-efficient than existing approaches. Experiments are performed on rotated MNIST and CIFAR, evaluating error, number of parameters and certain ablations.