Infusing Lattice Symmetry Priors in Neural Networks Using Soft Attention Masks
Authors: Mattia Atzeni
Mrinmaya Sachan
Andreas Loukas
Venue: Submitted to ICLR 2023 (2023)

Abstract: Infusing inductive biases and knowledge priors in artificial neural networks is a promising approach for achieving sample efficiency in current deep learning models. 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. As a step towards this direction, in this paper, we introduce LatFormer, a model that incorporates lattice geometry and topology priors in attention masks.
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. Our experiments on ARC and on synthetic visual reasoning tasks show that LatFormer requires 2-orders of magnitude fewer data than standard attention and transformers in these tasks. 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.

Under review as a conference paper at ICLR 2023
INFUSING LATTICE SYMMETRY PRIORS IN NEURAL NETWORKS USING SOFT ATTENTION MASKS
Anonymous authors Paper under double-blind review
ABSTRACT
Infusing inductive biases and knowledge priors in artificial neural networks is a promising approach for achieving sample efficiency in current deep learning models. 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. As a step towards this direction, in this paper, we introduce LATFORMER, a model that incorporates lattice geometry and topology priors in attention masks. 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 convolutional neural network. Our experiments and on synthetic visual reasoning tasks show that LATFORMER requires 2-orders of magnitude fewer data than standard attention and transformers in these tasks. 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.
ARC
INTRODUCTION
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 al., 2017; Lake & Baroni, 2018; Bahdanau et al., 2019). The Core Knowledge priors for human intelligence have been studied extensively developmental science (Spelke Kinzler, 2007). following the theory that humans endowed with a small number of separable systems of core knowledge, that new flexible skills and belief systems can build on these core foundations. Recent research in artificial intelligence (Al) has postulated the idea that the same priors should be incorporated in Al systems (Chollet, 2019), but is an open question how to incorporate these priors in neural networks.
Following this chain of thought, the Abstraction and Reasoning Corpus (ARC) (Chollet, 2019) was proposed as an Al benchmark built on top of the Core Knowledge priors from developmental science. Chollet (2019) posits that ARC "cannol be meaningfully approached by current machine learning techniques, including Deep Learning". Further, he argues that developing domain- specific approach based on the Core Knowledge priors is challenging first step and that "solving this specific subproblem is critical to general Al progress".
An important category of Core Knowledge priors includes geometry and topology priors. Indeed, significant attention has been devoted to incorporating such priors in deep learning architectures by rendering neural networks invariant (or equivariant)
transformations represented through group actions (Bronstein et al., 2021). Group invariant learning helps to build models that systematically ignore specific transformations applied to the input (such as translations or rotations).
We take 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). In particular, we focus on group actions that belong to the symmetry group of lattice. These transformations are pervasive in machine learning applications, as basic transformations of sequences, images, and other higher-dimensin dinensoannnnsionll regular grids fall in this category. While attention and transformers can in principle learn these kind of group actions, we show that they require significant amount of training data do

---

Under review as a conference paper at ICLR 2023
Training data
Test data
?
Figure 1: We consider problems that involve learning geometric transformation on the input data as a sub-problem. The displayed task (taken from ARC) entails learning to map, for each pair, the left to the right image. We investigate how to solve such tasks more sample-efficienty by imbuing self-attention with the ability to exploit lattice symmetry priors.
To address this sample complexity issue, we introduce LATFORMER, 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. 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. 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 mask generated by convolutional neural network (CNN).
Our paper focuses on ARC and its variants. We see the extension of LATFORMER to other tasks as 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). First, to probe the sample efficiency of our method, we compared its ability learn synthetic geometric transformations against Transformers and attention modules. Then, we annotated ARC tasks based on the knowledge priors they require, and we evaluated LATFORMER on the ARC and LARC tasks requiring geometric knowledge priors. Our results provide evidence that LATFORMER can learn geometric transformations with orders of magnitude fewer training data than transformers and attention. 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 deep learning models.
2 FORMALIZING THE GROUP-ACTION LEARNING PROBLEM
We are interested in helping neural networks to learn lattice transformations sample efficiently by infusing knowledge priors in the model. Motivated by ARC, we focus on learning geometric transformations that belong the symmetry group of lattice. This pertains to the more general problem of learning group actions given the input and the output of the transformation.
Concretely, we consider input-output transformations involving group element g taken from some known group G that can be expressed under the general formulation:
y f(g x:) for some
g g(x)) G
(group-action learning)
Above, E dd and y E Rd are input and output examples, f.g are unknown functions, and denotes the application of group action. As seen, the group element g can depend on the input data itself. More generally, the function ʃ may depend on more than one transformations of x based on elements belonging to various groups of interest.
A simple instance of the group-action learning problem is presented in Figure 1. The example task is borrowed from ARC (Chollet, 2019) and entails learning to fill out the yellow patches in the leftmost image (input) that the resulting image satisfies a 90" degree rotation symmetry. The learner is given only a small set of input- pairs (the ARC tasks have 3.3 training examples on average) and the prior knowledge of discrete two-dimenn -vo-dimensional point groups, of f-fold rotations C.. 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
of which is the cyclic group
2

---

Under review as a conference paper at ICLR 2023
neural network, group action, and non-linear activation: 1) mapping yellow zero, 2) rotating each image x by some g E C1. and 3) taking a pixel-wise max.
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):
y = fgg r)) for every g gE BEG
(invariant learning)
g y = [g x)) for every E G.
(equivariant learning)
Intuitively speaking, whereas in group-action learning 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 a systematic manner.
aims to learn functions that involve
ATTENTION MASKS FOR CORE GEOMETRY PRIORS
This section prepares some theoretical grounding for LATFORMER, our approach to learn the transformations for lattice symmetry groups in the form of attention masks. The section defines attention masks and explains how the former can be leveraged incorporate geometry priors when solving group action learning problems on sequences and images.
3.1 MODULATING ATTENTION WEIGHTS WITH SOFT MASKING
Consider the scaled dot-product attention mechanism as defined in Vaswani al. (2017). In our formulation, we consider real-valued masks M E 101]"*xm that rescale attention weights:
QK
A softmax
vd @
where Q E Rnexd is the query parameter of the attention mechanism, K E Rexxd is the key, the dimensionality of the model, ni and nk are the sizes of the sets encoded by the query and key matrices respectively, and @ is the Hadamard product. 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 al., 2017; Sartran et al., 2022). However, as aim to learn M, apply the mask after the softmax operation order to avoid squashing the gradient. Therefore, we rescale the attention weights to to I when calculating the output X of the attention mechanism:
A
Masked (Q,K, V;M)
V. Bx
A 1T
A
with 1, being vector of ones of size and V Rdx=x being the value parameter of the attention mechanism. Though masking can also be applied in cross-attention, in the following we primarily focus self-attention, where Q K =V X. For ease of notation, write Maskod X; M) whenever the query, key and value are the same matrix X.
3.2 EXISTENCE OF ATTENTION MASKS FOR LATTICE SYMMETRY ACTIONS
This section discusses group actions that be represented by attention masks. To develop intuition, let us first consider the simple example of translation in a one-dimensional lattice. Supposing that (Z,,...,,,,) a vector of elements, have:
0
MaskedAAttention(zz:.a M) (In, T1, In M
In-1)
Hence, when M is the circulant permutation matrix shown above, we have that the mask shifts the input x: by one element the right.
Beyond translation, it is natural to ask what kinds of group actions we can perform with attention masks on data with more digh-dimensional topological structure. The following theorem provides existence statements for data whose underlying topological space hypercubic lattice (such as sequences, images and higher-dimensional regular grids)
3

---

Table I: Fourier shifts for the transformations on the 1-dimensional and square lattices. We denote with ogg)), the k-tth component of the vector o(g;) E R", for I..... As stated in Theorem 2, attention masks for higher-dimensional lattices can be obtained by the Kronecker product of Under review as a conference paper at ICLR 2023
(a) Translation (1,
(b) Rotation by 90
(c) Vertical reflection
(d) Horiz. reflection
Figure 2: Examples of attention masks implementing transformations two dimensions, including: (a) translation by pixel on both axes, (b) rotation by 90º counterclockwise, (c) vertical reflection and (d) horizontal reflection around the center. White represents value and black 0.
Transformation
Fourier shift
Size of X
Identity
o(g:); ) 0
Translation (by 6)
ofgi) ): -5
11 1
Reflection
o(g:)1 (n 1), (B)) 0(g) -2
11
Rotation (90%)
olgi) = k. 04 - 1) [k - 1)/41
n = 11 X l2
Upscaling (by h)
og:)): (k: - mod h) + (h - 1) L(k- 10/--11 1)/h]
n 11
primitive masks defined over the 1-dimensional and square lattices. Composition of actions is given by matrix multiplication of the masks.
Theorem 1 (Existence). Let Gm be the symmetry group of the m-dimensional hypercubic lattice, including translational symmetry, 4-fold rotational symmetry and vertical, horizontal and diagonal reflections. Let X E Raxd be vectorized representation R x xlm, with n = 1 ..... lm. For any group action g E Gmo there exists an attention mask Mg (0,1)xxn, such that:
m-dimensional tensor X
MaskedAAtenntion X; M,) = gox X.
In other words, Theorem states that any translation, rotation or reflection can be expressed in terms of an attention mask. Figure 2 shows some examples of masks corresponding to translation, rotation and reflection operations on square lattices.
In the following, we adopt the convention of writing Mg to mean the mask that implements action For more details and for proof of Theorem 1, we refer the reader to Appendix D.
3.3 REPRESENTING ATTENTION MASKS FOR LATTICE TRANSFORMATIONS
To facilitate the learning of lattice symmetries, one needs to determine methods to parameterize the set of feasible group elements. Fortunately, precised in the following theorem, the attention masks considered Theorem can be expressed conveniently under the same general formulation.
Theorem 2 (Representation). Let Gms be the symmetry group of the m-dimenssiona hypercubic lattice and E Gm be an action on a tensor X E Rli glixx.xtmm Then, there exist some primitive atlention masks Mg. (0,1)m= x such that
2 Mg = X OM and = exp( o(gi) r.).))
where Mg E 10 is an attention mask implementing g, g, E Gm, for some m, E {1,2] 2} is action the one-dimensional or square lattice, & is the Kronecker product, F is the Fourier transform applied column-wise In, is the n, X n, identity matrix, j is the imaginary unit, Te. = (1,2....,n)) n..) and o(g.) is defined in Table I.
4

---

Under review as a conference paper at ICLR 2023
To obtain an intuitive understanding of Theorem 2, it helps to revisit the example of translation by 5 1 of a sequence a E R" on the 1-dimensional lattice (m = 1). Consulting Table 1, we find that ((g) is vector containing - at every position and know Mg is the permutation circulant matrix of Section 3.2. Indeed, by the time-shifting property of the Fourier transform, Mg can be obtained by shifting the rows of the identity by -1. In general, vector o(g) has 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. For instance, in the one-dimensional example of translation by one element to the right, each element attends to the one immediately before. Hence, we have (g)) 1 for any k ...
For higher- dineeroaimnnion lattices, attention masks can be expressed as the Kronecker product of the attention masks for lower-dimensional cases. For instance, on the square lattice, translation by pixel on both dimensions is the Kronecker product of the two circulant matrices corresponding to a translation by pixel on the one-dimensional lattice, as shown Figure 2a. On more than one dimension, we can additionally define 4-fold rotations, still following the same formulation, with (g.) defined as in Table
Although strictly not symmetry operation, seeraatoo,ssaalng 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. Therefore, for completeness, will consider scaling transformations well in our experiments.
Notice that Theorem 2 allows us to derive a way to calculate the attention masks. particular, we express attention masks convolution operation the identity, as stated below.
Corollary 1. Let Gm be the symmetry group of the m-dimensional hypercubic lattice and let Mg Raxa be an attention mask implementing action E Gm. Then:
2 n
=
F=
o(g) (lg).r]))es (.))) In
where * denotes the convolution operation.
In other words, we can represent any mask in our framework as convolution of the identity matrix with predefined kernels. This motivates to design convolutional neural network that produces our attention masks by successive convolutions of the identity.
4 THE LATFORMER ARCHITECTURE
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', Moreover, we are interested in learning unknown functions jointly with the transformation, which cannot be solved by searching on the space of group actions. Using a neural agent to search the space of possible actions would be a viable alternative, but this would make the problem non- non-differentiable and we would need to resort to reinforcement learning methods.
In this work, we aim to solve the problem in a disferentiable way. Inspired by the observations above, we introduce LATFORMER, which incorporates the insights of Section 3 into neural architecture that leverages our attention masks. We propose to use gated CNNs to parameterize the masks and we introduce an additional smoothing technique for casier optimization.
4.1 LATTICE MASK EXPERTS AS CONVOLUTIONAL NEURAL NETWORKS
Attention modules in neural networks usually include an attention mechanism with learnable linear transformations of the inputs² followed by feed-forward network (FFN), as in the Transformer encoder layer (Vaswani et al., 2017).
To infuse geometry priors in the attention module, we propose modulate the attention weights with a mask generated by an additional layer, as shown in Figure 3a. We refer to this layer as Lattice mask expert, it specializes towards specific transformations of the lattice. To understand the
'The size of the groups we consider grows with a polynomial of n and exponentially with m. For simplicity, we omitted linear transformations in the definition of BaskedAttentionn Section 3.1.
5

---

Under review as a conference paper at ICLR 2023
Lattice Transiation Expert
FFN
Cone M
Conu M
Come M
Cony M
Clonu M
Masked Attention
Cless M
Come M
Coner M
Clannr M
Conny
Latlice Mask Export
-
M
(a)
(b)
Figure 3: A LATFORMER layer (a) and an architecture for Lattice translation expert (b). The LATFORMER layer (a) is standard Transformer encoder layer augmented with Lattice mask expert constrained to generate attention masks corresponding to geometric transformation of the input. The Lattice translation expert (b) is particular instance of Latfice mask expert that produces translation masks. In the architecture above, every convolutional layer is meant to shift the input by power of 2 and can be skipped by gating function (denoted as σ).
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.
By Corollary 1, we know that each group action on the lattice can be represented by mask that is convolution of the identity and have analytical expression to calculate the kernels of the convolution. We can leverage this notion to design CNNs that produce attention masks corresponding to specific group actions by following the general formulation:
Mo = I and M on Conv Conv(M, Ki) + Mt)M for I 1 O.....,,
where Mr is the predicted mask, oy or X 0) = FFN, 0) is the output of gating function, O is learnable parameter, and K2 is the kernel of the 1-th convolutional layer whose weights determined based on Corollary and Table 1.
As an example, Figure 3b shows an architecture that generates translation masks. Following Theorem 2. the expert computes the translations along the two dimensions separately and then aggregates the resulting masks doing the Kronecker product. Hence, Lattice translation expert with L convolutional layers for each dimension can generate any translation up to 5 = 2 elements per dimension. At inference time, the values of the gates can be discretized, in such way that the generated mask provably performs meaningful group action.
Similarly to the expert in Figure 3b, we can define gated CNNs for transformations like rotation, reflection, and scaling. 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. For more details, refer the reader to Appendix A.
4.2 MASK SMOOTHING FOR EASIER TRAINING
The framework described so far parameterizes discrete transformations of lattice in a diffcrentiable manner. Nevertheless, to improve the training of LATFORMER, we found it beneficial to also apply a smoothing operation on the attention masks w.r.t. the intrinsic metric of the group in question.
Our approach entails defining an adjacency relation between group elements and applying graph convolution with a heat kernel on the corresponding graph. This encourages the optimizer to favor weight updates that change the masks in a smooth manner w.r.t. the geodesic distance implied by the graph. Concretely, we define the neighbors of each element g; on the lattice those g; e gi reachable by an application of a primitive action e, such as translation by single pixel in one dimension, rotation by 90", and vertical/horizontal reflection. The notion of neighborhood gives rise graph whose vertex set is the lattice group and that contains one edge for every pair of neighboring actions.
As before, it helps to consider different kinds of transformations separately. For instance, as shown in Figure for 2D rotations the underlying graph is cycle with 4 elements due to the underlying point group for 4-fold rotations being the cyclic group C.. Performing heat diffusion can be achieved by repeated neighborhood averaging the cycle and yields smoothed rotation

---

Under review as a conference paper at ICLR 2023
Smooth rotation mask
C4
M 21)
M
M (100
Figure 4: Rotational smoothing can be obtained by performing heat diffusion over cyclic graph where each node corresponds to a rotation mask.
mask that performs all rotations at the same time (rightmost image in Figure 4). We can extend the same approach to all lattice transformations: for instance, the case of translation, the underlying graph is grid and the smoothing operation is akin to convolution with a Gaussian kernel.
To train FORMER 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. The final loss the sum of two terms, one for the original prediction and one for the smoothed prediction. For smooth predictions, in all our experiments the original non-smooth predictions we use the cross-entropy or MSE loss depending on whether the output data categorical continuous.
utilize mean squared enror (MSE) loss, whereas for
5 EXPERIMENTS
To evaluate our method, we first developed set of synthetic tasks in order to compare FARFOR to attention modules and Transformers with respect to sample efficiency in learning basic geometric transformations. Then, annotated the ARC tasks based on the knowledge priors they require, and we assessed the performance our method on this challenging dataset. Finally, we experimented with the LARC (Acquaviva et al., 2021) dataset and compared our method stronger baselines based on neural program synthesis. We report additional experimental results in Appendix B.3.
5.1 SAMPLE EFFICIENCY ON GEOMETRIC TRANSFORMATIONS
As preliminary study, probed the ability of LATFORMER learn geometric transformations efficiently. To this end, we compared the performance of our model to transformer (Vaswani et al., 2017) and attention module (the same architecture on synthetic tasks with increasing number of examples. Inspired by ARC, we generated a set tasks where the model needs infer geometric transformation from input-output pairs. The input is grid taken from the ARC tasks and the output is either translation, rotation, reflection or scaling of the input. The specific transformation applied the input grid defines the task and consistent across all examples in the same task.
approach, without the mask expert)
We evaluated the models based on the mean accuracy across tasks. Figure 5 shows the accuracy of our model compared to the baselines and to version of LATFORMER without smoothing. The plots show that LATFORMER 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). Additionally, our results show that the smoothing operation described in Section 4.2 is helpful for larger groups. More details this experiment reported in Appendix B.
5.2 VISUAL REASONING ON ARC TASKS WITH GEOMETRIC KNOWLEDGE PRIORS
To assess the ability of our approach to learn efficiently on a more challenging use case, we focused on subset of the ARC dataset (Chollet, 2019) requiring geometric priors for which method could be a principled solution. 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 reference. Appendix B.2 provides more details about the annotation of ARC and Figure 7a the Appendix shows the knowledge priors that we considered and their distribution across the ARC tasks.

---

Under review as a conference paper at ICLR 2023
1.00
0.50
0.23
2t 2 2' 2 gigggge Number of examples
2 23 2' 2 pligh glight gloglagte
22 2 2- 27 2' git plaglags Number of examples (c)
24 21 27 2" 2i2 و Number of examples (d)
Number of examples
(a)
(b)
Altention pob))
Attention (rel. pos.)
Transformer
LafFormerr
LatFormer (wlo smoothing)
Figure 5: Sample efficiency of our method compared to the baselines on synthetic tasks on translation (a), rotation (b), reflection (c) and scaling (d). The y axis denotes the mean accuracy across tasks belonging to the same category, whereas the error shade is the standard deviation.
Translate
Rotate Translate
Reffect Translate
Scale Translate
CNN
0.019
0.000
0.000
0.000
Attention (abs, pos.)
0.019
0.000
0.023
0.000
Attention pos.)
0.000
0,023
PixcICNN
0.019
0.000
0.000
0.000
Transformer
0.000
Transformer data augmentation
0.200
0,184
0.091
LatFormer
0.800
0.591
Table 2: Performance ARC tasks that involve lattice symmetry priors.
We assessed the performance of our model on the tasks that require only knowledge priors corresponding the basic geometrical transformations that we addressed in this work, namely translation, rotation, reflection and scaling. Table 2 shows our results compared to neural baselines, including CNNs, attention with relative positional encodings (Shaw al., 2018), PixeICNN (Gul et al., 2019), and Transformers (Vaswani et al., 2017). We additionally compared to 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). Though we restrict to only a subset of the tasks and there is definitely room for improvement even these tasks, reach considerably better performance than the baselines. Therefore, we believe our results advocate for the applicability of end-to-end differentiable models even on problems requiring sample-efficient abstract reasoning. To the extent of knowledge, this is the first evidence of a neural network achieving this performance on ARC tasks.
5.3 COMPARISON WITH NEURAL PROGRAM SYNTHESIS ON LARC
Recently, Acquaviva et (2021) introduced the Language-complent 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 task. Acquaviva et al. (2021) evaluated several models based on neural program synthesis on LARC. All models generate symbolic programs from a carefully designed domain- specific language (DSL). LARC (IO) is model that has only access input-output pairs, as our LATFORMER. LARC (IO + NL) has access to the natural language descriptions as well and uses pre-trained T5 model (Raffel al., 2020) represent the text. LARC (I0 NL pseudo) uses pseudo-annotation to encourage the learning of compositional relationships between language and programs: during training, the model is given additional synthetic singuage-to-rrognan pairs generated by annotating primitive examples in the DSL with linguistic comments.
In order to compare to the work of Acquaviva al. (2021), we evaluated their models on the set of LARC tasks that correspond to ARC tasks in our subset requiring geometric knowledge priors. Additionally, following Acquaviva et al. (2021) we allowed LATFORMER to access textual descriptions by using pre-trained T5 model to generate representation of the text. This embedding is provided as input both to the Lattice Mask Expert and the FFN layers of LATFORMER. We refer to this model as LatFormer + NL. Table 3 shows the results of our experiments on the LARC dataset. The program- -ogrram-syntheis methods require a training stage on a portion of the tasks. Therefore, the LATFORMER models where only evaluated
the same testing tasks of LARC, using the same train-test split of Acquaviva et al. (2021). Overall, our results shows that LATFORMER performs better than program synthesis the subset of tasks requiring geometric priors, with need for
8

---

Under review as a conference paper at ICLR 2023
Translate
Rotate + Translate
Reflect Translate
Scale Translate
LARC (IO)
0.17
0.42
1.00
LARC (IO NL)
0.17
0.00
0.42
1,00
LARC (IO + NL pseudo)
0.25
0.00
0.42
1.00
Latmere
0.33
1.00
0.50
1.00
LatFormer NL
0.33
0.58
1.00
Table 3: Comparison of LATFORMER with neural program synthesis methods with access both input-output pairs and natural language descriptions on LARC
carefully designed DSL. This advantage comes the expense of being restricted to tasks involving geometric priors, whereas program synthesis approaches can be used a wider set of tasks. We also observe that the natural language descriptions marginally helped our model on one category of tasks. Our findings corroborate with Acquaviva al. (2021) in this remark.
6 RELATED WORK
Our work was inspired by 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). 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.
To the extent of our knowledge, the group-action learning problem has not been explicitly and generally formulated in previous work. That being said, many previous works have focused on specific instances, such as learning to sort (Graves et al., 2014; Reed & De Freitas, 2015; Li et al., 2020) by selecting an element of the permutation group Sn, docking/folding by roto-translating objects according to an action in the special Euclidean group SE(3) (Sverrisson et al., 2022; Stärk al., 2022; Jumper et al., 2021), and graph spectrum generation where the learned actions belong to the Stiefel manifold (Martinkus et al., 2022).
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). Many approaches based on attention mechanisms have been proposed in the past few years (Hudson Manning, 2018: 2019). Our work differentiates from the former in that aim to learn basic geometrie reasoning in sample-efficient way, rather than modeling relationships between high-level concepts.
Finally, some recent works came to same conclusion on the advantages of using attention masks to incorporate prior knowledge in neural networks. As an example, Yan et al. (2020) focus on the task of learning subroutines (e.g., sorting algorithms) and for a Transformer encoder. They show that learning the attention mask allows them to generalize longer sequences than the ones provided training time. Similarly, Sartran al. (2022) used precomputed attention masks to incorporate syntactic compositional biases in language models.
CNN to generate attention mask
7 CONCLUSION
This paper focused on how to help deep learning models to learn geometric transformations efficiently. Specifically, we proposed incorporate lattice symmetry biases into attention mechanisms by modulating the attention weights using learned soft masks. We have shown that attention masks implementing the actions of the symmetry group of hypereubic lattice exist, and we provided a way to represent these masks. This motivated us to introduce LATFORMER, model that generates attention masks corresponding to lattice symmetry priors using CNN. Our results on synthetic tasks show that our model can generalize better than the same attention modules without masking and Transformers. 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 Al research.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Samuel Acquaviva, Yewen Pu, Marta Kryven, Catherine Wong, Gabrielle E. Ecanow, Maxwell I. Nye. Theodoros Sechopoulos, Michael Henry Tessler, and Joshua B. Tenenbaum. Communicating natural programs to humans and machines. CoRR, abs/2106.07822 2021. URL
https://arxaa //rpivv org/ 0tts/216...088.
Jean-Mare Andreoli. Convolution, attention and structure embedding.
arXiv preprint
arXiv:1905.01289, 2019.
Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron C. Courville. Systematic generalization: What is required and can it be learned? In ICLR 2019, 2019.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonnaler, Vinícius Flores Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehre, Francis Song. Andrew Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. CoRR, hbs//800600221,, 2018.
Yoshua Bengio. The consciousness prior. CoRR, abs/1709.08 2017.
Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXfv:2104.13478, 2021.
François Chollet. On the measure of intelligence. CoRR, abs/1911.01544 2019. URL http: / //arxiv. ottpli/..r//b////ea/a////o /9911 01547.
dean-Baptisse Cordonnier, Andreas Loukas, and Martin Jaggi. On the relationship between self-attention and convolutional layers. Eighth International Conference on Learning Representations-IIRR 2020, number CONF, 2020.
Yash Goyal, Tejas Khot, Douglas Summers- Stay, Dhruv Batra, and Devi Parikh. Making the V in
VQA matter: Elevating the role of image understanding in visual question answering. 2017 IEEE Conference Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HII USA, July 21-26, 2017, pp. 6325-6334. IEEE Computer Society, 2017. doi: 2.1109/CVPR.2017.600
URL https::////.m.. doi. org/10/g///oooo 1119/1100PP 2017. 670.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014.
arXiv preprint
Muhammad Shahzeb Khan Gul, Michel Bätz, and Joachim Keinert. Pixel-wise confidences for stereo disparities using recurrent neural networks. In 30th British Machine Vision Conference 2019, BMVC 2019, Cardiff, UK, September 9-12, 2019, pp. 23. BMVA Press, 2019. URL https:: / www.c201l....m http://oo/ooooooooeetttotoo//ete copp-ot/tttuulllooo hatpllaas/pppppee / /027/424444 0274 saaper.pff.
Irina Higgins, Nicolas Sonnerat, Loic Matthey, Arka Pal, Christopher P. Burgess, Matko Bosnjak, Murray Shanahan, Matthew M. Botvinick, Demis Hassabis, and Alexander Lerchner. SCAN: learning hierarchical compositional visual concepts. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference net forum?id=
Track Proceedings. OpenReview.net, 2018. URL wttpttttss/p/eeeeeeeeeeseee rkN2 -1-2..
Drew A. Hudson and Christopher D. Manning. Compositional attention networks for machine reasoning. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
URL https: /wwppreevviiieoeeeoeee eww net/ forum? id=s1 Fuwz Rb...
Drew A. Hudson and Christopher D. Manning. Learning by abstraction: The neural state machine. In Hamna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems
10

---

Under review as a conference paper at ICLR 2023
2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 5901- 5914, 2019. URL https::/pproceedings. wtw.deeeeddeeinnneneo..eeeoee..e neurips cc/ httpppppppppppppppppptppppppppep hwpppppppll... /hash/ html.
Wei Jiang, Eduard Trulls, Jan Hosang. Andrea Tagliasacchi, and Kwang Moo Yi. COTR: correspondence transformer for matching across images. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pp. 6187-6197. IEEE, 2021. doi: 10.109//CCV4992..321.0000 10.1109/ ICCV48922 2021. 00615.
URL https: http://oooooooooooooooooooo...o. org/
Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, Lawrence Zitnick, and Ross B. Girshick. CLEVR: A diagnostic dataset for compositional language and elementary visual reasoning. In 2017 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2017, Honolulu, HI, USA, July 21-26, 2017, pp. 1988-1997. IEEE Computer Society, 2017a. doi: 10.1 10.11099CVPR.201.2255 URL https: / doi. , org 1wxg/1100 0. 111. /ww.1100////..222 Cc//////R..20 215.
Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Judy Hoffman, Li Fei-Fei, C. Lawrence Zitnck, and Ross B. Girshick. Inferring and executing programs for visual reasoning. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017, pp.
IEEE Computer Society, 201 7b. doi: 10.1 N..100//CCV22011.33355 URL https: //doi. .rr/orgg/// 0 11 2//9///////// 0..0....55.
John Jumper, Richard Evans, Alexander Pritzel, tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58)) 2021.
Brenden M. Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-fo-sequencc recurrent networks. In ICMI, 2018.
Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40:e253, 2017. doi:
10.10177/0100555X00000033
Yujia Li, Felix Gimeno, Pushmeet Kohli, and Oriol Vinyals. Strong generalization and efficiency in neural programs. CoRR, abs/2007.03622 2020. URL https: httpiv.iivvi.///////ee//aaaae.l/ org/ abs/ arg0 03629.
Joakim Lindblad and Natasa Sladoje. Linear time distances between fuzzy sets with applications to pattern matching and classification. IEEE Trans. Image Process., 2361):1261336, 2014. doi: 10.1 10.1109/TIP.211 2...I..201..288604.4444 2286904. URL https: http::///ooooogooee org 10. 1111911111PPP 2013 2286904.
D.G. Lowe. Object recognition from local scale-invariant features. In Proceedings of the Seventh IEEE International Conference on Computer Vision, volume 2, pp. 1150-1157 vol.2, 1999, doi: 10. 1911000C...11999.99999
Jiahao Lu, Johan Öfverstedt, Joakim Lindblad, and Natasa Sladoje. Is image-to-image translation the panacea for multimodal image registration? A comparative study. CoRR, 0ass2103.11266, 2021. URL https: //aaai/////vvv..//// org/ abs 22103.13.220 16262.
Jayuan Mao, Chuang Gan, Pushmeet Kohli, Joshua B. Tenenbaum, and Jiajun Wu. The neuro- symbolic concept learner: Interpreting scenes, words, and sentences from natural supervision. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans. LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https: /ttp::///peeeeeeeeeee www.wieweeeet.e.. www.nettffrrmme forum? d= rJgMlhRctm.
Karolis Martinkus, Andreas Loukas, Nathanaël Perraudin, and Roger Wattenhofer. Spectre: Spectral conditioning helps to overcome the expressivity limits of one-shot graph generators. In International Conference on Machine Learning, ICML, 2022.
Nicolas Pielawski, Elisabeth Wetzer, Johan Öfverstedt, Jiahao Lu, Carolina Wählby, Joakim Lindblad, and Natasa Sladoje. Comir: Contrastive multimodal image representation for registration. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual
11

---

Under review as a conference paper at ICLR 2023
Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL http:::/ppproceeiinsssss wwwp:/ccwweiiinn...iin...e. .c/ wttppappe//2///aaaaee.eeeee d6 2
html.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to- text transformer. J. Mach. Learn. Res., 21:140:1-140:67, 2020. URL http: /mlz.//m....///////// orgg
papers/v21/20-072 .ttm
Scott Reed and Nando De Freitas. Neural programmer-inteepreeerr. arXiv:1511.00279, 2015.
arXiv preprint
Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. Superglue:
Learning feature matching with graph neural networks. In 2020 IEEE/CVF Conference
Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp- 4937-4946. Computer Vision Foundation IEEE, 2020. doi: N0.1109/CVPR2220022222 0w00226000.2200044000 URL https: opwnaccesssssscceeppppppcc..e.se thecvf.ee.. com conteetttt _eteeeett_______eeeeeeeee CVPR 2020/htmm/seeeee SuperGlue Learning BettnninnnFFauuuueeeeeeeee Feature Mat With_Grapp Neural Neurra_Netwooke. CVPR_2020_pape html.
Laurent Sartran, Samuel Barrett, Adhiguna Kuncoro, Milos Stanojevic, Phil Blunsom, and Chris Dyer. Transformer grammars: Augmenting transformer language models with syntactic inductive biases scale. CoRR, 0ttp/2203.000333 2022. doi: 10.48550/arxiv.22000 2203.000....220033... URL https:
//doi.org/10
/0rg/11000 1..0..85550 1/88//////ii.../ www..iii............ 220..006333...3 60333.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Marilyn A. Walker. Heng and Amanda Stent (eds.), Proceedings of 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT, New Orleans, Louisiana, USA, June 1-6, 2018, Volume 2 (Short Papers), pp. 464-468. Association for Computational Linguistics, 2018. doi: 10.18653/v1/nI1 2/744- URL htw.h:tpp/t//...t. /doii orrel/orr/1// 11..1..1661111111.11
Elizabeth S. Spelke and Katherine D. Kinzler. Core knowledge. Developmental Science, 10 (1):89-96, 2007. doi: wwt//ooiabs//.........o/
URL https:// wi wiley. com doi wtt.s//11. 1ttp://1..1111/.1111..1........1 700777777772277..... 2000000069.......
Hannes Stärk, Octavian-Eugee Eugen Ganea, Lagnajit Pattanaik, Regina Barzilay, and Tommi Jaakkola. Equibind: Geometric deep learning for drug binding structure prediction. 2022.
Freyr Sverrisson, Jean Feydy, Joshua Southern, Michael M Bronstein, and Bruno Correia. Physics- informed deep neural network for rigid-body protein docking. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 2017. URL https: http:::://ppcceeiiggs.. neurips......................... ww/pappe///000 /hash// Abstract
3f5ee2435 232224444474e333 deeel www.dfddd33ccccccc oolee2/4aardeeeoooooooococcoece Abstsactceessssaaaaaraatttttt html.
Yujun Yan, Kevin Swersky, Danai Koutra, Parthasarathy Ranganathan, and Milad Hashemi. Neural execution engines: Learning execute subroutines. In Hugo Larochelle, Mare" Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL ht httpss://prrceeddnggg neurips ctt.cisscc..//e....eeeee cc/ pappr/2020/haaa/
wttppr/pss..cppppe//OO/// cb92afffsSaf77660006111e12222223 Atsstaaaat...................... html.
Abstract.
12

---

Under review as a conference paper at ICLR 2023
A ADDITIONAL DETAILS ON THE MODEL
This section describes the LATFORMER architecture providing additional details that were not covered in Section 4.1. As mentioned in Section 4.1, it is possible to design convolutional neural networks that perform considered transformations of the lattice. Figure 6 shows the architecture of the four expert models that generate translation, rotation, reflection and scaling masks.
Export
Scaling Expert
Cloree
Cless M
Classs M
Cloese M
My
K
ME
Cone M
Closu M
Clave M
Closeer M
Conp
Clna
M
8 (MM. Mg)
Cona M
Rotation Expert
Clomaa
M
+
Mr
M
Lattice Retlecion Expert
Cone
Clome M2
Figure 6: Model architecture of all the mask experts that considered.
models are CNNs applied to the identity matrix. the figure, we use the following notation:
M (6) denotes an attention mask implementing a translation by along one dimension;
M M denotes an attention mask implementing translation by 90%;
(90)
MM denotes an attention mask implementing reflection along one dimension;
M denotes attention mask implementing an upscaling by along one dimension.
(h)
Using Corollary 1, we can derive the kernels of the convolutional layers shown Figure 6. These kernels frozen at training time, the model only learns the gating function, denoted figure. Notice that all the models follow the same overall structure. However, for scaling, we also learn additional gate, denoted (Ms, M ) in the Figure 6. This gate allows the model to transpose the mask and serves the purpose of implementing down-scaling operations (down-scaling the transpose up-scaling).
in the
The composition of more actions can be obtained by combining different experts. This can be done either by chaining the experts by matrix multiplication of the masks. 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.
B ADDITIONAL EXPERIMENTS AND DETAILS ON THE EXPERIMENTAL SETUP
This section provides additional details the experimental setup of all our experiments, including further information on the generation of the synthetic tasks and the data annotation process for ARC.
B.1 EXPERIMENTS ON SYNTHETIC DATA
We considered four categories of tasks, namely translation, rotation, reflection and scaling. Each task 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. To each input grid, synthetic transformation is applied in order to obtain the corresponding output grid. For each task in each category, we generated 2048 training pairs and 100 test pairs.
For translation tasks, we have total of 900 possible translations a 30 x 30 grid. However, generating data and training models 900 tasks is computationally expensive, randomly
13

---

Under review as a conference paper at ICLR 2023
sampled 5 translations in the interval 29] [9]
[1, 29), obtaining total of 100 translation tasks. Rotation tasks include all 4-fold rotations except the identity. Similarly, reflection tasks involve horizontal, vertical and diagonal reflections. Scaling tasks include all possible up/down scaling transformations of the input grid by factors of [2, 5] x [2, 5) for a total of 32 scaling tasks.
The models evaluated based on the mean accuracy each category. For each task compute the accuracy on the test set based on how many of the predicted images match exactly the ground truth.
B.2 EXPERIMENTS ARC
Training data
Test data
Trunstaiion
Rotation
Refection
Sceling
Repeating pateer
Projectinn
Object cohosion
Object stance
Object
Moth
Countng
Gool droctedness
Calogrics
Tranaistion
None
Pg
Cor monit
Projecionn
Counting
Sooting
Object persislance
Object infuunco ea contact diectedress
Previous
Next
(b)
Figure 7: Distribution the considered core knowledge priors across the ARC tasks (a) and user interface built annotate the dataset (b).
In order to experiment with ARC, first performed annotation of the dataset to identify the underlying knowledge priors for each task. To this end, we built user interface where the annotator could browse the tasks and label them by selecting any combination of the available knowledge priors. Figure 7b shows the user interface provided to the annotator, whereas Figure 7a shows the distribution of knowledge priors across the ARC tasks. Most tasks follow in more than one of the categories represented in Figure 7a.
ARC can be regarded of unseen tasks to evaluate the performance of the model learned on the meta-training data. It is important to emphasize that we do not target this use case, we instead use the same setup in the synthetic data and learn each task from scratch using only its training set.
meta-learning benchmark, provides a training tasks and
Though simple and elegant, the ampervised-learnine formulation prevents our models from reusing knowledge that can be shared between different tasks. In order to mitigate this issue, we rely on lata-atggentttion strategy. At training time, for each model and every iteration, we augment each grid 10 times by mapping each color to different color (using the same mapping across training examples). The rationale behind this data- auttenaggmenttiion strategy is that (I) we assume that for tasks involving only geometric knowledge priors to be not affected by color mapping and (2) all models (including LATFORMER) need to learn a function from d-dimensional color representations categorical variables, hence it is beneficial if all colors represented in the training set.
14

---

Under review as a conference paper at ICLR 2023
Aerial data
Cytological data
a-AMD
SIFT
LatFormer
a-AMD
SIFT
LatFormer
CycleGAN (A + B)
5.3
67.2 - 16.8
68.3 4.5
74.2 + 3,8
30.2 = 4.2
68.3 2.2
CycleGAN (B - A)
65.7 6.7
84.0 + 2.5
86.1 + 3.1
21.3 + 1.8
18.2 + 3.5
24,2 + 3.3
DRIT++ (A + B)
35.3 2.1
38.1 8.1
38.2 + 5.9
50.4 + 12.1
24.2
DRIT++ B+
38.3 4.5
43.2
30.1 1..1
5.2
3.56
pixel2pixel (A-+ B)
8.0
98.7 0.4
2..
6.9
9.5
5.5
pixel2pixel (B - A)
6..2
4.03
8..7
0.2
4.0
4.2
StarGAN (A + B)
6.81
2.7
72.2
+ 6.3
60.2 + 12.2
12.2
12.2-22 2.0
59.5
+
StarGAN (B (B+A)
52.1 4.0
1.9
53.3 3 4.0
20.8 + 3,9
1.1
13,4
CoMIR
94.2
100.0 1 0.0
90.2 3.3
76.2 - 12.1
74.1 6.3
78.1 1 3,4
Table 4: Results of the experiment on image registration. The rows represent different models
trained to translate images from modality A to B (A + B) viceversa (B A).
All models are evaluated based on the ratio of solved tasks and a task is considered solved the model can predict the correct output grid for all examples in the test set.
B.3 ADDITIONAL. EXPERIMENTS ON IMAGE REGISTRATION
As an additional experiment, to assess the applicability of our LATFORMER on natural images, we performed experiments on multimodal image registration, namely the problem of spatially aligning images from different modalities. Image registration a well-studied problem in computer vision and we do not aim to establish state-of-the-ee performance. The main purpose of this experiment is giving a hint on the applicability our method natural images beyond ARC. We refer the reader to SuperGlue (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 model that converts one modality into the other (or to transfer both modalities the same representation proposed by Pielawski et al. (2020)), then they align the two images using traditional techniques. We follow the experimental setup of Lu et al. (2021) and experiment with two datasets, one containing aerial views of urban neighborhood and containing cytological images. 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. We use the code of the authors generate data involving only translations. Lu et al. (2021) additionally consider small rotations, but these transformation are not actions in the symmetry group of lattice, so we are not interested in resolving them.
We employ several state-of-the-art methods for modality translation and we compare our method a-AMD (Lindblad & Sladoje, 2014) and SIFT (Lowe, 1999) based on the success rate metric defined by Lu et al. (2021). A registration is considered successful if the relative registration eттоr (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%. 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. We the same models of Lu et al. (2021) for the modality translation task. Then, order to solve the image registration task with LATFORMER, we model to predict the translation from one
divide each image into 30 30 patches and we patch in an image to its counterpart in the corresponding image.
C LIMITATIONS AND FUTURE WORK
Although we believe our results are interesting and promising for learning group actions with neural networks, would like to point out some limitations of our approach.
First, our method is limited to actions on the symmetry group of the hypercubic lattice and it is not immediately extendable to other groups. For instance. though permutation matrices 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. Second, the model is
15

---

Under review as a conference paper at ICLR 2023
hard fine-tune: 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 be addressed by still keeping the same overall idea of modulating attention weights using soft attention masks, possibly with different parametrization of the masks. Future work will focus this research direction and extending work to cover a wider set the ARC tasks.
D DEFERRED PROOFS
We prove both Theorem 1 and 2 by induction on the dimensionality of the hypercubic lattice m.
D.1 BASE CASE FOR THEOREMS AND 2
First, it is useful to notice that whenever M E {0,1)aXm has exactly single 1 per row, in other words M In I.; then, for any X E Raxd
A
MaskedAttention X; M A 111
X
softmax
Xx vd
④
M
X
softmax
XX Vd O M 1
= -
In order to prove the theorem, we need to show that, for any action g E G1, including translations, reflections and rotations, there exists a mask Mg such that:
MuskedAttentionnXX X; M2) go X.
Let us consider different families of actions separately.
Translation. As mentioned in Section 3.2, in the I-dimensional case, a translation by one element is given by the circulant permutation matrix:
to the right for vector (æ1, T2.-..... ........................... Fn
0
0
1
0
0
M M
0
1
0/
This holds because M (1) 1, 1n, S0:
MaskedAttentiinn(;;: M()) = M x (In, Tnn... TTn11
In general, translation by å elements is given by the circulant matrix M (M yo Therefore, masks implementing translation operations exist in the 1-dimensional case and they are circulant permutation matrices. This is enough for base case for Theorem 1.
For Theorem 2, simply notice that:
2nj (8) M (6) = =F exp(- of rTT) where (8)
This comes directly from the time-shifting property of the Fourier transform.
16

---

Under review as a conference paper at ICLR 2023
Reflection.
In the 1-dimensional case, the reflection of vector (Tu, ,....,, R is:
MasstodAttention((; Mf)
Mr ,,
with
0
0
0
0 ...
0
MF
0
0
The attention mask ME can be obtained by shifting the rows of the identity matrix by:
n
T1
OF 71
Therefore, by the time-shifting property of the Fourier transform we have:
Mr MF=F-((F(I)) exp(- 2TJ OF TT).
n
Rotation. Rotation (4-fold) is not defined in one dimension, for a base we need consider the square lattice. Let X E R/11-2 be vectorized representaiton of a n = 11 x dimensional matrix. We need define a vector OR IR" such that:
M (90) = (F(,) exp( oRT,)) R
2n
is a rotation mask. Since rotation is permutation of the identity, we know the vector exists. As X vectorized, the oR (90) needs to take into account the size of the first dimension 4. For example, in order to perform rotation on vectorized representation, we need to map the first element of X to the position (4 1). The reader can check that the vector given by
(90) ) k (4 1) L(k 1)/4, (oR
satisfies the equation above.
Scaling, 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 the other transformations. We can take the I-dimensional lattice as a base case and consider vector æ (Io,21, .nn) Let h N be parameter specifying the filter size of the scaling operation. As example, for h we have:
MastedAAttettionzzzz M = M
x (x1, T1, 22, 22 .2,....,.11........... F[m/2))
where:
1 0 0
0
M (h)
0
0
0
This kind of matrix can also be obtained by shifting the rows of the identity as follows:
2nj (h)
M (h) F (F(() exp(
))
0s
where og (b) (Is mod h) + - 1) l(k 1)/bj.
17

---

Under review as a conference paper at ICLR 2023
D.2 INDUCTIVE STEP FOR THEOREMS 1 AND 2
Suppose that Mg E {0, 1 and Ma E {0, 1 are attention masks implementing
actions gi E Gmy and g2 Gm2 on some tensors X I and X2 E #1 = 4 ..... Lm and n2 = 14 ..... !'m Consider a tensor E x xti Xlmg and its vectorization X _2 with n1n2.
xiix-=xlm x22,, with
We have:
Attnnint X; Mg. @ M (Mgi @ Mg. X
(Mai @ (In, & M
In, @ Mgg)i @ Lnz).
Now notice that:
MaskedAtrentionn(:: In & Mg.) = (In, @ Mg.) X
= vec(Mgg vec (X) I,,)
vec(Mgs vec (X)),
and similarly
MaskodAAteetioonX* Mg 3 Ing) = (Mgi &n1)) X
= veo(In, vec (X) M 81
vec((Mg, vec (X)
TT)
Therefore, we conclude that performing masked attention with the mask Mg. @ Mg on X is equivalent applying gi on the first m, dimensions and ga on the last m2 dimensions of X. This provides way for building attention masks for higher-dimensional lattices using the primitive masks defined in Section D.1, proving both Theorem and
D.3 PROOF OF COROLLARY
The proof of Corollary follows immediately from Theorem 2 and from the property of the Fourier transform according to which multiplying the Fourier domain implements a convolution in the original domain.
18


Reviews:
[{"review_id": "gBOTL_NftFn", "reviewer": "Reviewer_vL2a", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, a LATFORMER model is introduced for learning effective attention. It incorporates both lattice geometry and topology priors for learning attention masks. The paper is demonstrated on ARC and on synthetic visual reasoning tasks, showing improved performances compared to standard attention and transformer on these tasks. ", "summary_of_the_review": "The experiments could be enhanced. More challenging tasks and datasets could be considered to better validate the proposed attention mechanism. The novelty is moderate. ", "strength_and_weaknesses": "Pros:\n+ The idea of introducing lattice geometry and topology priors for attention design is interesting\n+ The paper is clearly written.\n\nCons:\n- The experiments on ARC and synthetic datasets and tasks are not sufficient to demonstrate the effectiveness and generalization of the design. It would be much more interesting if the method can be demonstrated on main-stream vision and NLP tasks to show the effectiveness of the designed attention mechanism.\n- The ablation study should be extended to show the proposed attention mechanism with the lattice geometry and topology priors. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HHduiI0lka", "reviewer": "Reviewer_TgJP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper addresses problems that involve learning a geometric transformation (e.g., rotation) on the input data. Examples of the problem are Abstraction and Reasoning Corpus. The core idea of the proposed method is to inject lattice symmetry prior as soft masks to modulate the attention weights in a transformer. The paper shows that the various lattice symmetry actions can be realized with attention masks (such as translation, rotation, and reflection). The results on ARC tasks show that using the proposed method with infused lattice symmetry prior can perform well (Table 2 and Figure 5) and is sample-efficient. ", "summary_of_the_review": "I think this is an interesting paper that shows end-to-end trained models can perform well on abstract reasoning tasks. I don't know the practical implication of solving these lattice symmetry problems and how the ideas in the paper could be used to solve some other tasks. The paper did show an image registration task. But it does not provide sufficient details (e.g., are these only translations to rotation and translation? what's the task? what metrics were used to evaluate the success in Table 3?) for me to understand the significance. For registering images from different modalities, one should also try stronger descriptor-matching methods beyond SIFT, e.g., SuperGlue.\nGiven these concerns, I am thus on the fence about this paper but leaning slightly positive.", "strength_and_weaknesses": "Strength: \n+ Interesting idea of infusing lattice symmetry prior to a problem into an end-to-end trainable transformer model. \n+ The progress on ARC datasets using deep learning models is promising. \n+ The method is technically sound. The CNN-based mask prediction (Lattice Mask Expert) is interesting and supports compositions of multiple actions (e.g., reflection and rotation).\n\nWeakness:\n- From the examples in the paper and supplementary material, it seems that rotation is only 90 degrees. Not sure if this makes the problem trivial to solve (if we know the symmetric group of testing input data).\n- No visual examples and results on the cross-modal image registration tasks. I am not familiar with the tasks so it's a bit hard for me to judge the significance of the results and how the results validate the proposed method. \n- I wasn't able to follow some of the details in Section 3.3 and Section 4. But it may be because I don't have the background on related problems. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "B-ztCTPk8U", "reviewer": "Reviewer_YP2T", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on helping deep learning models learn fundamental geometric transformations efficiently. Specifically, LatFormer is proposed to incorporate lattice symmetry biases into attention mechanisms by modulating the attention weights using learned soft masks. Experiments are conducted on both synthetic tasks and image registration tasks to show that the proposed model can generalize better than the same attention modules without masking and Transformers. Besides, LatFormer also shows potential on ARC tasks that incorporates geometric priors. ", "summary_of_the_review": "This is an interesting paper about learning geometric transformations with deep neural networks, with experimental evidences from both synthetic tasks, visual reasoning tasks as well as applications on natural images.", "strength_and_weaknesses": "Strength:\n\n+ The proposed method presents an interesting way to infuse inductive biases and knowledge priors in deep neural networks, which is an important open problem for next step's developments of AI and neural networks. \n\n+ The idea of using a convolutional neural network to generate soft attention masks to scaled attention weights is reasonable. \n\n+ Experimental results on synthetic tasks and ARC tasks show that the proposed LatFormer can generalize better and from fewer\nexamples than transformers and unmasked attention modules. Experiments on image registration task further validates the applicability of LatFormer on natural images.\n\n\nConcerns:\n\n- LatFormer shows potential of inductive biases and knowledge priors. I'm wondering if LatFormer also performs better than learning transformations invariant representations by applying augmentations? Do we have such studies on the current experimental settings/tasks?\n\n-  For image registration experiments in Section 5.3, how about reuslts of using the baseline Transformer?\n\n- How about the application of the proposed LatFormer to other reasoning tasks, e.g. Physical Reasoning in [1] \n\n[1] Physical Reasoning Using Dynamics-Aware Models, ICML 2021", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6wUljBhIWg6", "reviewer": "Reviewer_ch2n", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents an approach to incorporate lattice priors into attention mask for deep learning. The theoretical results show this prior can be formulated as convolutions on an identity matrix. The proposed approach improves the sample efficiency of the model to solve the ARC tasks (few-shot geometric transformation learning) compared to vanilla transformers. \n", "summary_of_the_review": "The idea of the paper is reasonable. Theoretical results and experiments on synthetic data show the proposed approach is superior to vanilla transformer. However, the experiments on real data seem relatively weak and it’d be valuable to see more evidences of LatFormer outperforming transformers and more existing approaches than SIFT/alpha-AMD on image registration benchmarks.", "strength_and_weaknesses": "Strength\n* The idea of introducing structural priors in the form of attention masks for geometric transformation tasks is intuitive and well-motivated\n* The theoretical results show that such priors can expressed as a convolution on an identity matrix, which motivates the design of LatFormer\n* Empirical results show remarkable sample efficiency gains over vanilla transformers in ARC tasks and outperform baselines on real image registration\n\nWeakness\n* It’d be valuable to show more experiments on real data. How about showing the results on standard image registration/correspondence learning benchmarks?\n* In Table 3, are there other baselines to compare with? I’d imagine a vanilla transformer would be a reasonable one.\n* In Figure 5, it seems that the transformer is barely learning the task. I’m curious how many samples would it take for a transformer to learn this task perfectly (how far should we extend the x-axis).\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6WPY-uA_Oij", "reviewer": "Reviewer_9teR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes LAtFormer, which incorporates lattice geometry and topology priors in attention masks. The proposed architecture is a modification of the standard attention mechanism where attention weights are scaled using soft attention masks generated by a convolution neural net. Experiments on ARC and synthetic visual reasoning tasks show that LatFormer requires 2 orders of magnitude fewer data compared to standard attention and transformers.", "summary_of_the_review": "Please see the comments in the boxes above.", "strength_and_weaknesses": "Strengths:\n1. The paper is evaluated on a number of tasks.\n2. Results show that Latformer can learn geometric transformations with fewer training data.\n\nWeaknesses:\n1. Table 3 only compares the proposed method to alpha-AMD and SIFT which are old methods. Since this is an image registration task, the proposed method should be compared to more recent methods, e.g., COTR: Correspondence transformer for matching across images. ICCV 2021.\n2. The paper should also discuss failure modes, limitations of the proposed method and future work.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject