Paper: CodeIt: Abstract Reasoning with Iterative Policy-Guided Program Synthesis (ICLR-2024)

Reviewer comment (Evaluation):
There is limited insight into understanding the capabilities of the model. For example, the paper does not convey well what exactly the model is improving at over time, or how competent the model is at generating syntactically correct code or using the full range of DSL operators once the fine-tuning stage is completed.

Author response:
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.

Reviewer comment (Evaluation):
It is hard to have a good grasp of how well the approach actually worked — 40/400 tasks seems a bit low given the technique, and there's not much space in the paper devoted to understanding what exactly is going on with the model.

Author response:
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.

Reviewer comment (Reproducibility):
I don't understand how hindsight relabelling works with the mutation baseline in section 3.1 / appendix A.2 — for searching via mutation, it seems like you just need to store the program and the inputs into the buffer, no need to store the outputs?

Author response:
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.

Reviewer comment (Experiments):
Do you have any evidence that the pretrained model is comfortable using the DSL after pretraining? Or does it still struggle to use the correct input args, etc? Any information on this could help understand the model's capabilities better.

Author response:
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.

Reviewer comment (Writing):
You mention that including the mutated programs is a form of regularization, but I would prefer to call it data augmentation. Both imply preventing overfitting, but based on my understanding of the term, adding more data doesn't really regularize a model per se.

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

Reviewer comment (Evaluation):
What % of the time do the sampled programs execute successfully?

Author response:
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.

Reviewer comment (Presentation):
It might be good to include more examples of the DSL operators, so readers can have a rough sense of what your LM is generating without having to fully read Hodel's DSL explanation. Figure 4 is useful for this, but maybe you can show an example earlier on for this.

Author response:
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.

Reviewer comment (Experiments):
You have a "policy only" baseline in Figure 3, which should be described or at least listed in the "baselines" section. It is unclear whether "policy only" means that the initial fine-tuning doesn't have the mutated tasks, or if the "policy+mutation" means that you're searching for new task solutions at each iteration both with the policy and via random mutations.

Author response:
Other comments:
 - we added a description for the CodeIt variations in Figure 3, also see the main response above for details on the changes.

Reviewer comment (Writing):
Clarify that the solution programs were also written by Hodel — they deserve credit for that.

Author response:
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.

Reviewer comment (Experiments):
As the method still largely relies on the data it samples, I wouldn't be surprised to see it becomes better than previous methods. However, 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? From my perspective, your method might only be better than the mutation baseline, because your pretrained policy network serves as some smart prior, and to get exactly the programs your policy samples, random mutation might simply take more time (less efficient). However, using an unverified prior would also risk data coverage, meaning that it might not cover as much data as the random mutation method. In this sense, the random mutation, while inefficient, when taken to the extreme, could possibly be better.

Author response:
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.

Reviewer comment (Experiments):
As the method is basically data-driven, but smartly, I would be expecting to see 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 performance versus number of samples per task? If you decrease that number, your performance might not be better than the mutation baseline.

Author response:
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.

Reviewer comment (Evaluation):
Another problem regarding ARC in general is evaluation. There has been no equal footing as for the number of data one could use. I note that in your method, 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, which already sees quite a lot of data. In the extreme, one would like to sample the space as much as possible and feed them to a model, saving all the trouble in modeling.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Novelty):
One thing I've been thinking about ARC is that the community has been doing program search on a fixed DSL for a while. Would it be possible to jointly search over the DSL space, maybe starting simple such as using a mutation method for the DSL space?

Author response:
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.

Reviewer comment (Experiments):
Scalability is limited. In the discussion section, authors mention that “..via program mutation and hindsight relabeling, thus requiring only a small amount of expert programs to start training”. Although this is much better than requiring a large number of programs as data, the program mutation is only possible and effective if the DSL is efficiently designed by domain experts and also the program mutation algorithm is carefully designed by domain experts. Also for filtering, the execution engine for DSL is required as well. Overall, these requirements limit the scalability of the proposed approach, and it is important to study how well the proposed method will perform depending on the quality/availability of these prerequisites.

Author response:
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.

Reviewer comment (Experiments):
Comparison with baselines in common settings. Although I recognize that there seems no widely accepted common experiment setting exists, when it comes to comparing with other approaches, it would be more convincing if the proposed method and the compared methods are compared in a common setting. I do agree that the proposed DSL, grid representation, program mutation, and DSL execution engines are the important contributions of this work. However, still for evaluating the effectiveness of the proposed *learning framework* alone, including an apple-to-apple comparison result would greatly improve the paper’s significance. Also, comparing with other learning-based baselines would be great.

Author response:
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.

Reviewer comment (Writing):
Intuition behind recency-based sampling of input-output pairs. It would be helpful to provide intuition behind the proposed design choice such as recency-based sampling of input-output pairs.

Author response:
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.

Reviewer comment (Presentation):
In Figure 3, the meaning of x-axis is undefined and unclear. Section 3.3 mentions “across meta-iterations in Figure 3.” Does it mean “number of sampled programs” is the same as meta-iterations?

Author response:
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.

Reviewer comment (Writing):
The agents “CodeIt: mutated+policy” and “CodeIt: policy only” are not defined.

Author response:
We added more descriptive names for the methods in Figure 3. Also see the global response for more info on the updated Figure 3.

Reviewer comment (Experiments):
Suggestion: population-based policy searching. As indicated by the ablation A2 experiment, policy improvement plays an important role in program searching. It would be an interesting future direction to try population-based policy searching approaches; i.e., maintaining multiple population of policy networks Q to enable more efficient exploration in program search space.

Author response:
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.

Reviewer comment (Novelty):
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 [1].

Author response:
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.

Reviewer comment (Evaluation):
From Table 1, we can see that the full version of CodeIt solves only 1% more programs compared with the “mutation only” baseline (40/400 vs 36/400). This naive baseline simply samples programs from the training dataset perturbed by changing one line of code. This indicates the DSL and the original reference program for the training dataset play a major role in achieving SOTA performance. The author also mentions this point in the discussion, but it is important to expand on this.

Author response:
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.

Reviewer comment (Evaluation):
To better illustrate the performance of CodeIt, it is suggested that the authors 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?

Author response:
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.

Reviewer comment (Reproducibility):
The task mutation procedure plays a crucial part in the procedure, so it is important to expand the pseudocode to explain more details about it. There is no explanation about functions in the pseudocode, though some can be inferred from the names. How is this mutation procedure determined? How sensitive are the results w.r.t to the hyperparameter $\phi$?

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

Reviewer comment (Experiments):
How do authors determine the number of training epochs during continual learning? Will training for more than 1 epoch for each iteration degrade the performance?

Author response:
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.

Reviewer comment (Presentation):
In Figure 4, it is hard to interpret the figure because the task being considered is not shown and the meaning of each function is unknown.

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

Reviewer comment (Writing):
Why are the example input and output named S_0 and S_N?

Author response:
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$

Reviewer comment (Reproducibility):
It is encouraged for authors to include the DSL in the paper, as the original write-up of the DSL is not straightforward to find.

Author response:
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.