Title: ATTENTION ON ABSTRACT VISUAL REASONING

Abstract: Attention mechanisms have been boosting the performance of deep learning models on a wide range of applications, ranging from speech understanding to program induction. However, despite experiments from psychology which suggest that attention plays an essential role in visual reasoning, the full potential of attention mechanisms has so far not been explored to solve abstract cognitive tasks on image data. In this work, we propose a hybrid network architecture, grounded on selfattention and relational reasoning. We call this new model Attention Relation Network (ARNe). ARNe combines features from the recently introduced Transformer and the Wild Relation Network (WReN). We test ARNe on the Procedurally Generated Matrices (PGMs) datasets for abstract visual reasoning. ARNe excels the WReN model on this task by 11.28 ppt. Relational concepts between objects are efficiently learned demanding only 35% of the training samples to surpass reported accuracy of the base line model. Our proposed hybrid model, represents an alternative on learning abstract relations using self-attention and demonstrates that the Transformer network is also well suited for abstract visual reasoning.

1. INTRODUCTION:

Psychological models of human intelligence identify different manifestations of intellect. FLUID INTELLIGENCE describes the ability to adapt to new problems and situations without relating to previous learning outcomes (Hagemann et al., 2016) . This ability is considered as one of the most important aspects for learning and is essential for solving higher cognitive tasks (Jaeggi et al., 2008) . In order to excel in this type of intelligence, cognitive capabilies such as figural relations, memory span and inductive thinking are decisive. Fluid intelligence also paves the ground for ABSTRACT REASONING, the ability to use symbols instead of concrete objects. Furthermore, empirical data show that attention and fluid intelligence are strongly interlinked (Stankov, 1983; Schweizer, 2010; Ren et al., 2013) . Subjects that perform bad on attention tasks are also more likely to show deficits in abstract reasoning and fluid intelligence (Ren et al., 2012; 2013) . RAVEN'S PROGRESSIVE MATRICES (RPM) is an established test method for intelligence, especially fluid intelligence and abstract reasoning (Bilker et al., 2012) . It is a set of non-verbal tests showing several geometric objects arranged to a certain implicit rule (Figure 1a shows an example). The test subject has to complete the 3 × 3 matrix by picking an object that matches the implicit rule. The PROCEDURALLY GENERATED MATRICES (PGMs) (Santoro et al., 2018) dataset is motivated by RPMs and synthetically generated for training neural networks. Its core feature comprises several relations between objects and one of their attributes, e.g. the number of an object of a certain type in each matrix panel. A single matrix can contain up to four such relations simultaneously. A PGM's context and possible answers are shown in Figures 1a and 1b , respectively. Importantly, the PGM matrices are presented to a neural network as images such that the task requires a certain level of understanding of geometrical relations. The Wild Relation Network (WReN) uses different combinations of matrix elements as a mechanism for relational reasoning and is currently among the best performing solutions for the PGM dataset. However, despite the close relationship between attention and reasoning, the WReN is lacking a mechanism for attention.

Apart from the visual domain, it was recently shown that attention is also essential for another cognitive task that fundamentally relies on abstract reasoning: language understanding. The Transformer, an attention-based neural network, was introduced to improve machine translation and transduction (Vaswani et al., 2017) . The Transformer embodies a self-attention mechanism to relate parts of its input with each other for writing output. Since this form of attention and abstract reasoning is vital for fluid intelligence we hypothesize that a Transformer augmented with relational reasoning capabilities can perform well on the PGM task.

Another special trait of human intelligence is the ability to infer rules with little supervision. Often, only few samples suffice for people to grasp the idea of tasks such as visual reasoning. In contrast to that, current deep network architectures are trained on huge quantities of data, e.g. the PGM dataset contains more than 1.2 million labelled samples. To see if self-attention also enhances this capability of deep networks we investigate the sample efficacy by training our proposed model only on a fraction of the total available data samples.

Our contribution to the field of machine learning is three-fold:

• We introduce the Attention Relation Network (ARNe) that combines features from the WReN and the Transformer network and can be directly trained on visual reasoning tasks.

• We evaluate ARNe on the PGM task and show that it significantly outperforms the current state-of-the art (Steenbrugge et al., 2018) by 11.28 ppt.

• We demonstrate that ARNe is very sample-efficient and achieves its peak performance with only 35% of the full PGM dataset.

This paper is organized as follows. In Section 2 we discuss related approaches and training datasets.

In Section 3 we introduce the ARNe architecture and in Section 4 we show our simulation results. We conclude in Section 5.

2. RELATED WORK:

Language Modelling Language modelling, i.e. predicting the likelihood of future words given a set of contextual words, is a core task in natural language processing (NLP). While neural approaches to language modelling have some tradition (Bengio et al., 2003) , the success of the word2vec (Mikolov et al., 2013) has revived and boosted the interest in neural networks for this task. While word2vec defines fixed vectors for words, recently, several approaches were proposed that condition word representations on their context (Howard & Ruder, 2018; Devlin et al., 2018; Peters et al., 2018) . In particular attention-based models have emerged as the standard tool for language modelling as well as several downstream tasks, such as question answering (Rajpurkar et al., 2016) and translation (Vaswani et al., 2017) . Starting with the seminal Transformer (Vaswani et al., 2017) , follow-up models were successively improved (Devlin et al., 2018) and trained on larger text datasets achieving remarkable success across various NLP tasks (Radford et al., 2019) . Besides the work in language modelling, in concurrent work, Transformer-based models have been employed for visual question answering (Li et al., 2019) . In this work we take inspiration from the model Equation ( 2) Equation ( 3) Equation ( 1 design of the attention Transformer. However, instead of applying it in a language modelling task, we use it for abstract reasoning.

Abstract Reasoning Recently, there has been growing interest in abstract reasoning indicated by the introduction of datasets operating on textual data, e.g. babl (Weston et al., 2015) , as well as on visual data, involving CLEVR (Johnson et al., 2017) , FigureQA (Kahou et al., 2017) and NLVR (Suhr et al., 2017) .

While classic recurrent neural networks have often been used for these tasks, it has been found that novel architectures such as attention and memory are better suited to address abstract reasoning. Attention mechanisms are a method to simplify complex information by attending to different locations of the given data. This is an advantageous procedure to let a neural network determine a specific reasoning operation on these selected locations. Memory is needed to preserve information in tasks where multiple consecutive reasoning operations are required. In LSTM (Hochreiter & Schmidhuber, 1997) for example, the hidden state is protected by gates from being overwritten by new input. Models incorporating attention and memory involve the Neural Turing Machine (Graves et al., 2014) which evolved into the Differentiable Neural Computer (Graves et al., 2016) . Dynamic Memory Networks (Xiong et al., 2016) address both, visual and textual question answering. The recently introduced Memory, Attention and Composition (MAC) Network Hudson & Manning (2018) achieved excellent scores in the CLEVR benchmark by using multiple explicit reasoning steps. Recently, a reasoning module relying on attention Zambaldi et al. was designed to be employed in a reinforcement learning framework.

Our work differs by measuring fluid intelligence by understanding visual patterns rather than solving referential expressions (to which part of an image belongs a certain word) which is required in CLEVR, FigureQA, NLVR and many visual question answering tasks. The former task is considered in (Santoro et al., 2018) and (Steenbrugge et al., 2018) with which we compare our proposed method.

3. ATTENTION RELATION NETWORK (ARNE):

In this section we describe the Attention Relation Network (ARNe), which combines techniques from language modelling to abstract reasoning. It takes eight contextual panels and eight choice panels from the PGM, of which one fits to abstract relational rule implicitly determined by the context panels. The model output are logits of a corresponding fitness probability for each choice First, representations for context panels, denoted by x i , and representations for choice panels, denoted by c k , are generated using a shared Convolutional Neural Network (CNN). x i denotes the i-th context panel feature, c k denotes the k-th choice panel feature. This network has the same hyperparameters as the CNN in the Wild Relation Network (WReN). A one-hot positional encoding indicating one out of nine possible positions within the panel grid is concatenated to each extracted panel feature and the resulting vector is projected to obtain a final representation for each panel (see panel encoder and concatenation blocks in Figure 2 ). Sequences χ k of length N are composed of the context and choice panel representations. Optionally, a learnable deliminiter d which has the same number of dimensions as the panel representations can be included between contexts and choice. Thus, sequences of N = 9 or N = 10 elements are obtained:

χ k = (x 1 , x 2 , . . . , x 8 , d, c k ) .

(1)

We generate a sequence χ k for all eight choices. A multi-step attention network inspired by the encoder of the Transformer model t (Vaswani et al., 2017) , processes this sequence using the selfattention mechanism for abstract reasoning. All activations are accumulated before the network's output is generated by the MLP f .

EQUATION

where g denotes the output of the 4-layer MLP (see f and g decoder blocks in Figure 2 ). The vector p ∈ R 8 denotes the logits of the corresponding probability distribution over the choices and the matrix Â ∈ R 8×12 indicates the logits of the corresponding presence probabilities of rules imposed by the panels where âk ∈ R 12 is defined as a transposed row vector of this matrix. The construction of these meta-targets is explained subsequently.

The network returns a matrix o ∈ R 8×13 , which contains both, logits for the prediction of a choice panel and logits for the prediction of underlying patterns across the PGM Meta-targets In the PGM dataset, the layout and appearance of the nine panels follow an implicit rules. Each rule is represented by a triplet containing an object type, object attribute and relation type. An example rule could be (shape, size, progression) to describe a pattern of triangles with increasing size. Up to four different relational rules (i.e. four triplets) may occur simultaneously in a single PGM. Meta-targets are created by aggregating all binary encoded rules with an OR operation resulting in a 12 dimensional meta-target vector.

EQUATION

In order to identify correct panels, a successful model should be able to infer the underlying construction principle reliably. To enhance rule prediction in the model, we incorporate the auxiliary information encoded in the meta-targets in the loss function. Since each choice panel results in a prediction âk of a meta-target, the predictions need to be aggregated by a sum over choice panels, implying an OR relation between the elements (Santoro et al., 2018) .

Loss The loss L is defined by a weighted combination of finding the right choice p and detection of the correct logical pattern a (Santoro et al., 2018) :

L = CE(p, p) + β • BCE(a, 8 k=1 âk ) ,

with CE and BCE being cross entropy and binary cross entropy. The parameter β controls the influence of the meta-targets, i.e. for β = 0 the meta-targets are ignored. By enforcing correctness of the meta-targets we enable richer gradients that identify better parameters.

Implementation The learning is carried out using an adam Kingma & Ba (2014) optimizer with a batch size of 64 and an initial learning rate of 0.5 • 10 -4 using a learning rate scheduler with exponential decay. We apply early stopping (Goodfellow et al., 2016) with a patience of three epochs. Parameters of the individual components of our method are presented Table 1 . For further details we refer to the appendix or the implementation of our approach which is available here: http://hidden for blind review. Will be disclosed upon acceptance

4. EXPERIMENTS:

We evaluate ARNe on the PGM dataset (DeepMind, 2017) . Each PGM input sample consists of 160 × 160 pixel images. In total, there are 16 panels in every sample, eight of which define the context and the remaining eight define possible choices (see Figure 1 ). Similar to RPMs, a correct panel has to be chosen that matches the implicit relational rules encoded in the context panels. PGM includes meta-targets in the form of 12-bit feature vectors that denote relation, object and attribute types. The rules that underlie each sample are composed of 1 to 4 relational rules, chosen from the set (Progression, AND, OR, XOR, Consistent Union). Figure 1b, c shows an example of the Progression rule.

We trained ARNe on the 1. implementation of WReN also outperformed (Santoro et al., 2017) by 2.1 ppt despite careful code validation (see Table 2 and Appendix B).

In addition, we included a comparison with the Attention and Composition (MAC) Network (Hudson & Manning, 2018) . To the best of our knowledge this is the first time MAC is tested on the PGM dataset. To our surprise MAC did not perform significantly above chance level on this task. We also tested a version of MAC that uses a WReN at the input stage, similar to ARNe. We call this augmented variant WReN-MAC. Our first results with this model suggest that also WReN-MAC does not reach the performance of ARNe. We ran tests on a reduced PGM dataset that used only 20% of the full training set. WReN-MAC achieved 46.9% test accuracy, about 10% below that of ARNe on the same dataset size (see Figure 3 ). We are running simulations with WReN-MAC on the full dataset but by the time of the submission of this paper these experiments were not concluded. Additional details on the implementation of WReN-MAC are provided in Appendinx D.

To gain additional insights into the learning behavior of ARNe we clustered the test dataset by the number of relational dependencies. ARNe showed good performance on all types, performing best on samples with four (accuracy: 90.54%) and worst for samples with three (accuracy: 82.74%) relational rules. Next, we evaluated the ability of ARNe to predict the meta-targets. Table 3 shows individual performances for each of the 12 meta-targets. ARNe achieved high accuracy in all categories. Detection rates revealed an unconditional accuracy rate of above 90% consistently across all meta-target types.

Furthermore, we evaluate the generalization abilities of ARNe by the provided samples in the PGM dataset. Consistently with the scores on the neutral split, we find our method to outperform the wild relation network in Tab. 4. 

4.1. SAMPLE EFFICIENCY:

In order to evaluate the sample efficiency of ARNe, we trained the model subsequently with growing fractions of the training dataset. The validation and test sets were kept unchanged during the entire experiment. Prior to splitting of the training set, the dataset was shuffled and the random seed was kept fixed throughout all experiments. The results are presented in Figure 3 . ARNe achieved higher test accuracy across all dataset ratios ≥20%. Using a split size of 35%, the accuracy reads 87.64%, already close to the value when trained on the full training set and 8.65 ppt better than the WReN model at 100% of the training dataset. The progression of the loss shows a similar behavior (Figure 3b ).

4.2. SELF-ATTENTION FOR VISUAL REASONING:

In order to gain insights into how the self-attention mechanism of ARNe works internally we visualize the activations of the panel encoder for different encodings of the meta-targets in Figure 4 . It is important to note that these features captures all relevant information for reasoning about relational rules between objects. Therefore, recognizing objects of specific types is an important sub-task for deciding which objects are relevant for a particular relational rule. Figure 4 shows two examples for the object type meta-targets shape and line. For each object type, output channels of the panel encoder CNN were selected that strongly responded. The eight choices and the corresponding correct context panel in the bottom right, are shown.

We found that, after training, many panel encoder CNN output channels significantly encoded activation maps for a single meta-target type. The activations in Figure 4 most saliently select regions within the input panels that are part of line or shapes objects.

4.3. ABLATION:

We experimented with various configurations of ARNe. This involves dropout to reduce over-fitting, a different learning rate, the inclusion of the delimiter token when assembling sequences and the weighting of meta-targets β. Table 5 depicts the results of the subsequent and more detailed ablation studies of ARNe. Further values are as listed in the Appendix A. For small dataset dropout gives a substantial performance boost, while this is not the case when the full dataset is used. This is likely due to the additional samples preventing the network from over-fitting. Similarly, the addition of a deliminiter tends to improve performance for small datasets, thus it might act as a regularizer. Furthermore, we find that the method is robust to a changed initial learning rate.

An intriguing insight is that ARNe does not converge when no meta-targets are provided (β = 0). This indicates the relevance of the auxiliary signals which connect the panel choice with the underlying implicit relational rules during training (at test no meta-targets were presented). Beyond the findings of (2019) . We found it achieve a performance of 92.23% for 50,000 samples for each figure configuration. However, using Raven-10000, i.e. using only 10,000 samples, the performance dropped to 19.67%.

5. CONCLUSION:

In this work we introduced ARNe, a new deep learning model that combines features from the Wild Relation Network (WReN) and the Transformer network to discover patterns in progressive matrix panels using aspects of fluid intelligence. More precisely, ARNe builds upon the WReN and extends it with the attention mechanism of the Transformer, which originates from language modelling, making it to our knowledge the first deep learning approach that uses self-attention specifically for abstract visual reasoning. The learning is driven by an auxiliary loss that gives hints about the underlying patterns of the progressive matrix panels. In an extensive experimental comparison we find that ARNe outperforms state-of-the-art abstract reasoning methods on the PGM dataset by a large margin. Moreover, the analysis shows that our model is substantially more sample-efficient than competing approaches.

Our experiments also including the first application of MAC to the PGM dataset. But the MAC did not seem to be well suited for this task, and also showed significantly worse performance compared to baseline, when it was augmented with a WReN-based panel encoder. Additional experiments are needed to gain a deeper understanding on why MAC appears to fail here. First insights into the inner workings of ARNe are given in Figure 4 . We show that the panel encoder CNN generates meaningful features that aid the learning goal. This suggests that the WReN and Transformer network parts of ARNe learn to efficiently cooperate for abstract reasoning. of MAC on the CLEVER dataset can be seen in Figure 6 . Best value for validation regarding the accuracy is 97.83%. The training, validation demanded a computation effort of about half a week on a GeForce GTX TITAN X graphics card.

D DETAILS TO WREN-MAC IMPLEMENTATION:

The WReN-MAC model uses the same encoding mechanism as the WReN baseline model and includes 12 MAC-cells. In order to interface a MAC-cell, additional convolutional layers were appended to the adapted WReN convolutional network. The model is analogously sequentially aligned like WReN or WReN-Transformer whereas during each pass one knowledge base which encodes one distinct choice panel is used. The question vector q is computed in analogy to WReN where g θ and the subsequently applied sums were used. The computation of the model's loss also respected auxiliary structure set data of the corresponding PGMs.

We were not able to train WReN-MAC successfully on the PGM dataset. Early stopping finished learning after 33 epochs. The final accuracies and losses are 49. 26% and 3.04, 46.97% and 3.20, 46.89% and 3.23 for training, validation and testing respectively. Both accuracy and loss showed marginal improvements throughout all epochs. Due the recurrent nature of MAC-cells, we used gradient clipping (Goodfellow et al., 2016) . The learning rate scheduler was set but didn't diminished the the learning rate which was initially set to 1 • 10 -4 . Additional hyperparameter tuning did not show significant performance increase.

:

https://github.com/pytorch/pytorch/releases/tag/v1.2.0

Figures and Tables:

Figure 1: Datasets for visual reasoning. a: Sample from the RPM dataset (Bilker et al., 2012). Context: blue, choices: red. Correct choice: 1. Implicit rule: Subtraction along rows and columns. b, c: Sample from the PGM dataset (Santoro et al., 2018), contexts (b) and corresponding choices (c). The ? in (b) should be replaced by a correct choice from (c). Correct choice: The hexagon in the last tile. Implicit rule: progression of shape types (number of edges) in each column. Colors added to enhance visualization.

Figure 2: Overview of the ARNe network for abstract visual reasoning.

Figure 3: Sample efficiency: Test accuracy (a) and loss (b) after training on different numbers of samples (percent of original dataset size). Our implementation of WReN was used here (see Appendix B).

Figure 4: Sum of selected activations of the last layer of the panel encoder CNN, that encoded predominantly for shapes(maps 14, 24, 26, 29, 31)  and for lines(maps 1, 2, 3, 4, 8, 12). Matrices show complete sets of context panels and the correct choice panel highlighted in red.

Figure5: Accuracy and Loss of the self implemented WReN(Santoro et al., 2018) model for β = 10 on the neural PGM dataset. The development of both metrics include the progression of the panel choices, additional auxiliary data which describes the setup of the PG-Matrices. In addition, the loss shows the overall sum.

Modules and their corresponding parameters of ARNe.

Table content: <table><tr><td>Module</td></tr></table>

PGM accuracy by previous methods and our model (ARNe). Accuracy of WReN as reported in(Santoro et al., 2017)  and our implementation. MAC model included for comparison.

Table content: <table><tr><td>Model</td><td/><td>Accuracy [%]</td></tr><tr><td>MAC (our implementation)</td><td/><td>12.6</td></tr><tr><td colspan="2">VAE-WReN (Steenbrugge et al., 2018) β = 4</td><td>64.2</td></tr><tr><td>WReN (Santoro et al., 2017)</td><td>β = 10</td><td>76.9</td></tr><tr><td>WReN (our implementation)</td><td>β = 10</td><td>79.0</td></tr><tr><td>WReN-MAC</td><td/><td>79.6</td></tr><tr><td>ARNe (our implementation)</td><td>β = 10</td><td>88.2</td></tr></table>

Meta-target prediction performance.

Table content: <table><tr><td>Meta-target</td><td>Accuracy [%]</td><td>Precision [%]</td><td>Recall [%]</td><td>F1-Score [%]</td></tr><tr><td>Progression</td><td>96.53</td><td>94.07</td><td>90.52</td><td>92.26</td></tr><tr><td>AND</td><td>97.76</td><td>96.92</td><td>94.01</td><td>95.44</td></tr><tr><td>OR</td><td>98.15</td><td>98.53</td><td>93.99</td><td>96.21</td></tr><tr><td>XOR</td><td>97.59</td><td>99.07</td><td>91.23</td><td>94.99</td></tr><tr><td>Consistent Union</td><td>96.87</td><td>97.63</td><td>91.79</td><td>94.62</td></tr><tr><td>Shape</td><td>99.82</td><td>99.91</td><td>99.81</td><td>99.86</td></tr><tr><td>Line</td><td>99.98</td><td>99.98</td><td>99.98</td><td>99.98</td></tr><tr><td>Size</td><td>99.93</td><td>99.89</td><td>99.77</td><td>99.83</td></tr><tr><td>Type</td><td>99.95</td><td>99.91</td><td>99.98</td><td>99.94</td></tr><tr><td>Position</td><td>99.98</td><td>99.91</td><td>99.99</td><td>99.95</td></tr><tr><td>Number</td><td>99.97</td><td>99.92</td><td>99.84</td><td>99.88</td></tr><tr><td>Color</td><td>98.13</td><td>99.01</td><td>96.80</td><td>97.89</td></tr></table>

Generalisation results. No training results for WReN were available(Santoro et al., 2018). Values in the training, validation and test columns characterise the model's accuracy.

Table content: <table><tr><td colspan="2">WReN (Santoro et al., 2018) 10</td><td>-</td><td>93.60</td><td>15.50</td><td/></tr><tr><td>ARNe</td><td>10</td><td>99.43</td><td>98.93 min</td><td>17.76 activation</td><td>max</td></tr><tr><td>shape</td><td>line</td><td>shape</td><td/><td>line</td><td/></tr><tr><td>PGM sample 15817</td><td/><td colspan="3">PGM sample 141081</td><td/></tr></table>

Comparison of different model configurations. The highlighted lines refer to the best performing model of the given dataset ratio. Values in the training, validation and test columns characterize the model's accuracy. (lr = learning rate).

Table content: <table><tr><td colspan="2">Dropout [%] lr×10 -4 Delimiter</td><td>β</td><td colspan="4">Ratio [%] Train [%] Val. [%] Test [%]</td></tr><tr><td>10</td><td>0.5</td><td>10</td><td>35</td><td>81.83</td><td>78.67</td><td>78.32</td></tr><tr><td>10</td><td>0.5</td><td>10</td><td>35</td><td>82.21</td><td>79.39</td><td>78.69</td></tr><tr><td>17</td><td>0.5</td><td>10</td><td>35</td><td>93.74</td><td>85.10</td><td>84.35</td></tr><tr><td>17</td><td>0.5</td><td>10</td><td>35</td><td>93.83</td><td>87.65</td><td>87.64</td></tr><tr><td>10</td><td>0.5</td><td>0</td><td>100</td><td>12.50</td><td>12.65</td><td>12.55</td></tr><tr><td>17</td><td>0.5</td><td>10</td><td>100</td><td>86.27</td><td>87.72</td><td>87.11</td></tr><tr><td>10</td><td>1</td><td>10</td><td>100</td><td>88.86</td><td>88.66</td><td>87.95</td></tr><tr><td>17</td><td>0.5</td><td>10</td><td>100</td><td>88.23</td><td>88.42</td><td>88.04</td></tr><tr><td>10</td><td>0.5</td><td>10</td><td>100</td><td>89.06</td><td>88.77</td><td>88.18</td></tr><tr><td colspan="2">Feed-forward net instead of Encoder</td><td/><td/><td>32.02</td><td>35.87</td><td>35.06</td></tr><tr><td colspan="2">Feed-forward net instead of Self-attention</td><td/><td/><td>38.79</td><td>45.12</td><td>44.56</td></tr></table>

Table 5, we investigated improvements of components of WReN as well as the addition of FiLM Perez et al. (2018) layers but did not notice a performance improvement. However, by reducing the input image size from 160 × 160 to 80 × 80 pixels, accuracy decreased by about 15 ppt. This suggests that relevant small structures exist and are used by the model to solve the PGM.

Table content: <table/>

Peer review (iclr2020):
This paper describes a somewhat novel approach to abstract visual reasoning using transformers in the so-called *Attention Relation Network* (ARNe), which the authors show to improve on the *Wild Relation Network* (WReN). The Transformer is motivated by the role that attention may play in Human information processing - which sounds plausible, but the paper does not expand on this theme. The paper is well written and makes an interesting contribution, but I feel the results are not quite yet ready for publication. The authors are writing that they are still working on baseline results on the full dataset, which would provide interesting comparisons, and some details on the implementation (number of parameters, etc) are missing - or maybe I missed them. The learning curve in Figure 3 (sample efficiency, test accuracy) suggests that the ARNe training is not fully stable - why would the model deteriorate when going from ~40% of the training data to ~60%? Is the model potentially overfitting, and how does the size of the proposed model compare to the size of the baseline model(s)? It seems that the field is also moving towards the RAVEN dataset, which presents a more complex structure; it would be more convincing to present results on both datasets, to show that attention can indeed also improve results on more complex setups. The text in the *Acknowledgments* section should be removed for the camera ready version!