Solving Abstract Reasoning Problems With Neurosymbolic Program Synthesis and Task Generation

The Perception module comprises the raster encoder and the demonstration encoder.

As argued at the beginning of this section (and verified in our preliminary studies), typical convolutional layers are insufficient for capturing the information that is essential in ARC tasks. Therefore, we abandon purely convolutional processing in favor of encoding rasters as one-dimensional sequences of tokens that convey information about pixels. Prior to the sequential processing block, the data undergoes a preprocessing with two 2D convolutional layers. The layers preserve the spatial dimensions of the input tensor, ensuring that both height and width remain unchanged, which is essential due to extremely small rasters in the dataset. The information about the raster size itself, often crucial for solving the task, is also conveyed by the length of the sequence.

Our raster encoder is based on an attention module that allows processing rasters of different sizes, which is required when solving ARC tasks (see, e.g., the task shown in Figure 2). Each pixel resulting from the last convolutional layer gives rise to a separate token, tagged with the color and $(x,y)$ coordinates, the latter acting as input to a positional embedding layer, architecturally similar to those employed in transformer models Vaswani et al. (2017), but adapted for a 2D spatial context. This embedding strategy allows the model to explicitly incorporate spatial information into its representation of the input data. The resulting tokens are flattened into a one-dimensional sequence (see the left part of Figure 3), which is then processed by a series of eight self-attention blocks. Subsequently, a global average pooling operation reduces the final sequence to a single, fixed-size vector, yielding a compact representation of the raster. These raster representations are then passed to the demonstration encoder.

The raster encoder can be trained alongside the other TransCoder components; however, this turned out to be computationally inefficient. More importantly, when training the entire architecture end-to-end, it is hard to determine the component that should be blamed for the failure to deliver the correct answer. For these reasons, we pre-train the raster encoder within an autoencoder architecture, where the encoder is combined with a compatible decoder that can reproduce varying-length sequences of pixels (and thus the input raster) from a fixed-dimensionality latent. The decoder is shown in the right part of Figure 3 and comprises a sequence of steps that attempt to recover the spatial and color information from the fixed-dimensional latent.

The objective of pre-training is to develop an encoder capable of effectively encapsulating the essential information from the input data in the latent representation. To achieve this, we train the autoencoder in the self-supervision mode, by confronting it with a “fill-in” task. The decoder is tasked with predicting the correct pixel value (pixel color) at a given location, conditioned on the latent representation of the input raster from the encoder. By providing the decoder with the complete set of target pixel positions, the requirement for the decoder to determine the output raster is relaxed (compared to an alternative scenario in which one could require the decoder to restore the entire sequence, i.e., both colors and coordinates). The decoder then concatenates the embeddings of target positions with the reduced raster representation to prevent the encoder from overspecializing on the “fill-in” tasks. As the “query” position representation contains partial raster representation, the encoder is forced to reason in terms of dependencies between pixels and their colors and coordinates, rather than overfitting to specific locations and colors of pixels. The decoder then processes the tokens with a stack of cross-attention blocks. Eventually, a fully connected layer produces per-pixel predictions, which are confronted with the actual pixels from the input sequence using pixel-wise categorical cross-entropy loss on the color variable.

The autoencoder is pre-trained until convergence, using all rasters (inputs, outputs, and tests) collected from the training part of the ARC collection. In this process, we intend to make feature extraction invariant to permutations of colors, while accounting for the black color serving the special role of the background in a large fraction of tasks. To this end, the rasters derived from the ARC are treated as dynamic templates, which are being randomly re-colored on the fly, while preserving the distinctions of colors. This enforces permutation-invariance in the space of colors and improves generalization. As a result of pre-training, the raster encoder used in the experimental part of this study achieves the per-pixel reconstruction accuracy of 99.99% and the per-raster accuracy of 96% on the testing part of the original ARC. The resulting encoder is thus guaranteed to pass on almost the entirety of the information available in the input rasters.

Upon completion of pre-training, the raster encoder is prepended to the demonstration encoder in TransCoder. For each demonstration pair, the demonstration encoder invokes the raster encoder twice, separately for the input and output raster, concatenates the resulting latent vectors, and passes them through a two-layer multilayer perceptron (MLP). This is repeated for all demonstrations, and the output vectors produced by the MLP are chained into a sequence, which is then subject to processing with four consecutive self-attention blocks. The final sequence ${\mathbf{z}}_i$ is the representation of the demonstrations.

The sequential latent ${\mathbf{z}}_i$ produced by the perception task forms a representation of demonstrations of the given ARC task. As the raster encoder has been almost perfectly pre-trained via auto-association (see previous section), we expect this sequence to convey the entirety of information about the content of the task.⁴ However, this does not imply the capacity of solving the task. Therefore, the role of the Solver is to map ${\mathbf{z}}_i$ to a latent representation $\mathbf{z}$ of the program to be synthesized in the DSL, which is meant to act on the task and solve it.

Technically, the Solver is implemented as a two-layer MLP preceded by averaging over the sequence of ${\mathbf{z}}_i$s. In principle, the resulting latent vector $\mathbf{z}$ could be directly passed to the Program Generator (Section 3.5). However, as in most programming languages, in our DSL (detailed in Section 3.4), the relationship between the programs (their syntax) and their input–output behaviors (semantics) is many-to-one, that is, the same semantics can be implemented with more than one program. As a result, a given ARC task can be solved with more than one program, and this very program can be a solution to more than one ARC task.

To account for this many-to-many correspondence, we make the Solver stochastic by following the blueprint of the Variational Autoencoder Kingma and Welling (2022): the last layer of Solver’s MLP does not directly produce $\mathbf{z}$, but parameterizes the normal distribution $\mathcal{Z}({\mathbf{z}}_{\mu },{\mathbf{z}}_{\sigma })$ with two vectors ${\mathbf{z}}_{\mu }$ and ${\mathbf{z}}_{\sigma }$ and then calculates $\mathbf{z}={\mathbf{z}}_{\mu }+{\mathbf{z}}_{\sigma }N(0,1)$ where $N(0,1)$ comes from a random number generator. The resulting latent $\mathbf{z}$ is subsequently passed to the Program Generator, accompanied by the additional information about the task harvested from the Workspace.

The variational layer allows the Solver to propose alternative solutions to the same ARC task. Ideally, one would hope ${\mathbf{z}}_{\mu }$ to represent the centroid of the set of programs that solve a given task (and thus share in this sense a common semantics), while ${\mathbf{z}}_{\sigma }$ models the extent of that set in the latent space (i.e., the syntactic variations between those programs). Of course, such ideal convergence is not guaranteed, and the subsequent stages of processing do not rely on its emergence. However, as we found in preliminary experiments, the variational layer improves exploration in training.

The presence of sampling makes TransCoder non-deterministic and brings tangible benefits in training: as we found in preliminary experiments, it significantly improves exploration, helping the method to produce non-trivial programs, and is particularly important for generating synthetic tasks (Section 3.7. If desired, sampling can be disabled at test time, with ${\mathbf{z}}_{\mu }$ passed as $\mathbf{z}$, allowing for deterministic replication of TransCoder’s querying. Importantly, even though sampling may lead to different programs in response to the same ARC task, each of those programs has well-defined, deterministic semantics, and will always produce the same output when applied to a given input.

The Workspace is the core symbolic module of TransCoder, which is meant to provide the Program Generator (Section 3.5) with additional information, alongside the latent $\mathbf{z}$ resulting from the Solver. It supplies the Generator with two types of symbolic entities, that is, (i) the percepts representing parts of the rasters in demonstrations, and (ii) the elements of the DSL, that is, its instructions and constants. The access to both these types of elements is provided in a unified manner, by embedding them in the same space, which allows them to be retrieved on-demand by the Generator during program synthesis, composed in an arbitrary way (limited only by the syntax of the DSL), and ultimately “tried out” via program execution.

Concerning the access to percepts in the Workspace, we supplement the information provided by the Perception module with selected symbolic percepts inferred independently from the task via direct “symbolic parsing” of images. This supplementation is essential, as many ARC tasks involve qualitative and combinatorial concepts (e.g., counting of objects) that cannot be naturally conveyed by the neural nature of processing and representation provided by the Perception module. We provide two data structures for this purpose:The relation between these two storages can be likened to a dictionary, with the entries in the Workspace Template acting as keys that index specific values (realizations) in the Workspace of a specific demonstration. For instance, the Scene key in the Workspace Template may have a different value in each Workspace. For the task shown in Figure 2, the Workspace Template contains the Scene key, and its values in the first two Workspaces are different, because they feature different sets of colors:

The list of Workspace keys is predefined and includes constants (universal keys shared between all ARC tasks, e.g., zero and one), local invariants (values that repeat within a given task, across all demonstrations, e.g., the set of used colors), and local keys (information specific to a single demonstration pair, e.g., an input Region). For a complete list of available keys, see the Appendix.

The Workspaces require appropriate “neural presentation” for the Generator of DSL programs. For a given task, all keys available in the Workspace Template are first embedded in a Cartesian space, using a learnable embedding similar to those used in conventional deep learning. This representation is context-free, that is, each key is processed independently, and the embedding has in total as many entries as there are unique keys defined in DSL.

Concerning the access to the elements of the DSL, they are also embedded in the same Cartesian space so that the Generator can choose from them, that is, from the DSL instructions and constants from Tables 8 and 7, respectively, alongside the symbolic percepts. In this way, the elements of the DSL are presented to the Generator (on the neural level) in the context of the given task, and, among others, constants (such as “Red”) may have a different embedding depending on the perception result. This is expected to facilitate alternative interpretations of the roles of particular percepts; for instance, while the black pixels should often be interpreted as the background, some tasks are exceptions to this rule.

The DSL we devised for TransCoder is a typed, functional programming language, which allows conveniently representing programs as ASTs. The leaves of ASTs are responsible for fetching input data and constants, inner nodes are occupied by DSL instructions/functions, and the root of the tree produces the return value of the entire program. Each tree node has a well-defined type. The DSL features concrete data types (e.g., Int, Bool, and Region) and generic data types (e.g. List[T]). Tables 1 and 2 present the complete list of, respectively, types and instructions.

Within this DSL, a complete program defines a mapping from a set of input values to a corresponding output value of DSL types. Program execution is deterministic, that is, application of a given program to a given input will always produce the same output; however, the variational layer included in the solver (Section 3.2) enables TransCoder to respond with a different program in training and so provide exploration of the program space. Program execution leverages the Workspace for selective retrieval of necessary items using keys. When TransCoder is queried on an ARC problem, the same synthesized DSL program is executed independently for each test Workspace.

The DSL features 40 operations presented in Table 2 (and in Table 8 in more detail). These can be divided into:Our DSL features also higher-order functions, among them Map and Filter, which apply a subprogram to the elements of a compound data structure such as a List. The complete definition of the DSL can be found in the Appendix.

The type system in Table 1 together with the operations in Table 8 implicitly define the grammar of the DSL. We do not present it formally, as the type-parametricity (the generics) makes it context-dependent and quite complex. Crucially, the Program Generator detailed in the subsequent Section 3.5 traverses the combined tree of types and function signatures, and so guarantees the resulting programs to be syntactically correct and thus executable. For similar reasons, we do not attempt to formalize the semantics of the DSL, other than via our implementation of the DSL interpreter, which translates each instruction of the DSL from Table 2 into a snippet of Python code. This implementation and the natural language description in Table 8 provide informal, yet quite precise, denotational semantics of our DSL.

There is a number of arguments in favor of using a bespoke DSL rather than relying on generic languages such as Prolog or/and related approaches such as Inductive Logic Programming (Muggleton & Watanabe, 2014). First, our DSL is already equipped with adequate types for coordinates, regions of connected pixels, and generics (e.g., Pair). This is accompanied by a number of bespoke functions that embody “natural” operations on these data types. These elements of the DSL greatly facilitate learning, especially in the difficult initial phase, which is critical for the successful formation of a model. Secondly, a fair share of ARC puzzles are “operational” in nature, that is, require the input panel to be somehow transformed into the answer panel. It is thus natural to express such transformations as executable sequences (or other control structures) of steps that transform the initial state (as observed in the input panel) into a resulting state, represented by the output panel. Examples include moving objects, connecting pixels with segments, mirroring fragments of the input image, counting objects in the input image, and “expressing” the obtained number in the output raster, and more. Expressing such staged operations in, for example, Prolog, which is declarative rather than imperative, while hypothetically possible, would be quite cumbersome.

Figure 4 shows an example program and the process of its application to an input raster, including the effects of intermediate execution steps.

The Program Generator is a bespoke architecture based on the blueprint of the doubly recurrent neural network (DRNN); see Alvarez-Melis and Jaakkola (2017) for a simple variant of DRNN. The latent $\mathbf{z}$ obtained from the Solver (Section 3.2) determines the initial state of this network, which then recursively iterates over the AST nodes of the program being generated.

Before the generation loop, an embedding of each key presented by the Workspace Template is enriched with the information in the latent ${\mathbf{z}}_j$ produced by the Perception (Section 3.1) by applying two subsequent cross-attention blocks. This gives the Generator access to precise, symbolic information on demonstration representation, while enabling its “cross-linking” with the presumably less precise, but raster-wide (and thus in this sense global) representation provided by the neural Perception. The resulting vectors form the contextual embeddings ${\mathbf{z}}_{s_j}$ of the symbols (values in the Workspace), which are accessed by the Generator in the program generation process described below.

In each iteration, the Generator receives the data on the context, including the current tree size (the number of already generated AST nodes), the depth of the node in the AST, the parent of the current node, and the return type of the node. For the root node, the return type is Region; for other nodes, it is determined recursively by the types of arguments required by DSL functions picked in previous iterations. The Generator is also supplied with the set of symbols available in the Workspaces (including the elements of the DSL), via their embeddings. From this set, the Generator selects only the symbols that meet the requirements regarding the type and the maximum depth of the tree. Then, it applies an attention mechanism to the embedded representations of the selected symbols. The outcome of attention is the symbol to be “plugged” into the AST at the current location. Figure 5 illustrates a few iterations of this process.

By chaining the calls of DSL instructions via type consistency, the generated programs are guaranteed to be syntactically correct. This is an essential advantage of TransCoder, compared to alternative mechanisms based on, for example, language models. In those approaches, programs are stochastically generated sequences of tokens, which need to be checked for syntactic correctness before execution. By providing syntactic correctness by design, the generation mechanism proposed here is more elegant and computationally efficient.

The Generator also determines the hidden state of the DRNN to be passed to each of the to-be-generated children subtrees. This is achieved by merging the current state with a learnable embedding indexed with children’s indices, so that generation in deeper layers of the AST tree is informed about the node’s position in the sequence of parents’ children.

The generation process iterates recursively until the current node requests no children, which terminates the current branch of the AST (but not the others). It is also possible to enforce termination by narrowing down the set of available symbols.

To efficiently utilize computational resources, we adopted the dynamic batching mechanism presented in Bednarek et al. (2018).

TransCoder can be trained with reinforcement learning (RL) or supervised learning (SL). In RL, the program $p$ synthesized by the Generator is applied to the query raster and returns an output raster, which is then sent to the oracle. The oracle deems it correct or not, and that response determines the value of the reward, which is then used to update the Generator. The reward value is 1 when the oracle’s answer is true for all pairs of demonstration and test; otherwise, the value is 0. Unfortunately, the a priori odds for a generated program to solve the given task are minuscule. As a result, training TransCoder only with RL is usually ineffective and inefficient, especially in the early stages, when the generated programs are almost entirely random: most episodes lead to zero reward and, consequently, no parameter updates. Therefore, based on preliminary experimentation, we abandoned the idea of training TransCoder end-to-end with RL.

The ineffectiveness of RL motivates the SL mode, in which we use the program that solves the given task (target) to guide the Generator. We directly confront the actions of the Generator (i.e., the AST nodes it produces) with the target program node-by-node, and apply a loss function (categorical cross-entropy) that rewards choosing the right symbols at individual nodes and penalizes the incorrect choices. In SL, every prediction produces non-zero gradients, and thus non-zero updates for the model’s parameters (unless the target program has been perfectly generated). We employ the AdamW (Loshchilov & Hutter, 2017) optimizer with default parameter values provided in the TensorFlow (Abadi et al., 2016) library. Additionally, within each outer loop (consecutive training phases culminating in the transition to the exploration phase), the learning rate is adjusted according to the method introduced in Smith (2018) (with five warm-up and 25 decay inner cycles).

The prerequisite for SL is the availability of target programs; as those are not given in the ARC benchmark (in any form, not mentioning our DSL), we devise a method for producing them online during training, presented in Section 3.7. In general, the specific program used as the target will be one of many programs that map the input panels of the task’s demonstrations to the corresponding output panels (see Section 2). Deterministic models cannot realize one-to-many mappings, and the variational layer introduced in Section 3.2 is meant to address this limitation. Upon the ultimate perfect convergence of TransCoder’s training, we would expect the deterministic output of the Solver (corresponding to ${\mathbf{z}}_{\mu }$) to abstractly represent the common semantic of all programs that solve the presented task, and the variational layer to sample the latents that cause the Generator to produce concrete programs with that very semantics.

The programs produced by the Generator are syntactically correct by construction. Therefore, barring the occasional run-time errors (e.g., attempt to retrieve the head of an empty list), a generated program,⁵ applied to an input raster, will always produce some response output raster. By applying such a program $p$ to each of the input rasters $I$ of some task $T$, we obtain the corresponding list of responses $O$. We observe that the resulting raster pairs $(x,y)\in I\times O$ form another ARC task $T^{\prime }$, to which $p$ is the solution (usually one of possible solutions, to be precise). The resulting pair $(T^{\prime },p)$ forms thus a complete example that can be used to train TransCoder in SL mode, as explained in the previous section, where $T^{\prime }$ is the task to be solved and $p$ is the corresponding target program.

This observation allows us to learn from mistakes: whenever the Generator produces a program $p$ that fails the presented training task $T$, we pair the inputs with the actual outputs generated by $p$, add the synthetic task $(T^{\prime },p)$ formed in this way to the working collection $S$ of solved tasks, and subsequently use them for SL. Crucially, we expect $T^{\prime }$ to be on average easier than $T$, and thus provide the training process with a valuable “learning gradient.” By doing so, we intend to help the model make progress in the early stages of training, when its capabilities fall behind the difficulty of the original ARC tasks.

To model the many-to-many relation between tasks and programs, we implement $S$ as a relational database to facilitate the retrieval of all programs (known to date) that solve a given task, and vice versa—the retrieval of all tasks solved by a given program. We disallow duplicates in $S$.

The top-level learning process starts with $S=\mathrm{\varnothing }$ and the working training set $L$ filled with the original ARC tasks, and stages learning into cycles of Exploration, Training, and Reduction phases (Figure 6). As TransCoder spends most of its time iterating between Training and Reduction, this will be referred to as the inner loop of the method; Exploration happens more sparingly, and thus we consider it a part of the outer loop.

The specification of the DSL comprises:

Examples of programs creating new synthetic problems generated in subsequent phases of Exploration are presented in Table 9.