Diffusion On Syntax Trees For Program Synthesis

Thank you for your positive review, we are very encouraged by your positive feedback and comments. We also thank you for providing insightful discussion and thoughts about our work.

It is better to mention the size of the decoder model in the architecture section rather than in the appendix, so that readers with LLM background can quickly understand the edge of the model on this task.

This is a really good point, and we will state this in the main text. For reference, the model we are using is much smaller than typical LLMs, around ~10 million parameters.

It is better to discuss the number of steps in the diffusion procedure, and the model's potential ability limit in terms of output sequence length or number of symbols.

We wanted to ask for clarification to this point, is this asking us to measure performance conditioned on the number of symbols in the ground-truth test tasks? So, fewer symbols should be easier and more symbols should be harder? We discovered that increasing the number of symbols ended up making the test tasks easier than harder, since the probability we would subtract a big shape to make the image fully black was increased, so we went for the procedure in Appendix D.

Two highly related work should be cited and discussed …

Both of these works are very relevant, and we will cite these works. Thank you for bringing these to our attention.

Thank you for your feedback, your positive review, and insightful questions.

Thank you for your encouraging review. We are glad that you see the removal of the partial compiler and having to use no reinforcement learning while still using the execution output as strengths of our approach. We will answer your insightful questions here.,

In the ablations, there is one about training only on the last mutation step. It's illustrated by Fig. 12 by only showing the transition from z_5 to z_4. Do I understand correctly that the other reverse transitions (e.g., z_4 -> z_3, ..., z_1 -> z_0) are not used for training?

This is correct, the other transitions are not used for training. We considered using all the transitions at once as a training signal, but we were not actually sure how to do that, since the network must receive the execution output of the “current” program, which program would be the “current”?

Can you define the "Rollout" method, and the differences with the "Search" one?

This is our mistake for not stating, we will edit to add this description. Tree Diffusion Rollout is using the “policy” network’s actions (i.e. edit actions) being sampled from the model repeatedly, progressively editing an initial program, $z_0 \to z_1 \to \ldots$. Tree Diffusion Search is using this same policy, but instead sampling multiple possible “edits” and ranking them using the value network, as in beam search, which we described in Section 3.3.

What's the relationship between "Number of nodes expanded" and the "number of compilations needed"?

This is indeed a typo in Line 372, we will update this to be consistent and say “Number of nodes expanded” instead of “number of compilations needed”. Each node is a full program that is either edited through pure policy (rollouts), via search (policy+value), or by repeatedly re-sampling the node from the VLM baseline. We will fix this.

We thank you for your review, your feedback, and thoughtful discussion about our work.

Thank you for your review, and the strengths you identified and also the insightful feedback. We are glad you found our paper innovative, clear, reproducible and sound.

The authors should consider citing "Outline, Then Details: Syntactically Guided Coarse-To-Fine Code Generation" in the Neural program synthesis section since this work also takes multiple passes of the program and edits the program.

This is really good feedback, we will cite this paper and it is indeed very relevant.

The value network (vϕ) training and effectiveness aren't thoroughly evaluated. Alternative approaches to edit distance estimation, including direct calculation from syntax trees, are not explored or compared.

During inference time, we don’t have access to the ground truth syntax tree of the goal image, so we can’t directly compute the edit distance. To that end we need a value network to provide an estimate of progress to be used in tandem with search. Removing the value network from search would reduce it to doing just policy rollouts which we study in Figure 4. The value network is trained via supervised learning by directly calculating the edit distance between syntax trees like you mentioned. We agree that the specifics of the value network are not discussed in detail, to which we will make changes for clarity. We apologize if we misunderstood this point, and please let us know if you meant something else or if there is an ablation or experiment we could do to strengthen our results!

In the limitations section, you mention that your syntax tree currently supports only a limited set of operators. What are the bottlenecks in expanding support to other operators and generalizing to broader coding problems?

This is a really good question, and something we can expand on in the limitations section. There are a couple of challenges to applying this to general purpose languages, which we actively want to work on for future work:

For graphics programs, the execution output is self contained, an image with a predictable size, which makes it easier to study. For general purpose programs, not only do we need to provide the final execution output, we also need to provide debugging traces for variables within the program to effectively make use of the execution. Ideally we actually want to do these editing passes on computational graphs instead of trees. It is not entirely obvious how to perform random mutations on graphs that make the program executable, and also only change a small part of the program.

What is the cost of training the value network that predicts the edit distance?

The value network is a small 2-layer MLP with 128 units in each layer that takes in the output of the vision encoder (NF-ResNet 26), and outputs a single scalar predicting the edit distance in program space between two images. We train this for 100,000 steps. For efficiency, we could have trained this alongside the training of the policy. Again, this is a detail we missed in the paper, and we will make changes to clarify this.

(continued...)

(... continued)

If so, please state explicitly and explain how this metric is applied to the RGB images in TinySVG.

We accept an image if at least 99% of the pixels have every channel (RGB) within $1/256$ of the goal image. For each channel we have 256 possible values, and for a value to be considered the same they have to be within 1 step of each other. So an RGB tuple of $(76, 14, 243)$ is accepted with $(76, 13, 244)$ but not with $(88, 9, 10)$. We want at least 99% of the pixels to match in this way. Since these images are normalized between $0$ and $1$, this threshold is $0.005$.

The difference between tree diffusion search and tree diffusion rollouts is not explicitly stated or defined in section 4.

Again, this is our mistake for not stating, we will edit to add this description. Tree Diffusion Rollout is using the “policy” network’s actions (i.e. edit actions) being sampled from the model repeatedly, progressively editing an initial program, $z_0 \to z_1 \to \ldots$. Tree Diffusion Search is using this same policy, but instead sampling multiple possible “edits” and ranking them using the value network, as in beam search, which we described in Section 3.3.

There were references to computational demand and efficiency, yet no time-related metrics were reported to demonstrate gains in this regard, despite claims about improving performance efficiency.

All methods (including baselines) use the same neural architecture. All methods also use the same renderer (compiler). The cost of expanding a node is the cost of a forward pass through the neural architecture + the cost of rendering the program and checking. Therefore the X-axis in Figure 4 represents the total computation needed. For instance, increasing the model parameter count would scale the time taken by each method identically.

We hope we are able to clarify some of your concerns, and please let us know if we misunderstood any point. We really appreciate the feedback, both positive and negative, and the insightful discussions.

Thank you for your insightful review, we really appreciate the thorough reading alongside your detailed and engaged questions. We are also encouraged that you found our neurosymbolic design to be a strength of our work. We will attempt to address your mentioned weaknesses and concerns here.

the illustrative example shown in Figure 2, enabling the approach to modify node types rather than shape, the base syntax tree structure is governed by the initial generated program.

We will try to illustrate how our mutation scheme allows for not only changing node types, but also the shape and structure of the whole syntax tree here. Consider the tree in Figure 8, in Appendix A. Under our approach, all nodes in green are potential mutation targets. Node 44 is the width of the rectangle and has type Number. And it can be replaced with any other Number. This allows changing node types. For changing tree shape and structure, consider Node 33. Node 33 contains 2 primitives, the left Quad and the right Quad. A mutation target can replace this entire node with a single node that contains only 1 primitive. Let’s say we replace Node 33 with a Circle node. We have effectively changed the shape of the tree. With this Circle node, Node 32 is now also a potential mutation target, and can be replaced with a smaller subtree and so on. This way, we can actually delete the whole tree with our mutations. The reverse is also true, a node with just one primitive can be replaced with a node with two, which allows us to “grow” any tree using a series of small mutations.

We encourage the reviewer to look at our interactive demo for this (https://td-anon.github.io/). If you scroll down on the page and click the “Add Noise” button, it will run our code in the browser to show how our mutations can edit both the type and shape of the tree.

If this makes it clearer, we will include an illustrative description such as above to make this important point clearer.

Specifically, it is not clear how replacing the "(" denoted by the edit position token and replacing it with the grammar constrained autoregressive decoding yield valid syntax …

The figure is indeed confusing in this way. It looks like it’s replacing the opening “(“ with the whole expression, but your hunch is correct, every parenthesis actually uniquely refers to a node in the syntax tree (or alternatively every opening “(“ is implicitly matched with a closed “)” and the entire sub-expression is replaced). These details are actually very nicely handled by the fact that we are using a context-free grammar and the excellent Lark library for parsing, which automatically maps the node in the syntax tree to tokens in the S-expression representation. Our implementation is flexible such that only the grammar needs to be changed and it would work with other formats.

How are varying input lengths handled? … In addition, it is not clear what the purpose of "EOS" is in this context.

The Transformer architecture we’re using (akin to GPT-2) can handle varying input lengths (in the same way as large language models). The sequences are modeled autoregressively, and hence the <EOS> (i.e. “end of sequence”) token is used for the model to indicate that it has stopped generating. We could just stop decoding when all syntax constraints are met, however, using the <EOS> token allows us to re-use standard Python libraries that work with these Transformer models.

The fraction of problems solved by the method trained with "no reverse path" is nearly the same as that of the control after about 60 expanded nodes.

Again, great catch! We will remove the word “significantly” here. It does improve the results, but as you mentioned, not by much. We actually think this is a strength, since computing the optimal reverse path is not always possible, and the experiment shows that it is not strictly needed. We will restructure the description of this to be more along the lines of “it helps a little but not by much”.

Overall the presentation of section 3, especially 3.4, requires careful rework.

We are grateful for your feedback on Section 3.4 (based on your previous weaknesses and questions), and will make the issues you bring up clearer.

However, access to this ground truth may not be readily available in most cases.

The ground-truth mutations are actually just the reverse of the noise mutations. Every mutation can be inverted, so if we add noise mutations, the ground-truth are the reverse of these mutations. We add a variable number of mutations, so the ground-truth distance is the number of mutations we added.

(continued ...)

1. Yes, a detailed explanation via an example is a must. However, as I had mentioned earlier, it remains unclear (at least it has not been proven) that diffusion + base tree always yields the optimal syntax tree (a statement regarding suboptimal steps in section 4.3 is thus not justified).

2. Given the figure is confusing, which the authors agree with as well, this needs to be addressed and made clear.

3. Yes, declaiming the "significant" would eliminate the exageration. Yet, as I mentioned, there are no computational or time-related metrics presented which help put this into context.

4. Re: tree search vs rollouts - such clarification is warranted for the reader to understand the differences.

5. Re: time of computations, use of the same network architecture should be explicity stated for the reader to grasp this.

Thank you for your positive review, we are very encouraged that you thought our paper was simple but non-obvious, well-written, and easy to understand.

Fig 3 is confusing. v(x) is the value function or the pre-trained image encoder? if its pretrained, why is there a _phi subscript?

You are correct about Figure 3, this is a typo and we will make a correction to this. We were trying to use the same notation as in Tsimpoukelli et al. (2021), where they used $v_\phi$ for the image encoder, but just made the mistake of re-using the notation as the value function. Thank you for pointing this out!

Where do the initial problems come from? It seems like they are generated randomly, but how?

We (very) briefly touch on this in Section 3.1, Line 183. The initial problems come from our ConstrainedSample function, which follows Luke (2000) and Zeller et al. (2023), which randomly samples programs from the context-free grammar. You are correct that we don’t specify that this method is used to generate the problems, and we will definitely clarify this.

Do the fuzzing and edit algorithms generalize easily to non-context-free grammars (e.g. general programming languages)?

This is a really good question, and we believe it is an open research question for future work. Indeed, there have been fuzzers written for general programming languages (for instance, AFL can fuzz JavaScript code for systems like V8, etc.), but it is not obvious yet how to use that appropriately for an execution guided system like ours. We think this is an exciting future direction!

How many steps did you train for? I don't think this is covered in the appendix.

We trained all methods, including the baselines for 1 million steps. You are correct that it wasn't covered in the Appendix, and is also something we will add.

Again, we thank you for your encouraging review, insightful questions, and we will take the feedback into account and make the appropriate clarity changes.