SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bhishma Dedhia<sup>1</sup>, Michael Chang<sup>2</sup>, Jake C. Snell<sup>3</sup>, Thomas L. Griffiths<sup>3,4</sup>, Niraj K. Jha<sup>1</sup>
$^{1}$ Department of Electrical and Computer Engineering, Princeton University  $^{2}$ Department of Computer Science, University of California Berkeley  $^{3}$ Department of Computer Science, Princeton University  $^{4}$ Department of Psychology, Princeton University {bdedhia, js2523, tomg, jha} @princeton.edu, mbchang@berkeley.edu
Large language models are few-shot learners that can solve diverse tasks from a handful of demonstrations. This implicit understanding of tasks suggests that the attention mechanisms over word tokens may play a role in analogical reasoning. In this work, we investigate whether analogical reasoning can enable in-context composition over composable elements of visual stimuli. First, we introduce a suite of three benchmarks to test the generalization properties of a visual in-context learner. We formalize the notion of an analogy-based in-context learner and use it to design a meta-learning framework called Im-Promptu. Whereas the requisite token granularity for language is well established, the appropriate compositional granularity for enabling in-context generalization in visual stimuli is usually unspecified. To this end, we use Im-Promptu to train multiple agents with different levels of compositionality, including vector representations, patch representations, and object slots. Our experiments reveal tradeoffs between extrapolation abilities and the degree of compositionality, with non-compositional representations extending learned composition rules to unseen domains but performing poorly on combinatorial tasks. Patch-based representations require patches to contain entire objects for robust extrapolation. At the same time, object-centric tokenizers coupled with a cross-attention module generate consistent and high-fidelity solutions, with these inductive biases being particularly crucial for compositional generalization. Lastly, we demonstrate a use case of Im-Promptu as an intuitive programming interface for image generation.
No thought can be formed that isn't informed by the past; or, more precisely, we think only thanks to analogies that link our present to our past.
D. Hofstadter and E. Sander, Surfaces and Essences [1]
Humans represent complex concepts by combining and organizing simpler concepts [2, 3]. This endows us with an uncanny skill for constructing an unlimited number of concepts by composing a relatively small set of previously learned basic building blocks, an ability more formally called compositional generalization. For example, we can generate sentences by combining a dictionary of words in different ways and solve a wide range of mathematical problems using a small set of basic arithmetic operations and variables. In the case of the visual world, objects form naturally
composable entities and object-centric primitives can be flexibly combined using abstract syntactic rules to yield novel and counterfactual scenes.
The growing body of work in object-centric learning methods [4, 5, 6] enable the extraction of latent object representations from perceptual scenes. However, slot composition methods have been limited to ad-hoc composition rules based on hard-coding slots into concept libraries [7, 8]. A central challenge in learning a generative model that explicitly models the compositional grammar of the visual world is that the seemingly unlimited number of composition rules that undergird real-life visual entities prevents the explicit specification of each rule. For example, a warm spring day consists of blooming flowers and freshly verdant expanses, but come autumn, the same scene stipulates the presence of golden meadows carpeted with fallen leaves. State-of-the-art text-to-image models [9, 10, 11, 12] circumvent this bottleneck by grounding images in natural language. This aligns the representations of entities in an image to the compositional structure of language. While the visual quality of these models is stunning, they often require engineering complex language prompts to elicit the desired output. They also suffer from limited understanding of object relations [13]. Moreover, studies [14] have shown that infants as young as six months old can extract compositional structure from sequences of visual events, even before they acquire language. Such compositional abstraction independent of language faculty exhibited by humans remains elusive for contemporary models.
Figure 1: LLMs perform in-context understanding of a task from a few examples in an analogical manner (left). We use analogy-making to implicitly understand the composition rules over visual stimuli made of object-like entities (right).
Large language models (LLMs) [15, 16, 17, 18, 19, 20], on the other hand, demonstrate implicit composition skills. Pre-trained on large text corpora, these models have impressive emergent few-shot learning capabilities. They can solve novel tasks on the fly without gradient updates by following examples shown in the instruction prompts, a process known as in-context learning. Under the hood, in-context learning uses the inherently compositional nature of the input text to implicitly infer the task structure and how to apply the task structure to a query [21]. In the most simple form, the prompt-based few-shot learning paradigm can be formalized as analogy solving of the type  $A: B::C: D$ , where  $A$  represents an example task,  $B$  represents the example solution, and  $C$  denotes the query. The generated solution is the analogy completion  $D$ . Analogies over compositional elements, therefore, provide a direct isomorphism between in-
context learning for language and learning implicit composition rules for visual objects, as depicted in Fig. 1. While LLMs use pre-specified tokenizers [22, 23, 24], the compositional granularity for implicit visual understanding is unspecified. On the other hand, slot methods induce object-level representations but fail to learn composition rules. In this work, we then ask:
Can analogy-solving enable in-context generalization over object-centric compositional elements of visual entities?
While many linguistic corpora have been extracted from the Internet for text-based tasks, to the best of our knowledge, no datasets exist for measuring in-context composition capability from visual prompts. To this end, we introduce a suite of three diverse benchmarks in Section 3 built on top of rich combinatorial spaces to systematically test image-based in-context learning abilities: (a) 3D Shapes [25], (b) BitMoji Faces, and (c) CLEVr Objects [26]. If the attention over word tokens enables analogical understanding in LLMs, can we transfer this mechanism to visual analogies? We formulate a unifying framework in Section 4.1 to formalize this notion. We present a visual analogy-based meta-learning algorithm called Im-Promptu in Section 4.2 and use it to train agents with a cross-attention module. How does one tokenize visual inputs? To answer this, we model several generative agents in Section 5, encompassing a continuum ranging from strong compositional learners using object-centric abstractions to non-compositional monolithic representation learners. Extensive experiments presented in Section 6 show the dependence of robust in-context learning for visual stimuli on object-centric slots as visual tokenizers and cross-attention. We demonstrate the use
of our visual analogy framework in task extrapolation (Section 6.2), combinatorial generalization (Section 6.3), and generating counterfactual images (Section 6.5).
In this section, we discuss prior works and draw connections to relevant research areas.
2.1 Object-Centric Learning: This growing body of work targets unsupervised decomposition of scenes into a set of object-centric vectors from raw perception data. A common approach involves modeling an autoencoder consisting of a latent factorization of object slots that are independently decoded to reconstruct the scene [4, 5, 6, 27, 28, 29]. These methods effectively segment scenes into representations of underlying objects but do not learn the flexible combination of object representations. On the other hand, the authors in [7] replace the independent slot decoding bias with an Image-GPT [30] and demonstrate novel slot composition. While this work largely answers how to compose a given set of visual entities, their slot selection is performed manually through organization of slots into clustered libraries. In our work, we use a meta-learning approach to help the learner implicitly understand the composition task. Recently, the authors of [31] recast the iterative slot refinement problem as finding of fixed points through implicit differentiation to improve the stability and computational cost of the slot-attention (SA) technique. The authors of [32] use specific inductive biases to disentangle object contents from spatial information that enables relational reasoning. Our work produces a more generalized composition beyond spatial relations.
2.2 Analogical Reasoning: The ability to make analogies has been shown to have important implications in various domains, such as problem-solving, creativity, learning, and counterfactual thinking [1, 33, 34, 35, 36]. Various models have been proposed to explain the underlying mechanisms of analogy-making. Popular symbolic approaches include ARGUS [37], Transfer Frames [38], and Structural Mapping Engine [39], connectionist models include LISA [40] and Star-1 [41], and models that occupy a middle ground like Copycat [42]. Contemporary deep-learning models [43, 44, 45] attempt to solve Raven's Progressive Matrices and Bongard Problems, types of visual analogies present in IQ tests. The authors of [46] propose a generative deep visual analogy-making model based on a parallelogram regularizer on the latent space to disentangle variation factors and traverse transformation manifolds. Most similar to us, a recent work [47] reformulates the analogy-making task as Inpainting in an image grid and proposes an auto-encoder pre-training task to that end.
2.3 Compositional Generative Models: Prior works [48, 49] use energy-based models to represent composable factors of variation of a scene. However, these methods use concept labels or are grounded in text and are limited to a few concept composition operations. A follow-up work [50] describes unsupervised energy-based concepts; however, the re-composition is performed via manually-picked concepts. Recent works in text-to-image models use language grounding to induce compositionality. DALL-E [9] uses language modeling over the joint space of the text and image tokens to facilitate image generation. The latest text-to-image models [10, 11, 12] use diffusion modeling to guide the image generated from the text. A growing body of work [51, 52, 53] explores text-grounded controllable image editing via prompt engineering, where most parts of the image are preserved, but specific objects are re-composed.
2.4 Modular Representations and Attention: A swath of recent works proposes neural methods that combine modular primitives with recurrent mechanisms and cross-attention. The authors of [54] propose the Recurrent Independent Mechanisms (RIMs) to learn weakly coupled dynamical systems with independent compositional neural modules that communicate via a sparse attention bottleneck. Another work [55] uses object-centric abstractions as 'files' to store factorized declarative knowledge and the dynamics of the environment. A follow-up work [56] introduces neural abstractions of the classical Production Systems framework to learn rule-based templates, thereby facilitating compositional dynamics of the environment that can quickly adapt to novel stimuli. Most recently, the authors in [57] propose an attention-based architecture to learn in-context causal relationships between entities.
3 Benchmarks
Humans can contextualize past visual stimuli in the face of everyday experiences and even use them for imagining counterfactuals. Having seen an oak table, we can, with little difficulty, look at a metal armchair and understand the concept of an oak armchair. How do we then systematically test such
Figure 2: We use combinatorial image spaces to set up visual in-context learning tasks. Each task requires the generation of a solution from a query and supporting examples. The latent description of the task is described within the square parentheses. Left: Primitive tasks modify a composable element in isolation. Right: Composite tasks combine primitives to test the combinatorial generalization skills of the learner.
abilities of a visual in-context learner? In this work, we posit that an answer to this question lies in combinatorial visual spaces with controllable factors of variation, which provide a semantically rich, simple, scalable, and composable avenue for setting up such tasks. We create a suite of three benchmarks (Fig. 2) from compositional image creators that include (a) 3D Shapes [25], (b) BitMoji Faces, and (c) CLEVr Objects [26]. Each benchmark has a split of primitive training tasks and out-of-distribution test tasks. The primitive tasks (Fig. 2, left) constitute an analogy where a single composable element of the visual stimulus is modified in isolation. The description of each benchmark with its primitive set of tasks is given next:
3.1 3D Shapes Image Prompts: This dataset consists of static scenes of various objects lying on a colored floor in front of a colored wall viewed from different orientations. To set up primitive tasks, we consider the four properties  $P = \{ \text{object color}, \text{wall color}, \text{floor color}, \text{scene orientation} \}$ , where each property can take multiple values within a domain. For each such task, the latent relation  $r(.)$  for the analogy prompt takes a property  $p \in P$  and modifies it from a source domain value to a target value. This benchmark comprises 80000 such tasks with a maximum of four supporting image pairs for each instance and with a roughly equal split across properties.
3.2 BitMoji Faces Image Prompts: BitMoji is an avatar creator service for social media users that allows them to create intricate cartoon faces. We queried the BitMoji API [58] to collect faces with four underlying dynamic elements  $P = \{ \text{skin tone}, \text{facial hair style}, \text{hair style}, \text{eyewear} \}$ . Much like the previous case, primitive tasks are generated by intervening on a source value of a property and modifying it to a target value. We populate the training set with 80000 tasks, with four demonstrations available per task.
3.3 CLEVr Objects Image Prompts: CLEVr is a popular visual question-answering dataset with the visual component consisting of multiple objects lying in a scene. We use the CLEVr rendering engine to set up primitive tasks that include adding and deleting the same object across various scenes. In this fashion, we generate 60000 tasks with three examples per task.
A composable image space allows the concurrent modification of individual visual elements. One can construct composite tasks (Fig. 2, right) out of a combination of primitives shown in the training set. A  $k$ -composite task is denoted as  $R^k(.)=r_1 \circ r_2 \circ \cdots \circ r_k$  where  $r_1, \cdots, r_k \in \mathcal{R}$ . We provide details of the  $k$ -compositeness test of each of our agents in Section 6.3. Finer details about each
benchmark, along with visual examples from every task, are detailed in Appendix A. We will release the complete dataset of primitive and composite tasks upon the publication of this article.
To learn compositional structure over visual stimuli, we begin by formulating a unified mechanistic interpretation of an in-context compositional learner using insights from design elements of LLMs.
4.1 In-Context Learning as Analogy Completion: A general compositional in-context learner  $\mathcal{M}_{\phi, \alpha, \theta}(.)$  can be formalized via the following three key components:
- An encoder  $\mathcal{E}_{\phi}(.)$  that maps the input space  $\mathcal{X}$  to the compositional space  $\mathcal{C} \triangleq \mathcal{E}_{\phi}(.) : \mathcal{X} \longmapsto \mathcal{C}$
- An executor  $\mathcal{T}_{\alpha}(., ., .)$  that maps compositional entities from example tasks, example solutions, and a query task to the query solution  $\triangleq \mathcal{T}_{\alpha}(., ., .): \mathcal{C} \times \mathcal{C} \times \mathcal{C} \longmapsto \mathcal{C}$
- A decoder  $\mathcal{D}_{\theta}(.)$  that maps the compositional space back to the input space  $\triangleq \mathcal{D}_{\theta}(.) : \mathcal{C} \longmapsto \mathcal{X}$
Coupling the above functions together yields the learner:
$$
\mathcal {M} \triangleq \mathcal {D} _ {\theta} \left(\mathcal {T} _ {\alpha} \left(\varepsilon_ {\phi} (.) , \varepsilon_ {\phi} (.) , \varepsilon_ {\phi} (.)\right): \mathcal {X} \times \mathcal {X} \times \mathcal {X} \longmapsto \mathcal {X} \right. \tag {1}
$$
The inclusion of  $\mathcal{E}_{\phi}(.)$  and  $\mathcal{D}_{\theta}(.)$  is crucial for compositional abstraction since the input space may not be a priori composable as in the case of images. Beyond this, we do not place any strong parametric constraints on the model. Let  $\mathcal{R}$  denote a task set. For any task  $r \in \mathcal{R}$ , define the task analogy  $\zeta$  over input pairs  $x_{1}, x_{2}$  as  $x_{1}: r(x_{1}):: x_{2}: r(x_{2}) \triangleq A: B:: C: D$ . An effective in-context learner over  $\mathcal{R}$  then satisfies the following property [59]:
$$
\forall r \in \mathcal {R}, \mathbb {E} _ {\zeta \sim \mathcal {P} (r, \mathcal {X})} [ \mathcal {L} (D; \mathcal {M} _ {\phi , \alpha , \theta} (A, B, C)) ] \leq \epsilon \tag {2}
$$
Here,  $\mathcal{P}$  denotes the distribution of the analogies resulting from the input space  $\mathcal{X}$  and task  $r(.)$ ,  $\mathcal{L}(.)$  denotes a loss metric, and  $\epsilon$  is an error bound. The above property simply states that a good in-context learner solves analogies defined over a task set within some finite error bound. In Appendix B, we describe LLMs as an instantiation of this framework.
4.2 Im-Promptu Learning: Pre-training LLMs on large-scale text data leads to the emergence of highly transferable internal representations that allow them to adjust to novel tasks by simple priming on a few demonstrations. However, it is still unclear whether the underlying structure of the training data, architecture, or inductive biases instilled by choice of pre-training learning algorithms cause these intriguing properties to emerge. How can we encourage the emergence of such a general-purpose representation for visual data? In this work, we take an explicit approach based on the formalism presented in the previous section. Having established the desideratum (Eq. 2), we set up a learning procedure called Im-Promptu that trains the model to generate analogy completions from image prompts. To this end, given a task set  $\mathcal{R}$ , in each minibatch, we sample an analogy  $A:B::C:D$  that follows a latent primitive  $r(.)\in \mathcal{R}$ . Given a loss criterion,  $\mathcal{L}$ , we make the following stochastic gradient updates of the model parameters:
$$
\left\{\phi^ {t + 1}, \alpha^ {t + 1}, \theta^ {t + 1} \right\} = \left\{\phi^ {t}, \alpha^ {t}, \theta^ {t} \right\} - \frac {\alpha}{N} \sum_ {i = 1} ^ {N} \nabla_ {\theta^ {t}, \alpha^ {t}, \phi^ {t}} \mathcal {L} \left(D, \mathcal {M} _ {\phi^ {t}, \alpha^ {t}, \theta^ {t}} (A, B, C)\right) \tag {3}
$$
where the step size  $\alpha$  is a hyperparameter. The full Im-Promptu algorithm is laid out in Algorithm 1.
5 Learning Agents
In this section, we explore various modeling choices for a visual in-context learner ranging from simple pixel-space rules to object-centric learners (OCLs).
5.1 Pixel Baseline: This non-parametric baseline manipulates the pixel space using a simple addition rule. We denote the support prompt as  $A:B$ , the query as  $C$ , and the solution as  $\hat{D}$ . Then the Pixel baseline predicts  $\hat{D}$  as the outcome of a linear transformation,  $\hat{D} = C + (B - A)$ .
5.2 Monolithic Learner: Our first learning agent, inspired by [46], is entirely non-compositional and uses latent monolithic vectors to execute the image prompts as follows:
$$
\mathcal {V} = f _ {\alpha} \left(\left[ e _ {\phi} (A), e _ {\phi} (B) \right]\right), \hat {D} = d _ {\theta} \left(h _ {\beta} \left(\left[ \mathcal {V}, e _ {\phi} (C) \right]\right)\right) \tag {4}
$$
Algorithm 1 Im-Promptu Learning Algorithm
Require: Task set  $\mathcal{R}$  Step Size  $\alpha$    
Initialize parameters  $\theta ,\alpha ,\phi$  of model  $\mathcal{M}$    
while not done do Sample primitive task  $r(.)\sim \mathcal{R}$  Sample  $N$  input image pairs  $X_{1}^{1:N},X_{2}^{1:N}\sim \mathcal{P}_{train}$  for  $i$  in  $1\dots N$  do Sample analogy  $\zeta^i\triangleq x_1^i:r(x_1^i)::x_2^i:r(x_2^i)\triangleq A:B::C:D$  end for Update  $\{\theta ,\alpha ,\phi \} = \{\theta ,\alpha ,\phi \} -\frac{\alpha}{N}\nabla_{\theta ,\alpha ,\phi}\sum_i\mathcal{L}_\mathcal{M}(\zeta^i)$    
end while
The encoder  $e_{\phi}(.)$  and decoder  $d_{\theta}(.)$  are convolutional and deconvolutional networks, respectively. The inference network  $f_{\alpha}(.)$  and executor network  $h_{\beta}(.)$  are modeled as multi-layered perceptrons. Here, [] denotes the concatenation operator. This architecture is detailed in Appendix D.1.
5.3 Inpainting Model: Inspired by recent work on inpainting for In-Context learning [47], this agent (Inpaint., visualized in Appendix D.2) is modeled via an autoencoder architecture that fills in missing patches of an image. The architecture is composed of a discrete variational autoencoder (dVAE) that discretizes the input patches and a Vision Transformer [60] auto-encoder to impute missing values. While the monolithic agent separately encodes the components of the visual analogy, this model jointly represents them as a  $2 \times 2$  image grid. The model is pre-trained using the masked autoencoder reconstruction task [61]. Subsequently, at inference time, the analogy solution is represented as a missing image in the grid. Details have been given in Appendix D.2.
5.4 Object-Centric Learner (OCL): This agent, as depicted in Fig. 3, is a strongly compositional learner that uses object-centric inductive biases. The encoder and decoder of the framework in Section 4.1 are parametrized using Slot Attention Transformer (SLATE) [7] (see Appendix C.2) as follows:
$$
\mathcal {E} _ {\phi} (.) \triangleq S A _ {\phi} \left(f _ {\phi} ^ {\mathrm {d V A E}} (.)\right), \mathcal {D} _ {\theta} (.) \triangleq g _ {\theta} ^ {d V A E} \left(g _ {\theta} ^ {G P T} (.)\right) \tag {5}
$$
A Context Encoder Transformer (CET)  $\triangleq$ $\mathcal{T}_{\alpha}(. . .)$  with interleaved layers of:  $(a)$  cross-attention on the context slots and  $(b)$  self-attention, induces sparse modifications on the query slots to complete the analogy.
$$
S _ {1: N} ^ {A} = \mathcal {E} _ {\phi} (A), S _ {1: N} ^ {B} = \mathcal {E} _ {\phi} (B), S _ {1: N} ^ {C} = \mathcal {E} _ {\phi} (C) \tag {6}
$$
$$
S _ {1: N} ^ {D} = C E T \left(\text {q u e r y} = S _ {1: N} ^ {C}, \text {k e y s , v a l u e s} = S _ {1: N} ^ {A}, S _ {1: N} ^ {B}\right) \tag {7}
$$
The latent sequence  $\hat{Z}_D$  predicted by Image-GPT is mapped by the dVAE back to the pixel space  $\hat{D}$ .
$$
\hat {Z} _ {D} = g _ {\theta} ^ {G P T} \left(S _ {1: N} ^ {D}\right), \hat {D} = g _ {\theta} ^ {d V A E} \left(\hat {Z} _ {D}\right) \tag {8}
$$
The CET forward pass is shown in Appendix D.3.
5.5 Sequential Prompt: This (Seq., visualized in Appendix D.4) is an ablation over the OCL that replaces CET with an LLM-like sequence prompt obtained from a simple concatenation of slots of images  $A$ ,  $B$ , and  $C$ , injected with position embeddings  $p_i$ .
$$
\mathcal {T} _ {\alpha} (.,..) \triangleq \left[ S _ {1: N} ^ {A}, S _ {1: N} ^ {B}, S _ {1: N} ^ {C} \right] + p _ {i} \tag {9}
$$
While the OCL explicitly encodes the entire context into a fixed-length slot sequence via the CET, the ablated agent models a longer concatenated sequence.