Group Equivariant Subsampling
Year: 2021

Abstract:
Subsampling is used in convolutional neural networks (CNNs) in the form of pooling or strided convolutions, to reduce the spatial dimensions of feature maps and to allow the receptive fields to grow exponentially with depth. However, it is known that such subsampling operations are not translation equivariant, unlike convolutions that are translation equivariant. Here, we first introduce translation equivariant subsampling/upsampling layers that can be used to construct exact translation equivariant CNNs. We then generalise these layers beyond translations to general groups, thus proposing group equivariant subsampling/upsampling. We use these layers to construct group equivariant autoencoders (GAEs) that allow us to learn low-dimensional equivariant representations. We empirically verify on images that the representations are indeed equivariant to input translations and rotations, and thus generalise well to unseen positions and orientations. We further use GAEs in models that learn object-centric representations on multiobject datasets, and show improved data efficiency and decomposition compared to non-equivariant baselines.

heading: 1 Introduction
text: 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.

heading: 2 Equivariant Subsampling and Upsampling

heading: 2.1 Translation Equivariant Subsampling for CNNs
text: 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.

heading: 2.2 Group Equivariant Subsampling and Upsampling
text: 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 .

heading: 2.3 Constructing Φ
text: We use the following simple construction of the equivariant mapping Φ : IG → G/K for subsampling/upsampling operations, although any equivariant mapping would suffice. For an input feature map f ∈ IG, we define
pK = Φ(f) := (argmax g∈G ∥f(g)∥1)K (7)
Provided that the argmax is unique, it is easy to show that (up) · K = Φ(π(u)f), hence Φ is equivariant. In practice one can insert arbitrary equivariant layers to f before and after we take the norm ∥ · ∥1 to avoid a non-unique argmax (see Appendix F). Note that the argmax function alone may not be noise-robust. In Appendix E.2, we empirically show that applying smoothing equivariant layers before taking the argmax would improve the stability of the output sampling indices.
Non-unique argmax case When the input feature map f ∈ IG has inherent symmetries, i.e. there exists u ∈ G, u ̸= e, such that f = π(u)f , one cannot avoid a non-unique argmax in Equation (7). That is because if there is a unique argmax g∗ such that g∗ = argmaxg∈G ∥f(g)∥1, we would have:
f(u−1g∗) = f(g∗) = max g∈G ∥f(g)∥1
Therefore u−1g∗ is also a valid argmax, hence the argmax is not unique. For symmetric inputs, the equivariant map Φ would give a set of sampling indices (cosets) rather than a single one. If we instead consider including this set of sampling indices in zeq, and let group acts on this set, it can be shown that the exact equivariance would still hold. In practice, we uniformly sample a sampling index from this set to perform subsampling, and the subsampled feature maps will be the same for all sampling indices from this set because the inputs are symmetric. This complexity is unavoidable because an equivariant map that maps the feature map to a single coset does not exist in this case. However, perfectly symmetric inputs are very rare for real-world applications and we only encounter this problem for synthetic data.

heading: 3 Application: Group Equivariant Autoencoders
text: Group equivariant autoencoders (GAEs) are composed of alternating G-convolutional layers and equivariant subsampling/upsampling operations for the encoder/decoder. One important property
of GAEs is that the final subsampling layer of the encoder subsamples to a feature map defined on the trivial group {e}, outputting a vector (instead of a feature map) that is invariant. For the 1D-translation case, suppose the input to the final subsampling layer is a feature map f defined on Z. Then the final layer produces an invariant vector fb(0) = f(iL) where iL = argmaxx∈Z ∥f(x)∥1. Note that there is no scale factor cL here. Intuitively we can think of this as setting the scale factor cL = ∞. Hence the encoder of the GAE outputs a representation that is disentangled into an invariant part zinv = fb(0) (the vector output by the final subsampling layer) and an equivariant part zeq = (i1, ..., iL).
For the general group case, instead of specifying scale factors as in Section 2.1, we specify a sequence of nested subgroups G = G0 ≥ G1 ≥ · · · ≥ GL = {e}, where the feature map for layer l is defined on subgroup GL. For example, for the p4 group G = Z ⋊ C4, we can use the following sequence for subsampling: Z⋊C4 ≥ 2Z⋊C4 ≥ 4Z⋊C4 ≥ 8Z⋊C2 ≥ {e}. Note that for the final two layers of this example, we are subsampling translations and rotations jointly.
We lift the input defined on the homogeneous input space to IG (see Appendix A.3 for details on homogeneous spaces and lifting), and treat f0 ∈ IG as inputs to the autoencoders. The group equivariant encoder ENC can be described as follows:
[fl, plGl] = Sb↓ Gl−1 Gl (G-CNNEl−1(fl−1); Φl)
[zinv, zeq] = [fL(e), (p1G1, p2G2, . . . , pLGL)] (8) where l = 1, . . . , L and G-CNNl(·) denotes G-convolutional layers before the lth subsampling layer. The decoder DEC simply goes in the opposite direction, and can be written formally as:
fL is defined on GL = {e} and fL(e) = zinv fl−1 = G-CNNDl−1(Su↑ Gl−1 Gl (fl, plGl)) (9)
where l = L, . . . , 1 and f̂ = f0 gives the final reconstruction.
Recall from Section 2.1 that the tuple (i1, . . . , iL) can be expressed equivalently as a single integer. Similarly, the tuple (p1G1, p2G2, . . . , pLGL) can be expressed as a single group element in G. We show in Appendix B.2 that the action implicitly defined on the tuple via Equation (6) simplifies elegantly to the left-action on the single group element in G.
We now have the following properties for the learned representations (see Appendix D for a proof): Proposition 3.1. When ENC and DEC are given by Equations (8) and (9), and the group actions are specified as in Equation (3) and Equation (6), for any g ∈ G and f ∈ IG, we have
[zinv, g · zeq] = ENC(π(g)f) π(g)f̂ = DEC(zinv, g · zeq)

heading: 4 Related Work
text: Group equivariant neural networks The equivariant subsampling/upsampling that we propose deals with feature maps (functions) defined on the space of the group G or its subgroups K, which transform under the regular representation with the group action. Hence our equivariant subsampling/upsampling is compatible with lifting-based group equivariant neural networks defined on discrete groups (Cohen and Welling, 2016; Hoogeboom et al., 2018; Romero and Hoogendoorn, 2020; Romero et al., 2020) that define a mapping between feature maps on G. We also discuss the extension of group equivariant subsampling to be compatible with those defined on continuous/Lie groups (Cohen et al., 2018a; Esteves et al., 2018; Finzi et al., 2020; Bekkers, 2020; Hutchinson et al., 2021) in Section 6. This is in contrast to group equivariant neural networks that do not use lifting and use irreducible representations, defining mappings between feature maps on the input space X. (Cohen and Welling, 2017; Worrall et al., 2017; Thomas et al., 2018; Kondor et al., 2018; Weiler et al., 2018b,a; Weiler and Cesa, 2019a; Esteves et al., 2020; Fuchs et al., 2020).
Coset pooling In particular, Cohen and Welling (2016) propose coset pooling, which is also a method for equivariant subsampling. Here a feature map f on G is mapped onto a feature map Φ(f) on G/K (as opposed to K, for our equivariant subsampling) as follows:
Φ(f)(gK) = POOLk∈Kf(gk) (10)
such that the feature values on the coset gK are pooled. For the 1D-translation case, where G = Z,K = cZ, this amounts to pooling over every cth pixel, which disrupts the locality of features as opposed to our equivariant subsampling that preserves locality, and hence is more suitable to use with convolutions for translation equivariance. See Figure 2 for a visual comparison. As such, the p4-CNNs in Cohen and Welling (2016) use standard max pooling with stride=2 rather than coset pooling for Z2, and coset-pooling is only used in the final layer to pool over feature maps across 90-degree rotations, achieving exact rotation equivariance but imperfect translation equivariance. In our work, we use translation equivariant subsampling in the earlier layers and rotation equivariant subsampling in the final layers to achieve exact roto-translation equivariance.
Unsupervised disentangling and object discovery GAEs produce equivariant (zeq) and invariant (zinv) representations, effectively separating position and pose information with other semantic information. This relates to unsupervised disentangling (Higgins et al., 2017; Chen et al., 2018; Kim and Mnih, 2018; Zhao et al., 2017) where different factors of variation in the data are separated in different dimensions of a low-dimensional representation. However unlike equivariant subsampling, there is no guarantee of any equivariance in the low-dimensional representation, making the resulting disentangled representations less interpretable. Works on unsupervised object discovery (Burgess et al., 2019; Greff et al., 2019; Engelcke et al., 2020; Locatello et al., 2020) learn object-centric representations, and we showcase GAEs in MONet (Burgess et al., 2019) where we replace their VAE with a V-GAE in order to separate position and pose information and learn more interpretable representations of objects in a data-efficient manner.
Shift-invariance in CNNs As early as Simoncelli et al. (1992), it has been discussed that shiftinvariance cannot hold for conventional subsampling. Although standard subsampling operations such as pooling or strided convolutions are not exactly shift invariant, they do not prevent strong performance on classification tasks (Scherer et al., 2010). Nonetheless, Zhang (2019) integrates antialiasing to improve shift-invariance, showing that it leads to better performance and generalisation on classification tasks. Chaman and Dokmanić (2020) explore a similar strategy to our equivariant subsampling by partitioning feature maps into polyphase components and select the component with the highest norm. However, unlike the proposed group equivariant subsampling/upsampling which tackle general equivariance for arbitrary discrete groups, both works focus only on translation invariance.
Equivariant/invariant autoencoders GAEs learn exact low-dimensional equivariant representations under the autoencoding framework, and this has also been explored in previous work. Locatello et al. (2019) constructs autoencoders that are equivariant to the D12 group (30◦ rotations and reflections) but not to translations, and the dimension of their learned representation grows with the group size. Lohit and Trivedi (2020) explores a rotation-invariant encoder on spheres with a global pooling layer and a rotation-invariant loss function, without inserting subsampling layers between convolutional layers. Moreover, unlike the work above that focuses on group equivariant neural networks, Hinton et al. (2011); Sabour et al. (2017); Kosiorek et al. (2019) learn equivariant representations with capsule networks. However, general capsule networks do not come guaranteed exact equivariances or invariances.

heading: 5 Experiments
text: In this section, we compare the performance of GAEs with equivariant subsampling to their nonequivariant counterparts that use standard subsampling/upsampling in object-centric representation learning. We show that GAEs give rise to more interpretable representations that show better sample complexity and generalisation than their non-equivariant counterparts. In Appendix E.1, we show that we can also observe generalisation performance gains when using group equivariant subsampling for classification tasks.
Models and Baselines (G-)Convolutional autoencoders (G)ConvAE are composed of alternating (G-)convolutional layers and subsampling/upsampling operations with a final MLPs applied to the flattened feature maps. We categorize models by the types of equivariance preserved by the convolutional layers. We consider three different discrete symmetry groups: p1 (only translations), p4 (composition of translations and 90 degree rotations), p4m (composition of translations, 90 degree rotations and mirror reflection). The baseline models are: ConvAE-p1 (standard convolutional autoencoders), GConvAE-p4, GConvAE-p4m, where the corresponding equivariance is preserved in the (G-)convolutional layers but not in the subsampling/upsampling operations. The equivariant counterparts of these baseline models are GAE-p1, GAE-p4, GAE-p4m, where the subsampling/upsampling operations are also equivariant. For baseline models, we use a scale factor of 2 for all subsampling/upsampling layers. For GAEs, we subsample first the translations, then rotations, followed by reflections, all with scale factor 2. e.g. for GAE-p4m, the feature maps at each layer are defined on the following chain of nested subgroups: Z2⋊ (C4⋊C2) ≥ (2Z)2⋊ (C4⋊C2) ≥ (4Z)2⋊ (C4⋊C2) ≥ (8Z)2 ⋊ (C4 ⋊ C2) ≥ (16Z)2 ⋊ (C2 ⋊ C2) ≥ {e}. As in Cohen and Welling (2016), we rescale the number of channels such that the total number of parameters of these models roughly match each other.
Data To demonstrate basic properties of GAEs and compare sample complexity under the single object scenario, we use Colored-dSprite (Matthey et al., 2017) and a modification of FashionMNIST (Xiao et al., 2017), where we first apply zero-padding to reach a size of 64× 64, followed by random shifts, rotations and coloring. For multi-object datasets, we use Multi-dSprites (Kabra et al., 2019) and CLEVR6 which is a variant of CLEVR (Johnson et al., 2017) with up to 6 objects. All input images are resized to a resolution of 64× 64. See Appendix F and our reference implementation 3 for more details on hyperparameters and data preprocessing. Our implementation is built upon open source projects Harris et al. (2020); Paszke et al. (2019); Yadan (2019); Weiler and Cesa (2019b); Engelcke et al. (2020); Hunter (2007); Waskom (2021).

heading: 5.1 Basic Properties: Equivariance, Disentanglement and Out-of-Distribution Generalization
text: Equivariance The encoder-decoder pipeline in GAEs is exactly equivariant. In Figure 3, we train GAE-p4m on 6400 examples from Colored-dSprites, and we show how to manipulate reconstructions by manipulating the equivariant representation zeq (left). If an image x is encoded into [zinv, zeq], then decoding [zinv, g · zeq] will give g · x̂ where x̂ is the reconstruction of x. When the input has perfect symmetries (e.g. squares, ellipses in Figure 3), zeq is obtained by sampling from a set of sampling indices but different sampling indices in this set would give the same reconstruction (see Section 2.3).
Disentanglement The learned representations in GAEs are disentangled into an invariant part zinv and an equivariant part zeq. In Figure 3 (left), we vary the equivariant part while the invariant part remains the same. In Figure 3 (right), we show the frames of a movie of a heart, and show its reconstruction after replacing zinv representing a heart with that of an ellipse. Note that the ellipse shape undergoes the same sequence of transformations as the heart.
3https://github.com/jinxu06/gsubsampling
Out-of-distribution generalisation GAEs can generalise to data with unseen object locations and poses. We train an GAE-p4 on 6400 constrained training examples, where we only use examples with locations in the top-left quarter and orientations within [0, 90] degrees, as shown in Figure 4. During test time, we evaluate mean squared error (MSE) of reconstructions on unfiltered test data to see how models generalise to unseen location and poses. Both ConvAE-p1 and GConvAE-p4 cannot generalise well to object poses out of their training distribution. In contrast, GAE-p1 generalise to any locations without performance degradation but not to unseen orientations, while GAE-p4, which encodes both translation and rotation equivariance, generalises well to all locations and orientations. We only use heart shapes for evaluation, because the square and ellipse have inherent symmetries (see Section 2.3).

heading: 5.2 Single Object
text: Since GAEs are fully equivariant and can generalize to unseen object poses, it is natural to conjecture that such models can significantly improve data efficiency when symmetry-transformed data points are also plausible samples from the data distribution. We test this hypothesis on Colored-dSprites and transformed FashionMNIST, and the results are shown in Figure 5. On both datasets, equivariant autoencoders significantly outperform their non-equivariant counterparts for all considered training set sizes. In fact, as shown in the figure, equivariant models trained with a smaller training set size is often comparable to baseline models trained on a larger training set. Furthermore, the results demonstrate that it is beneficial to consider symmetries beyond translations in these problems: for both non-equivariant and equivariant models, variants that encode rotation and reflection symmetries consistently show better performance compared to models that only consider the translation symmetry.

heading: 5.3 Multiple Objects
text: In multi-object scenes, it is often more interesting to consider local symmetries associated with objects rather than the global symmetry for the whole image. To exploit object symmetries in
image data, one needs to first discover objects and separate them from the background, which is a challenging problem on its own. Currently, GAEs do not have inherent capability to solve these problems. In order to investigate whether our models could improve data efficiency in multi-object settings, we rely on recent work on unsupervised object discovery and only use GAEs to model object components. More specifically, we explored replacing component VAEs in MONet (Burgess et al., 2019) with V-GAEs (probabilistic version of our GAEs, where a standard Gaussian prior is put on zinv and zeq remains deterministic), and train models end-to-end. Again we study the low data regime to show results on data efficiency.
We train models on Multi-dSprites and CLEVR6 with training set sizes 3200, 6400 and 12800. We consider two evaluation metrics: mean squared error (MSE) to measure the overall reconstruction quality, and adjusted rand index (ARI), which is a clustering similarity measure ranging from 0 (random) to 1 (perfect) to measure object segmentation. As in Burgess et al. (2019), we only use foreground pixels to compute ARI. Component VAEs in MONet use spatial broadcast decoders (Watters et al., 2019) that broadcast the latent representation to a full scale feature map before feeding them into the decoders, and the decoders therefore do not need upsampling. It has the implicit effect of encouraging the smoothness of the decoder outputs. To encourage similar behaviour, we add average pooling layers with stride 1 and kernel size 3 to our equivariant decoders. As shown in Table 1, using GAEs to model object components significantly improves reconstruction quality, which is consistent with our findings in single-object scenario. As shown in Table 2, using GAEs to model object components also leads to better object discovery in the low data regimes, but this advantage seems to diminish as the dataset becomes sufficiently large.

heading: 6 Conclusions, Limitations and Future Work
text: Conclusions We have proposed subsampling/upsampling operations that exactly preserve translation equivariance, and generalised them to define exact group equivariant subsampling/upsampling for discrete groups. We have used these layers in GAEs that allow learning low-dimensional representations that can be used to reliably manipulate pose and position of objects, and further showed how GAEs can be used to improve data efficiency in multi-object representation learning models.
Limitations and Future work Although the equivariance properties of subsampling layers also hold for Lie groups, we have not discussed the practical complexities that arise with the continuous case, where feature maps are only defined on a finite subset of the group rather than the whole group. We leave this as important future work, as well as application of equivariant subsampling for tasks other than representation learning where equivariance/invariance is desirable e.g. object classification, localization (See Appendix E.1 for a preliminary exploration of classification tasks). Another limitation is that our work focuses on global equivariance, like most other works in the literature. An important direction is to extend to the case of local equivariances e.g. object-specific symmetries for multi-object scenes.
Acknowledgments and Disclosure of Funding
We would like to thank Adam R. Kosiorek for valuable discussion. We also thank Lewis Smith, Desi Ivanova, Sheheryar Zaidi, Neil Band, Fabian Fuchs, Ning Miao, and Matthew Willetts for providing feedback on earlier versions of the paper, and anonymous reviewers for their constructive suggestions during the review process. JX gratefully acknowledges funding from Tencent AI Labs through the Oxford-Tencent Collaboration on Large Scale Machine Learning.

Review summaries:
This paper introduces a group-theoretic framework for equivariant subsampling. Based on a similar approach to ref [3] in the case of translations, the idea is to subsample according to intrinsic landmarks and keeping track of the shifts involved. The framework introduced can handle this procedure for general symmetry groups by seeing the subsampled space as a subgroup, and the shifts are stored as elements of the quotient of the original group and the subgroup. The method is evaluated in the context of VAEs, where it is shown that an architecture using equivariant subsampling and upsampling is capable of fully disentangling covariant from invariant features, and, depending on the group used, is shown to generalize perfectly to unseen object positions.

The paper proposes and evaluates equivariant subsampling and upsampling operations for group convolutional networks, which include vanilla CNNs as a special case. These operations address the issue that conventional subsampling and upsampling operations are only equivariant under the action of a subgroup of symmetries. The models have the property to disentangle the appearance of objects from their G-pose, which allows to build disentangled group equivariant autoencoders.  As a simple introduction, the authors formulate their method initially for the special case of $G = \mathbb{Z}$ being the 1-dimensional translation group. In contrast to conventional subsampling operations, which sample the values on a grid $c\mathbb{Z}$ with spacing $c$, the equivariant subsampling operation samples features on a shifted grid $c\mathbb{Z} + i$, where $i = 0, \dots, c-1$ is an integer offset that is predicted from the signal itself via some equivariant function. The tuple of subsampled features and the integer offset are shown to transform in an equivariant way under the action of the full symmetry group $\mathbb{Z}$. If multiple subsampling layers are successively applied in a deep network, each of them adds a new offset value to a sequence of offsets. When eventually pooling to the trivial group $G=\{e\}$, the resulting feature is position invariant and all position information is encoded in the offset sequence. Upsampling layers are accordingly taking a low-resolution signal and an offset and embed it with this offset into a higher-resolution grid.  These constructions are in the obvious way generalized to feature maps on arbitrary discrete groups $G$: The lower resolution sampling grid is here given by a coset $pK \in G/K$ of a subgroup $K<G$ in $G$. The group equivariant subsampling operation is therefore producing a tuple which consists of a subsampled signal on the coset and the coset (offset) itself. To upsample such a signal, it is embedded in the coset of the supergroup. Lemma 2.1 and Proposition 2.2 assert the full $G$-equivariance of these operations.  The integer offset, or, more generally, coset, needs to be computed via any equivariant map from the input signal to the space of cosets. A simple choice, made by the authors, is that coset $gK$ that consist to the argmax group element $g$ of the signal's l1-norm.  The equivariant subsampling layers are used to construct various group equivariant autoencoders (GEAs), which are evaluated against baseline autoencoders that are constructed from group convolutions but use conventional subsampling and upsampling layers. GEAs are empirically shown to be exactly equivariant and to generalize perfectly out of distribution. This is not the case for the baseline models, whose conventional subsampling and upsampling layers break the models' overall equivariance. The disentanglement of object appearance and $G$-pose in the autoencoders' latent space is in figure 3 (left) shown to allow for an independent manipulation of these factors. GEAs perform on single and multi-object reconstruction tasks significantly better than the baselines.

The paper proposes a method to equip discrete-group convolutional neural networks (GCNNs) with a perfectly equivariant subsampling layer. The method relies on a non-fixed sampling grid which is dynamically chosen as an equivariant function of the current input. This method enables encoder-decoder architectures where each layer is equivariant to the global action of the discrete group on the input of the model. Experimental results confirm the accurate equivariance and the improved generalization of the models.

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.