Original: I NFUSING L ATTICE S YMMETRY P RIORS IN N EURAL N ETWORKS U SING S OFT A TTENTION M ASKS
Revised: I NFUSING L ATTICE S YMMETRY P RIORS IN N EURAL N ETWORKS U SING S OFT A TTENTION M ASKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Infusing inductive biases and knowledge priors in artificial neural networks is a promising approach for achieving sample efficiency in current deep learning models.
Revised: Infusing inductive biases and knowledge priors in artificial neural networks is a promising approach for achieving sample efficiency in current deep learning models.

Original: Core knowledge priors of human intelligence have been studied extensively in developmental science and recent work has postulated the idea that research on artificial intelligence should revolve around the same basic priors.
Revised: Core knowledge priors of human intelligence have been studied extensively in developmental science and recent work has postulated the idea that research on artificial intelligence should revolve around the same basic priors.

Original: As a step towards this direction, in this paper, we introduce L AT F ORMER , a model that incorporates lattice geometry and topology priors in attention masks.
Revised: As a step towards this direction, in this paper, we introduce L AT F ORMER , a model that incorporates lattice geometry and topology priors in attention masks.

Original: Our study of the properties of these masks motivates a modification to the standard attention mechanism, where attention weights are scaled using soft attention masks generated by a convolutional neural network.
Revised: Our study of the properties of these masks motivates a modification to the standard attention mechanism, where attention weights are scaled using soft attention masks generated by a convolutional neural network.

Original: Our experiments on ARC and on synthetic visual reasoning tasks show that L AT F ORMER requires 2-orders of magnitude fewer data than standard attention and transformers in these tasks.
Revised: Our experiments on ARC and on synthetic visual reasoning tasks show that L AT F ORMER requires 2-orders of magnitude fewer data than standard attention and transformers in these tasks.

Original: Moreover, our results on ARC tasks that incorporate geometric priors provide preliminary evidence that deep learning can tackle this complex dataset, which is widely viewed as an important open challenge for AI research.
Revised: Moreover, our results on ARC tasks that incorporate geometric priors provide preliminary evidence that deep learning can tackle this complex dataset, which is widely viewed as an important open challenge for AI research.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Infusing inductive biases and knowledge priors in neural networks is regarded as a critical step to improve their sample efficiency (Battaglia et al., 2018; Bengio, 2017; Lake et al., 2017; Lake &
Revised: Infusing inductive biases and knowledge priors in neural networks is regarded as a critical step to improve their sample efficiency (Battaglia et al., 2018; Bengio, 2017; Lake et al., 2017; Lake &

Original: Baroni, 2018; Bahdanau et al., 2019).
Revised: Baroni, 2018; Bahdanau et al., 2019).

Original: The Core Knowledge priors for human intelligence have been studied extensively in developmental science (Spelke & Kinzler, 2007), following the theory that humans are endowed with a small number of separable systems of core knowledge, so that new flexible skills and belief systems can build on these core foundations.
Revised: The Core Knowledge priors for human intelligence have been studied extensively in developmental science (Spelke & Kinzler, 2007), following the theory that humans are endowed with a small number of separable systems of core knowledge, so that new flexible skills and belief systems can build on these core foundations.

Original: Recent research in artificial intelligence (AI) has postulated the idea that the same priors should be incorporated in AI systems (Chollet, 2019), but it is an open question how to incorporate these priors in neural networks.
Revised: Recent research in artificial intelligence (AI) has postulated the idea that the same priors should be incorporated in AI systems (Chollet, 2019), but it is an open question how to incorporate these priors in neural networks.

Original: Following this chain of thought, the Abstraction and Reasoning Corpus (ARC) (Chollet, 2019) was proposed as an AI benchmark built on top of the Core Knowledge priors from developmental science.
Revised: Following this chain of thought, the Abstraction and Reasoning Corpus (ARC) (Chollet, 2019) was proposed as an AI benchmark built on top of the Core Knowledge priors from developmental science.

Original: Chollet (2019) posits that ARC “ cannot be meaningfully approached by current machine learning techniques, including Deep Learning ”.
Revised: Chollet (2019) posits that ARC “ cannot be meaningfully approached by current machine learning techniques, including Deep Learning ”.

Original: Further, he argues that developing a domainspecific approach based on the Core Knowledge priors is a challenging first step and that “ solving this specific subproblem is critical to general AI progress ”.
Revised: Further, he argues that developing a domainspecific approach based on the Core Knowledge priors is a challenging first step and that “ solving this specific subproblem is critical to general AI progress ”.

Original: An important category of core Knowledge priors includes geometry and topology priors .
Revised: An important category of Core Knowledge priors includes geometry and topology priors .

Original: Indeed, significant attention has been devoted to incorporating such priors in deep learning architectures by rendering neural networks invariant (or equivariant) to transformations represented through group actions (Bronstein et al., 2021).
Revised: Indeed, significant attention has been devoted to incorporating such priors in deep learning architectures by rendering neural networks invariant (or equivariant) to transformations represented through group actions (Bronstein et al., 2021).

Original: Group invariant learning helps to build models that systematically ignore specific transformations applied to the input (such as translations or rotations).
Revised: Group invariant learning helps to build models that systematically ignore specific transformations applied to the input (such as translations or rotations).

Original: We take a complementary perspective and aim to help neural networks to learn functions that incorporate geometric transformations of their input (rather than to be invariant to such transformations).
Revised: We take a complementary perspective and aim to help neural networks to learn functions that incorporate geometric transformations of their input (rather than to be invariant to such transformations).

Original: In particular, we focus on group actions that belong to the symmetry group of a lattice.
Revised: In particular, we focus on group actions that belong to the symmetry group of a lattice.

Original: These transformations are pervasive in machine learning applications, as basic transformations of sequences, images, and other higher-dimensional regular grids fall in this category.
Revised: These transformations are pervasive in machine learning applications, as basic transformations of sequences, images, and other higher-dimensional regular grids fall in this category.

Original: While attention and transformers can in principle learn these kind of group actions, we show that they require a significant amount of training data to do so.
Revised: While attention and transformers can in principle learn these kind of group actions, we show that they require a significant amount of training data to do so.

Original: To address the aforementioned sample complexity issue, we introduce L AT F ORMER , a model that relies on attention masks in order to learn actions belonging to the symmetry group of a lattice, such as translation, rotation, reflection, and scaling, in a differentiable manner.
Revised: To address this sample complexity issue, we introduce L AT F ORMER , a model that relies on attention masks in order to learn actions belonging to the symmetry group of a lattice, such as translation, rotation, reflection, and scaling, in a differentiable manner.

Original: We show that, for any such action, there exists an attention mask such that an untrained self-attention mechanism initialized to the identity function performs that action.
Revised: We show that, for any such action, there exists an attention mask such that an untrained self-attention mechanism initialized to the identity function performs that action.

Original: We further prove that these attention masks can be expressed as convolutions of the identity, which motivates a modification to the standard attention module where the attention weights are modulated by a mask generated by a convolutional  network (CNN).
Revised: We further prove that these attention masks can be expressed as convolutions of the identity, which motivates a modification to the standard attention module where the attention weights are modulated by a mask generated by a convolutional neural network (CNN).

Original:  We  conducted anextensive evaluation of our approach  on a diverse set of tasks.
Revised: Our paper focuses on ARC and its variants. We see the extension of L AT F ORMER to other tasks as a promising avenue for future research. Therefore, we conducted an evaluation of our approach based on synthetic tasks, ARC and the recently proposed LARC (Acquaviva et al., 2021).

Original: First, to probe the sample efficiency of our method, we compared it on synthetic tasks against Transformers and attention modules.
Revised: First, to probe the sample efficiency of our method, we compared its ability to learn synthetic geometric transformations against Transformers and attention modules.

Original: Then, wemanually annotated ARC tasks based on the knowledge prior they require, and we evaluated L AT F ORMER on the  tasks requiring geometric knowledge priors.
Revised: Then, we annotated ARC tasks based on the knowledge priors they require, and we evaluated L AT F ORMER on the ARC and LARC tasks requiring geometric knowledge priors.

Original: Finally, we tested our method on natural images on the task of image registration (spatially aligning two images).

Original: Our results provide evidence that L AT F ORMER can learn geometric transformations with 2 orders of magnitude fewer training data than transformers and attention mechanisms.
Revised: Our results provide evidence that L AT F ORMER can learn geometric transformations with 2 orders of magnitude fewer training data than transformers and attention.

Original: We also provide the first neural network reaching good performance on a subset of ARC, suggesting that this kind of problems do not lie out of the reach of deep learning models.
Revised: We also provide the first neural network reaching good performance on a subset of ARC, suggesting that this kind of problem does not lie out of the reach of deep learning models.

Original:  F ORMALIZING THE GROUP - ACTION LEARNING PROBLEM
Revised:  F ORMALIZING THE GROUP - ACTION LEARNING PROBLEM

Original: We are interested in helping neural networks to learn lattice transformations sample efficiently by infusing knowledge priors in the model.
Revised: We are interested in helping neural networks to learn lattice transformations sample efficiently by infusing knowledge priors in the model.

Original: Motivated by ARC, we focus on learning geometric transformations that belong to the symmetry group of a lattice.
Revised: Motivated by ARC, we focus on learning geometric transformations that belong to the symmetry group of a lattice.

Original: This pertains to the more general problem of learning group actions given the input and the output of the transformation.
Revised: This pertains to the more general problem of learning group actions given the input and the output of the transformation.

Original: Concretely, we consider input-output transformations involving a group element g taken from some known group G that can be expressed under the general formulation:
Revised: Concretely, we consider input-output transformations involving a group element g taken from some known group G that can be expressed under the general formulation:

Original: [Equation]
Revised: [Equation]

Original: Above, x ∈ R d in and y ∈ R d out are input and output examples, f, g are unknown functions, and ◦ denotes the application of a group action.
Revised: Above, x ∈ R d in and y ∈ R d out are input and output examples, f, g are unknown functions, and ◦ denotes the application of a group action.

Original: As seen, the group element g can depend on the input data itself.
Revised: As seen, the group element g can depend on the input data itself.

Original: More generally, the function f may depend on more than one transformations of x based on elements belonging to various groups of interest.
Revised: More generally, the function f may depend on more than one transformations of x based on elements belonging to various groups of interest.

Original: A simple instance of the group-action learning problem is presented in Figure 1.
Revised: A simple instance of the group-action learning problem is presented in Figure 1.

Original: The example task is borrowed from ARC (Chollet, 2019) and entails learning to fill out the yellow patches in the leftmost image (input) so that the resulting image satisfies a 90 ◦ degree rotation symmetry.
Revised: The example task is borrowed from ARC (Chollet, 2019) and entails learning to fill out the yellow patches in the leftmost image (input) so that the resulting image satisfies a 90 ◦ degree rotation symmetry.

Original: The learner is given only a small set of input-output pairs (the ARC tasks have 3.3 training examples on average) and the prior knowledge of discrete two-dimensional point groups, one of which is the cyclic group of 4-fold rotations C 4 .
Revised: The learner is given only a small set of input-output pairs (the ARC tasks have 3.3 training examples on average) and the prior knowledge of discrete two-dimensional point groups, one of which is the cyclic group of 4-fold rotations C 4 .

Original: Though the task is challenging for a general neural network (due to the small number of samples), under the rotation prior it can be easily solved by the composition of a shallow neural network, a group action, and a non-linear activation: 1) mapping yellow to zero, 2) rotating each image x by some g ∈ C 4 , and 3) taking a pixel-wise max.
Revised: Though the task is challenging for a general neural network (due to the small number of samples), under the rotation prior it can be easily solved by the composition of a shallow neural network, a group action, and a non-linear activation: 1) mapping yellow to zero, 2) rotating each image x by some g ∈ C 4 , and 3) taking a pixel-wise max.

Original: It is important to stress that group-action learning is the exact antithesis of the typical group invariant and equivariant learning problems (Bronstein et al., 2021):
Revised: It is important to stress that group-action learning is the exact antithesis of the typical group invariant and equivariant learning problems (Bronstein et al., 2021):

Original: [Equation]
Revised: [Equation]

Original: Intuitively speaking, whereas in group-action learning one aims to learn functions that involve specific (and data-dependent) transformations of our data by actions of the group, in in/equivariant learning the goal is to build models that are oblivious to such group actions in a systematic manner.
Revised: Intuitively speaking, whereas in group-action learning one aims to learn functions that involve specific (and data-dependent) transformations of our data by actions of the group, in in/equivariant learning the goal is to build models that are oblivious to such group actions in a systematic manner.

Original:  A TTENTION MASKS FOR CORE GEOMETRY PRIORS
Revised:  A TTENTION MASKS FOR CORE GEOMETRY PRIORS

Original: This section prepares some theoretical grounding for L AT F ORMER , our approach to learn the transformations for lattice symmetry groups in the form of attention masks.
Revised: This section prepares some theoretical grounding for L AT F ORMER , our approach to learn the transformations for lattice symmetry groups in the form of attention masks.

Original: The section defines attention masks and explains how the former can be leveraged to incorporate geometry priors when solving group action learning problems on sequences and images.
Revised: The section defines attention masks and explains how the former can be leveraged to incorporate geometry priors when solving group action learning problems on sequences and images.

Original: 3. M ODULATING ATTENTION WEIGHTS WITH SOFT MASKING
Revised: 3. M ODULATING ATTENTION WEIGHTS WITH SOFT MASKING

Original: Consider the scaled dot-product attention mechanism as defined in Vaswani et al.
Revised: Consider the scaled dot-product attention mechanism as defined in Vaswani et al.

Original: In our formulation, we consider real-valued masks M ∈ [0 , 1] n Q × n K that rescale attention weights:
Revised: In our formulation, we consider real-valued masks M ∈ [0 , 1] n Q × n K that rescale attention weights:

Original: [Equation]
Revised: [Equation]

Original: where Q ∈ R n Q × d is the query parameter of the attention mechanism, K ∈ R n K ×
Revised: where Q ∈ R n Q × d is the query parameter of the attention mechanism, K ∈ R n K ×

Original: d is the key, d is the dimensionality of the model, n Q and n K are the sizes of the sets encoded by the query and key matrices respectively, and ⊙ is the Hadamard product.
Revised: d is the key, d is the dimensionality of the model, n Q and n K are the sizes of the sets encoded by the query and key matrices respectively, and ⊙ is the Hadamard product.

Original: Attention masks have been widely used to constrain the values of the attention weights and are usually binary masks applied before the softmax activation (Vaswani et al., 2017; Sartran et al., 2022).
Revised: Attention masks have been widely used to constrain the values of the attention weights and are usually binary masks applied before the softmax activation (Vaswani et al., 2017; Sartran et al., 2022).

Original: However, as we aim to learn M , we apply the mask after the softmax operation in order to avoid squashing the gradient.
Revised: However, as we aim to learn M , we apply the mask after the softmax operation in order to avoid squashing the gradient.

Original: Therefore, we rescale the attention weights to sum to 1 when calculating the output X of the attention mechanism:
Revised: Therefore, we rescale the attention weights to sum to 1 when calculating the output X of the attention mechanism:

Original: MaskedAttention
Revised: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: with 1 n being a vector of ones of size n and V ∈ R d × n K being the value parameter of the attention mechanism.
Revised: with 1 n being a vector of ones of size n and V ∈ R d × n K being the value parameter of the attention mechanism.

Original: Though masking can also be applied in cross-attention, in the following we primarily focus on self-attention , where Q = K = V = X .
Revised: Though masking can also be applied in cross-attention, in the following we primarily focus on self-attention , where Q = K = V = X .

Original: For ease of notation, we write
Revised: For ease of notation, we write

Original: MaskedAttention( X ; M ) whenever the query, key and value are the same matrix X .
Revised: MaskedAttention( X ; M ) whenever the query, key and value are the same matrix X .

Original: 3. E XISTENCE OF ATTENTION MASKS FOR LATTICE SYMMETRY ACTIONS
Revised: 3. E XISTENCE OF ATTENTION MASKS FOR LATTICE SYMMETRY ACTIONS

Original: This section discusses group actions that can be represented by attention masks.
Revised: This section discusses group actions that can be represented by attention masks.

Original: To develop intuition, let us first consider the simple example of translation in a one-dimensional lattice.
Revised: To develop intuition, let us first consider the simple example of translation in a one-dimensional lattice.

Original: Supposing that x = ( x 1 , . . .
Revised: Supposing that x = ( x 1 , . . .

Original: , x n ) ⊤ is a vector of n elements, we have:
Revised: , x n ) ⊤ is a vector of n elements, we have:

Original: [Equation]
Revised: [Equation]

Original: Hence, when M is the circulant permutation matrix shown above, we have that the mask shifts the input x by one element to the right.
Revised: Hence, when M is the circulant permutation matrix shown above, we have that the mask shifts the input x by one element to the right.

Original: Beyond translation, it is natural to ask what kinds of group actions we can perform with attention masks on data with a more high-dimensional topological structure.
Revised: Beyond translation, it is natural to ask what kinds of group actions we can perform with attention masks on data with a more high-dimensional topological structure.

Original: The following theorem provides existence statements for data whose underlying topological space is a hypercubic lattice (such as sequences, images and higher-dimensional regular grids).
Revised: The following theorem provides existence statements for data whose underlying topological space is a hypercubic lattice (such as sequences, images and higher-dimensional regular grids).

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Identity
Revised: Identity

Original: [Equation]
Revised: [Equation]

Original: Translation (by δ )
Revised: Translation (by δ )

Original: [Equation]
Revised: [Equation]

Original: Reflection
Revised: Reflection

Original: [Equation]
Revised: [Equation]

Original: Rotation
Revised: Rotation

Original: [Equation]
Revised: [Equation]

Original: Upscaling (by
Revised: Upscaling (by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Theorem 1 (Existence) .
Revised: Theorem 1 (Existence) .

Original: Let G m be the symmetry group of the m -dimensional hypercubic lattice, including translational symmetry, 4-fold rotational symmetry and vertical, horizontal and diagonal reflections.
Revised: Let G m be the symmetry group of the m -dimensional hypercubic lattice, including translational symmetry, 4-fold rotational symmetry and vertical, horizontal and diagonal reflections.

Original: Let X ∈ R n × d be a vectorized representation of an m -dimensional tensor X ∈ R l 1 ×···× l m , with n = l 1 · . . .
Revised: Let X ∈ R n × d be a vectorized representation of an m -dimensional tensor X ∈ R l 1 ×···× l m , with n = l 1 · . . .

Original: · l m .
Revised: · l m .

Original: For any group action g ∈ G m , there exists an attention mask
Revised: For any group action g ∈ G m , there exists an attention mask

Original: M g ∈ { 0 , 1 } n × n , such that:
Revised: M g ∈ { 0 , 1 } n × n , such that:

Original: [Equation]
Revised: [Equation]

Original: In other words, Theorem 1 states that any translation, rotation or reflection can be expressed in terms of an attention mask.
Revised: In other words, Theorem 1 states that any translation, rotation or reflection can be expressed in terms of an attention mask.

Original: Figure 2 shows some examples of masks corresponding to translation, rotation and reflection operations on square lattices.
Revised: Figure 2 shows some examples of masks corresponding to translation, rotation and reflection operations on square lattices.

Original: In the following, we adopt the convention of writing M g to mean the mask that implements action g .
Revised: In the following, we adopt the convention of writing M g to mean the mask that implements action g .

Original: For more details and for a proof of Theorem 1, we refer the reader to Appendix C.
Revised: For more details and for a proof of Theorem 1, we refer the reader to Appendix D.

Original: 3. R EPRESENTING ATTENTION MASKS FOR LATTICE TRANSFORMATIONS
Revised: 3. R EPRESENTING ATTENTION MASKS FOR LATTICE TRANSFORMATIONS

Original: To facilitate the learning of lattice symmetries, one needs to determine methods to parameterize the set of feasible group elements.
Revised: To facilitate the learning of lattice symmetries, one needs to determine methods to parameterize the set of feasible group elements.

Original: Fortunately, as precised in the following theorem, the attention masks considered in Theorem 1 can be expressed conveniently under the same general formulation.
Revised: Fortunately, as precised in the following theorem, the attention masks considered in Theorem 1 can be expressed conveniently under the same general formulation.

Original: Theorem 2 (Representation) .
Revised: Theorem 2 (Representation) .

Original: Let G m be the symmetry group of the m -dimensional hypercubic lattice and g ∈
Revised: Let G m be the symmetry group of the m -dimensional hypercubic lattice and g ∈

Original: G m be an action on a tensor X ∈ R l 1 ×···× l m .
Revised: G m be an action on a tensor X ∈ R l 1 ×···× l m .

Original: Then, there exist some primitive attention masks M g i ∈ { 0 , 1 } n i × n i such that
Revised: Then, there exist some primitive attention masks M g i ∈ { 0 , 1 } n i × n i such that

Original: [Equation]
Revised: [Equation]

Original: where M g ∈ { 0 , 1 } n × n is an attention mask implementing g , g i ∈
Revised: where M g ∈ { 0 , 1 } n × n is an attention mask implementing g , g i ∈

Original: G m i for some m i ∈ { 1 , 2 } is an action on the one-dimensional or square lattice, ⊗ is the Kronecker product, F is the Fourier transform applied column-wise, I n i is the n i × n i identity matrix, j is the imaginary unit, r n i = (1 , 2 , . . . , n i ) ⊤ , and o ( g i ) is defined as in Table 1.
Revised: G m i for some m i ∈ { 1 , 2 } is an action on the one-dimensional or square lattice, ⊗ is the Kronecker product, F is the Fourier transform applied column-wise, I n i is the n i × n i identity matrix, j is the imaginary unit, r n i = (1 , 2 , . . . , n i ) ⊤ , and o ( g i ) is defined as in Table 1.

Original: To obtain an intuitive understanding of Theorem 2, it helps to revisit the example of translation by δ = 1 of a sequence x ∈ R n on the 1 -dimensional lattice ( m = 1 ).
Revised: To obtain an intuitive understanding of Theorem 2, it helps to revisit the example of translation by δ = 1 of a sequence x ∈ R n on the 1 -dimensional lattice ( m = 1 ).

Original: Consulting Table 1, we find that o ( g ) is a vector containing − 1 at every position
Revised: Consulting Table 1, we find that o ( g ) is a vector containing − 1 at every position

Original: and we know M g is the permutation circulant matrix of Section 3.2.
Revised: and we know M g is the permutation circulant matrix of Section 3.2.

Original: Indeed, by the time-shifting property of the Fourier transform, M g can be obtained by shifting the rows of the identity by − 1 .
Revised: Indeed, by the time-shifting property of the Fourier transform, M g can be obtained by shifting the rows of the identity by − 1 .

Original: In general, vector o ( g ) has a convenient intuitive interpretation as its k -th component represents the relative position (with respect to k ) of the element that the k -th row of X attends to.
Revised: In general, vector o ( g ) has a convenient intuitive interpretation as its k -th component represents the relative position (with respect to k ) of the element that the k -th row of X attends to.

Original: For instance, in the one-dimensional example of translation by one element to the right, each element attends to the one immediately before.
Revised: For instance, in the one-dimensional example of translation by one element to the right, each element attends to the one immediately before.

Original: Hence, we have o ( g ) k
Revised: Hence, we have o ( g ) k

Original: = − 1 for any k = 1 , . . . , n .
Revised: = − 1 for any k = 1 , . . . , n .

Original: For higher-dimensional lattices, attention masks can be expressed as the Kronecker product of the attention masks for lower-dimensional cases or, equivalently, by applying each mask independently along each dimension of the input tensor.
Revised: For higher-dimensional lattices, attention masks can be expressed as the Kronecker product of the attention masks for lower-dimensional cases.

Original: For instance, on the square lattice, a translation bypixel on both dimensions is the Kronecker product of the two circulant matrices corresponding to a translation by 1 pixel on the one-dimensional lattice, as shown in Figure 2a.
Revised: For instance, on the square lattice, a translation bypixel on both dimensions is the Kronecker product of the two circulant matrices corresponding to a translation by 1 pixel on the one-dimensional lattice, as shown in Figure 2a.

Original: On more than one dimension, we can additionally define 4-fold rotations, still following the same formulation, with o ( g i ) defined as in Table 1.
Revised: On more than one dimension, we can additionally define 4-fold rotations, still following the same formulation, with o ( g i ) defined as in Table 1.

Original: Although strictly not a symmetry operation, scaling transformations of the lattice can also be defined in terms of attention masks under the same general formulation of Theorem 2, as reported in Table 1.
Revised: Although strictly not a symmetry operation, scaling transformations of the lattice can also be defined in terms of attention masks under the same general formulation of Theorem 2, as reported in Table 1.

Original: Therefore, for completeness, we will consider scaling transformations as well in our experiments.
Revised: Therefore, for completeness, we will consider scaling transformations as well in our experiments.

Original: Notice that Theorem 2 allows us to derive a way to calculate the attention masks.
Revised: Notice that Theorem 2 allows us to derive a way to calculate the attention masks.

Original: In particular, we can express our attention masks as a convolution operation on the identity, as stated below.
Revised: In particular, we can express our attention masks as a convolution operation on the identity, as stated below.

Original: Corollary 1.
Revised: Corollary 1.

Original: Let G m be the symmetry group of the m -dimensional hypercubic lattice and let M g ∈ R n × n be an attention mask implementing action g ∈ G m .
Revised: Let G m be the symmetry group of the m -dimensional hypercubic lattice and let M g ∈ R n × n be an attention mask implementing action g ∈ G m .

Original: Then:
Revised: Then:

Original: [Equation]
Revised: [Equation]

Original: where ⋆ denotes the convolution operation.
Revised: where ⋆ denotes the convolution operation.

Original: In other words, we can represent any mask in our framework as a convolution of the identity matrix with predefined kernels.
Revised: In other words, we can represent any mask in our framework as a convolution of the identity matrix with predefined kernels.

Original: This motivates us to design a convolutional neural network that produces our attention masks by successive convolutions of the identity.
Revised: This motivates us to design a convolutional neural network that produces our attention masks by successive convolutions of the identity.

Original:  T HE L AT F ORMER ARCHITECTURE
Revised:  T HE L AT F ORMER ARCHITECTURE

Original: While in principle the problem of inferring group actions from input-output pairs can be solved via search over finite groups, in practice the size of the group for lattice symmetry actions makes this approach unfeasible 1 .
Revised: While in principle the problem of inferring group actions from input-output pairs can be solved via search over finite groups, in practice the size of the group for lattice symmetry actions makes this approach unfeasible 1 .

Original: Moreover, we are interested in learning unknown functions jointly with the transformation, which cannot be solved by searching on the space of group actions.
Revised: Moreover, we are interested in learning unknown functions jointly with the transformation, which cannot be solved by searching on the space of group actions.

Original: Using a neural agent to search the space of possible actions would be a viable alternative, but this would make the problem non-differentiable and we would need to resort to reinforcement learning methods.
Revised: Using a neural agent to search the space of possible actions would be a viable alternative, but this would make the problem non-differentiable and we would need to resort to reinforcement learning methods.

Original: In this work, we aim to solve the problem in a differentiable way.
Revised: In this work, we aim to solve the problem in a differentiable way.

Original: Inspired by the observations above, we introduce L AT F ORMER , which incorporates the insights of Section 3 into a neural architecture that leverages our attention masks.
Revised: Inspired by the observations above, we introduce L AT F ORMER , which incorporates the insights of Section 3 into a neural architecture that leverages our attention masks.

Original: We propose to use gated CNNs to parameterize the masks and we introduce an additional smoothing technique for easier optimization.
Revised: We propose to use gated CNNs to parameterize the masks and we introduce an additional smoothing technique for easier optimization.

Original: 4. L ATTICE MASK EXPERTS AS CONVOLUTIONAL NEURAL NETWORKS
Revised: 4. L ATTICE MASK EXPERTS AS CONVOLUTIONAL NEURAL NETWORKS

Original: Attention modules in neural networks usually include an attention mechanism with learnable linear transformations of the inputs 2 followed by a feed-forward network (FFN), as in the Transformer encoder layer (Vaswani et al., 2017).
Revised: Attention modules in neural networks usually include an attention mechanism with learnable linear transformations of the inputs 2 followed by a feed-forward network (FFN), as in the Transformer encoder layer (Vaswani et al., 2017).

Original: To infuse core geometry priors in the attention module, we propose to modulate the attention weights with a mask generated by an additional layer, as shown in Figure 3a.
Revised: To infuse core geometry priors in the attention module, we propose to modulate the attention weights with a mask generated by an additional layer, as shown in Figure 3a.

Original: We refer to this layer as Lattice 
Revised: We refer to this layer as Lattice mask expert , as it specializes towards specific transformations of the lattice. To understand the

Original: [Figure]
Revised: [Figure]

Original: mask expert , as it specializes towards specific transformations of the lattice. To understand the purpose of this layer, it is useful to remember that, by the analysis conducted in Section 3, even if the attention and FFN layers are initialized to the identity function, the mask expert can generate attention masks that produce precise geometric transformations of the input.
Revised: purpose of this layer, it is useful to remember that, by the analysis conducted in Section 3, even if the attention and FFN layers are initialized to the identity function, the mask expert can generate attention masks that produce precise geometric transformations of the input.

Original: By Corollary 1, we know that each group action on the lattice can be represented by a mask that is a convolution of the identity and we have an analytical expression to calculate the kernels of the convolution.
Revised: By Corollary 1, we know that each group action on the lattice can be represented by a mask that is a convolution of the identity and we have an analytical expression to calculate the kernels of the convolution.

Original: We can leverage this notion to design CNNs that produce attention masks corresponding to specific group actions by following the general formulation:
Revised: We can leverage this notion to design CNNs that produce attention masks corresponding to specific group actions by following the general formulation:

Original: [Equation]{+M 
 =+}
Revised: {+M 
 =+}

Original:  M l for l = 0 , . . ., L [Equation]
Revised: I and M l +1 = α l Conv ( M l , K l ) + (1 − α l ) M l for l = 0 , . . .

Original:  where M L is the predicted mask, α l = σ l ( X ; θ ) =
Revised: , L − 1 , where M L is the predicted mask, α l = σ l ( X ; θ ) =

Original: FFN l ( X , θ ) is the output of a gating function, θ is a learnable parameter, and K l is the kernel of the l -th convolutional layer whose weights are determined based on Corollary 1 and Table 1.
Revised: FFN l ( X , θ ) is the output of a gating function, θ is a learnable parameter, and K l is the kernel of the l -th convolutional layer whose weights are determined based on Corollary 1 and Table 1.

Original: As an example, Figure 3b shows an architecture that generates translation masks.
Revised: As an example, Figure 3b shows an architecture that generates translation masks.

Original: Following
Revised: Following

Original: Theorem 2, the expert computes the translations along the two dimensions separately and then aggregates the resulting masks doing the Kronecker product.
Revised: Theorem 2, the expert computes the translations along the two dimensions separately and then aggregates the resulting masks doing the Kronecker product.

Original: Hence, a Lattice translation expert with L convolutional layers for each dimension can generate any translation up to δ = 2 L −elements per dimension.
Revised: Hence, a Lattice translation expert with L convolutional layers for each dimension can generate any translation up to δ = 2 L −elements per dimension.

Original: At inference time, the values of the gates can be discretized, in such a way that the generated mask provably performs a meaningful group action.
Revised: At inference time, the values of the gates can be discretized, in such a way that the generated mask provably performs a meaningful group action.

Original: Similarly to the expert in Figure 3b, we can define gated CNNs for transformations like rotation, reflection, and scaling.
Revised: Similarly to the expert in Figure 3b, we can define gated CNNs for transformations like rotation, reflection, and scaling.

Original: The product of experts (i.e., the combination of more actions) can be obtained by either chaining the experts or multiplying the attention masks generated by different experts.
Revised: The product of experts (i.e., the combination of more actions) can be obtained by either chaining the experts or multiplying the attention masks generated by different experts.

Original: For more details, we refer the reader to Appendix A.
Revised: For more details, we refer the reader to Appendix A.

Original: 4. M ASK SMOOTHING FOR EASIER TRAINING
Revised: 4. M ASK SMOOTHING FOR EASIER TRAINING

Original: The framework described so far parameterizes discrete transformations of a lattice in a differentiable manner.
Revised: The framework described so far parameterizes discrete transformations of a lattice in a differentiable manner.

Original: Nevertheless, to improve the training of L AT F ORMER , we found it beneficial to also apply a smoothing operation on the attention masks w.r.t.
Revised: Nevertheless, to improve the training of L AT F ORMER , we found it beneficial to also apply a smoothing operation on the attention masks w.r.t.

Original: the intrinsic metric of the group in question.
Revised: the intrinsic metric of the group in question.

Original: Our approach entails defining an adjacency relation between group elements and applying graph convolution with a heat kernel on the corresponding graph.
Revised: Our approach entails defining an adjacency relation between group elements and applying graph convolution with a heat kernel on the corresponding graph.

Original: This encourages the optimizer to favor weight updates that change the masks in a smooth manner w.r.t.
Revised: This encourages the optimizer to favor weight updates that change the masks in a smooth manner w.r.t.

Original: the geodesic distance implied by the graph.
Revised: the geodesic distance implied by the graph.

Original: Concretely, we define the neighbors of each element g i on the lattice as those g j = e ◦ g i reachable by an application of a primitive action e , such as translation by a single pixel in one dimension, rotation by 90 ◦ , and vertical/horizontal reflection.
Revised: Concretely, we define the neighbors of each element g i on the lattice as those g j = e ◦ g i reachable by an application of a primitive action e , such as translation by a single pixel in one dimension, rotation by 90 ◦ , and vertical/horizontal reflection.

Original: The notion of neighborhood gives rise to a graph whose vertex set is the lattice group and that contains one edge for every pair of neighboring actions.
Revised: The notion of neighborhood gives rise to a graph whose vertex set is the lattice group and that contains one edge for every pair of neighboring actions.

Original: As before, it helps to consider different kinds of transformations separately.
Revised: As before, it helps to consider different kinds of transformations separately.

Original: For instance, as shown in Figure 4, for 2D rotations the underlying graph is a cycle with 4 elements due to the underlying point group for 4-fold rotations being the cyclic group C 4 .
Revised: For instance, as shown in Figure 4, for 2D rotations the underlying graph is a cycle with 4 elements due to the underlying point group for 4-fold rotations being the cyclic group C 4 .

Original: Performing heat diffusion 
Revised: Performing heat diffusion can be achieved by repeated neighborhood averaging over the cycle and yields a smoothed rotation

Original: [Figure]
Revised: [Figure]

Original: can be achieved by repeated neighborhood averaging over the cycle and yields a smoothed rotation mask that performs all rotations at the same time (rightmost image in Figure 4).
Revised: mask that performs all rotations at the same time (rightmost image in Figure 4).

Original: We can extend the same approach to all lattice transformations: for instance, in the case of translation, the underlying graph is a grid and the smoothing operation is akin to convolution with a Gaussian kernel.
Revised: We can extend the same approach to all lattice transformations: for instance, in the case of translation, the underlying graph is a grid and the smoothing operation is akin to convolution with a Gaussian kernel.

Original: To train L AT F ORMER with smoothed masks, we compute two predictions: one with the non-smooth mask predicted by the model and one with a smoothed version of the same mask.
Revised: To train L AT F ORMER with smoothed masks, we compute two predictions: one with the non-smooth mask predicted by the model and one with a smoothed version of the same mask.

Original: The final loss is the sum of two terms, one for the original prediction and one for the smoothed prediction.
Revised: The final loss is the sum of two terms, one for the original prediction and one for the smoothed prediction.

Original: For smooth predictions, in all our experiments we utilize a mean squared error (MSE) loss, whereas for the original non-smooth predictions we use the cross-entropy or MSE loss depending on whether the output data are categorical or continuous.
Revised: For smooth predictions, in all our experiments we utilize a mean squared error (MSE) loss, whereas for the original non-smooth predictions we use the cross-entropy or MSE loss depending on whether the output data are categorical or continuous.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluated our method on three categories of tasks. First, we  developed a set of synthetic tasks in order to compare L AT F ORMER to attention modules and Transformers with respect to sample efficiency in learning basic geometric transformations.
Revised: To evaluate our method, we first developed a set of synthetic tasks in order to compare L AT F ORMER to attention modules and Transformers with respect to sample efficiency in learning basic geometric transformations.

Original: Then, we annotated the ARC tasks based on the knowledge priors they require, and we assessed the performance of our method on this challenging dataset.
Revised: Then, we annotated the ARC tasks based on the knowledge priors they require, and we assessed the performance of our method on this challenging dataset.

Original: Finally,in order to test L AT F ORMER on natural images, we investigated its application to the task of image registration (Lu et al., 2021).
Revised: Finally, we experimented with the LARC (Acquaviva et al., 2021) dataset and compared our method to stronger baselines based on neural program synthesis. We report additional experimental results in Appendix B.3.

Original: 5. S AMPLE EFFICIENCY ON GEOMETRIC TRANSFORMATIONS
Revised: 5. S AMPLE EFFICIENCY ON GEOMETRIC TRANSFORMATIONS

Original: As a preliminary study, we probed the ability of L AT F ORMER to learn geometric transformations efficiently.
Revised: As a preliminary study, we probed the ability of L AT F ORMER to learn geometric transformations efficiently.

Original: To this end, we compared the performance of our model to a transformer  and an attention module (the same architecture as our approach, without the mask expert) on synthetic tasks with increasing number of examples.
Revised: To this end, we compared the performance of our model to a transformer (Vaswani et al., 2017) and an attention module (the same architecture as our approach, without the mask expert) on synthetic tasks with increasing number of examples.

Original: Inspired by ARC, we generated a set tasks where the model needs to infer a geometric transformation from input-output pairs.
Revised: Inspired by ARC, we generated a set tasks where the model needs to infer a geometric transformation from input-output pairs.

Original: The input is a grid taken from the ARC tasks and the output is either a translation, rotation, reflection or scaling of the input.
Revised: The input is a grid taken from the ARC tasks and the output is either a translation, rotation, reflection or scaling of the input.

Original: The specific transformation applied to the input grid defines the task and is consistent across all examples in the same task.
Revised: The specific transformation applied to the input grid defines the task and is consistent across all examples in the same task.

Original: We evaluate the models based on the mean accuracy across tasks.
Revised: We evaluated the models based on the mean accuracy across tasks.

Original: Figure 5 shows the accuracy of our model compared to the baselines and to a version of L AT F ORMER without smoothing.
Revised: Figure 5 shows the accuracy of our model compared to the baselines and to a version of L AT F ORMER without smoothing.

Original: The plots show thatour L AT F ORMER can generalize better and from fewer examples than transformers andunmasked attention modules.
Revised: The plots show that L AT F ORMER can generalize better and from fewer examples than transformers and attention modules both with absolute positional encodings and relative positional encodings (Shaw et al., 2018).

Original: Additionally, our results provide evidence that the smoothing operation described in Section 4.2 is helpful in groups of larger size.
Revised: Additionally, our results show that the smoothing operation described in Section 4.2 is helpful for larger groups.

Original: For more details on this experiment, we refer the reader to Appendix B.1.
Revised: More details on this experiment are reported in Appendix B.1.

Original: 5. V ISUAL REASONING ON ARC TASKS WITH GEOMETRIC KNOWLEDGE PRIORS
Revised: 5. V ISUAL REASONING ON ARC TASKS WITH GEOMETRIC KNOWLEDGE PRIORS

Original: To assess the ability of our approach to learn efficiently on a more challenging use case, we focused on a subset of the ARC dataset (Chollet, 2019) requiring geometric priors for which our method could be a principled solution.
Revised: To assess the ability of our approach to learn efficiently on a more challenging use case, we focused on a subset of the ARC dataset (Chollet, 2019) requiring geometric priors for which our method could be a principled solution.

Original: To this end, we annotated the ARC tasks based on the knowledge priors they require, using the list of priors provided by Chollet (2019) as a reference.
Revised: To this end, we annotated the ARC tasks based on the knowledge priors they require, using the list of priors provided by Chollet (2019) as a reference.

Original: Appendix
Revised: Appendix

Original: B.2 provides more details about the annotation of ARC and Figure 7a in the Appendix shows the knowledge priors that we considered and their distribution across the ARC tasks.
Revised: B.2 provides more details about the annotation of ARC and Figure 7a in the Appendix shows the knowledge priors that we considered and their distribution across the ARC tasks.

Original: [Figure]
Revised: [Figure]

Original: Translate

Original: Rotate +

Original: [Figure]
Revised: [Figure]

Original: Reflect

Original: [Figure]

Original: [Table]

Original: [Figure] [Table] [Figure] We assessed the performance of our model on the tasks that require only knowledge priors corresponding to the basic geometrical transformations that we addressed in this work, namely translation, rotation, reflection and scaling.
Revised: We assessed the performance of our model on the tasks that require only knowledge priors corresponding to the basic geometrical transformations that we addressed in this work, namely translation, rotation, reflection and scaling.

Original: Table 2 shows our results compared to neural baselines.
Revised: Table 2 shows our results compared to neural baselines, including CNNs, attention with relative positional encodings (Shaw et al., 2018), PixelCNN (Gul et al., 2019), and Transformers (Vaswani et al., 2017). We additionally compared to a Transformer model that has access to precomputed transformations of the input ( Transformer + data augmentation ). Precomputing all group actions is only feasible for smaller groups (rotation, reflection and scaling).

Original: Though we restrict to only a subset of the tasks and there is definitely room for improvement even on these tasks, we reach considerably better performance than other neural baselines.
Revised: Though we restrict to only a subset of the tasks and there is definitely room for improvement even on these tasks, we reach considerably better performance than the baselines.

Original: Therefore, we believe our results advocate for the applicability of end-to-end differentiable models even on problems requiring sample-efficient abstract reasoning.
Revised: Therefore, we believe our results advocate for the applicability of end-to-end differentiable models even on problems requiring sample-efficient abstract reasoning.

Original: To the extent of our knowledge, this is the first evidence of a neural network achieving this performance on ARC tasks.
Revised: To the extent of our knowledge, this is the first evidence of a neural network achieving this performance on ARC tasks.

Original: 5. I MAGE REGISTRATION
Revised: 5. C OMPARISON WITH NEURAL PROGRAM SYNTHESIS ON LARC

Original: As a final experiment, to assess the applicability of our L AT F ORMER on natural images, we performed experiments on multimodal image registration , namely the problem of spatially aligning images from different modalities.
Revised: Recently, Acquaviva et al. (2021) introduced the Language-complete Abstraction and Reasoning Corpus (LARC), which provides natural language descriptions of 88% of the ARC tasks, generated by human participants who where asked to communicate to other humans a set of precise instructions to solve a task.

Original: Popular approaches to multimodal image registrations work in two stages: first, they learn a model that converts one modality into the other (or to transfer both modalities in the same representation as proposed by Pielawski et al. 
Revised: Acquaviva et al. (2021) evaluated several models based on neural program synthesis on LARC. All models generate symbolic programs from a carefully designed domainspecific language (DSL).

Original: , then they align the two images using traditional techniques.
Revised: LARC (IO) is a model that has only access to input-output pairs, as our L AT F ORMER . LARC (IO + NL) has access to the natural language descriptions as well and uses a pre-trained T5 model (Raffel et al., 2020) to represent the text. LARC (IO + NL pseudo) uses pseudo-annotation to encourage the learning of compositional relationships between language and programs: during training, the model is given additional synthetic language-to-program pairs generated by annotating primitive examples in the DSL with linguistic comments.

Original: We follow the experimental setup of Lu et al. (2021) and experiment on two datasets, one containing aerial views of a urban neighborhood and one containing cytological images. The images we employ
Revised: In order to compare to the work of Acquaviva et al. (2021), we evaluated their models on the set of LARC tasks that correspond to ARC tasks in our subset requiring geometric knowledge priors.

Original: [Table]
Revised: Additionally, following Acquaviva et al. we allowed L AT F ORMER to access the textual descriptions by using a pre-trained T5 model to generate a representation of the text. This embedding is provided as input both to the Lattice Mask Expert and the FFN layers of L AT F ORMER . We refer to this model as LatFormer + NL .

Original: CycleGAN ( [Table]
Revised: Table 3 shows the results of our experiments on the LARC dataset.

Original: −→ A )
Revised: The program-synthesis methods require a training stage on a portion of the tasks.

Original: [Table]
Revised: Therefore, the

Original: DRIT+ [Table] [Table]
Revised: L AT F ORMER models where only evaluated on the same testing tasks of LARC, using the same train-test split of Acquaviva et al.

Original: −→ B )
Revised: Overall, our results shows that L AT F ORMER performs better than program synthesis on the subset of tasks requiring geometric priors, with no need for a

Original: [Table]
Revised: [Table]

Original: A )
Revised: carefully designed DSL.

Original: [Table] [Table] [Table] are views of the same scene, but they are taken with different modalities and they are translated with respect to one another.
Revised: This advantage comes to the expense of being restricted to tasks involving geometric priors, whereas program synthesis approaches can be used on a wider set of tasks.

Original: We employ several state-of-the-art methods for modality translation and we compare our method to α -AMD (Lindblad & Sladoje, 2014) and SIFT (Lowe, 1999) based on the success rate metric defined of Lu et al. Table 3 reports our results on the image registration tasks and shows that our approach performs well on both datasets coupled with different methods for modality translation.
Revised: We also observe that the natural language descriptions marginally helped our model on one category of tasks. Our findings corroborate with Acquaviva et al. (2021) in this remark.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Our work was inspired by a previous investigation of self-attention layers which identified sufficient conditions such that they can perform convolution when equipped with relative positional encodings (Cordonnier et al., 2020; Andreoli, 2019).
Revised: Our work was inspired by a previous investigation of self-attention layers which identified sufficient conditions such that they can perform convolution when equipped with relative positional encodings (Cordonnier et al., 2020; Andreoli, 2019).

Original: Rather than relying on relative encodings, we here show how soft-masking can be used to learn sample efficiently more general input transformations, such as rotation, reflection, and scaling.
Revised: Rather than relying on relative encodings, we here show how soft-masking can be used to learn sample efficiently more general input transformations, such as rotation, reflection, and scaling.

Original: To the extent of our knowledge, the group-action learning problem has not been explicitly and generally formulated in previous work.
Revised: To the extent of our knowledge, the group-action learning problem has not been explicitly and generally formulated in previous work.

Original: That being said, many previous works have focused on specific instances, such as learning to sort (Graves et al., 2014;
Revised: That being said, many previous works have focused on specific instances, such as learning to sort (Graves et al., 2014;

Original: Reed & De Freitas, 2015; Li et al., 2020) by selecting an element of the permutation group S n , docking/folding by roto-translating objects according to an action in the special Euclidean group SE (3) (Sverrisson et al., 2022; Stärk et al., 2022; Jumper et al., 2021), and graph spectrum generation where the learned actions belong to the Stiefel manifold (Martinkus et al., 2022).
Revised: Reed & De Freitas, 2015; Li et al., 2020) by selecting an element of the permutation group S n , docking/folding by roto-translating objects according to an action in the special Euclidean group SE (3) (Sverrisson et al., 2022; Stärk et al., 2022; Jumper et al., 2021), and graph spectrum generation where the learned actions belong to the Stiefel manifold (Martinkus et al., 2022).

Original: Our work is similar in spirit to recent efforts in neuro-symbolic visual reasoning (Johnson et al., 2017b;a; Goyal et al., 2017; Mao et al., 2019; Higgins et al., 2018).
Revised: Our work is similar in spirit to recent efforts in neuro-symbolic visual reasoning (Johnson et al., 2017b;a; Goyal et al., 2017; Mao et al., 2019; Higgins et al., 2018).

Original: Many approaches based on attention mechanisms have been proposed in the past few years (Hudson & Manning, 2018; 2019).
Revised: Many approaches based on attention mechanisms have been proposed in the past few years (Hudson & Manning, 2018; 2019).

Original: Our work differentiates from the former in that we aim to learn basic geometric reasoning in a sample-efficient way, rather than modeling relationships between high-level concepts.
Revised: Our work differentiates from the former in that we aim to learn basic geometric reasoning in a sample-efficient way, rather than modeling relationships between high-level concepts.

Original: Recently, the ARC dataset has been augmented with natural language descriptions of the tasks by

Original: Acquaviva et al.

Original: (2021), who introduced LARC ( Language-complete ARC ).

Original: The dataset provides precise natural-language instructions to solve an ARC task, in such a way that an human would be able to solve the task even with only access to the instructions.

Original: The authors use a pre-trained T model (Raffel et al., 2020) to construct an approach that has access to both the input-output ARC grids and natural language instructions.

Original: Their results show that the tasks are challenging enough that even including natural language instructions does not allow neural networks to solve them.

Original: Finally, some recent works came to our same conclusion on the advantages of using attention masks to incorporate prior knowledge in neural networks.
Revised: Finally, some recent works came to our same conclusion on the advantages of using attention masks to incorporate prior knowledge in neural networks.

Original: As an example, Yan et al.
Revised: As an example, Yan et al.

Original: (2020) focus on the task of learning subroutines (e.g., sorting algorithms) and use a CNN to generate an attention mask for a Transformer encoder.
Revised: (2020) focus on the task of learning subroutines (e.g., sorting algorithms) and use a CNN to generate an attention mask for a Transformer encoder.

Original: They show that learning the attention mask allows them to generalize to longer sequences than the ones provided at training time.
Revised: They show that learning the attention mask allows them to generalize to longer sequences than the ones provided at training time.

Original: Similarly, Sartran et al.
Revised: Similarly, Sartran et al.

Original: (2022) used precomputed attention masks to incorporate syntactic compositional biases in language models.
Revised: (2022) used precomputed attention masks to incorporate syntactic compositional biases in language models.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Motivated by the long-term ambitious goal of infusing core knowledge priors in neural networks, this paper focused on how to help deep learning models to learnfundamental geometric transformations efficiently.
Revised: This paper focused on how to help deep learning models to learn geometric transformations efficiently.

Original: Specifically, we proposed to incorporate lattice symmetry biases into attention mechanisms by modulating the attention weights using learned soft masks.
Revised: Specifically, we proposed to incorporate lattice symmetry biases into attention mechanisms by modulating the attention weights using learned soft masks.

Original: We have shown that attention masks implementing the actions of the symmetry group of a hypercubic lattice exist, and we provided a way to represent these masks as convolutions of the identity.
Revised: We have shown that attention masks implementing the actions of the symmetry group of a hypercubic lattice exist, and we provided a way to represent these masks.

Original: This motivated us to introduce L AT F ORMER , a modification to the standard Transformer encoder layer that generates attention masks corresponding to lattice symmetry priors using a CNN.
Revised: This motivated us to introduce L AT F ORMER , a model that generates attention masks corresponding to lattice symmetry priors using a CNN.

Original: Our results on synthetic tasks show that our model can generalize better than the same attention modules without masking and Transformers.
Revised: Our results on synthetic tasks show that our model can generalize better than the same attention modules without masking and Transformers.

Original: Moreover, the performance of our method on a subset of ARC provides the first evidence that deep learning can be used on this dataset, which is widely considered as an important challenge for AI research.
Revised: Moreover, the performance of our method on a subset of ARC provides the first evidence that deep learning can be used on this dataset, which is widely considered as an important challenge for AI research.

Original: Experimental details for reproducibility of our results are reported in Appendix B. Both the code to reproduce our experiments and the dataset with our annotations of the ARC tasks will be made publicly available upon acceptance. E THICS S TATEMENT
Revised: R

Original: We foresee no major ethical considerations for this work.

Original: The datasets used in this work do not contain any sensitive information to the best of our knowledge

Original: A A DDITIONAL DETAILS ON THE MODEL
Revised: A A DDITIONAL DETAILS ON THE MODEL

Original: This section describes the L AT F ORMER architecture providing additional details that were not covered in Section 4.1.
Revised: This section describes the L AT F ORMER architecture providing additional details that were not covered in Section 4.1.

Original: As mentioned in Section 4.1, it is possible to design convolutional neural networks that perform all considered transformations of the lattice.
Revised: As mentioned in Section 4.1, it is possible to design convolutional neural networks that perform all considered transformations of the lattice.

Original: Figure 6 shows the architecture of the four expert models that generate translation , rotation , reflection and scaling masks.
Revised: Figure 6 shows the architecture of the four expert models that generate translation , rotation , reflection and scaling masks.

Original: [Figure]
Revised: [Figure]

Original: All models are CNNs applied to the identity matrix.
Revised: All models are CNNs applied to the identity matrix.

Original: In the figure, we use the following notation:
Revised: In the figure, we use the following notation:

Original: • M ( δ ) T denotes an attention mask implementing a translation by δ along one dimension; • M (90) R denotes an attention mask implementing a translation by 90 ◦ ; • M F denotes an attention mask implementing a reflection along one dimension; • M ( h ) S denotes an attention mask implementing an upscaling by h along one dimension.
Revised: • M ( δ ) T denotes an attention mask implementing a translation by δ along one dimension; • M (90) R denotes an attention mask implementing a translation by 90 ◦ ; • M F denotes an attention mask implementing a reflection along one dimension; • M ( h ) S denotes an attention mask implementing an upscaling by h along one dimension.

Original: Using Corollary 1, we can derive the kernels of the convolutional layers shown in Figure 6.
Revised: Using Corollary 1, we can derive the kernels of the convolutional layers shown in Figure 6.

Original: These kernels are frozen at training time, the model only learns the gating function, denoted as σ in the figure.
Revised: These kernels are frozen at training time, the model only learns the gating function, denoted as σ in the figure.

Original: Notice that all the models follow the same overall structure.
Revised: Notice that all the models follow the same overall structure.

Original: However, for scaling, we also learn an additional gate, denoted as σ ( M S , M ⊤ S ) in the Figure 6.
Revised: However, for scaling, we also learn an additional gate, denoted as σ ( M S , M ⊤ S ) in the Figure 6.

Original: This gate allows the model to transpose the mask and serves the purpose of implementing down-scaling operations (down-scaling is the transpose of up-scaling).
Revised: This gate allows the model to transpose the mask and serves the purpose of implementing down-scaling operations (down-scaling is the transpose of up-scaling).

Original: The composition of more actions can be obtained by combining different experts.
Revised: The composition of more actions can be obtained by combining different experts.

Original: This can be done either by chaining the experts or by matrix multiplication of the masks.
Revised: This can be done either by chaining the experts or by matrix multiplication of the masks.

Original: In preliminary experiments, we did not notice any significant difference in performance between the two options and we rely on the latter in our implementation.
Revised: In preliminary experiments, we did not notice any significant difference in performance between the two options and we rely on the latter in our implementation.

Original: B D ETAILS ON THE EXPERIMENTAL SETUP
Revised: B A DDITIONAL EXPERIMENTS AND DETAILS ON THE EXPERIMENTAL SETUP

Original: This section provides additional details on the experimental setup of all our experiments, including further information on the generation of the synthetic tasks and the data annotation process for ARC
Revised: This section provides additional details on the experimental setup of all our experiments, including further information on the generation of the synthetic tasks and the data annotation process for ARC

Original: . B. E XPERIMENTS ON SYNTHETIC DATA
Revised: . B. E XPERIMENTS ON SYNTHETIC DATA

Original: We considered four categories of tasks, namely translation , rotation , reflection and scaling .
Revised: We considered four categories of tasks, namely translation , rotation , reflection and scaling .

Original: Each task is defined in terms of input-output pairs, which are sampled from the set of all ARC grids and padded to the size of 30 × 30 cells.
Revised: Each task is defined in terms of input-output pairs, which are sampled from the set of all ARC grids and padded to the size of 30 × 30 cells.

Original: To each input grid, a synthetic transformation is applied in order to obtain the corresponding output grid.
Revised: To each input grid, a synthetic transformation is applied in order to obtain the corresponding output grid.

Original: For each task in each category, we generated 2048 training pairs and 100 test pairs.
Revised: For each task in each category, we generated 2048 training pairs and 100 test pairs.

Original: For translation tasks, we have a total of 900 possible translations in a 30 × 30 grid.
Revised: For translation tasks, we have a total of 900 possible translations in a 30 × 30 grid.

Original: However, generating data and training models on 900 tasks is computationally expensive, so we randomly[Figure] sampled 100 translations in the interval [1 , 29] × [1 , 29] , obtaining a total of 100 translation tasks.
Revised: However, generating data and training models on 900 tasks is computationally expensive, so we randomly sampled 5 translations in the interval [1 , 29] × [1 , 29] , obtaining a total of 100 translation tasks.

Original: Rotation tasks include all 4-fold rotations except the identity.
Revised: Rotation tasks include all 4-fold rotations except the identity.

Original: Similarly, reflection tasks involve horizontal, vertical and diagonal reflections.
Revised: Similarly, reflection tasks involve horizontal, vertical and diagonal reflections.

Original: Scaling tasks include all possible up/down scaling transformations of the input grid by factors of [2 , 5] ×
Revised: Scaling tasks include all possible up/down scaling transformations of the input grid by factors of [2 , 5] ×

Original: [2 , 5] for a total of 32 scaling tasks.
Revised: [2 , 5] for a total of 32 scaling tasks.

Original: The models are evaluated based on the mean accuracy on each category.
Revised: The models are evaluated based on the mean accuracy on each category.

Original: For each task we compute the accuracy on the test set based on how many of the predicted images match exactly the ground truth.
Revised: For each task we compute the accuracy on the test set based on how many of the predicted images match exactly the ground truth.

Original: B. E XPERIMENTS ON ARC
Revised: B. E XPERIMENTS ON ARC

Original:  In order to experiment with ARC, we first performed an annotation of the dataset to identify the underlying knowledge priors for each task.
Revised: [Figure] In order to experiment with ARC, we first performed an annotation of the dataset to identify the underlying knowledge priors for each task.

Original: To this end, we built a user interface where the annotator could browse the tasks and label them by selecting any combination of the available knowledge priors.
Revised: To this end, we built a user interface where the annotator could browse the tasks and label them by selecting any combination of the available knowledge priors.

Original: Figure 7b shows the user interface provided to the annotator, whereas Figure 7a shows the distribution of knowledge priors across the ARC tasks.
Revised: Figure 7b shows the user interface provided to the annotator, whereas Figure 7a shows the distribution of knowledge priors across the ARC tasks.

Original: Most tasks follow in more than one of the categories represented in Figure 7a.
Revised: Most tasks follow in more than one of the categories represented in Figure 7a.

Original: ARC can be regarded as a meta-learning benchmark, as it provides a set of training tasks and a set of unseen tasks to evaluate the performance of the model learned on the meta-training data.
Revised: ARC can be regarded as a meta-learning benchmark, as it provides a set of training tasks and a set of unseen tasks to evaluate the performance of the model learned on the meta-training data.

Original: It is important to emphasize that we do not target this use case, as we instead use the same setup as in the synthetic data and learn each task from scratch using only its training set.
Revised: It is important to emphasize that we do not target this use case, as we instead use the same setup as in the synthetic data and learn each task from scratch using only its training set.

Original: Though simple and elegant, the supervised-learning formulation prevents our models from reusing knowledge that can be shared between different tasks.
Revised: Though simple and elegant, the supervised-learning formulation prevents our models from reusing knowledge that can be shared between different tasks.

Original: In order to mitigate this issue, we rely on a data-augmentation strategy.
Revised: In order to mitigate this issue, we rely on a data-augmentation strategy.

Original: At training time, for each model and every iteration, we augment each grid 10 times by mapping each color to a different color (using the same mapping across training examples).
Revised: At training time, for each model and every iteration, we augment each grid 10 times by mapping each color to a different color (using the same mapping across training examples).

Original: The rationale behind this data-augmentation strategy is that (1) we assume that for tasks involving only geometric knowledge priors to be not affected by color mapping and (2) all models (including L AT F ORMER ) need to learn a function from d -dimensional color representations to categorical variables, hence it is beneficial if all colors are represented in the training set.
Revised: The rationale behind this data-augmentation strategy is that (1) we assume that for tasks involving only geometric knowledge priors to be not affected by color mapping and (2) all models (including L AT F ORMER ) need to learn a function from d -dimensional color representations to categorical variables, hence it is beneficial if all colors are represented in the training set.

Original: All models are evaluated based on the ratio of solved tasks and a task is considered solved if the model can predict the correct output grid for all examples in the test set.
Revised: [Table]

Original: B. E XPERIMENTS ON IMAGE REGISTRATION
Revised: B. A DDITIONAL EXPERIMENTS ON IMAGE REGISTRATION

Revised: As an additional experiment, to assess the applicability of our L AT F ORMER on natural images, we performed experiments on multimodal image registration , namely the problem of spatially aligning images from different modalities.

Revised: Image registration is a well-studied problem in computer vision and we do not aim to establish state-of-the-art performance.

Revised: The main purpose of this experiment is giving a hint on the applicability of our method to natural images beyond ARC.

Revised: We refer the reader to SuperGlue

Original: For the experiments on image registration, we use the experimental setup of (Lu et al., 2021).
Revised: (Sarlin et al., 2020) and COTR (Jiang et al., 2021) to have a sense of approaches specifically designed for this task. Popular approaches to multimodal image registrations work in two stages: first, they learn a model that converts one modality into the other (or to transfer both modalities in the same representation as proposed by Pielawski et al. , then they align the two images using traditional techniques. We follow the experimental setup of Lu et al.

Revised: (2021) and experiment with two datasets, one containing aerial views of a urban neighborhood and one containing cytological images.

Revised: The images we employ are views of the same scene, but they are taken with different modalities and they are translated with respect to one another.

Original: We use the code of the authors to generate data involving only translations.
Revised: We use the code of the authors to generate data involving only translations.

Original: Lu et al.
Revised: Lu et al.

Original: (2021) additionally consider small rotations, but these transformations are not actions in the symmetry group of a lattice, so we are not interested in resolving them.
Revised: (2021) additionally consider small rotations, but these transformations are not actions in the symmetry group of a lattice, so we are not interested in resolving them.

Revised: We employ several state-of-the-art methods for modality translation and we compare our method to α -AMD (Lindblad & Sladoje, 2014) and SIFT (Lowe, 1999) based on the success rate metric defined by Lu et al.

Revised: A registration is considered successful if the relative registration error (i.e., the residual distance between the reference patch and the transformed patch after registration normalized by the height and width of the patch) is below 2% .

Revised: Table 4 reports our results on the image registration tasks and shows that our approach performs well on both datasets coupled with different methods for modality translation.

Original: We use the same models of Lu et al.
Revised: We use the same models of Lu et al.

Original: (2021) for the modality translation task.
Revised: (2021) for the modality translation task.

Original: Then, in order to solve the image registration task with L AT F ORMER , we divide each image into 30 × 30 patches
Revised: Then, in order to solve the image registration task with L AT F ORMER , we divide each image into 30 × 30 patches

Original: and we run our model to predict the translation from one patch in an image to its counterpart in the corresponding image.
Revised: and we run our model to predict the translation from one patch in an image to its counterpart in the corresponding image.

Revised: C L IMITATIONS AND FUTURE WORK

Revised: Although we believe our results are interesting and promising for learning group actions with neural networks, we would like to point out some limitations of our approach.

Revised: First, our method is limited to actions on the symmetry group of the hypercubic lattice and it is not immediately extendable to other groups.

Revised: For instance, though permutation matrices are still convolutions of the identity and they can be generated by a CNN, providing an architecture with predefined kernels that can compute any permutation matrix is not feasible.

Original: C D EFERRED P ROOFS
Revised: Second, the model is hard to fine-tune: we noticed that once the gates of the CNN have been trained, it is hard for the model to adapt to different actions. We believe that both limitations can be addressed by still keeping the same overall idea of modulating attention weights using soft attention masks, possibly with a different parametrization of the masks. Future work will focus on this research direction and on extending our work to cover a wider set of the ARC tasks. D D EFERRED P ROOFS

Original: We prove both Theorem 1 and 2 by induction on the dimensionality of the hypercubic lattice m .
Revised: We prove both Theorem 1 and 2 by induction on the dimensionality of the hypercubic lattice m .

Original: C. B ASE CASE FOR THEOREMS 1 AND
Revised: D. B ASE CASE FOR THEOREMS 1 AND

Original: First, it is useful to notice that whenever M ∈ { 0 , 1 } n × n has exactly a single 1 per row, in other words M · 1 n = 1 n , then, for any X ∈ R n × d
Revised: First, it is useful to notice that whenever M ∈ { 0 , 1 } n × n has exactly a single 1 per row, in other words M · 1 n = 1 n , then, for any X ∈ R n × d

Original: MaskedAttention( X ;
Revised: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: In order to prove the theorem, we need to show that, for any action g ∈ G 1 , including translations, reflections and rotations, there exists a mask M g such that:
Revised: In order to prove the theorem, we need to show that, for any action g ∈ G 1 , including translations, reflections and rotations, there exists a mask M g such that:

Original: [Equation]
Revised: [Equation]

Original: Let us consider different families of actions separately.
Revised: Let us consider different families of actions separately.

Original: Translation.
Revised: Translation.

Original: As mentioned in Section 3.2, in the 1-dimensional case, a translation by one element to the right for a vector x = ( x 1 , x 2 , . .
Revised: As mentioned in Section 3.2, in the 1-dimensional case, a translation by one element to the right for a vector x = ( x 1 , x 2 , . .

Original: , x n ) ⊤ is given by the circulant permutation matrix:
Revised: , x n ) ⊤ is given by the circulant permutation matrix:

Original: [Equation]    · [Equation]
Revised: M = M (1) T =     · · ·   · · ·   · · ·...

Original: 0 · · ·     .
Revised: 0 · · ·     .

Original: This holds because M (1) T · 1 n = 1 n , so:
Revised: This holds because M (1) T · 1 n = 1 n , so:

Original: MaskedAttention
Revised: MaskedAttention(

Original: [Equation]
Revised: [Equation]

Original: In general, a translation by δ elements is given by the circulant matrix M ( δ ) T =
Revised: In general, a translation by δ elements is given by the circulant matrix M ( δ ) T =

Original: ( M (1) T ) δ .
Revised: ( M (1) T ) δ .

Original: Therefore, masks implementing translation operations exist in the 1 -dimensional case and they are circulant permutation matrices.
Revised: Therefore, masks implementing translation operations exist in the 1 -dimensional case and they are circulant permutation matrices.

Original: This is enough for a base case for Theorem 1.
Revised: This is enough for a base case for Theorem 1.

Original: For Theorem 2, simply notice that:
Revised: For Theorem 2, simply notice that:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where o ( δ ) T =   − δ − δ ...

Original: − δ   .

Original: This comes directly from the time-shifting property of the Fourier transform.
Revised: This comes directly from the time-shifting property of the Fourier transform.

Original: Reflection.
Revised: Reflection.

Original: In the 1 -dimensional case, the reflection of a vector x = ( x 0 , x 1 , . . .
Revised: In the 1 -dimensional case, the reflection of a vector x = ( x 0 , x 1 , . . .

Original: , x n ) ⊤ is:
Revised: , x n ) ⊤ is:

Original: MaskedAttention( x ; M F ) =
Revised: MaskedAttention( x ; M F ) =

Original: M F · x = ( x n , x n − 1 , . . .
Revised: M F · x = ( x n , x n − 1 , . . .

Original: , x 2 , x 1 ) ⊤ with
Revised: , x 2 , x 1 ) ⊤ with

Original: [Equation]
Revised: M F =   0 · · · 0 00 · · · 0 10 · · · 1 0... 1 · · · 0 0  .

Original: The attention mask M F can be obtained by shifting the rows of the identity matrix by: 
Revised: The attention mask M F can be obtained by shifting the rows of the identity matrix by: o F =   n −n −n

Original: [Equation]
Revised: [Equation]

Revised:   .

Original: Therefore, by the time-shifting property of the Fourier transform we have:
Revised: Therefore, by the time-shifting property of the Fourier transform we have:

Original: [Equation]
Revised: [Equation]

Original: Rotation.
Revised: Rotation.

Original: Rotation (4-fold) is not defined in one dimension, so for a base case we need to consider the square lattice.
Revised: Rotation (4-fold) is not defined in one dimension, so for a base case we need to consider the square lattice.

Original: Let X ∈ R l 1 · l 2 be a vectorized representaiton of a n = l 1 × l 2 dimensional matrix.
Revised: Let X ∈ R l 1 · l 2 be a vectorized representaiton of a n = l 1 × l 2 dimensional matrix.

Original: We need to define a vector o R ∈ R n such that:
Revised: We need to define a vector o R ∈ R n such that:

Original: [Equation]
Revised: [Equation]

Original: is a rotation mask.
Revised: is a rotation mask.

Original: Since rotation is a permutation of the identity, we know the vector exists.
Revised: Since rotation is a permutation of the identity, we know the vector exists.

Original: As X is vectorized, the o (90) R needs to take into account the size of the first dimension l 1 .
Revised: As X is vectorized, the o (90) R needs to take into account the size of the first dimension l 1 .

Original: For example, in order to perform a rotation on a vectorized representation, we need to map the first element of X to the position ( l 1 − 1) .
Revised: For example, in order to perform a rotation on a vectorized representation, we need to map the first element of X to the position ( l 1 − 1) .

Original: The reader can check that the vector given by
Revised: The reader can check that the vector given by

Original: [Equation]
Revised: [Equation]

Original: satisfies the equation above.
Revised: satisfies the equation above.

Original: Scaling.
Revised: Scaling.

Original: Although scaling is not a group action of the symmetry group of the lattice, we pointed out that it still can be defined within the same general formulation as the other transformations.
Revised: Although scaling is not a group action of the symmetry group of the lattice, we pointed out that it still can be defined within the same general formulation as the other transformations.

Original: We can take the 1-dimensional lattice as a base case and consider a vector x = ( x 0 , x 1 , . . .
Revised: We can take the 1-dimensional lattice as a base case and consider a vector x = ( x 0 , x 1 , . . .

Original: , x n ) ⊤ .
Revised: , x n ) ⊤ .

Original: Let h ∈ N be a parameter specifying the filter size of the scaling operation.
Revised: Let h ∈ N be a parameter specifying the filter size of the scaling operation.

Original: As an example, for h =we have:
Revised: As an example, for h =we have:

Original: MaskedAttention
Revised: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: where:
Revised: where:

Original: [Equation]
Revised: [Equation]

Original: This kind of matrix can also be obtained by shifting the rows of the identity as follows:
Revised: This kind of matrix can also be obtained by shifting the rows of the identity as follows:

Original: [Equation]
Revised: [Equation]

Original: where o ( h ) S =
Revised: where o ( h ) S =

Original: ( k − 1 mod h )
Revised: ( k − 1 mod h )

Original: + ( h − 1) ·
Revised: + ( h − 1) ·

Original: ⌊ ( k − 1) /h
Revised: ⌊ ( k − 1) /h

Original: C
Revised: D. I NDUCTIVE STEP FOR THEOREMS 1 AND

Original: I NDUCTIVE STEP FOR THEOREMS 1 AND

Original: Suppose that M g 1 ∈ { 0 , 1 } n 1 × n 1 and M g 2 ∈ { 0 , 1 } n 2 × n 2 are attention masks implementing actions g 1 ∈ G m 1 and g 2 ∈ G m 2 on some tensors X 1 ∈ R l 1 ×···× l m 1 and X 2 ∈ R l ′ 1 ×···× l ′ m 2 , with n 1 = l 1 · . . .
Revised: Suppose that M g 1 ∈ { 0 , 1 } n 1 × n 1 and M g 2 ∈ { 0 , 1 } n 2 × n 2 are attention masks implementing actions g 1 ∈ G m 1 and g 2 ∈ G m 2 on some tensors X 1 ∈ R l 1 ×···× l m 1 and X 2 ∈ R l ′ 1 ×···× l ′ m 2 , with n 1 = l 1 · . . .

Original: · l m 1 and n 2 = l ′ 1 · . . .
Revised: · l m 1 and n 2 = l ′ 1 · . . .

Original: · l ′ m 2 .
Revised: · l ′ m 2 .

Original: Consider a tensor X ∈ R l 1 ×···× l m 1 × l ′ 1 ×···× l ′ m 2 and its vectorization X ∈ R n with n = n 1 n 2 .
Revised: Consider a tensor X ∈ R l 1 ×···× l m 1 × l ′ 1 ×···× l ′ m 2 and its vectorization X ∈ R n with n = n 1 n 2 .

Original: We have:
Revised: We have:

Original: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: MaskedAttention

Original: [Equation]
Revised: Now notice that:

Original: MaskedAttention
Revised: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: g [Equation] M g 1 ⊗ In 2 ) . Now notice that:
Revised: I

Original: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: and similarly
Revised: and similarly

Original: MaskedAttention( X ;
Revised: MaskedAttention

Original: [Equation]
Revised: [Equation]

Original: Therefore, we conclude that performing masked attention with the mask M g 1 ⊗ M g 2 on X is equivalent to applying g 1 on the first m 1 dimensions and g 2 on the last m 2 dimensions of X .
Revised: Therefore, we conclude that performing masked attention with the mask M g 1 ⊗ M g 2 on X is equivalent to applying g 1 on the first m 1 dimensions and g 2 on the last m 2 dimensions of X .

Original: This provides a way for building attention masks for higher-dimensional lattices using the primitive masks defined in Section C.1, proving both Theorem 1 and 2.
Revised: This provides a way for building attention masks for higher-dimensional lattices using the primitive masks defined in Section D.1, proving both Theorem 1 and 2.

Original: C. P ROOF OF C OROLLARY
Revised: D. P ROOF OF C OROLLARY

Original: The proof of Corollary 1 follows immediately from Theorem 2 and from the property of the Fourier transform according to which multiplying in the Fourier domain implements a convolution in the original domain.
Revised: The proof of Corollary 1 follows immediately from Theorem 2 and from the property of the Fourier transform according to which multiplying in the Fourier domain implements a convolution in the original domain.