Title: No title found

Abstract: Attention mechanisms, and most prominently self-attention, are a powerful building block for processing not only text but also images. These provide a parameter efficient method for aggregating inputs. We focus on self-attention in vision models, and we combine it with convolution, which as far as we know, are the first to do. What emerges is a convolution with data dependent filters. We call this an Affine Self Convolution. While this is applied differently at each spatial location, we show that it is translation equivariant. We also modify the Squeeze and Excitation variant of attention, extending both variants of attention to the roto-translation group. We evaluate these new models on CIFAR10 and CIFAR100 and show an improvement in the number of parameters, while reaching comparable or higher accuracy at test time against self-trained baselines.

1. INTRODUCTION:

Computer vision has seen great success thanks to the use of the convolution operation in Convolutional Neural Networks (CNNs) (LeCun et al., 1989; Krizhevsky et al., 2012; He et al., 2017; Mnih et al., 2013) . This operation takes advantage of the translational symmetry in visual perception tasks such as image classification. Meanwhile, in tasks that require sequence processing, attention (Chorowski et al., 2015; Bahdanau et al., 2014) and self-attention (Vaswani et al., 2017) have emerged as a powerful technique.

One of the peculiarities of CNNs is that filters are defined independently of the data. At the same time, self-attention is data dependent, but does not provide a template matching scheme, as does the convolution operation, since it merely reweights neighborhoods. While there is work towards using attention in CNNs, the current models use them independently, sequentially, or in parallel.

We unify convolution and self-attention, taking the best of both worlds. Our method provides a translationally equivariant convolution, where the filters are also dependent on the input. These data dependent filters more efficiently describe the relations present in the input. Moreover, by formulating it as a special convolution, it can be extended to be equivariant to other groups of transformations. As a result, we apply the rich literature on group equivariant neural networks (Cohen & Welling, 2016) and develop the roto-translation equivariant counterpart. This module can be used as a replacement for standard convolutional layers and we call it an Affine Self Convolution.

Another variant of attention in computer vision is Squeeze and Excitation (Hu et al., 2018) . This provides global attention and we extend it to the roto-translation variant in order to compare it to our module. We plan to release code for all the experiments soon.

The contributions of this work are:

• We introduce the Affine Self Convolution (ASC), merging convolution and self-attention.

• We prove ASC is translation equivariant.

• We extend ASC to roto-translation equivariant ASC.

• We develop group Squeeze and Excitation.

• We evaluate these modules on CIFAR10 and CIFAR100.

2. BACKGROUND:

In order to combine convolution and self-attention we first look at the group convolutions (Cohen et al., 2018a ) and then at the self-attention mechanism (Vaswani et al., 2017; Parmar et al., 2018) .

2.1. GROUP CONVOLUTION:

Group equivariant convolutional neural networks extend the operation of convolution. Cohen & Welling (2016) show that we can consider the convolution operation to be defined on a group and that this allows for a natural generalization to other objects that have a group structure.

Translation equivariant convolution:

The set of points in Z 2 with the operation of vector addition forms a group. For each value x in this group, the translation operator L (x) translates the domain of a function: L (x) [f ](y) = f (-x + y). Therefore, given two functions f, ψ : Z 2 → R, the convolution between the image f and the filter ψ can be written in terms of translations by elements of the group (Z 2 , +) (we depict the convolution in Figure 1 ):

EQUATION

Where the convolution operation Z 2 is indexed by the domain of the input and the filter. Most importantly, the convolution is equivariant to translations:

L (z) [f ] Z 2 ψ = L (z) [f Z 2 ψ].

This property connects a transformation of the input image L (z) [f ] with a precise transformation of the activations

L (z) [f Z 2 ψ]

. This is desirable because a model that implements such an operation can also be invariant to translations by taking a max over spatial positions at the end.

Roto-translation equivariant convolution A natural extension of the group of translations is the group of planar rotations and translations (Weiler et al., 2018b; Diaconu & Worrall, 2019) . The elements of this group have 2 components, a (proper) rotation R and a translation y:

EQUATION

Where I is the identity matrix of order 2 and (I, 0) is the identity element of SE(2). Throughout this work we will use P, R, S to denote rotation elements and x, y, z translation elements of the group SE(2). From here, we denote the G = SE(2). Similarly to Diaconu & Worrall (2019), the operator L (P,x) inversely transforms the coordinates of a function f with domain R 2 by L (P,x) [f ](y) = f (P -1 (y -x)). As a result, we can evaluate the planar convolution at points on G, [f R 2 ψ](P, x). This operation is called a lifting layer and outputs activations with domain G.

The operator L (P,x) transforms such functions by L (P,x) [f ](R, y) = f (P -1 R, P -1 (y-x)). We can preserve G equivariance of these functions by replacing the planar convolution with G convolutions between f, ψ : G → R:

EQUATION

In practice we are limited by the sampling grid, as real world images are functions on Z 2 . This has as symmetries translations by integer values, which form the group Z 2 and rotations by multiples of 90 • , which form the group C 4 . By replacing the groups R 2 and SO(2) in the definition of SE(2) with Z 2 and C 4 we get the group p4. This is a subgroup of SE(2), and all the relations in this work hold for both.

2.2. SELF-ATTENTION IN COMPUTER VISION:



Simplified form:

We start with a simplified form of local self-attention to highlight each part and to make notation clearer. For an input image f : Z 2 → R, self-attention is defined by 2 parts: 1) a score function between center pixel f (x) and a neighbor pixel f (y). This is usually the dot product, score(f (x), f (y)) = f (x)f (y). Moreover, such that the score values in a neighborhood sum to 1, the scores are normalized with softmax, α(x, y) = softmax y (score(f (x), f (y)), 2) an aggregation of the neighbors y based on the score function: The scores and the input image neighborhood are then multiplied at each spatial position. This is shown over α(x, •)f (•) by sharpening the color on the building and whitening or removing the color from the sky or trees. By summing over the spatial dimension, we aggregate the neighborhood scaled by how similar the center is to its neighbors.

EQUATION

This operation defines a neighborhood dependent weighting. We depict this in Figure 2 . The dot product score function, and therefore Equation 4, does not take account of the relative position between x and y. Simply, this encodes no spatial information. To take advantage of the regular grid in images, recent methods use positional embeddings β : Z 2 → R. These can be added to neighbors at y, prior to computing the score, and is parametrized by the relative position from the center x, as done in Ramachandran et al. (2019) , α(x, y) = softmax y (score(f (x), f (y) + β(y -x))). It can also be added to the neighbors score, after computing the score function, as done in Bello et al. (2019) ; Hu et al. (2019) , α(x, y) = softmax y (score(f (x), f (y)) + β(y -x))

Self-attention In practice, self-attention uses three sets of parameters (Parmar et al., 2018; Wang et al., 2018) , W Q , W K , W V , where Q, K, V stands for Query, Key, Value. We turn our attention to an input function (image) f : Z 2 → R din , with d in input channels and parameters

W Q , W K : Z 2 → R d k ×din and W V : Z 2 → R dout×din . W Q , W K , W V ,

which are implemented as 1×1 convolutions, have the purpose of mapping to three separate embeddings, Q, K, V and not to process spatial information. These have d k or d out channels. The spatial information is weighted and mixed by the attention mechanism. In this more general setting, where we denote an arbitrary channel with c, selfattention is defined with: 1) three linear mappings of the input f (defined analogously for K and V ):

Q c (x) = [f Z 2 W Q c ](x)

2) a normalized score function, which can use positional embeddings: α(x, y) = softmax y (score(Q(x), K(y))) and 3) an aggregation of the V embeddings based on the score function:

EQUATION

It is also common to use the multi head mechanism from Transformer (Vaswani et al., 2017) alongside attention. This means that Q, K, V are split along the channel dimension and the self-attention mechanism is evaluated independently for each element of the partition (each head). These are then concatenated and passed through a linear layer.

3. METHOD:

In the following section we look at the local relative type of attention, self-attention and we formulate it similarly to convolution. This allows us to merge the two and then develop the roto-translational 

(x) [β](•), which results in A (x) [f ](•).

If we were to sum now over spatial positions, this would be an affine convolution. Nonetheless, the affinely transformed input is now passed through the self-attention mechanism. This consists of evaluating a similarity score between the center of the neighborhood and the neighbors α(x, y) and then multiplying the input with the score function. As a result, the ASC module can focus on the parts of the image that contain information relevant to the center. Similarly to both previous methods, the neighbors' information is aggregated as the final step. Intuitively, what we get is a data dependent convolution between for the image around x.

variant. We also derive the roto-translational variant of the Squeeze and Excite module, which defines global attention.

3.1. AFFINE SELF CONVOLUTION (ASC):

Affine map We have seen how simplified self-attention in Equation 4 is applied to an input image and that it was further developed using relative positional embeddings β in order to include spatial information. We note that these are additive terms and extend them by a multiplicative term ψ that is also spatially dependent. This results in a affine map A, which depends on β, ψ : Z 2 → R. We define the map A relative to a center x: A (x) : (Z 2 , R) → R and we apply it using the translation operator

L (x) as A (x) (y, r) = rL (x) [ψ](y) + L (x) [β](y)

. Furthermore, we can describe the affine map as acting on an image f : Z 2 → R:

EQUATION

Simplified form Applying the affine map A, then simplified self-attention from Equation 4 to f :

Attention[A (x) [f ]](x) = y α(x, y)A (x) [f ](y)

is how we define the simplified ASC. We denote this with Z 2 . This is applied to two functions, an input and an affine filter, and we index it by the domain of the two functions, similarly to the convolution operation:

Z 2 [f, A](x) := y α(x, y)A (x) [f ](y) = y∈Z 2 α(x, y) score (f (y) L (x) [ψ](y) filter + L (x) [β](y) positional emb ) (7)

This uses the self-attention score from Equation 4, the convolutional filter from Equation 1 and also adds positional embeddings. Moreover, we can distribute α and view α(x, y)L (x) [ψ](y) and α(x, y)L (x) [β](y) as the parameters of a normalized affine map. Intuitively, this not only performs template matching through ψ, which is independent of the information in the image, but scales the template relative to what is in the image through α(x, y). By unifying the convolution and self-attention, these data dependent filters can more efficiently describe the relations in the image because they are applied differently at each location. We call this an Affine Self Convolution and we depict it in Figure 3 . It is possible to recover the usual convolution by setting the scaling coefficients α(x, y) to 1 and β to 0. By setting ψ to 1, we recover a simplified self-attention, where positional embeddings are used for computing the score and for the aggregated term.

Translation equivariance We prove that this is translation equivariant in the Appendix 21 and therefore:

Z 2 [L (z) [f ], A](x) = L (z) [ Z 2 [f, A]](x).

This means that this model can detect objects regardless of their position in the image, as does the standard convolution.

Affine Self Convolution Similarly to the general form of self-attention in Equation 5, we use three sets of parameters,

(W Q , A Q ), (W K , A K ), (W V , A V ) for an input f : Z 2 → R din

. This is depicted in Appendix Figure 5 . By contrast to the simple variant, we now also have the affine maps. These process spatial information and are 0 outside the extent, which is controlled by a hyperparameter: kernel size. Moreover, each affine map, can be implemented with ψ Q , ψ K : Z 2 → R d k ×d k , mapping from all the channels in the input to all the channels in the output or ψ Q , ψ K : Z 2 → R d k , mapping from one channel in the input to the same channel in the output. The same is true for ψ V , if we replace d k with d out . Our experiments use the latter, due to computational constraints. The additive term of the map is defined as

β Q , β K : Z 2 → R d k , where for β V , d k is replaced with d out .

We denote an arbitrary channel with c and we define ASC as:

1) three linear mappings of the input f (defined analogously for K and V ):

EQUATION

2) three affine maps A Q , A K , A V for the Q, K, V terms (defined analogously for K and V ):

A Q (x) [Q] c (y) = Q c (y)L (x) [ψ Q ] c (y) + L (x) [β Q ] c (y) (9)

3) a score function between center x and neighbor y, which is then normalized with softmax:

α(x, y) = softmax y (score(A Q (x) [Q](x), A K (x) [K](y))) (10)

4) an aggregation of the V embeddings based on the score function:

EQUATION

We also use the multi head mechanism. We note that now we have a separate set of parameters for

A Q , A K , A V , and that A Q is only evaluated at:A Q (x) [Q](x) = Q(x)ψ Q + β Q . Therefore, we can learn ψ Q , β Q for only one spatial index 0.

By comparison, the positional embeddings in Ramachandran et al. (2019) are represented by the term β K . The positional embeddings in Hu et al. (2019) are equivalent to learning the product of the embeddings β Q β K , which arises when multiplying

A Q x [Q](x)A K x [K](y).

A difference between this work and their work is that we directly learn these parameters, while Ramachandran et al. (2019) ; Hu et al. (2019) use a separate network to learn β.

3.2. ROTO-TRANSLATION AFFINE SELF CONVOLUTION:

We now use the machinery of this new operation and the group theoretic background to develop the roto-translation ASC. Similarly to the standard ASC, we first define a simplified form based on an affine map, then we prove roto-translation equivariance, and finally, we present the general form.

Affine map:

We now turn to functions on SE(2). We will denote G := SE(2), H := SO(2), R 2 = G/H. Similarly to the relative affine map for functions defined on Z 2 in equation 6, we define a relative affine map A for functions on G. This map uses affine parameters ψ, β : G → R and is defined as A (P,x) : (G, R) → R. This acts on a function f with domain G as:

EQUATION

This affine map transforms a function f relative to a center (P, x).

Simplified form To define the score function, we notice that we can split the sum over the group G in Equation 3 into two sums, y∈G/H R∈H A (P,x) [f ](R, y). In this form, for each P ∈ H the convolution can be seen as a weighted sum of the neighbors y, relative to a center x. We can scale the affine group convolution based on this intuition. Precisely, we add a score function that is based on each center (P, x) and each neighbor y:

EQUATION

We note that it might also be possible to define a score function α((P, x), (R, y)). Nonetheless, this could be computationally prohibitive and is not required in order to preserve roto-translation equivariance. Using the score in Equation 13we define the Simple ASC on G:

EQUATION

We show in the Appendix 24 that we can replace R∈H β(R, x) with β(x). Therefore, we can learn β as a function with domain G/H.

Roto-translation equivariance:

We verify that ASC on groups is equivariant to actions of the group G, by checking the equivariance relation in the Appendix 25.

Roto-translation ASC In practice, the input is discrete and we turn to functions on p4, G := p4, H := C 4 , Z 2 = G/H. For the general roto-translation ASC all quantities are defined analogously to the ASC on Z 2 in Equation 11, using the score function in Equation 13, the aggregation in Equation 14 and replacing the domain Z 2 with G. The details can be found in Appendix D.

3.3. GROUP SQUEEZE AND EXCITE:

Interactions between filters that are spatially far apart is also tackled in Squeeze and Excitation (SE) (Hu et al., 2018) . The SE module proposes to rescale each feature map based on a global aggregation of the spatial dimension. An intuition for this is that, by allowing for channel interactions at all the spatial locations, this effectively enlarges the receptive field maximally. This is also a parameter efficient method for increasing the receptive field.

For a general group G (with g, h, s ∈ G), the squeeze term takes an average of a function f : G → R din over the group. This is invariant to transformations by actions of the group G, which we show in the Appendix 28. The average is then passed through a one hidden layer MLP (W 1 ∈ R din×(din/r) , W 2 ∈ R (din/r)×din , where d in is the number of channels of the input and r is the reduction ratio). A sigmoid unit (σ) is then used on the activation.

EQUATION

These are then broadcasted across the domain with an element-wise multiplication, f cin (g) = f cin (g)squeeze cin (f ). Therefore, multiplying a function f (g) by squeeze(f ) preserves the group structure of f :

EQUATION

This is added as the last operation in any bottleneck ResNet. In our experiments, the group is either the group of integer translations Z 2 (which is the original operation is Hu et al. (2018) and for which |G| is the height×width of the image) or the discrete roto-translation group p4 (for which |G| is height×width×4, since there are 4 rotation in p4).

4. EXPERIMENTS:

In this section we describe the dataset used and motivate our baseline architecture. The results are then divided into models that use convolution only and models that use self-attention, including ASC. We then present the overall trends. We specify various hyperparameters in the Appendix C.

Dataset:

In our experiments we test the models' performance on the CIFAR10 and CIFAR100 datasets (Krizhevsky et al., 2009) . These consist of 60k images each, 50k for training and 10k for testing. We further split the training set into 45k images for training and we leave 5k for validation. CIFAR10 consists of 10 classes and CIFAR100 consists of 100 classes.

Backbone ResNets (He et al., 2016) In our experiments we use a variant of the ResNet that is appropriate for CIFAR and also uses bottleneck residual blocks. We take this approach because the models using self-attention in the literature use it as a replacement for the 3 × 3 layer inside the bottleneck residual block and we do the same. As a result, from the standard ResNet20, which is an example of a ResNet with basic residual block (6n + 2 layers, with n = 3), we arrive at ResNet29, which is an example of a ResNet with bottleneck residual block (9n + 2 layers, with n = 3).

Models We present the results on CIFAR10 in Figure 4 . The models are divided into convolution only models and self-attention models, including ASC. Convolution models include squeeze and excite models (SE), since these do not change the convolution operation, but adds the SE module at the end of each bottleneck block. We show the baseline ResNet29, the roto-translational baseline p4ResNet29, together with these models with SE. We now turn to models that use self-attention, for which we replace the 3 × 3 convolutions in ResNet29s bottleneck with self-attention. This means that we still leave a convolution, the convolution in the stem (this is the first layer in all ResNets). We experiment with several models. We replicate the strategy for positional embeddings in Ramachandran et al. ( 2019) and we only learn a factorized variant of β K , instead of learning an affine map for each of the Q, K, V terms. These parameters are factorized over the spatial dimensions. This means that for each head, half of the positional embeddings β K is invariant to horizontal translations, while the other half is invariant to vertical translations. Models using these parametrization of self-attention are denoted with +SASA. Self-attention models includes ASC models. We use Simple ASC based on Equation 7. This does not use three separate pairs of parameters, just one. With +ASC we denote the general form of ASC as described in equation 11. We also include ResNet29+ASC+SE as one of the models. This is because ASC and SE are not mutually exclusive and we can add the SE module as for a standard ResNet. We also train a roto-translation ASC, which we denote with p4ResNet29+ASC. This uses the general form of roto-translation ASC as described in equation 20. 2 .

Results:

The overall trade-off between accuracy and parameter count in Figure 4 shows that all the self-attention models use fewer parameters than convolution models, while reaching an accuracy in the same range. This confirms the intuition that data dependent filters are more flexible and a powerful modeling choice. We see that +SE adds a small improvement to convolutional models, but no improvement to +ASC models. This indicates that local self-attention provides enough spatial context. In terms of the self-attention models, +SASA is the only one which does not manage to compete with the baseline. We conclude that a relative affine map is required, additive positional embeddings are not enough. Most insightful, +Simple ASC drastically decreases the number of parameters (≈ 30%) and reaches accuracy within one standard deviation of the baseline. The three models with the highest accuracy are the roto-translational counterparts to standard models. They show a clear constant increase in performance with little to no extra parameters. This confirms that the theory was applied consistently and that the group theoretic approach benefits attention mechanisms. We ran a subset of the models on Cifar100, results are shown in Table 1 . The table shows that both ASC and roto-translation equivariance are more beneficial on this dataset. This indicates that these models generalize better when there is less data available.

The experiments show that self-attention is competitive when including the affine map as done in ASC and that roto-translational equivariance is a robust improvement in all models.

5. RELATED WORK:



Group equivariant CNNs:

The theoretically founded approach of group equivariant neural networks has motivated several advances. These works are presented under a unifying framework of group equivariant convolutional networks in Cohen et al. (2018a) . Thomas et al. (2018) . We note that this work would benefit from extensions to semigroups (Worrall & Welling, 2019) or curved manifolds (Cohen et al., 2019) . Other relevant works include Kondor & Trivedi (2018) Self-attention Various forms of self-attention for CNNs have been introduced based on non-local means (Buades et al., 2005) or on the Transformer (Vaswani et al., 2017) . These models are generally trained for classification/segmentation (Wang et al., 2018; Hu et al., 2019) , but have also been used for generative tasks (Parmar et al., 2018; Zhang et al., 2018) . Some works (Hu et al., 2019) show that locality of the self-attention mechanism together with softmax helps and that relative positional embeddings are essential. In parallel, Ramachandran et al. (2019) also show improvements using local self-attention with local relative positional embeddings over Bello et al. (2019) which use global self-attention with global positional embeddings. We describe in more detail how our model is related to Ramachandran et al. (2019) ; Hu et al. (2019) at the end of Section 3.1. It is also worth mentioning that several of these works compare convolutional models with attention based models and show that convolutional models require more floating point operations per second (FLOPS).

Self-attention has also been applied to graphs (Veličković et al., 2017) . Regardless, they are faster than the attention based models. 

Data dependent filters:



6. CONCLUSION:

In this work we show that there is a mixture of convolution and self-attention that can be used to replace spatial convolutions in CNNs. This module can be described as a convolution with data dependent filters. By retaining all the benefits of self-attention and convolution, what emerges are filters that are translationally equivariant, while being applied differently for each location in the input. The results show that this method is able to achieve comparable if not better performance than the convolutional models, while using fewer parameters and a bigger receptive field. Under simplifying assumptions, we can recover both self-attention and convolution, which allows us to incorporate the group theoretic approach of Group equivariant CNNs. Therefore, we prove the translational equivariance of ASC and we also develop the roto-translation equivariant ASC. The latter, is more robust to transformations of the input while surpassing the other models in accuracy.

We expect the most fruitful directions for future work to be: an efficient implementation (because self-attention is slower), efficient parametrization (order and shape of W s and As), merge selfattention and convolution for NLP/graphs, and equivariant ASC for manifolds (equivariant selfattention could score transport methods without assuming a predefined geometry).

A QKV ASC PICTURE Figure 5 : Affine Self Convolution with extent 3×3 evaluated at x. Yellow depicts the cropped region. For clarity, the affine parameters are not depicted, just the output of the affine map. As opposed to Simple ASC in Figure 3 , here we have 3 linear transformations W Q , W K , W V , which do not process spatial information, just output 3 separate embeddings Q, K, V for the same input f . Moreover, three local affine maps A Q , A K , A V are applied to each neighborhood relative to a center x in order to process spatial information. The score α(x, y) is evaluated between the center taken from Q and the neighbors taken from K. The score is then used to aggregate taken from V . This depicts how self-attention is applied in practice. 

B RESULTS:



C HYPERPARAMETERS, INITIALIZATION AND TRAINING SCHEDULE:

We normalize the images and use the standard data augmentation technique of random horizontal flips and random crops of 32×32 from the zero padded 36×36 images.

We initialize convolutional layers using He initialization (He et al., 2015) (for the p4 variant, the number of channels is multiplied by 4 in the He initialization) and we initialize batchnorms scaling coefficient γ to 1 and shifting coefficient β to 0. The reduction ratio in SE is 16, while in p4 models, the reduction ratio in SE is 4.

In the self-attention models we replace the spatial convolutions in the bottleneck layers with selfattention layers. Where the baseline ResNet uses a stride of 2, the self-attention models applies selfattention then an average pooling layer with kernel size 2×2 and stride 2. The self-attention layers use the multi head mechanism with 8 heads and a kernel size of 5. We set d k = d out . In all examples we initialize ψ, β ∼ N (0, 1). We use the dot product score normalized by 1 d k . These models are more unstable in the first couple of epochs. Therefore, we initialize the scaling coefficient γ of the last batchnorm in each residual block to 0 as done in Goyal et al. (2017) . Moreover, we warmup (per epoch) the learning rate for the first 10 epochs, up to the learning rate of 0.1. The models are trained using Nesterov accelerated gradient with momentum 0.9 and weight decay of 0.0001. The ResNet29 models and its variants are trained for 100 epochs, where the learning rate was divided by 10 at epochs 50 and 75 and each model was trained for 100 epochs. We also include examples of ResNet83. These bigger models we trained for 200 epochs and divided the learning rate by 10 at epochs 100 and 150. We do this for all the models, for a fair comparison. Throughout our experiments we used Pytorch (Paszke et al., 2017) . We have also taken inspiration from the GitHub repositories: Smery; Idelbayev.

When we use roto-translation layers instead of standard layers, we divide the number of channels in each stage of the ResNet by |p4| = 2. This preserves a similar number of parameters between the roto-translation models and the standard models.

D ROTO-TRANSLATION ASC:

In practice, the input is discrete and we turn to functions on p4, G := p4, H := C 4 , Z 2 = G/H. For the general roto-translation ASC all quantities are defined analogously to the ASC on Z 2 in Equation 11, using the score function in Equation 13and replacing the domain Z 2 with G. Therefore, we define roto-translation ASC with: 1) three linear mappings of the input f (defined analogously for K and V ):

EQUATION

2) three affine maps A Q , A K , A V for the Q, K, V terms (defined analogously for K and V ):

EQUATION

3) a score function between center (P, x) and neighbor y, which is then normalized with softmax: 

Similarly to ASC, for roto-translation ASC we learn ψ Q , β Q for only one spatial index 0.

E PROOFS ASC translation equivariance proof:

Claim:

EQUATION

Proof:

Expanding from the left hand side: 

Proof:

For the ASC on groups, the map A and therefore, α and β, are always used inside R∈H . This leads to a more parameter efficient parametrization for β: 

:



Figures and Tables:

Figure 1: Convolution between an input f and a 3×3 filter ψ evaluated at x. We denote with (•) all the positions y in the neighborhood of x. This begins by cropping the 3×3 neighborhood in f centered at x. The filter is colored brown for clarity, but it can have different values. The filter is centered in the same neighborhood by translating it L (x) [ψ](•). The two quantities are then multiplied at each spatial position y, independently.Finally, the sum over the spatial positions aggregates the response of the convolution, concluding how similar the input neighborhood around x is to the filter.

Figure 2: Simple self-attention with extent 3×3 for an input f evaluated at x. Similarly to the convolution (•) denotes the neighborhood of x. The 3×3 region around x is cropped. The score function then compares the center with each neighbor. The normalized score is α(x, y).The scores and the input image neighborhood are then multiplied at each spatial position. This is shown over α(x, •)f (•) by sharpening the color on the building and whitening or removing the color from the sky or trees. By summing over the spatial dimension, we aggregate the neighborhood scaled by how similar the center is to its neighbors.

Figure 3: Simple ASC for an input f and an affine map A with extent 3×3 evaluated at x. As for Figures 1 and 2, (•) denotes the neighborhood of x. The neighborhood is first extracted from the input, then the affine map parameters are centered in the neighborhood of x. For each spatial position, f is transformed affinely by L(x) [ψ](•) and L (x) [β](•), which results in A (x) [f ](•).If we were to sum now over spatial positions, this would be an affine convolution. Nonetheless, the affinely transformed input is now passed through the self-attention mechanism. This consists of evaluating a similarity score between the center of the neighborhood and the neighbors α(x, y) and then multiplying the input with the score function. As a result, the ASC module can focus on the parts of the image that contain information relevant to the center. Similarly to both previous methods, the neighbors' information is aggregated as the final step. Intuitively, what we get is a data dependent convolution between for the image around x.

Figure 4: Comparing the model size (in thousands of parameters) to accuracy of several variants of ResNet29 on CIFAR10. Error bars represent 1 standard deviation from 8 runs. Precise numbers are in Appendix Table 2.

Closely related to our work are the developments in planar Euclidean groups in Worrall et al. (2017); Weiler et al. (2018b); Diaconu & Worrall (2019) and the discrete variants of these in LeCun et al. (1989); Cohen & Welling (2016); Dieleman et al. (2016); Hoogeboom et al. (2018) and 3D in Kondor (2018); Cohen et al. (2018b); Esteves et al. (2018a); Worrall & Brostow (2018); Weiler et al. (2018a); Winkels & Cohen (2018);

; Esteves et al. (2019; 2018b); Marcos et al. (2017); Zhou et al. (2017); Bekkers et al. (2018); Jacobsen et al. (2017).

Other works which approach the problem of learning to apply filters differently at each spatial positions are Stanley et al. (2019); Jia et al. (2016); Ha et al. (2016); Jaderberg et al. (2015); Sabour et al. (2017); Kosiorek et al. (2019).

aggregation of the V embeddings based on the score function: The final aggregation step completely describes the roto-translation ASC operation:G [V, A V ] c (P, x) = y∈Z 2 α((P, x), y) R∈H A V (P,x) [V ] c (R, y)

[L (z) [f ], A](x) = y∈Z 2 softmax y (score(A (x) [L (z) [f ]](x), A (x) [L (z) [f ]](y))) A (x) [L (z) [f ]](y) Using the substitution: y → y + z Using: y → y + z ⇒ A (x) [L (z) [f ]](y) = A (-z+x) [f ](y), which we prove in the Appendix 22. Using: A (x) [L (z) [f ]](x) = A (-z+x) [f ](-z + x), which we prove in the Appendix 23. = y∈Z 2 softmax y (score(A (-z+x) [f ](-z + x), A (-z+x) [f ](y))) A (-z+x) [f ](y) = Z 2 [f, A](-z + x) =L (z) [ Z 2 [f, A]](x) Proof Claim: y → y + z ⇒ A (x) [L (z) [f ]](y) = A (-z+x) [f ](y) (22) Proof: A (x) [L (z) [f ]](y) =L (z) [f ](y)ψ(y -x) + β(y -x) =f (y -z)ψ(y -x) + β(y -x) Using the substitution: y → y + z =f (y)ψ(y + z -x) + β(y + z -x) =A (-z+x) [f ](y) Proof Claim: A (x) [L (z) [f ]](x) = A (-z+x) [f ](x -z) (23) Proof: A (x) [L (z) [f ]](x) =L (z) [f ](x)ψ(x -x) + β(x -x) =f (x -z)ψ((x -z) -(x -z)) + β((x -z) -(x -z)) =A (x-z) [f ](x -z)Proof β is a function on G/H Claim: R∈H A (P,x) [f ](R, y) = R∈H f (R, y)L (P,x) [ψ](R, y) + L (P,x) [β](y)

(P,x) [f ](R, y) = R∈H f (R, y)L (P,x) [ψ](R, y) + L (P,x) [β](R, y) = R∈H f (R, y)L (P,x) [ψ](R, y) + R∈H L (P,x) [β](R, y) R∈H L (P,x) [β](R, y) = R∈H β(P -1 R, P -1 (y -x))Using the substitution: R → P R = R∈H β(1, P -1 (y -x))

the final layer. As a result we can choose basic block ResNets with 6n + 2 layers or bottleneck block ResNets with 9n + 2 layers.

Table content: <table/>

We show accuracy mean and 1 standard deviation from 8 runs while varying random seed of various models on CIFAR100. Plot variant in Appendix Figure4

Table content: <table><tr><td/><td>Model</td><td>Accuracy</td><td>#Parameters</td></tr><tr><td>translation</td><td>ResNet29</td><td>68.34 ± 0.38</td><td>336k</td></tr><tr><td>equivariant</td><td colspan="2">ResNet29+Simple ASC (ours) 68.40 ± 0.78</td><td>240k</td></tr><tr><td/><td>ResNet29+ASC (ours)</td><td>68.68 ± 0.77</td><td>291k</td></tr><tr><td>roto-translation</td><td>p4ResNet29</td><td>72.03 ± 0.46</td><td>321k</td></tr><tr><td>equivariant</td><td>p4ResNet29+SE</td><td>72.03 ± 0.45</td><td>354k</td></tr><tr><td/><td>p4ResNet29+ASC (ours)</td><td>72.71 ± 0.51</td><td>283k</td></tr></table>

We show accuracy mean and 1 standard deviation from 8 runs while varying random seed of various models on CIFAR10. With bold we outline the highest accuracy second highest with blue. We do the same for the most parameter efficient models. Plot variant in Figure4.We also run preliminary experiments with ResNet83 (bottleneck block 9n + 2, with n = 9), which are presented in Figure7. These show a similar trend to the ResNet29 models, but they also indicate that attention, either as SE or ASC might be more rewarding in bigger models.

Table content: <table><tr><td/><td>Model</td><td>Accuracy</td><td>#Parameters</td></tr><tr><td/><td>ResNet29</td><td>91.18 ± 0.34</td><td>313k</td></tr><tr><td>convolution</td><td>ResNet29 + SE</td><td>91.39 ± 0.15</td><td>347k</td></tr><tr><td>models</td><td>p4ResNet29</td><td>93.02 ± 0.16</td><td>310k</td></tr><tr><td/><td>p4ResNet29 + SE</td><td>93.24 ± 0.17</td><td>342k</td></tr><tr><td/><td>ResNet29+SASA</td><td>89.43 ± 0.40</td><td>235k</td></tr><tr><td/><td colspan="2">ResNet29+Simple ASC (ours) 91.01 ± 0.24</td><td>217k</td></tr><tr><td>self-attention</td><td>ResNet29+ASC (ours)</td><td>91.24 ± 0.32</td><td>268k</td></tr><tr><td>models</td><td>ResNet29+ASC+SE (ours)</td><td>91.21 ± 0.11</td><td>301k</td></tr><tr><td/><td>p4ResNet29+ASC (ours)</td><td>93.03 ± 0.25</td><td>272k</td></tr></table>

Peer review (iclr2020):
This paper introduces the Affine Self Convolution that combines local connectivity, relative positional embeddings, and a convolutional filter. This form of local self-attention is spatially equivariant and is further extended to include rotational symmetries. Finally, a generalized form of Squeeze and Excite is derived that can be implied in a setting where the symmetries of some group need to be preserved. The experiments demonstrate that networks using affine self convolution use less parameters while attaining similar accuracies compared to using convolutional models. The paper describes how to merge self attention with (depthwise) convolutions. This merge is executed in a sensible way and the derivations seems correct to me. The theory justifies that combined convolutions and attention could be beneficial and the experiments do verify that it is possible to train networks with affine self convolution to a reasonable accuracy. However, I recommend to reject this paper because: 1) The paper does not clearly explain its contributions. 2) The experiments do not conclusively demonstrate a significant practical advantage that justify the additional complexity. 3) A significant part of the paper is dedicated to generalizing spatial equivariance to arbitrary groups, which seems unrelated to the contributions. Main argument 1) The paper itself is not enough to understand how the proposed method compares to other local self-attention methods or adaptive filter methods. In the related work section various self-attention and data dependent filter methods are mentioned. It is unclear how the performance of those methods compares to the attention mechanism proposed in this paper. The convolutional operation that is merged into the attention has a special form of sparsity because it is factorized in the feature dimension (depthwise convolutions). In convolutional models, depthwise convolutions have been used to increase performance while reducing parameter count. The paper only compares attention models with depthwise convolutions to CNN models with normal convolutions. It is therefore unclear what the increased parameter efficiently should be attributed to and whether the attention adds any benefit on top of the depthwise convolutions. Finally, it is unclear what the contribution is in the section on Group Squeeze and Excite (Group SE). It seems obvious that if an average over all spatial dimensions is taken as an input, the function will be invariant to any transformation of the spatial coordinate like translation and rotation. Therefore, the Group SE does not seem to differ in any significant way from normal SE. 2) The reported accuracies on CIFAR are on par with the convolutional baseline. The authors argue that affine self convolution is more efficient in terms of parameters. However, as mentioned before this is not an apples to apples comparison because a method with sparse convolutions is compared to a model with dense convolutions. Other metrics that could give more information about the potential efficiency of the method (flops, wall-clock time, etc.) are lacking. There is also no details about how the affine self convolution can be implemented on modern accelerators which seems non-trivial. 3) Due to the structure of the paper it is difficult recognize the contributions of the authors. Local self-attention is derived within the context of an arbitrary group rather than just spatial transformations. However, it is unclear how generalizing the derivation to groups leads to any new insights. In fact, it seems obvious that any function that acts solely on a patch that is extracted using either spatial or rotational transformations will be invariant to those transformations. There is no special structure in the proposed attention mechanism that creates this symmetry or makes it easier to compute the function at every patch. Possible improvements: 1) Re-organize the paper such that all group equivariance material is in its own section. 2) Explain how the group equivariance affects the algorithm apart from operating on a different set of patches. 3) Perform an ablation study. This might be somewhat complicated by the fact that two different methods are merged into one. Both the ablation towards vanilla self-attention as well as depthwise convolutions would be valuable. 4) Provide more detailed information about the practical implementation of the attention mechanism: How efficiently can this be implemented on modern accelerators? How fast could it be in theory? What is the observed wall-clock time compared to CNNs?