SlowGuess/NeurIPS_2023_OCR huggingface dataset:
to-sample and identity rules. Problems were formed from objects of three shapes (cube, sphere, and cylinder), three sizes (small, medium, and large), eight colors (gray, red, blue, green, brown, purple, cyan, and yellow) and two textures (rubber and metal). To test systematic generalization, the training set consisted of problems using only small and medium sized rubber cubes in four colors (cyan, brown, green, and gray) whereas the test set consisted only of problems using large sized metal spheres and cylinders in four distinct colors (yellow, purple, blue, and red). The training and test sets thus involved completely disjoint perceptual features, similar to the original ART.
We compared our model to results from a set of baseline models evaluated by Vaishnav and Serre [43], including their GAMR architecture, ResNet50 [17], and a version of ResNet that employed self-attention (Attn-ResNet) [42]. To the best of our knowledge, these baselines represent the best performing models on the tasks that we investigated. We also investigated an additional set of baselines that combined our pretrained slot attention module with various reasoning architectures, including GAMR [43], ESBN [51], Relation Net (RN) [35], Transformer [45], Interaction Network (IN) [48], and CoRelNet [22]. These baselines benefited from the same degree of pre-training and object-centric processing as our proposed model, but employed different reasoning mechanisms. Finally, we investigated a self-supervised baseline (a Masked Autoencoder (MAE) [16]) to assess the extent to which a generative objective might encourage systematic generalization (see Tables S13-S14 for results). All baselines were evaluated using multi-object displays. See Supplementary Section S3 for more details on baselines.
4.3 Experimental Details
4.3.1 Pre-training slot attention
We pre-trained slot attention using a slot-based autoencoder framework. For ART, we generated a dataset of random multi-object displays (see Figure S2), utilizing a publicly available repository of unicode character images. We eliminated any characters that were identical to, or closely resembled, the 100 characters present in the ART dataset, resulting in a training set of 180k images, and a validation set of 20k images. For SVRT, we pre-trained two versions of slot attention on all tasks, one using 500 and the other using 1000 training examples from each of the 23 tasks. For CLEVR-ART, we generated a dataset of randomly placed 3D shapes from CLEVR, and pre-trained slot attention using a training set of 90k images, and validation set of 10k images.
To pre-train slot attention, we used a simple reconstruction objective and a slot-based decoder. Specifically, we used a spatial broadcast decoder [47], which generates both a reconstructed image and a mask for each slot. We then generated a combined reconstruction, normalizing the masks across slots using a softmax, and using the normalized masks to compute a weighted average of the slot-specific reconstructions.
After pre-training, the slot attention parameters were frozen for training on the downstream reasoning tasks. For ART, we found that despite not having been trained on any of the objects in the dataset, slot attention was nevertheless able to generate focused attention maps that corresponded closely to the location of the objects (see Figures S6, S9).
4.3.2 Training details and hyperparameters
Images were resized to  $128 \times 128$ , and pixels were normalized to the range [0, 1]. For SVRT, we converted the images to grayscale, and also applied random online augmentations during training, including horizontal and vertical flips. Tables S2 and S3 describes the hyperparameters for the convolutional encoder and spatial broadcast decoder respectively. For slot attention, we used  $K = 6$  ( $K = 7$  for CLEVR-ART) slots,  $T = 3$  attention iterations per image, and a dimensionality of  $D = 64$ . Training details are described in Table S4
Table S5 gives the hyperparameters for the transformer. To solve the SVRT tasks, and the same/different ART task, we used a CLS token (analogous to CLS token in [9]) together with a sigmoidal output layer to generate a binary label, and trained the model with binary cross-entropy
$^{2}$ https://github.com/bbvanexttechnologies/unicode-images-database
Figure 3: ART results. Results on the Abstract Reasoning Tasks dataset [51]. OCRA outperformed all baselines for all tasks, with especially strong performance in the most difficult generalization regime ( $m = 95$ ). Results reflect an average across 10 runs for each model  $\pm$  standard error. The X axis denotes the number of objects (out of 100) withheld during training, with no objects withheld in the leftmost condition ( $m = 0$ ), and almost all objects withheld in the rightmost condition ( $m = 95$ ).
loss. To solve the RMTS, distribution-of-3, and identity rules tasks, we inserted each candidate answer choice into the problem, and used a CLS token together with a linear output layer to generate a score. The scores for all answer choices were passed through a softmax and the model was trained with cross-entropy loss. For ART and CLEVR-ART, we used a default learning rate of  $8e - 5$ , and a batch size of 16. The number of training epochs is displayed for each ART task in Table S6, and for each CLEVR-ART task in Table S10. The training details for some baseline models were modified as discussed in Section S3 and detailed in Tables S7-S9. For SVRT, we used a learning rate of  $4e - 5$ , a batch size of 32, and trained for 2000 epochs on each task. We used the ADAM optimizer [23], and implementation was done using the Pytorch library [31].
4.3.3 Model selection
After pre-training slot attention for a fixed number of epochs (number of epochs for each dataset listed in Table S4), we selected the version of the model from the epoch with the lowest validation loss (on the unsupervised pre-training task) to use in the visual reasoning tasks. We used a validation set of 20k images for ART (again excluding the objects used in the ART tasks), 4k images per task for SVRT, and 10k images for CLEVR-ART.
When training on the primary reasoning tasks, for SVRT, we trained for 2000 epochs on each task, and selected the version of the model from the epoch with the highest accuracy on the validation set. This is consistent with the approach taken by Vaishnav and Serre [43], who used validation accuracy as a criterion for early stopping. For ART and CLEVR-ART, we did not perform any model (or epoch) selection when training on the primary reasoning tasks.
Figure 3 shows the results on the ART dataset (all results are presented in tabular form in Tables S11-S17). OCRA achieved state-of-the-art results across all tasks and generalization regimes, outperforming both the previous state-of-the-art model (GAMR), and other strong baselines such as
Figure 4: CLEVR-ART. Results for systematic generalization on two CLEVR-ART tasks. Results reflect test accuracy averaged over 5 trained networks  $\pm$  standard error. SVRT. Results on the Synthetic Visual Reasoning Test (SVRT) [13]. OCRA and GAMR showed comparable overall performance, with OCRA performing marginally better on problems involving same / different relations, and GAMR performing marginally better on problems involving spatial relations.
the slot-transformer. This disparity was especially striking in the most difficult generalization regime  $(m = 95)$ , in which models were trained on only 5 objects, and tested on a heldout set of 95 objects. Importantly, no component of OCRA (including both the slot attention and relational reasoning components) had been exposed to these heldout objects. OCRA outperformed the next-best baselines in this regime (GAMR and slot-transformer) by as much as  $25\%$ , demonstrating a significant degree of systematic generalization. We also analyzed OCRA's learned relational embeddings, finding that they captured relational information even in this out-of-distribution setting (Figure S5).
Many of the baseline models that we evaluated did not contain a relational bottleneck (Transformer, IN, RN, GAMR, ResNet), limiting their ability to generalize relational patterns to novel objects (i.e., higher values of  $m$ ). The comparison with these models thus demonstrates the importance of the relational bottleneck, which enables systematic generalization of learned relational patterns. The ESBN and CoRelNet architectures did include a relational bottleneck, but they were not designed with multi-object inputs in mind, and therefore performed poorly on these tasks in all regimes (even when combined with our slot attention module). This is due primarily to the random permutation of slots in slot attention, which motivated our positional embedding variable-binding scheme (allowing OCRA to keep track of which relations correspond to which pairs of objects).
On SVRT, OCRA showed comparable overall performance with GAMR, displaying marginally better performance on tasks characterized by same/different relations, while GAMR performed marginally better on tasks involving spatial relations (Figure 4, right). OCRA also outperformed GAMR on CLEVR-ART (Figure 4, left), demonstrating that its capacity for systematic generalization can be extended to problems involving more realistic visual inputs. Overall, the results on ART, SVRT, and CLEVR-ART demonstrate that OCRA is capable of solving a diverse range of relational reasoning problems from complex multi-object displays.
We performed an ablation study targeting each of OCRA's major components, focusing on the most difficult regime of the relational match-to-sample and identity rules tasks (Table I). To evaluate the importance of OCRA's capacity for object-centric processing, we tested a version of the model without slot attention ('w/o Slot Attention'), in which the visual feature map was instead divided into a  $4 \times 4$  grid, treating each image patch as though it were an 'object'. We also tested a version in which slot attention was only trained end-to-end on the visual reasoning tasks ('w/o Pre-training'), to test the importance of prior exposure to a diverse range of multi-object displays. To evaluate the importance of OCRA's relational embedding method, we tested five different ablation models. In one model ('w/o Relational Embeddings'), slot embeddings were passed directly to the transformer. In another model ('w/o Factorized Representation'), we did not compute separate feature and position embeddings for each object, and instead applied the relational operator directly to the slot embeddings. In a third model ('w/o Relational Bottleneck'), relational embeddings were computed simply by applying an MLP to each pair of slot embeddings (allowing the model to overfit to the perceptual
Table 1: Ablation study on  $m = {95}$  generalization regime for two ART tasks. Results reflect test accuracy averaged over 10 trained networks  $\pm$  standard error.
RMTS	ID
OCRA (full model)	85.31±2.0	92.80±0.3
w/o Slot Attention	56.58±3.1	84.42±1.1
w/o Pre-training	58.69±3.6	53.48±4.8
w/o Relational Embeddings	74.79±1.8	81.21±1.1
w/o Factorized Representation	50.26±0.2	75.40±0.6
w/o Relational Bottleneck	63.62±1.0	76.56±1.8
w/o Inner Product	50.70±0.4	48.07±2.6
Replace Inner Product w/ Outer Product	62.84±1.6	69.58±1.1
w/o Transformer	51.29±0.2	44.18±0.6
details of these embeddings), thus removing OCRA's inductive bias for relational abstraction. In a fourth model ('w/o Inner Product') the dot product in the relational operator (Equation 3) was replaced with a learned bottleneck, in which each pair of slot embeddings was passed to a learned linear layer with a one-dimensional output. This model tested the extent to which the inherently relational properties of the dot product are important for relational abstraction (as opposed to resulting merely from compression to a single dimension). In a fifth model ('Replace Inner Product w/ Outer Product'), the dot product was replaced with an outer product (resulting in a  $D \times D$  matrix, which was then flattened and projected via a learned linear layer to an embedding of size  $D$ ). Finally, to evaluate the importance of OCRA's capacity to process higher-order relations, instead of passing relational embeddings to a transformer, we instead summed all relation embeddings elementwise and passed the resulting vector to an MLP, as is done in the Relation Net ('w/o Transformer'). All ablation models performed significantly worse than the complete version of OCRA, demonstrating that all three elements in our model - object-centric processing, relational abstraction, and higher-order relations - play an essential role in enabling systematic generalization.
6 Limitations and Future Directions
In the present work, we have presented a model that integrates object-centric visual processing mechanisms (providing the ability to operate over complex multi-object visual inputs) with a relational bottleneck (providing a strong inductive bias for learning relational abstractions that enable human-like systematic generalization of learned abstract rules). Though this is a promising step forward, and a clear advance relative to previous models of abstract visual reasoning, it is likely that scaling the present approach to real-world settings will pose additional challenges. First, real-world images are especially challenging for object-centric methods due to a relative lack of consistent segmentation cues. However, there has recently been significant progress in this direction [36], in particular by utilizing representations from large-scale self-supervised methods [54, 7], and it should be possible to integrate these advances with our proposed framework. Second, the current approach assumes a fixed number of slot representations, which may not be ideal for modeling real-world scenes with a highly variable number of objects [55]. Though we did not find that this was an issue in the present work, in future work it would be desirable to develop a method for dynamically modifying the number of slots. Third, OCRA's relational operator is applied to all possible pairwise object comparisons, an approach that may not be scalable to scenes that contain many objects. In future work, this may be addressed by replacing the transformer component of our model with architectures that are better designed for high-dimensional inputs [18]. Finally, it is unclear how our proposed approach may fare in settings that involve more complex real-world relations, and settings that require the discrimination of multiple relations at once. It may be beneficial in future work to investigate a 'multi-head' version of our proposed model (analogous to multi-head attention), in which multiple distinct relations are processed in parallel. A major challenge for future work is to develop models that can match the human capacity for structured visual reasoning in the context of complex, real-world visual inputs.
Shanka Subhra Mondal was supported by Office of Naval Research grant N00014-22-1-2002 during the duration of this work. We would like to thank Princeton Research Computing, especially William G. Wischer and Josko Plazonic, for their help with scheduling training jobs on the Princeton University Della cluster.
[1] A. Altabaa, T. Webb, J. Cohen, and J. Lafferty. Abstractors: Transformer modules for symbolic message passing and relational reasoning. arXiv preprint arXiv:2304.00195, 2023.
[2] V. Bapat, A. Sanchez-Gonzalez, C. Doersch, K. Stachenfeld, P. Kohli, P. Battaglia, and J. Hamrick. Structured agents for physical construction. In International conference on machine learning, pages 464-474. PMLR, 2019.
[3] D. Barrett, F. Hill, A. Santoro, A. Morcos, and T. Lillicrap. Measuring abstract reasoning in neural networks. In International conference on machine learning, pages 511-520. PMLR, 2018.
[4] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261, 2018.
[5] L. Bergen, T. O'Donnell, and D. Bahdanau. Systematic generalization with edge transformers. Advances in Neural Information Processing Systems, 34:1390-1402, 2021.
[6] C. P. Burgess, L. Matthey, N. Watters, R. Kabra, I. Higgins, M. Botvinick, and A. Lerchner. Monet: Unsupervised scene decomposition and representation. arXiv preprint arXiv:1901.11390, 2019.
[7] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9650-9660, 2021.
[8] K. Cho, B. Van Merrienboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078, 2014.
[9] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[10] D. Ding, F. Hill, A. Santoro, M. Reynolds, and M. Botvinick. Attention over learned object embeddings enables complex visual reasoning. Advances in neural information processing systems, 34:9112-9124, 2021.
[11] A. Dittadi, S. Papa, M. De Vita, B. Schölkopf, O. Winther, and F. Locatello. Generalization and robustness implications in object-centric learning. arXiv preprint arXiv:2107.00637, 2021.
[12] M. Engelcke, O. Parker Jones, and I. Posner. Genesis-v2: Inferring unordered object representations without iterative refinement. Advances in Neural Information Processing Systems, 34:8085-8094, 2021.
[13] F. Fleuret, T. Li, C. Dubout, E. K. Wampler, S. Yantis, and D. Geman. Comparing machines and humans on a visual categorization test. Proceedings of the National Academy of Sciences, 108(43):17621-17625, 2011.
[14] A. Gopalakrishnan, S. van Steenkiste, and J. Schmidhuber. Unsupervised object keypoint learning using local spatial predictability. arXiv preprint arXiv:2011.12930, 2020.
[15] K. Greff, R. L. Kaufman, R. Kabra, N. Watters, C. Burgess, D. Zoran, L. Matthey, M. Botvinick, and A. Lerchner. Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, pages 2424-2433. PMLR, 2019.
[53] V. Zambaldi, D. Raposo, A. Santoro, V. Bapst, Y. Li, I. Babuschkin, K. Tuyls, D. Reichert, T. Lillicrap, E. Lockhart, et al. Relational deep reinforcement learning. arXiv preprint arXiv:1806.01830, 2018.
[54] J. Zhou, C. Wei, H. Wang, W. Shen, C. Xie, A. Yuille, and T. Kong. ibot: Image bert pre-training with online tokenizer. arXiv preprint arXiv:2111.07832, 2021.
[55] R. S. Zimmermann, S. van Steenkiste, M. S. Sajjadi, T. Kipf, and K. Greff. Sensitivity of slot-based object-centric models to their number of slots. arXiv preprint arXiv:2305.18890, 2023.