GAMR: A G UIDED A TTENTION M ODEL FOR ( VISUAL ) R EASONING
Anonymous authors
Paper under double-blind review
A BSTRACT
Humans continue to outperform modern AI systems in their ability to flexibly parse and understand complex visual scenes.
Here, we present a novel transformer-based module for visual reasoning, the Guided Attention Model for (visual) Reasoning ( GAMR ), which instantiates an active vision theory – positing that the brain solves complex visual reasoning problems dynamically – via sequences of attention shifts to select and route task-relevant visual information into memory.
Experiments on an array of visual reasoning tasks and datasets demonstrate GAMR’s ability to learn visual routines in a robust and sample-efficient manner.
In addition, GAMR is shown to be capable of zero-shot generalization on completely novel reasoning tasks. Overall, our work provides computational support for cognitive theories that postulate the need for a critical interplay between attention and memory to dynamically maintain and manipulate task-relevant visual information to solve complex visual reasoning tasks
I NTRODUCTION
Intelligence is not only the ability to reason but also to find relevant material in memory and deploy attention when needed.
Abstract reasoning refers to our ability to analyze information and discover rules to solve arbitrary tasks, and it is fundamental to general intelligence in human and non-human animals (Gentner & Markman, 1997; Lovett & Forbus, 2017).
It is considered a critical component for the development of artificial intelligence (AI) systems and has rapidly started to gain attention.
A growing body of literature suggests that current neural architectures exhibit significant limitations in their ability to solve relatively simple visual cognitive tasks in comparison to humans (see Ricci et al.
(2021) for review).
Given the vast superiority of animals over state-of-the-art AI systems, it makes sense to turn to brain sciences to find inspiration to leverage brain-like mechanisms to improve the ability of modern deep neural networks to solve complex visual reasoning tasks.
Indeed, a recent human EEG study has shown that attention and memory processes are needed to solve same-different visual reasoning tasks (Alamia et al., 2021).
It is thus not surprising that deep neural networks which lack an attention and/or memory system fail to robustly solve visual reasoning problems that involve such same-different judgments (Kim et al., 2018).
Recent computer vision work (Messina et al., 2021a; Vaishnav et al., 2022) has provided further computational evidence for the benefits of attention mechanisms in solving a variety of visual reasoning tasks.
Interestingly, in both aforementioned studies, a transformer module was used to implement a form of attention known as self-attention (Cheng et al., 2016; Parikh et al., 2016). In such a static module, attention mechanisms are deployed in parallel across an entire visual scene.
By contrast, modern cognitive theories of active vision postulate that the visual system explores the environment dynamically via sequences of attention shifts to select and route task-relevant information to memory (Ullman, 1984; 1987). Psychophysics experiments (Hayhoe, 2000) on overt visual attention have shown that eye movement patterns are driven according to task-dependent routines.
Inspired b active vision theories, we describe a dynamic extension of the self-attention mechanisms popularised by the transformer module, which we call guided attention . Our proposed transformerbased Guided Attention Module for (visual) Reasoning (GAMR) learns to shift attention dynamically, in a task-dependent manner, based on queries internally generated by an LSTM executive controller. Through extensive experiments on the two main visual reasoning challenges, the Synthetic Visual
Reasoning Test (SVRT) (Fleuret et al., 2011) and the Abstract Reasoning Task (ART) (Webb et al., 2021), we demonstrate that our neural architecture is capable of learning complex compositions of relational rules in a data-efficient manner and performs better than other state-of-the-art neural architectures for visual reasoning.
Using explainability methods, we further characterize the visual strategies leveraged by the model in order to solve representative reasoning tasks.
We demonstrate that our model is compositional – in that it is able to generalize to novel tasks efficiently and learn novel visual routines by re-composing previously learned elementary operations.
Contributions Our contributions are as follows:
• We present a novel transformer-based end-to-end trainable guided-attention module to learn to solve visual reasoning challenges in a data-efficient manner.
• We show that our multi-head transformer-based guided-attention module learns to shift attention to task-relevant locations and gate relevant visual elements into a memory bank; such a multi-head transformer is also shown to perform significantly better than the popular self-attention mechanisms present in state-of-the-art transformer networks (Vaswani et al., 2017).
• We show that our architecture learns compositionally and is capable of learning efficiently by re-arranging previously-learned elementary operations stored within a reasoning module.
• Our architecture sets new benchmarks on SVRT (Fleuret et al., 2011) and ART (Webb et al., 2021), the two main visual reasoning challenges.
P ROPOSED APPROACH
shift attention, and selectively route task-relevant object representations ( z t ) to be stored in a memory bank ( M ).
The recurrent controller ( f s ) generates a query vector ( query t ) at each time step to guide the next shift of attention based on the current fixation.
After a few shifts of attention, a reasoning module ( r θ ) learns to identify the relationships between objects stored in memory.
Our model can be divided into three components: an encoder, a controller, and a relational module (see Fig.
1 for an overview).
The encoder module includes a feature extraction block ( f e ) for an image ( x in ) which is composed of five convolutional blocks (SI Figure S6).
The output of the module is denoted as z img ∈ R (128 ,hw ) (with h height and w width).
We applied Temporal Context Normalization (TCN) as done in (Webb et al., 2020). TCN is a simple inductive bias, implemented similar to batch normalization but applied at the task-relevant temporal window.
Using TCN helps to preserve the relational information between the objects within the window, resulting in better learnability and generalization.
TCN is applied over z img before passing it to the controller for further processing.
k r t =1 ← 0 h t =1 ← 0 M t =1 ← {} z img ← f e ( x in ) for t in 1...
T do out, g, query t , h t ← f s ( h t − 1 , k r t − 1 ) w k ← ( guided attention ( z img , query t )) .sum
( axis = 1) z t ← ( z img ∗ w k ) .sum
( axis = 1) if t is 1 then k r t ← 0 else w k t ← w k .sum
( axis = 2)
k r t ← g ∗ ( M t − 1 ∗ w k t ) end if M t ← { M t − 1 , z t } end for all obj ← r θ ( (cid:80) Ti,j =1 ( M v i , M v j )) ˆ y ← f ϕ ( all obj || out ) After z img is built, a guidedattention block routes visual information from relevant image location at each time step ( t ).
This block builds on the now classic multi-head attention (MHA) transformer module used in Natural Language Processing (Vaswani et al., 2017), which differs substantially from the selfattention transformer module more commonly used in vision.
An MHA block works as a retrieval block to extract relevant information based on 3 attention variables: key ( k ), query ( q ), and value ( v ).
The query ( q ) is used to compute a similarity score with the key ( k ), which is then multiplied by the values ( v ).
In our implementation, MHA accepts keys and values from a controller ( f s ) (discussed in the next paragraph) and queries from the encoder block ( f e ) to generate a weight vector ( w k ∈ R 128 ) corresponding to the key (k).
The vectors w k is then used to modulate the feature vector z img to yield a new context vector ( z t ∈ R 128 ).
A similar approach was used in Visual Question Answering ( V QA ) (Yu et al., 2019) to compute the similarities between questions and images.
The context vector ( z t ) is then stored in the memory bank ( M ) to be subsequently accessed again later by a reasoning module.
This memory bank is inspired by the differential memory used in Webb et al.
The controller module is responsible for generating a query ( query t ) in response to a task-specific goal in order to guide attention in the transformer module.
The controller module ( f s ) uses a Long Short-Term Memory (LSTM) to provide a query vector ( query t ∈ R 128 ) as input to the guided attention module for the current time step t .
The controller also generates a gate vector ( g ∈ R 128 ) and output vector ( out ∈ R 512 ).
The gate vector g is used to generate the next input to the controller based on prior relevant features stored in the memory.
The gate ( g ) is later used to shift attention to the next task-relevant feature based on the features previously stored in M .
On the other hand, the decision layer uses the output vector ( out ) to produce the system classification output (SI Figure S4).
The relational module is where the reasoning takes place over the context vector ( z t ) stored in the memory bank ( M ).
This module is composed of a multilayer perceptron (MLP) layer ( r θ ) which produces a relational vector ( all obj ) similar to the relational network (Santoro et al., 2017).
As we will show in section 6, r θ learns elementary operations associated with basic relational judgments between context vectors ( z t ) stored in the memory ( M ).
It is concatenated with the output ( out ) of the controller ( f s ) at the last time step (t=T) and passed through decision layer ( f ϕ ) to predict the output ( ˆ y ) for a particular task.
We have summarized the steps in Algorithm 1.
M ETHOD
Datasets We used two datasets for our experiments, SVRT and ART.
Experiments in section 4, 5, 6, 7 is carried out with SVRT dataset and in section 9 with ART dataset.
SVRT dataset is composed of 23 different binary classification challenges, each representing either a single rule or a composition of multiple rules.
A complete list of tasks with sample images from each category is shown in SI (Fig. S19, S20).
We formed four different datasets with 0.5k, 1k, 5k, and 10k training samples to train our model.
We used unique sets of 4k and 40k samples for validation and test purposes.
Classes are balanced for all the analyses.
We trained the model for a maximum of 100 epochs with a stopping criterion of 99% accuracy on the validation set.
The model was trained using Adam (Kingma & Ba, 2014) optimizer and a binary cross-entropy loss.
All the models were trained from scratch. We used a hyperparameter optimization framework Optuna (Akiba et al., 2019) to get the best learning rates and weight decays for these tasks and reports the test accuracy for the models which gave the best validation scores.
Webb et al. (2021) proposed four visual reasoning tasks (Fig. S22) that we will henceforth refer to as the Abstract Reasoning Task (ART): (1) a same-different ( SD ) discrimination task, (2) a relation match to sample task ( RMTS ), (3) a distribution of three tasks ( Dist 3 ) and (4) an identity rule task ( ID ). These four tasks utilize shapes from a set of 100 unique Unicode character images 1 .
They are divided into training and test sets into four generalization regimes using different holdout character sets (m = 0, 50, 85, and 95) from 100 characters. We have described training and test samples and different hyperparameters for all four tasks in SI section S6.
Baselines For the baselines in this dataset, we compared our architecture performance to a Relational
Network ( RN ), a popular architecture for reasoning in VQA.
The RN uses the same CNN backbone as GAMR with feature maps of dimension R 128 ,hw where h = 8 and w = 8 .
We consider each spatial location of the encoded feature representation as an object (i.e., N = 8 × 8 = 64 object representations).
We computed all pairwise combinations between all 64 representations using a shared MLP between all the possible pairs (4096 pairs).
These combinations are then averaged and processed through another MLP to compute a relational feature vector before the final prediction layer ( f ϕ ).
In a sense, GAMR is a special case of an RN network endowed with the ability to attend to a task-relevant subset ( N = 4) of these representations with the help of a controller instead of exhaustively computing all 4,096 possible relations – thus reducing the computing and memory requirements of the architecture very significantly.
As an additional baseline model we used ResNet-50 (He et al., 2016)
( ResNet ) and its transformerbased self-attention network ( Attn-ResNet ) introduced in Vaishnav et al. (2022).
These have been previously evaluated on SVRT tasks (Funke et al., 2021; Vaishnav et al., 2022; Messina et al., 2021b;a).
It serves as a powerful baseline because of more free parameters and a self-attention module to compare the proposed active attention component of GAMR .
In our proposed method, the controller shifts attention heads sequentially to individual task-relevant locations against a standard self-attention module where all task-relevant locations are attended to simultaneously.
We also evaluated ESBN (Webb et al., 2021) in which we used a similar encoder to that of GAMR and passed the images in sequential order with each shape as a single stimulus and the number of time steps as the number of shapes present in the SVRT task.
In order to train these models we used images of dimension 128 × 128 for architectures such as RN , ESBN , GAMR and 256 × 256 for ResNet, Attn-ResNet (consistent with previous work).
B ENCHMARKING GUIDED ATTENTION
We evaluated our guided-attention module ( GAMR ) and compared it with alternative systems with comparable base architecture but endowed with self-attention ( With-SA ) or no attention and/or memory ( GAMR w/o Atn (RN) ) over 23 SVRT tasks.
As a side note, GAMR-SA turns out to be similar to ARNe (Hahne et al., 2019) used for solving Raven’s tasks. We found that, on average, our Guided Attention Model’s relative performance is 11.1% better than its SA counterpart and 35.6% than a comparable system lacking attention (or memory) for SD tasks; similarly, relative improvements for SR tasks are 4.5% and 10.4%.
It shows that GAMR is computationally efficient as it yields a higher performance for the same number (1k) of training samples. Results are shown in Figure 3 (right).  B ENCHMARKING THE SYSTEM
All twenty-three tasks in the SVRT dataset can be broadly divided into two categories, same-different (SD) and spatial relations (SR), based on the identification of relations.
Same-different ( SD ) tasks ( 7, 21, 5, 19, 6, 17, 20, 1, 13, 22, 16 ) have been found to be harder for neural networks (Ellis et al., 2015;
Kim et al., 2018; Stabinger et al., 2016; 2021; Puebla & Bowers, 2021; Messina et al.,
Vaishnav et al., 2022) compared to spatial relations ( SR ) tasks ( 12, 15, 14, 9, 23, 8, 10, 18, 4, 3, 11, 2 ).
We analyzed an array of architectures and found that, on average, GAMR achieves at least 15% better test accuracy score on SD tasks for 500 samples. In contrast, for SR tasks, average accuracy has already reached perfection.
We find a similar trend for other architectures when trained with different dataset sizes.
Overall, RN (GAMR minus attention) and ESBN struggled to solve SVRT tasks even with 10k training samples, pointing towards the lack of an essential component, such as attention.
On the other hand, Attn-ResNet architecture demonstrated the second-best performance, which shows its importance in visual reasoning.
Results are summarized in Figure 2.  L EARNING C OMPOSITIONALITY
Below, we provide evidence that GAMR is capable of harnessing compositionality. We looked for triplets of tasks ( x, y, z ) such that z would be a composition of tasks x and y .
We systematically looked for all such available triplets in the SVRT dataset and found three of them ( 15, 1, 10 ), ( 18, 16, 10 ) and ( 21, 19, 25 ).
All these tasks and their associated triplets are described in SI section S1.5. We study the ability of the network to learn to compose a new relation with very few training samples, given that it had previously learned the individual rules.
We first trained the model with tasks x and y so that the rules are learned with the help of the reasoning module r θ . The first layer learns the elementary operation over the context vectors stored in the memory block ( M ), and the second layer learns to combine those operations for the tasks z . We freeze the model after training with tasks x , y and only fine-tuned: (i) a layer to learn to combine elementary operations ( r θ ) and (ii) a decision layer ( f ϕ ) on tasks z with ten samples per category and 100 epochs total. Results are shown in Fig.(left).
Z ERO - SHOT GENERALIZATION
Training
Test
Test Accuracy
Task Task
GAMR
Attn-ResNet ResNet  72.
We hypothesize that if a model has learned the abstract rule underlying a given task, it should be able to re-use its knowledge of this task on other novel tasks which share a similar rule.
To verify that GAMR is indeed able to generalize across tasks that share similar rules, we searched for pairs of tasks in SVRT which were composed of at least one common elementary relation (Vaishnav et al., 2022) between them.
For example, in pair ( 1 , 22 ), task 1 involves the identification of two similar shapes in category 1 and task 22 involves the identification of three similar shapes in category 1.
In the selected pair, the category that judges the similar rule should belong to the same class (let us say categoryin the above example) so that we test for the right learnability.
We systematically identified a set x of tasks 1, 5, 7, 21, 23 representing elementary relations such as identifying same-different ( 1, 5 ), grouping ( 7 ), learning transformation like scaling and rotation ( 21 ) and learning insideness ( 23 ).
Then we paired them with other tasks sharing similar relations.
These pairs are task 1 with each of 5, 15 and 22 , task 5 with each of 1,and 22 .
Similarly other pairs of tasks are ( 7, 22 ), ( 21, 15 ) and ( 23, 8 ).
We have explained how all these pairs of tasks form a group in SI section S1.3.
We separately trained the model on the set x and tested the same model on their respective pairs without finetuning further with any samples from the test set (zero-shot classification).
We observed that GAMR could easily generalize from one task to another without re-training.
On the contrary, a chance level accuracy by ResNet-50 (ResNet) shows the network’s rote memorization of task-dependent features.
In comparison, GAMR exhibits far greater abstraction abilities – demonstrating an ability to comprehend rules in unseen tasks without any training at all.
We further explored the strategies learned by GAMR using attribution methods for all the tasks (see SI section S4).
These attribution methods confirm that GAMR does indeed use a similar visual routine between the original task for which it was trained and the new task for which it was never trained.
Table 1 summarizes all the results.
A BLATION S TUDY
We now proceed to study what role different components of the proposed architecture play in GAMR ’s ability to learn reasoning tasks.
In our first set of experiments, we selected essential building blocks of the model, such as the relational vector ( all obj ), feature gate vector weighing step t to how relevant are the elements stored in memory in or der to get the next query, the of controller’s output ( out ) in the final decision layer ( f ϕ ) and temporal context normalization ( tcn ) on the encoded representation.
We studied the effect of these components on SD and SR categories.
Our lesioning study revealed that TCN plays a vital role in the model’s reasoning capability even for learning simple rules, as in SR tasks.
We also found that for SD tasks, excluding out from the decision-making process is detrimental.
Finally, g and w k t have their visible role for tasks containing complex rules (task 7 , SI Figure S3).
We have summarized the results in Figure 4 (top) and added additional analyses in SI section S1.
We also plot the saliency maps of the model in Figure 4 (bottom) at each time step and show the way in which the model attends to task-dependent features while learning the rules.
A DDITIONAL E XPERIMENT
Baseline models As a baseline, we chose the ESBN (Webb et al., 2021) along with the two other prevalent reasoning architectures, the Transformer (Vaswani et al., 2017) and Relation Network (RN) (Santoro et al., 2017).
These three share a similar encoder backbone as in GAMR .
In order to make our baselines stronger, we evaluated these models in their natural order, i.e., by passing a single image at a time. We added a random translation (jittering) for the shapes in the area of ± 5 pixels around the center to prevent these architectures from performing template matching. For GAMR we present task-relevant images together as a single stimulus (SI, Fig. S23) while jittering each shape.
We have also added ART results where each image is centered and put together in a single stimulus in SI section S7. In order to make our architecture choose one option from multiple stimuli ( RMTS : 2, Dist 3 or ID : 4), we concatenate the relational vector ( all obj ) for every stimulus and pass them to a linear layer for final decision.
Results We observed a near-chance level accuracy for all the baseline models and in different generalization scenarios for the SD and RMTS tasks (Figure 5). Whereas, when we trained the networks like ESBN with the images centered in the stimulus in a similar scenario, they resulted in perfect accuracy.
However, our proposed architecture is robust to handle this jittering, as shown in SI Figure S24 where we compare its performance when images are not jittered.
For the other two tasks, Dist 3 and ID , baseline models performed better than the chance level (25%).
ESBN showed an increasing trend in accuracy for progressively easier generalization conditions approachingholdouts. This points toward the fact that the first three shapes in both these tasks allow ESBN to consider a translation factor while comparing the next three shapes, letting it choose the correct option accordingly. RN and Transformer still struggled to generalize. ESBN (memory-based model) performance on SD tasks in both the visual reasoning dataset shows the importance of attention needed for reasoning.
R ELATED W ORK
Previous studies (Alamia et al., 2021) have shown that mechanisms associated with attention and memory are involved in solving same-different visual reasoning tasks.
Multiple datasets have been used to assess the visual reasoning ability of neural networks.
One of the first challenges included the SVRT.
Recently introduced Raven’s Progressive Matrices (RPM) dataset (Barrett et al., 2018;
Zhang et al., 2019) focuses on seven unique relations to be learned.
However, it was found that the dataset was seriously flawed as it was later found that neural architectures could solve tasks by leveraging shortcuts (Hu et al., 2020; Spratley et al., 2020).
Prior work on SVRT studies has focused on the role of attention in solving some of these more challenging tasks.
In SVRT, some of the tasks are significantly more challenging for computer vision algorithms than others.
In particular, tasks that involve same-different (SD) judgements appear to be significantly harder for neural networks to learn compared to those involving spatial relation judgement (SR) (Stabinger et al., 2016; Yihe et al., 2019;
Kim et al., 2018)
(see Ricci et al.
(2021) for a recent review).
Motivated by neuroscience principles, Vaishnav et al.
(2022) studied how the addition of feature-based and spatial attention mechanisms differentially affects the learnability of the tasks.
These authors found that SVRT tasks could be further taxonomized according to their differential demands for these two types of attention.
In another attempt to leverage a transformer network to incorporate attention mechanisms for visual reasoning, Messina et al.
(2021a) proposed a recurrent extension of the classic Vision Transformer block (R-ViT).
Spatial attention and feedback connections helped the transformer to learn visual relations better.
The authors compared the accuracy of four same-different (SVRT) tasks (tasks 1,5,20,21 ) to demonstrate the efficacy of their model.
They also showed that, even with 400k samples available for training, neither a Relational Network (Santoro et al., 2017) nor a Vision Transformer (Dosovitskiy et al., 2020) were capable of learning these tasks.
While a recent work Webb et al.
(2021) has explored the role of memory in ART reasoning tasks.
With the introduction of transformer networks, attention mechanisms started gaining popularity in computer vision.
They can either complement (Bello et al., 2019; Vaishnav et al., 2022; d’Ascoli et al., 2021) or completely replace existing CNN architectures (Ramachandran et al., 2019; Touvron et al., 2021; Dosovitskiy et al., 2020).
Augmenting the attention networks with the convolution architectures helps them explore the best of both and train relatively faster.
In contrast, stand-alone attention architecture takes time to develop similar inductive biases as CNN.
As initially introduced by Vaswani et al.
(2017), transformer attention uses a key (k), query (q), and value (v) attention mechanisms for NLP.
Since the images are not like language, all the existing architectures for image recognition explore self-attention variants in which the key, query, and values are the same as an input image.
We used a similar attentional system, but instead of using it as a self-attention module, we called it a guided attention module.
This system internally generates a query to guide the attention module to the location essential for the task.
Since there could be more than one location where the model will attend, we then implemented a memory bank.
We took inspiration for the memory bank from Webb et al.
(2021), where mechanisms for variable binding and indirection were introduced in architecture for visual reasoning with the help of external memory.
Variable binding is the ability to bind two representations, and indirection is the mechanism involved in retrieving one representation to refer to the other.
These authors also introduce Temporal
Context Normalization (TCN) (Webb et al., 2020), which is found beneficial for out-of-distribution generalization for relational reasoning tasks.
However, the model exhibits significant limitations:
It assumes an object-centric image representation whereby objects are presented individually in a sequence.
We cannot evaluate such an architecture on the SVRT challenge because images in each task contain multiple objects which require individuation.
There are also some relations, like “touching”, which this individuation cannot represent (or any object-centric architecture).
ESBN also lacks an attentional mechanism and works best in a scenario where hard attention at the pre-processing level helps to simplify the tasks.
We tested this template matching behavior of the architecture by training it in the presence of Gaussian noise. It led to a chance-level performance (refer to SI section S7 for more details).
Here, we build on this work and describe an end-to-end trainable model that learns to individuate task-relevant scenes and store their representations in memory to allow the judging of complex relations between these objects.
Finally, our relational mechanism is inspired by the work in Santoro et al.
(2017) that introduced a plug-and-play model for computing relations between object-like representations in a network.
C ONCLUSION AND LIMITATIONS
In this paper, we described a novel transformer-based Guided Attention Module for (visual) Reasoning ( GAMR ) to try to start bridging the gap between the reasoning abilities of humans and machines.
Inspired by the cognitive science literature, our module learns to dynamically allocate attention to task-relevant image locations and to store relevant information in memory.
Our proposed guidedattention mechanism is shown to outperform the self-attention mechanisms commonly used in vision transformers.
Our ablation study demonstrated that an interplay between attention and memory was critical to achieve robust abstract visual reasoning.
Furthermore, we demonstrated that the resulting systems efficiently is capable to solve novel tasks without limited training – by simply rearranging the elemental processing steps to learn the rules without involving any training.
We demonstrated GAMR’s versatility, robustness, ability to generalize compositionality through an array of experiments.
In the process, we achieved state-of-the-art accuracy for the two main visual reasoning challenges.
One limitation of the current approach is that it currently only deals with a fixed number of time steps (t=4).
Training the model with four time steps was sufficient to efficiently solve all SVRT and ART tasks.
However, a more flexible approach is needed to allow the model to automatically allocate a number of time steps according to the computational demand of the task.
Supplementary Information S I N DEPTH ANALYSIS OF GAMR S1. G UIDED A TTENTION
Transformer (Vaswani et al., 2017) architectures used self-attention mechanisms to draw global dependencies between input and output.
In self-attention, input interacts with itself to estimate where more attention has to be paid with the help of key, query and value dot product.
Self-attention is composed of three steps (i) dot product similarity score, (ii) normalizing scores to obtain weights, and (iii) re-weighting the original embeddings using weights.
[Equation]
where ‘d’ is the dimension of K and V .
To calculate the Q , K and V from an input X , there are trainable weight matrices W q , W k and W v .
[Equation]
However, in guided-attention , we take query as z img , key and value vectors as query t from the controller without any trainable matrix.
So the attention score becomes
[Equation]
as the dimensionality of the vector is R 128 .
This guided-attention (GA ∈ R hw × 128 ) score is summed across the spatial dimension ( h × w ) to re-weight the feature channels of the z img and form z t vector.
This module acts like a feature-based attention module because of its ability to modulate the channels and only store the feature vector in the memory.
We also illustrated the multi-head attention in
Fig. S1.
S1.
C HARACTERIZING THE LEARNED STRATEGIES WITH EXPLAINABILITY METHODS
We used three attribution methods: Integrated-gradients (Sundararajan et al., 2017), Saliency (Simonyan et al., 2014) and SmoothGrad (Smilkov et al., 2017).
Common to these methods is the use of the gradient of the classification units w.r.t.
the input image.
These methods help characterize the visual features that are diagnostic for the network prediction (e.g., SI section S3).
S1. L EARNING RULES
Rules involving Same-Different (SD) identification In task 1, the challenge is to differentiate between the two shapes and recognize if these two shapes are the same or different.
When we trained the model on task 1 , it (1) learned a mechanism to attend to the two shapes presented in an image; (2) learned to judge the similarity between the two shapes.
To test the models’ generalization ability, we selected tasks demanding similar judgment to decide.
There are three of those from the set of 23 satisfying this criterion, 5, 15 and 22 .
In task 15 , four shapes form a square; in category 1, all these shapes are identical.
Compositions involved in this task are identifying the same shape and counting them to make a correct decision.
GAMR has shown its ability to attend to all the identical shapes (Fig. S13) and make a correct decision with 92.53% accuracy.
The model learns a similar ability when trained with task 21 .
It has two shapes that are in rotation, scaled, or translated form.
The reason task 21 could solve task 15 is same as before with 1 .
In task 22 , there are three shapes aligned in a row.
All of them are identical in category 1.
Again, the model must count the number of identical shapes and make the correct decision if the count reaches 3.
In this task, our model achieves 84.91% accuracy.
Another task taking the decision based on counting is 5 .
The model needs to identify if two pairs of identical shapes are present in the image.
However, when we evaluate ResNet-50 on these three tasks, we believe that the model must have found a shortcut to make a decision.
ResNet-50 has learned a strategy to count only one pair of identical shapes, and it takes the decision as soon as it finds that identical pair in tasks 5, 15, 22 and performs better than the chance level.
Next, we try to understand the learning of the model trained on task 5 .
The model trained with 5 has learned to identify the two pairs of identical shapes from a group of four.
We found that the network is exploring a shortcut to solve this task.
Rather than learning to find both pairs of identical shapes, it learned to attend to only one pair to take a judgment (Fig. S7).
Still, the model’s routine helps extend this mechanism and attends to more than two shapes if shown in the same stimulus.
At the same time,
ResNet uses its template matching technique to explore one pair of identical shapes, which helps it to solve task 1 but does not help on tasks 15 and 22 .
In task 7 , the model learns to identify and count identical shapes in an image.
There are three pairs of two identical shapes vs. two pairs of three identical shapes, as seen in Fig.
S7.
We tested this model on task 22 and obtained 83.8% accuracy.
Here the rule is to identify if three identical shapes form a row or not.
However, a model trained on task 22 cannot solve task 7 because it explores a shortcut in solving task 22 .
It compares two shapes and arrives at a decision (Fig. S8), whereas, in task 7 , grouping and counting are required to make a decision.
ResNet, on the other hand, might still be looking for only one pair of identical shapes (there are three pairs in total) to make a decision.
Thus fail to adapt to a novel task.
Rules involving Spatial Relations In addition to solving SD tasks, we show that GAMR also learns the routines when the judgment requires spatial reasoning.
In tasks 11 and 2 , the challenge is to inspect if the two shapes are touching/around the corner or not.
One of the smaller shapes in these tasks is towards the boundary from inside (task 2 ) or touching from outside (task 11 ).
Training one and testing the other yields > 95% accuracy.
In Fig.
S16, we can see from the attribution maps that the model is focusing on the area where the two shapes are touching to reach a decision.
In task 23 , the challenge is to learn insideness.
Task 23 has two small and one large shape; both the small shapes are either inside or outside in one category vs. one shape inside and one outside in category 2.
We test the model on task 8 , which has two shapes, one inside the other and identical in one category vs. outside in category 2.
When we train a model on task 23 , it learns the judgment of insideness and also counts the number of shapes present inside.
If we evaluate such a model on task 8 , it quickly figures out instances when the shape is inside or outside (probably not identical).
For all these experiments, we take 10k samples for training the network for tasks 1, 5, 7, 21, 23, 11,and then perform the zero-shot generalization on a novel set of tasks.
S
[Equation] C URRICULUM
In this part of the experiment, we go beyond making a decision based on the learned rule and analyze
GAMR ’s ability to parse those rules individually and evaluate tasks containing simpler rules.
We trained GAMR on a complex task and tested it on its elementary subset.
For example, in task 19 , the rules involved are to learn the scale invariance between two shapes in a stimulus.
After training the model on task 19 , we tested for the case where the scale factor is 1 (not seen during training), i.e., task 1 , and obtained 94.05% accuracy.
Similarly, we took another task 21 , where there are two elementary operations involved, scaling and rotation together (non-zero values).
To see if GAMR has the ability to parse this complex relation we tested the model trained on task 21 on task 19 (scaling = x and rotation = 0) and task 1 (scaling = 1 and rotation = 0) and obtain 86.95% and 90.86% respectively.
It shows GAMR ’s unique ability to parse a complex set of relations into independent simple elementary unseen relations.
S1. L EARNING C OMPOSITION
We selected group corresponding to each tasks ( 15, 18, 21 ) used for composition.
Task 15 has four shapes which are forming a square and are also identical.
It can be composed of task 1 helping to identify the same shapes and task 10 which helps to learn if the four shapes are forming a square.
In task 18 , a rule is needed to be learned related to symmetry along the perpendicular bisector of the image.
It can be taken as a composition of task 16 which requires learning mirror reflection of the image along the perpendicular bisector of the image and task 10 where in which symmetry could be discovered in between 4 shapes (forming a square).
At last, we took task 21 which involves both scaling and rotation between two shapes in an image.
As its compositional elements, we designed a variant where there is only rotation and no scaling and represented it with 25 and combined it with another counterpart of 21 where there is scaling and no rotation, i.e., task 19 .
S1. A BLATION S TUDY
As different SVRT tasks vary in terms of complexity, we studied the effect of these components on a subset of representative tasks instead of averaging over all task complexities.
We performed our quantitative analysis with a model trained on 1,000 samples on four different tasks ( 1, 5, 7, 2 ).
These are selected based on the factors like rules involved in their formulation, reasonable accuracy (when trained with 1k samples), and a representative from both SD or SR categories.
Overall, we found that the weighting factor at each time step ( w k t ) has an insignificant drop in the overall performance.
Still, we made a place for this variable in the final architecture to further investigate its effect on more complex tasks.
However, guided attention ( w k ) has a substantial role in learning the correct rule.
S1. A BSTRACT V ARIABLE
As seen in our ablation study (section 8), one of the fundamental components responsible for learning complex rules is out , and abstaining from it makes the model struggle.
We did a t-SNE analysis on this variable using the first 20 principal components from its 512-dimensional vector as these components contribute to 95% of the variance.
We analyzed 200 samples from each task of a class-balanced validation dataset.
None of these were included in the training.
These components encode the abstract variable for each task and have separable boundaries from each other, as seen in Fig.
S4.
S1. R OBUSTNESS
We did a study to inspect the robustness of the model.
We evaluated the model on two factors: robustness to scale and robustness to the noise.
We found that GAMR is able to extract the correct rules in both of these scenarios.
Scale We analyzed the performance of the model with changing scale.
We train the model with images of resolution 128 × 128 and test them at scale 1.5 times the original, i.e., 192 × 192 .
We found that the model can generalize at this scale to various tasks such as 1, 9, 11, 14, 15, 17 .
In
Table S1, we compared the performances of the model trained on tasks at scale (s=1) and tested on task with s=1 with tested on task with s=1.5 for two different training sets, 5k and 10k.
We did not use any augmentation relating to changing scale during the training, which might push the model to learn such scale-invariant representations.
In Fig.
S5, we demonstrate an example of task 1 highlighting the region considered for taking decision at s=1.5.
Noise We even checked the robustness of GAMR to noise.
We added Gaussian white noise to the dataset to perform this experiment with 0 mean and .05 variance.
We found a negligible difference in the test accuracy of the model with noise when compared to the case without any noise (Fig.
S27, S28).
We even went further to change the type of the noise to salt and pepper while testing the same model.
Nevertheless, the model learned the rules delivering nearly the same test accuracy as before.
S V ISUAL R OUTINES
Every instant vast amount of visual information is presented to the sensory systems in the brain, where it extracts the relevant information.
This ability to visually detect, recognize, search or obtain descriptions from the scene helps us execute various daily life tasks.
Studies have also shown the insensitivity to changes in the visual scenes during eye movement, also called “changing blindness.”
This phenomenon indicates that a scene is represented with only a small amount of T EST A CCURACY T RAINING 5 K SAMPLES [Table]
K SAMPLES [Table] available information. For doing this, some mechanism should exist in order to select the relevant information based on the current circumstances.
Despite the complexities of these tasks, the underlying principles of visual processing in the brain are relatively simple.
For example, even if the visual system has the capability to analyze curves and contours, it is impossible to represent all the possible combinations of these just by using feature detectors (Ullman, 1984).
A similar problem is posed in the visual search scenario.
Due to the combinatorial explosion of possible features involved in the formation of a scene, the feedforward detection mechanism finds it hard to understand (Tsotsos, 1990).
It shows us that mechanisms exist in our visual system other than simple feature detectors that aid in solving such tasks from the versatility of possible scenes.
To implement such task-dependent visual processing, (Ullman, 1984; 1987) proposed visual routines theory.
It is one such method to process visual information beyond creating representations, thereby helping in tasks like object recognition, manipulation and abstract reasoning.
A primitive set of operations are defined in Ullman (1987) that are applied to obtain relevant information and spatial relations from a scene.
Such a composition of operations is called a routine .
The visual routine theory is heavily influenced by the theory of vision by Marr (1982).
Ullman and Marr focussed on extensive analysis and used a functional approach as a starting point.
This type of analysis provides a framework that, on the one hand, can be used in the theory of human cognition while, on the other hand, equally applicable in designing computer vision systems.
This theory is based on similar principles that visual systems can solve tasks that feedforward feature detectors could not do.
Tasks are solved with the help of some mechanisms applied on top of these representations obtained by feature detectors (known as base representations).
Ullman showed that a model could solve complex tasks if a finite set of simple operations helping the model to reason are sequenced properly (forming a routine).
Another similar idea was presented in a utilitarian theory of perception (Ramachandran, 1985).
This theory mentions that perception is similar to various parts of the human body, which are collections of ad-hoc pieces, working in their own way as well as together.
Another piece of evidence presented by Roelfsema et al.
(2000) showed how the elemental operations required for routines are implemented in the brain.
When an image is presented, feedforward processes in the brain lead to the activity pattern distributed across visual areas.
Afterward, the elemental operation comes into play as firing rate modulations.
These modulations enforce the grouping of neural responses into coherent object representations.
Later, Hayhoe (2000) discussed evidence of routine in vision, indicating that only a small part of the information of a scene is represented in a brain at each moment.
A cognitive blackboard theory is postulated in Roelfsema & de Lange (2016), explaining how the activity of neurons in the early visual cortex helps infer the visual world.
The writing operations on the blackboard are compared to the feedback signals from higher cortical areas in the early visual cortex.
These feedback signals are read by cortical regions and subcortical structures in later processing steps, thereby supporting an exchange of intermediate computational steps.
S2. R ELATING VISUAL ROUTINE THEORY WITH GAMR :
Visual processing by Ullman can be divided into three stages.
First is the creation of base representation .
It is a bottom-up process in which base representations are generated from the visual input.
Once these representations are constructed, a visual processor is used to apply sequences of elemental operations to them in the second step.
It extracts the desired information from the base representations in a top-down manner.
In the third stage, this extracted information is used to reason in any given task.
In humans, the bottom-up process corresponds to the retinotopic maps to capture properties like edge, color, speed and direction of motion.
These base representations are dependent on a fixed set of operations uniformly applied over the entire input without any object/task-specific knowledge.
Such sequences of operations to extract information from the base representations are known as visual routines .
These elemental operations are task-specific that are used to create incremental representations containing information to be used in the next operation.
Thus, by applying these routines, incremental representations are constructed using sequences of elemental operations, which are finally used to provide a solution particular to the task.
Some examples of these elemental operations are as follows: shifting the processing focus, indexing, marking an object or location for future reference, tracing boundaries, and spreading activation over an area delimited by boundaries.
In the shift of processing focus, the spotlight of selective attention helps to shift the processing focus on the image from one part to the other.
Indexing enlists all locations where processing focus can shift.
Routines are generally generated in response to some internally generated queries by assembly mechanisms in response to some specific goals.
Aligned with this theory, our proposed model has different building blocks.
These are, base representations ( z img ), incremental representations ( z t ), memory block to save these incremental representations ( M ), guided attention and finally reasoning module ( r θ ).
In our model, base representation ( z img ) is created using a convolutional encoder ( f e ) which is a pure bottom-up process.
This representation is fixed and calculated using a feedforward network for any particular input.
z img is used to create an incremental representation ( z t ) using a visual routine processor represented by our guided attention module.
At each time step, the guided attention module takes the input as the base representation ( z img ) and query vector coming from a recurrent module ( f s ).
This guided attention module gives the desired task-relevant feature vector to create z t , which is stored in the memory bank (M) and later used for reasoning.
[Figure]
S E XPLANATION OF ROUTINES LEARNED AND CORRESPONDING ATTRIBUTION MAPS
This section explains the different mechanisms our model learns when we train it with a task and shows how it uses those learning mechanisms to understand the rule while testing on a novel set of tasks.
In Fig.
S10, we see the explainability maps corresponding to the model trained with tasks 5, 19,&
and we test those models on task 1 .
Fig.
S7 shows that in task 5 , the model learned a shortcut to select only one identical pair from a group of two pairs to take a decision.
The model learns this shortcut because it was sufficient to make the correct decision.
At the same time, to decide on images from category 2, the model stops searching as soon as it finds a non-identical region between two shapes.
In task 19, the model learns scale invariance.
The model chooses one shape and examines a similar part around the second shape to check the scale-invariant representation.
A broader version of transformations is learned in task 21 , where the model understands rotation, scaling, and translation concurrently.
So the model uniformly attends to the two shapes present in the image and compares them.
In task 20 , the model learns one of the three types of similarity transformation, i.e., reflection.
If a model learns to distinguish similarity, it can solve any reasoning tasks associated with same-different identification.
Moving on to Fig.
S11, we test the rules learned by tasks 1, 7 & 20 and test on task 22 where the challenge is to identify if the three shapes arranged in a row are identical or not.
As seen from the previous discussions, task 1 learns to identify if two shapes are the same or not.
As there are three shapes existing, the model tries to compare two shapes in pairs before deciding.
So we can see the high heat map values on the two shapes for both tasks.
As discussed before, task 20 learned a similarity transformation between two shapes, and it makes a pairwise comparison between three shapes to take a decision.
Glancing at the performance of task 7 on task 22 , we claim the model learned to compare the shapes as well as count them to make a decision.
Similar to category 2 of task 7 , here again, we have three shapes responsible for such a good performance.
Next, we look at the strategy used to solve task 15 in Fig.
S13.
The rule to formulate the task involved identifying if the four shapes required to form a square are identical.
The model trained on task 1, 5, 21 and 22 starts by picking up one shape and stop comparing when it comes to the fourth shape.
This is the reason all of the four shapes in category 2, which are similar, are seen to be involved in decision-making.
[Figure]
S H YPERPARAMETERS
Number of heads : In order to decide the number of heads of the guided-attention module, we ran experiments on the same-different differentiation task of the ART Dataset.
For all four holdout values, we evaluated GAMR with a varied number of heads in the multi-head attention module (no. of head = 1, 4, 8, 16).
We found that head = 4 (Fig. S1) works best for all the holdout scenarios [Table] and we fixed this value as 4 for all our experiments.
[Table]
Holdout set For example, holdout 0 represents a generalization regime in which the test sets contain the same characters as those used during training.
At the other extreme, in holdout 95 , the training set contains a minimal number of characters, most of which are actually used for tests.
Hence, it is necessary to learn the abstract rule in order to generalize to characters in this regime.
[Table]
SVRT
This dataset can be generated with the code 2 provided by the SVRT authors with images of dimension 128 × 128 .
No augmentation technique was used for training other than normalization and randomly flipping the image horizontally or vertically, as is customary for this challenge (Vaishnav et al., 2022).
We trained the model for a maximum of 100 epochs with a stopping criterion of 99% accuracy on the validation set.
The model was trained using Adam (Kingma & Ba, 2014) optimizer and a binary cross-entropy loss.
All the models were trained from scratch.
We used a hyperparameter optimization framework Optuna (Akiba et al., 2019) to get the best learning rates and weight decays for these tasks and reports the test accuracy for the models which gave the best validation scores.
Learning compositionality One of the triplets ( 21, 19, 25 ) involves rotation which needed more than 1,000 samples to learn.
So, we selected 5,000 samples each from the tasks x and y to pre-train the network.
This pre-training is carried out for 100 epochs for both tasks.
Once the model is trained, we fine-tune it on the novel unseen task z .
We confirmed that GAMR was able to learn the new rule with as few as ten samples per category – hence demonstrating an ability to harness compositionality.
S ESBN VS GAMR
We have identified some limitations of the ESBN architecture, such as its sensitivity to translation and noise.
It is also incapable in scenarios where there are multiple shapes presented together in the stimulus or if there exist complex relations (Fig. S21).
To compare our architecture with that of ESBN, we analyzed GAMR on ART tasks and analyzed
ESBN on SVRT tasks.
At first, we evaluated GAMR on the ART dataset as shown in Fig.
S22.
In the previous work Webb et al.
(2021), each image is passed as an individual image to the ESBN model, which is stored in the external memory bank for contemplating the rule, whereas GAMR is smart enough to figure out the rules from a single stimulus containing all the characters (Fig. S23).
This time we added translation to every character before passing.
To take into account all these complexities introduced, like numerous shapes, relations, and stimuli altogether, we let our model run for 2 additional time steps (t=6) for tasks Dist3 and ID.
For these two tasks, for our model to hold up the same level of performance as without translation, it was difficult without increasing the number of time steps.
We compared the accuracy of our architecture and found that even in the hardest generalization scenario (m=95), it performs relatively well, where 95% of the shapes were held during training.
Our model achieves at least 60% test accuracy when trained from scratch using ten different random seeds.
The holdout of 95% is relatively hard for GAMR as it takes some time for it to develop an attention mechanism.
We have already seen in section 5 that ESBN architecture has difficulty contemplating the rules for the SD and RMTS tasks.
To ensure that the updated backbone is not an issue, we ran a similar experiment where we used the same encoder module ( f e ) as the one proposed in Webb et al.
and obtained a consistent chance level accuracy.
As our next step we moved on to the part where we analyzed the performance of ESBN on SVRT tasks.
To make SVRT tasks compatible with the ESBN architecture we disintegrated the shapes of the dataset in different image frames such that their location is intact.
We did not center shapes in the dataset as it might alter the rules imbibed in the task because of their spatial relational structure.
A representative example of this method can be seen in Fig.
S25.
Still, two of the tasks ( 11 and 2 ) in
SVRT cannot be accustomed to this paradigm as they involve the concept of touching.
So we skipped those tasks from our analysis because of the complex preprocessing steps involved in separating two touching contours into separate channels.
After this pre-processing step we pass these individual images to the ESBN model and plot the performance in Fig.
S26.
We found that ESBN even struggles with relatively simple tasks that were learned by ResNet50 and other non-attentive models.
Such tasks can be solved with approximately 500 samples, whereas ESBN did not learn to solve these tasks ( 14, 9, 23, 10, 2 ) even after training with 10k samples.
We believe that the inability to formulate the spatial relations between shapes presented as
S7. N OISE R OBUSTNESS
In this experiment we tested the robustness of these two architectures against noise for two different datasets.
We selected the most basic task of same-different discrimination (Fig. S22) and used both the architecture in their naturalistic training paradigm.
During training and testing, we added
Gaussian noise to the images.
We found that ESBN is not able to pass chance level accuracy even for the 0-holdout case where all the symbols are presented during the training time, whereas
GAMR experienced an insignificant drop in accuracy (Fig. S27).
We observed a similar trend in the performance when we repeated the same experiment with a complex variant of same-different tasks, i e., SVRT and plotted the results in Fig.
S28.
We matched the number of parameters in the encoder block ( f e ) in both these experiments.