Group Equivariant Subsampling

1 Introduction . Convolutional Neural Networks ( CNNs ) are known to be more data efficient and show better generalisation on perceptual tasks than fully-connected networks , due to translation equivariance encoded in the convolutions : when the input image/feature map is translated , the output feature map also translates by the same amount . In typical CNNs , convolutions are used in conjunction with subsampling operations , in the form of pooling or strided convolutions , to reduce the spatial dimensions of feature maps and to allow receptive field to grow exponentially with depth . Subsampling/upsampling operations are especially necessary for convolutional autoencoders ( ConvAEs ) ( Masci et al. , 2011 ) because they allow efficient dimensionality reduction . However , it is known that subsampling operations implicit in strided convolutions or pooling layers are not translation equivariant ( Zhang , 2019 ) , hence CNNs that use these components are also not translation invariant . Therefore such CNNs and ConvAEs are not guaranteed to generalise to arbitrarily translated inputs despite their convolutional layers being translation equivariant . Previous work , such as Zhang ( 2019 ) ; Chaman and Dokmanić ( 2020 ) , has investigated how to enforce translation invariance on CNNs , but does not study equivariance with respect to symmetries beyond translations , such as rotations or reflections . In this work , we first describe subsampling/upsampling operations that preserve exact translation equivariance . The main idea is to sample feature maps on an input-dependent grid rather than a fixed one as in pooling or strided convolutions , and the grid is chosen according to a sampling index computed from the inputs ( see Figure 1 ) . Simply replacing the subsampling/upsampling in standard CNNs with such translation equivariant subsampling/upsampling operations leads to CNNs and transposed CNNs that can map between spatial inputs and lowdimensional representations in a translation equivariant manner . ∗Corresponding author : < jin.xu @ stats.ox.ac.uk > 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We further generalise the proposed subsampling/upsampling operations from translations to arbitrary groups , proposing group equivariant subsampling/upsampling . In particular we identify subsampling as mapping features on groups G to features on subgroups K ( vice versa for upsampling ) , and identify the sampling index as a coset in the quotient space G/K . See Appendix A for a primer on group theory that is needed to describe this generalisation . We note that group equivariant subsampling is different to coset pooling introduced in Cohen and Welling ( 2016 ) , which instead gives features on the quotient space G/K , and discuss differences in detail in Section 4 . Similar to the translation equivariant subsampling/upsampling , group equivariant subsampling/upsampling can be used with group equivariant convolutions to produce group equivariant CNNs . Using such group equvariant CNNs we can construct group equivariant autoencoders ( GAEs ) that separate representations into an invariant part and an equivariant part . While there is a growing body of literature on group equivariant CNNs ( G-CNNs ) ( Cohen and Welling , 2016 , 2017 ; Worrall et al. , 2017 ; Weiler et al. , 2018b , a ; Thomas et al. , 2018 ; Weiler and Cesa , 2019a ) , such equivariant convolutions usually preserve the spatial dimensions of the inputs ( or lift them to even higher dimensions ) until the final invariant pooling layer . There is a lack of exploration on how to reduce the spatial dimensions of such feature maps while preserving exact equivariance , to produce low-dimensional equivariant representations . This work attempts to fill in this gap . Such lowdimensional equivariant representations can be employed in representation learning methods , allowing various advantages such as interpretability , out-of-distribution generalisation , and better sample complexity . When using such learned representations in downstream tasks such as abstract reasoning , reinforcement learning , video modelling , scene understanding , it is especially important for representations to be equivariant rather than invariant in these tasks , because transformations and how they act on feature spaces are critical information , rather than nuisance as in image classification problems . In summary , we make the following contributions : ( i ) We propose subsampling/upsampling operations that preserve translational equivariance . ( ii ) We generalise the proposed subsampling/upsampling operations to arbitrary symmetry groups . ( iii ) We use equivariant subsampling/upsampling operations to construct GAEs that gives low-dimensional equivariant representations . ( iv ) We empirically show that representations learned by GAEs enjoys many advantages such as interpretability , out-of-distribution generalisation , and better sample complexity . 2 Equivariant Subsampling and Upsampling . 2.1 Translation Equivariant Subsampling for CNNs . In this section we describe the proposed translation equivariant subsampling scheme for feature maps in standard CNNs . Later in Section 2.2 , we describe how this can be generalised to group equivariant subsampling for feature maps on arbitrary groups . Standard subsampling Feature maps in CNNs can be seen as functions defined on the integer grid , e.g . Z for 1D feature maps , and Z2 for 2D . Hence we represent feature maps as f : Z → Rd , where d is the number of feature map channels . For simplicity , we start with 1D and move on to the 2D case . Typically , subsampling in CNNs is implemented as either strided convolution or ( max ) pooling , and they can be decomposed as CONVck = SUBSAMPLING c ◦ CONV1k MAXPOOLck = SUBSAMPLING c ◦ MAXPOOL1k where subscripts denote kernel sizes and superscripts indicate strides . c ∈ N is the scale factor for SUBSAMPLING , and this operation simply restricts the input domain of the feature map from Z to cZ , without changing the corresponding function values . Translation equivariant subsampling In our equivariant subsampling scheme , we instead restrict the input domain to cZ + i , the integers ≡ i mod c , where i is a sampling index determined by the input feature map . The key idea is to choose i such that it is shifts by t ( modc ) when the input is translated by t , to ensure that the same features are subsampled upon translation . Let i be given by the mapping Φc : IZ → Z/cZ . IZ denotes the space of vector functions on Z and Z/cZ is the space of remainders upon division by c. i = Φc ( f ) = mod ( argmax x∈Z ∥f ( x ) ∥1 , c ) ( 1 ) where ∥ · ∥1 denotes L1-norm ( other choices of norm are equally valid ) . Other choices for Φc are equally valid as long as they satisfy translation equivariance , ensuring that the same features are subsampled upon translation of the input : Φc ( f ( · − t ) ) = mod ( Φc ( f ) + t , c ) . ( 2 ) Note that this holds for Equation ( 1 ) provided the argmax is unique , which we assume for now ( see Appendix B.1 for a discussion of the non-unique case ) . We can decompose the subsampled feature map defined on cZ + i into its values and the offset index i , expressing it as [ fb , i ] ∈ ( IcZ , Z/cZ ) , where fb is the translated output feature map such that fb ( cx ) = f ( cx+ i ) for x ∈ Z . The subsampling operation described above , which maps from IZ to ( IcZ , Z/cZ ) is translation equivariant : when the feature map f is translated to the right by t ∈ Z , one can verify that fb will be translated to the right by c⌊ i+tc ⌋ , and the sampling index for the translated inputs would become mod ( i+ t , c ) . We provide an illustration for c = 2 in Figure 1 , and describe formal statements and proofs later for the general cases in Section 2.2 . Multi-layer case For the subsequent layers , the feature map fb is fed into the next convolution , and the sampling index i is appended to a list of outputs . When the above translation equivariant subsampling scheme is interleaved with convolutions in this way , we obtain an exactly translation equivariant CNN , where each subsampling layer with scale factor ck produces a sampling index ik ∈ Z/ckZ . Hence the equivariant representation output by the CNN with L subsampling layers is a final feature map fL and a L-tuple of sampling indices ( i1 , . . . , iL ) . This tuple can in fact be expressed equivalently as a single integer by treating the tuple as mixed radix notation and converting to decimal notation . We provide details of this multi-layer case in Appendix B.2 , including a rigorous formulation and its equivariance properties . Translation equivariant upsampling As a counterpart to subsampling , upsampling operations increase the spatial dimensions of feature maps . We propose an equivariant upsampling operation that takes in a feature map f ∈ IcZ and a sampling index i ∈ Z/cZ , and outputs a feature map fu ∈ IZ , where we set fu ( cx + i ) = f ( cx ) and 0 everywhere else . This works well enough in practice , although in conventional upsampling the output feature map is often a smooth interpolation of the input feature map . To achieve this with equivariant upsampling , we can additionally apply average pooling with stride 1 and kernel size > 1 . 2D Translation equivariant subsampling When feature maps are 2D , they can be represented as functions on Z2 . The sampling index becomes a 2-element tuple given by : ( x∗ , y∗ ) = argmax ( x , y ) ∈Z2∥f ( x ) ∥1 ( i , j ) = ( mod ( x∗ , c ) , mod ( y∗ , c ) ) and we subsample feature maps by restricting the input domain to cZ2 + ( i , j ) . The multi-layer construction and upsampling is analogous to the 1D-case . 2.2 Group Equivariant Subsampling and Upsampling . In this section , we propose group equivariant subsampling by starting off with the 1D-translation case in Section 2.1 , and provide intuition for how each component of this special case generalises to arbitrary discrete groups G. We then proceed to mathematically formulate group equivariant subsampling , and prove that it is indeed G-equivariant . Feature maps on groups First recall that the feature maps for the 1D-translation case were defined as functions on Z , or f ∈ IZ for short . To extend this to the general case , we consider feature maps f as functions on a group G , i.e . f ∈ IG = { f : G → V } 2 where V is a vector space , as is done in e.g . group equivariant CNNs ( G-CNNs ) ( Cohen and Welling , 2016 ) . Note that translating feature maps f on Z by displacement u is effectively defining a new feature map f ′ ( · ) = f ( · − u ) . In the general case , we say that the group action on the feature space is given by [ π ( u ) f ] ( g ) = f ( u−1g ) ( 3 ) where π is a group representation describing how u ∈ G acts on the feature space . Recap : translation equivariant subsampling Recall that standard subsampling that occurs in pooling or strided convolutions for 1D translations amounts to restricting the domain of the feature map from Z to cZ , whereas equivariant subsampling also produces a sampling index i ∈ Z/cZ , an integer mod c , and that this is equivalent to restricting the input domain to cZ + i. i is given by the translation equivariant mapping Φc : IZ → Z/cZ . We can translate the input domain back to cZ , and represent the output of subsampling as [ fb , i ] ∈ ( IcZ , Z/cZ ) , where fb is the translated output feature map and fb ( cx ) = f ( cx+ i ) for x ∈ Z . Group equivariant subsampling Similarly in the general case , for a feature map f ∈ IG , standard subsampling can be seen as restricting the domain from the group G to a subgroup K , whereas equivariant subsampling additionally produces a sampling index pK ∈ G/K , where the quotient space G/K = { gK : g ∈ G } is the set of ( left ) cosets of K in G. Note that we have rewritten i as p to distinguish between the 1D translation case and the general group case . This is equivalent to restricting the f to the coset pK . The choice of the coset pK is given by equivariant map Φ : IG → G/K ( the action of G on G/K is given by u ( gK ) = ( ug ) K for u , g ∈ G ) , such that pK = Φ ( f ) . This restriction of f to pK can also be thought of as having an output feature map fb on K and choosing a coset representative element p̄ ∈ pK , such that fb ( k ) = f ( p̄k ) . This choice of coset representative is described by a function s : G/K → G , such that p̄ = s ( pK ) . The function s is called a section and should satisfy s ( pK ) K = pK . Now let us formulate subsampling and upsampling operations Sb↓GK and Su↑ G K mathematically and prove its G-equivariance . Let IK = { f : K → V ′ } be the space of feature map on K. Sb↓GK takes in a feature map f ∈ IG and produces a feature map fb ∈ IK and a coset in G/K . In reverse , the upsampling operation Su↑GK takes in a feature map in IK , a coset in G/K , and produces a feature map in IG . We use a section s : G/K → G to represent a coset with a representative element in G , and point out that equivariance holds for any choice of s. Formally , given an equivariant map Φ : IG → G/K ( we will discuss how to construct such a map in Section 2.3 ) , and a fixed section s : G/K → G such that p̄ = s ( pK ) , the subsampling operation Sb↓GK : IG → IK ×G/K is defined as : pK = Φ ( f ) , fb ( k ) = f ( p̄k ) for k ∈ K [ fb , pK ] = Sb↓GK ( f ; Φ ) , ( 4 ) while the upsampling operation Su↑GK : IK ×G/K → IG is defined as : fu ( g ) = f ( p̄ −1g ) if g ∈ K else 0 fu = Su↑GK ( f , pK ) . ( 5 ) 2This is not to be confused with the space of Mackey functions in , e.g. , Cohen et al . ( 2019 ) , and rather it is the space of unconstrained functions on G. To make the output of the upsampling dense rather than sparse , one can apply arbitrary equivariant smoothing functions such as average pooling with stride 1 and kernel size > 1 , to compensate for the fact that we extend with 0s rather than values close to their neighbours . In practice , we observe that upsampling without any smoothing function works well enough . The statement on the equivariance of Sb↓GK and Su↑ G K requires we specify the action of G on the space IK ×G/K , which we denote as π′ . For any u ∈ G , p′K = upK , f ′b = π ( p̄ ′−1up̄ ) fb [ f ′b , p ′K ] = π′ ( u ) [ fb , pK ] ( 6 ) Lemma 2.1. π′ defines a valid group action of G on the space IK ×G/K . We can now state the following equivariance property ( See Appendix D for a proof ) : Proposition 2.2 . If the action of group G on the space IG and IK × G/K are specified by π , π′ ( as defined in Equations ( 3 ) and ( 6 ) ) , and Φ : IG → G/K is an equivariant map , then the operations Sb↓GK and Su↑ G K as defined in Equations ( 4 ) and ( 5 ) are equivariant maps between IG and IK ×G/K . In fact , we can also prove the converse ( See Appendix D ) : Proposition 2.3 . If Sb↓GK : IG → IK ×G/K ( as defined in Equation ( 4 ) ) is an equivariant map , then the corresponding Φ : IG → G/K must be equivariant . The above implies that Φ must depend on the input feature map f .

Summary: The authors propose subsampling and upsampling layers that are equivariant to common transformation groups. These constructions are used to design autoencoders consisting of layers that alternate between group equivariant convolutions and group equivariant subsampling/upsampling layers. The result is that the entire autoencoder is *exactly* group equivariant. Experiments on translations, translations+rotations, translations+rotations+reflections for simple images with single/multiple objects clearly show the improved performance over subsampling methods that are not group-equivariant.