CodeIt: Abstract Reasoning with Iterative Policy-Guided Program Synthesis

We thank the reviewer for their in-depth comments and constructive questions. We highlight updates to the paper based on this review in purple.

There is limited insight into understanding the capabilities of the model.

We included additional analysis and results to provide more intuition for why the method works. For example, we include a comparison between solutions identified by the mutation baseline and CodeIt in Appendix D.1 and D.2, and analyze how CodeIt solutions change over time in Appendix D.3. We observe, for example, that in 81% of tasks for which CodeIt finds a solution, a shorter solution is identified in a later meta-iteration. This indicates that the method may use insights gained in some tasks to improve its solution for other tasks.

Why don't you evaluate on the hidden test set?

Although this evaluation can provide another datapoint, this set is not available to the public, making it difficult to reproduce results. Not many published works report this performance, and it can only be obtained by submitting code to Kaggle or the ARC challenge. We did consider it however, and may do so at a later stage.

I don't understand how hindsight relabelling works with the mutation baseline in section 3.1 / appendix A.2

For the mutation baseline, we perform task relabeling, an exhaustive form of hindsight relabeling similar to [Gauthier 2022]. We run each program produced by mutation on all tasks, not just the task corresponding to the expert program we're mutating. Mutating is semi-random, so it sometimes results in a lucky hit, where we unexpectedly solve a different task. We do not perform this task relabeling procedure for CodeIt.

Do you have any evidence that the pretrained model is comfortable using the DSL after pretraining?

We do not test this explicitly, but we can infer this from the ablations in Figure 3, where we see that CodeIt without pretraining is unable to improve performance.

I'm wondering if a grid array number encoding works better

We picked our encoding scheme to minimize encoding length, as this would enable fitting this in the (limited) LLM context window. There is also some evidence that object-based representations of ARC tasks are easier to work with for LLMs [Xu 2023]. But it is possible that the best encoding for ARC tasks indeed depends on the base model and its pretraining data.

You mention that including the mutated programs is a form of regularization, but I would prefer to call it data augmentation.

We agree, this is a better choice of words. We updated the text.

What % of the time do the sampled programs execute successfully?

We did not track this number, as we filtered both programs that did not execute correctly and duplicate programs, so we have no way to distinguish between the two. After training we get about 5,000 programs for 24*400 = 9,600 sampled tasks, so it is likely that more than 50% of programs execute successfully.

I recall the original ARC kaggle competition winner saying they solved at least a hundred on the evaluation set using their DSL

This is correct, this method should give 129 correct programs. However, the DSL was designed by looking at the eval set. This means a human designer explicitly built priors into the DSL that would help with solving ARC evaluation programs. Michael Hodel first and foremost looked at the training set when designing the DSL, and thus stayed more true to the spirit of the ARC challenge as described by [Chollet 2019].

It might be good to include more examples of the DSL operators .. Clarify that the solution programs were also written by Hodel — they deserve credit for that!

Fully agreed, we did acknowledge this in Section 3. We added more details about the DSL and links to the GitHub repo of Hodel in Appendix B.4 - and explicitly mention that he built the 400 solution programs we used to kickstart CodeIt.

Other comments:

• we added more details about the hyperparameters and how we tuned these.
• we added a description for the CodeIt variations in Figure 3, also see the main response above for details on the changes.

References:

• [Gauthier 2022]: "Learning Program Synthesis for Integer Sequences from Scratch", Gauthier & Urban, 2022
• [Xu 2023]: "LLMs and the Abstraction and Reasoning Corpus: Successes, Failures, and the Importance of Object-based Representations", Xu et al., 2023

We thank the reviewer for their positive assessment of our work and their insights. Changes to the paper related to his review are listed in orange.

I’m also interested in hearing where the limit of the simple mutation baseline is: if you more extensively sample mutations and in the extreme case, the mutations cover all your newly sampled data, would your method becomes inferior?

Our experiments were indeed designed to maintain data-parity. If one were to increase the number of sampled mutations, one should also increase the number of policy samples to ensure fair comparison. It would be unfair to allow the mutation to sample new programs indefinitely and not allow CodeIt do to likewise.

ablation on the amount of sampled data, rather than the context length. How would the model perform if you reduce the number of samples, or even better, can you show the curve of perf vs. num of samples per task?

There are two ways to interpret this question. First, if the reviewer asks about the effect of the sampled programs per task $n_\rho$ per meta-iteration: we chose this number based on experiments on a custom validation set. We did not ablate this choice here, as doing so would require tuning other parameters as well (learning rate, batch size and/or number of iterations of continual learning).

Second, if the reviewer asks about performance over time, but shown on a different axis: we sample equally many programs per task during search, $n_\rho$. This means that dividing the number of sampled programs on the $x$-axis in Figure 3 by the number of tasks (400) will give the performance vs the number of sampled programs per task.

you have incorporated quite a lot of sampled programs, and I seriously doubt what would happen when the newly sampled programs are used by other methods. Also a pretrained CodeT5+ is used

Likely, an approach in which a capable enough LM is trained on the programs sampled by CodeIt would be capable of imitating CodeIt’s performance; but this approach would amount to nothing more than a behavioral clone of CodeIt.

We are unable to control for the amount of data used to train CodeT5+, since this information is not in the public domain (only the number of files used is reported, but not the amount of tokens in each). Our ablations show that using pre-trained weights does play an important role (Figure 3, right side).

A recent work [Gendron 2023] reports a performance of 11.9% of GPT-4 on the ARC eval set. While this is hard to compare directly to our result due to methodological differences (we do train on ARC, but see much less data), it is telling that we match one of the most capable generalist LLMs in performance, despite using a much smaller CodeT5+ model.

Experiments only on ARC are not necessarily sufficient to show the superiority of the method. I knew of the Raven matrices that also stress abstract reasoning, and the RAVEN dataset has similarly structured data and a much simplified program structure. Would it be possible to show similarly improved performance on this task?

Thanks for pointing out this benchmark, we were not aware of it. We chose the ARC benchmark because of its difficulty (as also remarked by reviewer NHnw), the availability of neural baselines to compare with, and the possibility to efficiently represent the grids. While the RAVEN tests are similar in spirit, choosing how to represent its tasks is an open question, meaning that applying our method would be a new research effort. We consider this beyond the scope of this work, but will consider it for future work.

Would it be possible to jointly search over the DSL space, maybe starting simple such as using a mutation method for the DSL space?

We see learning or refining the DSL as an interesting area for future research. Previous work [Ellis 2020] proposes an approach called DreamCoder that can perform DSL refinement (referred to as abstrac tion), and work [Banburski 2020] applies this method to ARC; however, this scaled poorly to the size of most ARC problems. We therefore deemed it to be beyond the scope of the present work, and focused on an approach that scales well for a given DSL.

Thank you for reading our response. Apart from the question on additional benchmark tasks and refining of the DSL, which questions do you believe remain unanswered? We will do our best to respond now and/or address these in the paper.

We thank the reviewer for highlighting the simplicity of our method, as well as its performance and the exhaustiveness of our ablation studies. Updates to the paper related to this review are listed in green.

Technical novelty is limited. As the related work mentions, iterative policy improvement and hindsight relabeling are not new ideas. In the regime of program synthesis, the idea sampling from the policy, filtering by execution, and then retraining is also explored in Haluptzok et al. (2022)

Thanks for pointing out reference [Haluptzok 2022], we have included a comparison in the related work section. The key differences between that work and ours are the way tasks are proposed and how output programs are used. In their work, synthetic tasks are proposed by a language model, and programs that solve those puzzle correctly are saved. In our work we use real task inputs only, but use hindsight relabeling to augment the outputs: we save all programs that can be executed and their outputs as new, synthetic tasks. Whether these techniques are complementary is an interesting question for future work.

The effectiveness of the proposed pipeline is not convincing:

CodeIt (full eval) solves 48 tasks compared to 41 tasks in our Mutation $d_1$ baseline. While this is only a 7 task difference, this represents an increase of 17%. Further, we designed the mutation procedure as a strong baseline that makes effective use of the DSL, and tuned its hyperparameters – it is not a naive random search. Although deployment it is not a focus of this work, also note that running inference on a never seen-before task is hard for the mutation baseline: the best one can do is execute-and-check all found programs. On the other hand, CodeIt distills knowledge about the ARC tasks into model parameters, which enables running a forward pass on the new task. CodeIt’s better performance demonstrates that distilling knowledge about the tasks ultimately trumps brute-force search of a solution, even with a DSL explicitly designed to enable search. While CodeIt (full eval) significantly outperforms Ainooson et al. (2023) and Mirchandani et al. (2023), it also performs in line with GPT-4 Gendron et al. (2023). Achieving similar performance here speaks to the efficiency of the proposed pipeline since GPT-4 is a much larger model, with a substantially larger context window, and trained on a much larger amount of data.

I strongly encourage the authors to provide the following results: 1. how many solved tasks are in common for CodeIt and the mutation baseline? 2. For tasks that both methods solve, are there any differences between the solutions of each method? Will CodeIt generate a more succinct program?

We thank the reviewer for this suggestion. We provide an analysis in this spirit in appendix D.2 of the revised paper, and we refer to it in the global reply. In short: there are differences between the solutions, CodeIt programs tend to be shorter, and some tasks are only solved by one of the two methods.

provide quantitative results on how well CodeIt compresses programs like the example in Figure 4

We carried out this analysis and we provide it Appendix D.3 of the revised paper. We observe that for 80% of tasks where a solution is found, CodeIt finds a shorter solution at a later point during the search.

expand the pseudocode [of the mutation procedure]

We have expanded and better outlined the mutation algorithm as requested.

How do authors determine the number of training epochs during continual learning

We tuned all of our hypeparameters on a custom validation set as described in Appendix B.2. On this set, we observed that training longer can lead to overfitting and is ultimately not beneficial.

Figure 4, it’s hard to interpret the figure because the task being considered is not shown and the meaning of each function is unknown

We updated the figure and figure caption. We hope that the figure is easier to parse now.

why the example input and output are named $S_0$ and $S_N$

A CodeIt solution can be thought of as a sequence of $S$(tate) transformations, going from the 0th input grid to the desired output grid. This structure is effectively promoted by the design of the DSL. Since the number of transformation steps is unknown, we use a placeholder $N$ for the output grid, $S_N$

It is encouraged for authors to include the DSL in the paper, the original write-up of the DSL is not straightforward to find.

We provide additional info and a link to the original write-up on the DSL in Appendix B.4. As the full DSL would effectively be larger than the current paper+appendix, we chose not to include it in our manuscript, but provide a link to Hodel's GitHub instead.

References:

• [Haluptzok 2022]: "Language models can teach themselves to program better", ICLR 2023

We thank the reviewer for their comments. Updates to the paper related to this review are listed in brown.

Limited novelty; the proposed idea shares the main idea with the iterative policy improvement works

We respond to this point in the global response above, as novelty was raised by reviewer WdyM as well.

Scalability is limited: it is important to study how well the proposed method will perform depending on the quality/availability of [the DSL, mutation procedure, expert programs]

It is true that our method relies on access to some expert programs and a DSL amenable to mutation. We investigate the effect of the mutation procedure in Figure 3, and through additional ablations in Appendix C, also see the global response. We observed that mutation is an essential ingredient for fast training - it is not impossible to train CodeIt without it, but it will take longer. For many program synthesis works, access to an interpreter is a necessary requirement, and as such we do not see this as a limitation.

Comparison with baselines in common settings

We have added comparisons to additional neural baselines to the revised paper, in particular LLM-based methods that solve a variation of the ARC eval set.

Intuition behind recency-based sampling of input-output pairs

We saw some evidence that as we ran CodeIt, programs would become shorter over time. This likely meant those solutions were of higher quality, and that they should therefore be sampled more often. We include quantitative proof of this phenomenon in Appendix D.3, where we show that for 81% of tasks for which CodeIt finds a solution, a shorter solution is found in a later meta-iteration. We also include qualitative examples in Appendix D.2.

Population-based policy searching

Thanks for the suggestion to explore population-based search, this could indeed be a way to increase data diversity during exploration and prevent stagnation. For example, [Jung 2020] shows that population-based learning can be an effective way to improve off-policy RL performance. Given the limited rebuttal period, we decided to leave more advanced search approaches for future work.

Other comments:

• We added the definition of $n_\rho$ (number of sampled programs per task per meta-iteration) and $n_{tasks}$ (number of tasks) to the text.
• The x-axis of Figure 3 shows the total number of sampled programs. As we sample a fixed number of programs per task per meta-iteration, dividing this number by $( n_\rho \times n_{tasks})$ gives the total number of meta-iterations.
• We added more descriptive names for the methods in Figure 3. Also see the global response for more info on the updated Figure 3.

References:

• [Jung 2020] "Population-Guided Parallel Policy Search for Reinforcement Learning", Jung et al., ICLR 2020.

Analyzing programs (WdyM, JPeY, xY18)

Reviewers ask to see example programs, and ask whether we can show qualitative differences between programs found by CodeIt versus the mutation baseline (which also uses the DSL). In Appendix D, we provide in-depth analysis of the programs produced by CodeIt and the baseline. We provide a brief summary here.

The main questions we tried to answer were:

1. Is there a quantitative difference between tasks solved by CodeIt and those solved by mutation? (Appendix D.1)
2. When a task is solved by both CodeIt and mutation, is there any qualitative difference between the solution programs? (Appendix D.2)
3. Does CodeIt learn to make programs shorter as training progresses? (Appendix D.3)

The answer to all three questions is "yes".

Compared to the mutation baseline, we observe that CodeIt tends to produce shorter solutions on average. We also note that CodeIt more often solves tasks with short task representations (i.e. smaller grids). Since the task representations need to be provided as context in the encoder, whose context length is limited, this is perhaps not surprising.

When both methods solve a task, CodeIt more often than not produces a shorter solution, sometimes significantly so. We report one such case as an example in the following table, more in Appendix D Table 5; note that one cannot obtain the CodeIt solution from the mutation one just by removing redundant lines: the two solutions use different DSL primitives. Nevertheless, in some cases CodeIt's solution is longer than the mutation one, and for fairness we report such cases in Table 6 in the appendix as well.

Example comparison of CodeIt and mutation solutions for the same task:

We show an example CodeIt solution (left column) and a mutation solution (right column) for the same task. The CodeIt solution is much shorter than the one obtained via mutation.

Finally, we confirm that CodeIt does compress programs as training progresses: in 80% of cases, a solution will get shorter in subsequent meta-iterations after it's first discovered. We show more examples of this phenomenon in Table 7 in the Appendix. We do not explicitly reward for program length, but do favor shorter programs during training (sampling shorter programs more often) and evaluation (if two solutions have equal performance on demonstration examples, we take the shorter ones).

Example of program compression through meta-iterations:

Selection of longest (left) and shortest programs (right) for ARC evaluation tasks solved by CodeIt.