SlowGuess/NeurIPS_2023_OCR huggingface dataset:
5.6 Patch Learner: This agent straddles the compositionality extremes of Monolithic Learners and OCLs, and uses image patch abstractions. The SA module from OCL is ablated to get a discrete sequence of  $4 \times 4$  patch latents as the degree of compositional granularity  $\mathcal{E}_{\phi}(.) \triangleq f_{\phi}^{\mathrm{dVAE}}(.)$ . A larger image implies longer sequences for this modeling choice.
In this section, we set up experiments to answer (1) whether analogical reasoning enables in-context generalization, (2) does such a generalization require key mechanisms from LLMs, namely compositionality and attention, (3) if compositionality is a key ingredient, what is the correct granularity for visual tokens, and (4) what are the limits to the generalization.
Figure 3: The OCL uses compositional object vectors to answer image prompts. (1) The encoder  $\mathcal{E}_{\phi}(.)$  is a pre-trained SLATE [7] encoder that uses a dVAE to obtain a latent discrete sequence for the input image and runs SA [6] over the sequence to extract object-centric slots. (2) The CET modifies query slots  $S_{1:N}^{C}$  via cross-attention over the support prompt slots  $S_{1:N}^{A}, S_{1:N}^{B}$ . (3) The modified slots are used to prompt an Image-GPT [30] that outputs the latent discrete sequence of the task solution  $D$  that is decoded by the dVAE decoder.
6.1 Training Setup: We trained each agent on the primitive set of tasks across the three benchmarks using Im-Promptu (Section 4.2). The SA-based object-centric encoders used in the OCL and Sequential Prompter were pre-trained as a SLATE autoencoder, as described by the authors in [7]. On the other hand, the Patch Learner was trained completely from scratch. The loss metric  $\mathcal{L}$  used to train the agents was cross-entropy (CE) loss between the true latent sequence  $Z_{D}$  and the predicted solution  $\hat{Z}_{D}$  obtained from Image-GPT, i.e.,  $\mathcal{L}_{\text{impromptu}} = CE(Z_{D}, \hat{Z}_{D})$ .
In addition to the above loss, the dVAE was trained using the mean-squared error (MSE) loss over the raw pixel space to yield the full loss function  $\mathcal{L} = \mathcal{L}_{impromptu} + MSE(D,\hat{D})$ . For inference, answers of the transformer-based agents were sampled from the Image-GPT decoder using top- $k$  nucleus sampling [62]. Hyperparameters for training and inference have been laid out in Appendix E.
6.2 Primitive Task Extrapolation: In the first out-of-distribution paradigm, we tested the ability of agents to apply learned rules to different domain pairs. In order to do this, we held out  $20\%$  of source-target pairs (see Section 3 for the definition of source and target) from the training primitives and tested the ability of agents to generalize from learned rules to unseen source-target pairs. For example, the object color property in the 3D Shapes benchmark can take 10 unique values and  $10 \times 9 = 90$  source-target combinations. Thus, an agent was trained on only  $90 \times 0.8 = 72$  pairs for object-hue primitives. We made sure that each value in the domain set was shown at least once as either the target or the source.
Fig. 4(a) plots scores of different agents across benchmarks against two key metrics: (1) MSE (lower is better) that quantitatively compares the construction against the ground truth and (2) Fréchet inception distance (FID, lower is better) score to measure the perceptual quality of the composition. We make several interesting observations:
(R1.1) Simple pixel manipulation produces implausible outputs. The pixel baseline has poor MSE scores ( $\sim 138$  for 3D Shapes,  $\sim 895$  for BitMoji) with clearly apparent artifacts in the output. While linear pixel transformations are sufficient when an object is independent of the remaining entities, it is unable to model complex dependencies (see Fig. 4(b),(d) under 'Pixel').
(R1.2) Inpainting model struggles to generalize beyond i.i.d. inputs. We observed that the Inpainting model is able to cogently complete analogies on the i.i.d. validation set (see Appendix F.1 for outputs). However, when tested on extrapolated primitives, it only fills in the high-level structure (see Fig. 4 (b), (c)) and, in the case of more complex inputs, produces entirely incoherent outputs (Fig. 4 (d)). We posit that a top-down method like inpainting pre-ordains larger datasets for improved generalization and, as such, suffers from sample inefficiency.
(R1.3) Monolithic Learners are reasonably effective at generalizing from learned rules on structured stimuli. These agents have the lowest MSE scores on 3D Shapes and BitMoji Faces, both spatially consistent datasets. However, the output deteriorated over CLEVr Objects where the
Figure 4: Primitive task extrapolation: (a) Plot showing scores where the left y-axis and the bars represent the MSE scores, while the right y-axis and the dotted line denote the FID scores. (b)-(d) Comparison of agent solutions with the dotted red circle representing anomalies with respect to the ground truth.
latent scene configuration is diverse. The edges of the generated shapes and lighting details were significantly blurred, and occluded objects were often fused together (FID  $\sim 92$ , worse than Pixel baseline, Fig. 4(d) under 'Monolithic').
(R1.4) Patch Learner and Sequential Prompt benefit from additional examples. For these agents, the addition of context examples was markedly apparent from the significant MSE drops for  $>1$  shot (see Fig. 4(a)). We posit that since these models learn from longer contexts, additional examples are particularly effective in regularizing the output sequence manifold of Image-GPT.
6.3 Composite Task Extrapolation: In this generalization paradigm, we probed the compositional generalization capabilities of the agents. A  $k$ -composite task ( $k \geq 2$ , see Section 3) is an unordered composite function of  $k$  primitives. The value of  $k$  is varied across each benchmark, where higher values of  $k$  yield visually complex contexts. To solve such a task, the agent must identify the underlying primitive relations and combinatorially modify the query.
In Fig. 5(a), we note (R1.1) and (R1.3) from the previous setup in this extrapolation paradigm as well, but we make additional key observations next:
(R2.1) The effect of object-centric biases is strong. This is indicated by the widening MSE gap between the OCL and monolithic agent for increasing values of  $k$  across all three benchmarks (see Fig. 5(a)-(c)). Moreover, the OCL and the Sequential Prompt generate outputs with consistent perceptual quality (lower FID scores).
(R2.2) Monolithic and Patch Learners apply shortcuts. The monolithic agent tends to produce a Frankenstein-like image, as if a superposition of context images (see Fig. 5(d)-(f) under 'Monolithic'). In the case of the CLEVr Objects dataset, both the Monolithic and Patch Learners often simply filled in 'blobs' agnostic to the color and shape of objects (Fig. 5(f), Appendix F.4), with missing shadows leading to lower MSE scores but poor visual quality.
6.4 Analysis: Next, we discuss key elements required for in-context composition.
(A1) In-context generalization critically hinges on compositional granularity. As evident from (R1.2), monolithic representations suffice for reliably learning primitive composition rules over simple contexts. However, they suffer from fidelity issues when the visual space becomes more complex. Beyond primitive generalization, patch abstractions are effective when the underlying space
Figure 5: Composite task extrapolation: (a)-(c) Plots showing scores across benchmarks where the left y-axis and the bars represent the MSE scores, while the right y-axis and the dotted line denote the FID scores.  $k$  denotes the level of compositeness. (d)-(f) Comparison of agent solutions with the dotted red circle representing anomalies with respect to the ground truth.
has object-patch equivalence, as evident from the CLEVr Objects dataset (Fig. 5(c)). However, even with object-patch equivalence, the Patch Learner omits the inter-object dependencies (e.g., object occlusions and shadows). It further fails to learn complex and diverse objects spanning multiple patches reliably (does poorly on the BitMoji Faces dataset with diverse facial components). Slot-based compositionality with the CET not only enables the generalization of the learned primitives but also enables combinatorial generalization to composite tasks that require abstracting out the composition rule over each entity in isolation as well as composing the resultant interplay between them to generate a globally consistent output.
(A2) Encoding the context via cross-attention leads to sample efficiency. Longer slot sequences in the Sequential Prompt require multi-shot contexts across all benchmarks. The sample inefficiency was particularly observed on the CLEVr Objects dataset that uses six slots and, hence, much longer sequences. Cross-attention via CET is essential for implicit inference, enabling the learner to encode specific objects from the context.
6.5 'Counterfactual Prompt Engineering' with Im-Promptu: LLMs have spurred interest in several priming techniques [63, 64] that can 'engineer' the model into following user instructions. While this natural-language-based programming language provides an intuitive interface for tasks that have a fundamental degree of language abstraction (writing code, summarization, symbolic
Figure 6: Scene creation from object components: (a) OCL trained via Im-Promptu is used to generate a scene of objects. Image-Prompts engineered using existing assets reliably generate the scene object by object. (b) The same scene is generated via language grounding and in-painting features of DALL-E [9]. Object properties are distorted and specifying the location of objects is tedious via language.
manipulation, etc.), it is not a priori obvious that textual descriptions should also scaffold image generators. Coming up with natural-language prompts for complex visual entities is often tedious and unintuitive. Instead, explaining an image is naturally easier from its entities and composition over them<sup>1</sup>. To this end, Im-Promptu provides an avenue for generating counterfactuals via image prompts. In Fig. 6(a), we demonstrate the creation of a scene from its object components via the OCL. As a point of comparison, in Fig. 6(b), using image inpainting and natural language prompts with DALL-E [9] yields an unreliable output with distorted objects. More practically, the OCL reduces human effort significantly in rendering images with desired properties by using existing visual assets. The procedure to generate scenes from OCL can be found in Appendix G.
This work investigated whether analogy-solving can enable in-context compositional generalization over visual entities. We first designed a new test suite. We probed if attention and compositionality, key ingredients in LLMs, also play a crucial role in the visual domain. We transported these ingredients from language to model several agents. Our experiments showed that non-compositional agents do not generalize well beyond primitive task extrapolation. Even with compositionality baked in, the performance of patch-based learners depends on idiosyncrasies of the visual structure. However, we found that object-centric biases consistently facilitate an implicit understanding of composition rules and generate outputs with global semantic consistency. Our ablation of the CET suggested that cross-attention plays a crucial role, just like language models, which posits the importance of analogy and attention in understanding the underpinnings of in-context learning. Future research challenges include collecting real-world primitives from photo-realistic graphic engines or human labels, improving object-centric inductive biases for real-life entities, and designing better prompting techniques. We hope that our work will spur further research on visual in-context learning.
We would like to thank Sreejan Kumar for helpful discussions and insights throughout the course of this work. The experiments reported in this article were performed on the computational resources managed and supported by Princeton Research Computing at Princeton University. This project was funded by NSF Grant No. CCF-2203399 and ONR Grant No. N00014-18-1-2873. JS gratefully acknowledges financial support from the Schmidt DataX Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.
No one is an artist unless he carries his picture in his head before painting it, and is sure of his method and composition" – Claude Monet
[60] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. CoRR, abs/2010.11929, 2021.
[61] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dólar, and Ross Girshick. Masked Autoencoders Are Scalable Vision Learners. CoRR, abs/2111.06377, 2021.
[62] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The Curious Case of Neural Text Degeneration. CoRR, abs/1904.09751, 2020.
[63] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing. CoRR, abs/2107.13586, 2021.
[64] Simon Willison. In Defense of Prompt Engineering. https://simonwillison.net/2023/Feb/21/in-defense-of-prompt-engineering/, February 2023.
[65] Jason Tyler Rolfe. Discrete Variational Autoencoders. CoRR, abs/1609.02200, 2016.
[66] Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learning. CoRR, abs/1711.00937, 2017.