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: The paper describes a generalization of group equivariant neural networks that typically rely on linear operators such as the group convolution, or pseudo linear operators such as attentive group convolution. The authors take a look at the aggregation function of the convolution operator (which is usually linear via a kernel times feature values followed by sum aggregation) and replaces it with a non-linear function that generally could depend not just on relative positions, but also on the central and neighbor feature values. The result is a more non-linear operator than group convolutions, which are a special case of this framework.  The paper is accurate and has an appropriate experimental section with decent ablation studies. The proposed work compares favorably compared to the baselines both in terms of accuracy, parameters and nr of operations (FLOPs).