Diffusion On Syntax Trees For Program Synthesis

## Abstract

Abstract Large language models generate code one token at a time. Their autoregressive generation process lacks the feedback of observing the program’s output. Training LLMs to suggest edits directly can be challenging due to the scarcity of rich edit data. To address these problems, we propose neural diffusion models that operate on syntax trees of any context-free grammar. Similar to image diffusion models, our method also inverts “noise” applied to syntax trees. Rather than generating code sequentially, we iteratively edit it while preserving syntactic validity, which makes it easy to combine this neural model with search. We apply our approach to inverse graphics tasks, where our model learns to convert images into programs that produce those images. Combined with search, our model is able to write graphics programs, see the execution result, and debug them to meet the required specifications. We additionally show how our system can write graphics programs for hand-drawn sketches. Video results can be found at https://tree-diffusion.github.io .

## 1 Introduction

Large language models (LLMs) have made remarkable progress in code generation, but their autoregressive nature presents a fundamental challenge: they generate code token by token, without access to the program’s runtime output from the previously generated tokens. This makes it difficult to correct errors, as the model lacks the feedback loop of seeing the program’s output and adjusting accordingly. While LLMs can be trained to suggest edits to existing code [, , ], acquiring sufficient training data for this task is difficult.

In this paper, we introduce a new approach to program synthesis using *neural diffusion* models that operate directly on syntax trees. Diffusion models have previously been used to great success in image generation [, , ]. By leveraging diffusion, we let the model learn to iteratively refine programs while ensuring syntactic validity. Crucially, our approach allows the model to observe the program’s output at each step, effectively enabling a debugging process.

In the spirit of systems like AlphaZero [], the iterative nature of diffusion naturally lends itself to search-based program synthesis. By training a value model alongside our diffusion model, we can guide the denoising process toward programs that are likely to achieve the desired output. This allows us to efficiently explore the program space, making more informed decisions at each step of the generation process.

We implement our approach for inverse graphics tasks, where we posit domain-specific languages for drawing images. Inverse graphics tasks are naturally suitable for our approach since small changes in the code produce semantically meaningful changes in the rendered image. For example, a misplaced shape on the image can be easily seen and fixed in program space.

Our main contributions for this work are (a) a novel approach to program synthesis using diffusion on syntax trees and (b) an implementation of our approach for inverse graphics tasks that significantly outperforms previous methods.

Figure: Figure 1: Examples of programs recovered by our system. The top row shows a hand-drawn sketch of an icon (left), the recovered program (middle), and the compilation of the recovered program (right). The top two rows are for the constructive solid geometry language (CSG2D-Sketch). The last row is an example output from our TinySVG environment that learns to invert hierarchical programs of shapes and colors. Video examples can be found at [https://tree-diffusion.github.io](https://tree-diffusion.github.io).
Refer to caption: extracted/5633576/figures/recovery.png

## 2 Background & Related Work

##### Neural program synthesis

Neural program synthesis is a prominent area of research, in which neural networks generate programs from input-output examples. Early work, such as Parisotto et al. [], demonstrated the feasibility of this approach. While modern language models can be directly applied to program synthesis, combining neural networks with search strategies often yields better results and guarantees. In this paradigm, the neural network guides the search process by providing proposal distributions or scoring candidate programs. Examples of such hybrid methods include Balog et al. [], Ellis et al. [], and Devlin et al. []. A key difference from our work is that these methods construct programs incrementally, exploring a vast space of partial programs. Our approach, in contrast, focuses on *editing* programs, allowing us to both grow programs from scratch and make corrections based on the program execution.

##### Neural diffusion

Neural diffusion models, a class of generative models, have demonstrated impressive results for modeling high-dimensional data, such as images [, , ].
A neural diffusion model takes samples from the data distribution (e.g. real-world images), incrementally corrupts the data by adding noise, and trains a neural network to incrementally remove the noise. To generate new samples, we can start with random noise and iteratively apply the neural network to denoise the input.

##### Diffusion for discrete data

Recent work extends diffusion to discrete and structured data like graphs [], with applications in areas such as molecule design [, , ]. Notably, Lou et al. [] proposed a discrete diffusion model using a novel score-matching objective for language modeling. Another promising line of work for generative modeling on structured data is generative flow networks (GFlowNets) [], where neural models construct structured data one atom at a time.

##### Diffusion for code generation

Singh et al. [] use a diffusion model for code generation. However, their approach is to first embed text into a continuous latent space, train a *continuous* diffusion model on that space, and then unembed at the end. This means that intermediate stages of the latent representation are not trained to correspond to actual code. The embedding tokens latch to the nearest embeddings during the last few steps.

Direct code editing using neural models has also been explored. Chakraborty et al. [] use a graph neural network for code editing, trained on a dataset of real-world code patches. Similarly, Zhang et al. [] train a language model to edit code by modifying or inserting [MASK] tokens or deleting existing tokens. They further fine-tune their model on real-world comments and patches. Unlike these methods, our approach avoids the need for extensive code edit datasets and inherently guarantees syntactic validity through our pretraining task.

##### Program synthesis for inverse graphics

We are inspired by previous work by Sharma et al. [], Ellis et al. [, ], which also uses the CSG2D language. Sharma et al. [] propose a convolutional encoder and a recurrent model to go from images to programs. Ellis et al. [] propose a method to provide a neural model with the intermediate program execution output in a read–eval–print loop (REPL). Unlike our method, the ability to execute partial graphics programs is a key requirement for their work. Our system operates on complete programs and does not require a custom partial compiler. As mentioned in their work, their policies are also brittle. Once the policy proposes an object, it cannot undo that proposal. Hence, these systems require a large number of particles in a Sequential Monte-Carlo (SMC) sampler to make the system less brittle to mistakes.

## 3 Method

Figure: Figure 2: An overview of our method. Analogously to adding noise in image diffusion, we randomly make small mutations to the syntax trees of programs. We then train a conditional neural model to invert these small mutations. In the above example, we operate in a domain-specific language (DSL) for creating 2D graphics using a constructive solid geometry language. The leftmost panel ($z_{0}$) shows the target image (bottom) alongside its program as a syntax tree (top). The $y$ value of the circle gets mutated from $16$ to $10$ in the second panel, making the black circle "jump" a little higher. Between $z_{1}$ and $z_{2}$, we see that we can mutate the Subtract ($-$) node to a Circle node, effectively deleting it.
Refer to caption: extracted/5633576/figures/overview.png

The main idea behind our method is to develop a form of denoising diffusion models analogous to image diffusion models for syntax trees.

Consider the example task from Ellis et al. [] of generating a constructive solid geometry (CSG2D) program from an image. In CSG2D, we can combine simple primitives like circles and quadrilaterals using boolean operations like addition and subtraction to create more complex shapes, with the context-free grammar (CFG),

$$ $\displaystyle\texttt{S}\to\texttt{S}+\texttt{S}\ |\ \texttt{S}-\texttt{S}\ |\ \texttt{Circle}^{r}_{x,y}\ |\ \texttt{Quad}_{x,y,\theta}^{w,h}.$ $$

In Figure 2, $z_{0}$ is our *target program*, and $x_{0}$ is the rendered version of $z_{0}$. Our task is to invert $x_{0}$ to recover $z_{0}$. Our noising process randomly mutates y=16 to y=10. It then mutates the whole $\stackMath\mathbin{\stackinset{c}{0ex}{c}{0ex}{-}{\bigcirc}}$ sub-tree with two shapes with a new sub-tree with just one shape. Conditioned on the image $x_{0}$, and starting at $z_{3},x_{3}$, we would like to train a neural network to reverse this noising process to get to $z_{0}$.

In the following sections, we will first describe how “noise” is added to syntax trees. Then, we will detail how we train a neural network to reverse this noise. Finally, we will describe how we use this neural network for search.

### 3.1 Sampling Small Mutations

Let $z_{t}$ be a program at time $t$. Let $p_{\mathcal{N}}(z_{t+1}|z_{t})$ be the distribution over randomly mutating program $z_{t}$ to get $z_{t+1}$. We want $p_{\mathcal{N}}$ mutations to be: (1) small and (2) produce syntactically valid $z_{t+1}$’s.

To this end, we turn to the rich computer security literature on grammar-based fuzzing [, , , ]. To ensure the mutations are small, we first define a function $\sigma(z)$ that gives us the “size” of program $z$. For all our experiments, we define a set of terminals in our CFG to be primitives. As an example, the primitives in our CSG2D language are $\{\texttt{Quad},\texttt{Circle}\}$. In that language, we use $\sigma(z)=\sigma_{\text{primitive}}(z)$, which counts the number of primitives. Other generic options for $\sigma(z)$ could be the depth, number of nodes, etc.

We then follow Luke [] and Zeller et al. [] to randomly sample programs from our CFG under exact constraints, $\sigma_{\min}<\sigma(z)\leq\sigma_{\max}$. We call this function $\texttt{ConstrainedSample}(\sigma_{\min},\sigma_{\max})$. Setting a small value for $\sigma_{\max}$ allows us to sample small programs randomly. We set $\sigma_{\max}=\sigma_{\text{small}}$ when generating small mutations.

To mutate a given program $z$, we first generate a set of candidate nodes in its tree under some $\sigma_{\text{small}}$,

$$ $\displaystyle\mathcal{C}$ $\displaystyle=\{n\in\texttt{SyntaxTree}(z)\mid\sigma(n)\leq\sigma_{\text{small }}\}.$ $$

Then, we uniformly sample a mutation node from this set,

$$ $\displaystyle m$ $\displaystyle\sim\text{Uniform}[\mathcal{C}].$ $$

Since we have access to the full syntax tree and the CFG, we know which production rule produced $m$, and can thus ensure syntactically valid mutations. For example, if $m$ were a number, we know to replace it with a number. If $m$ were a general subexpression, we know we can replace it with any general subexpression. Therefore, we sample $m^{\prime}$, which is $m$’s replacement as,

$$ $\displaystyle m^{\prime}\sim\texttt{ConstrainedSample}(\texttt{ProductionRule} (m),\sigma_{\text{small}}).$ $$

### 3.2 Policy

#### 3.2.1 Forward Process

We cast the program synthesis problem as an inference problem. Let $p(x|z)$ be our observation model, where $x$ can be any kind of observation. For example, we will later use images $x$ produced by our program, but $x$ could also be an execution trace, a version of the program compiled to bytecode, or simply a syntactic property. Our task is to invert this observation model, i.e. produce a program $z$ given some observation $x$.

We first take some program $z_{0}$, either from a dataset, $\mathcal{D}=\{z^{0},z^{1},\ldots\}$, or in our case, a randomly sampled program from our CFG. We sample $z_{0}$’s such that $\sigma(z_{0})\leq\sigma_{\max}$. We then add noise to $z_{0}$ for $s\sim\text{Uniform}[1,s_{\max}]$, steps, where $s_{\max}$ is a hyper-parameter, using,

$$ $\displaystyle z_{t+1}\sim p_{\mathcal{N}}(z_{t+1}|z_{t}).$ $$

We then train a conditional neural network that models the distribution,

$$ $\displaystyle q_{\phi}(z_{t-1}|z_{t},x_{t};x_{0}),$ $$

where $\phi$ are the parameters of the neural network, $z_{t}$ is the current program, $x_{t}$ is the current output of the program, and $x_{0}$ is the target output we are solving for.

#### 3.2.2 Reverse Mutation Paths

Since we have access to the ground-truth mutations, we can generate targets to train a neural network by simply reversing the sampled trajectory through the forward process Markov-Chain, $z_{0}\to z_{1}\to\ldots$. At first glance, this may seem a reasonable choice. However, training to simply invert the last mutation can potentially create a much noisier signal for the neural network.

Consider the case where, within a much larger syntax tree, a color was mutated as,

$$ $\displaystyle{\color[rgb]{1,0,0}\texttt{Red}}\to{\color[rgb]{0,0,1}\texttt{ Blue}}\to{\color[rgb]{0.23046875,0.58984375,0.48828125}\texttt{Green}}.$ $$

The color in our target image, $x_{0}$, is Red, while the color in our mutated image, $x_{2}$, is Green. If we naively teach the model to invert the above Markov chain, we are training the network to turn the Green to a Blue, even though we could have directly trained the network to go from Green to a Red.

Therefore, to create a better training signal, we compute an *edit path* between the target tree and the mutated tree. We use a tree edit path algorithm loosely based on the tree edit distance introduced by Pawlik and Augsten [, ]. The general tree edit distance problem allows for the insertion, deletion, and replacement of any node. Unlike them, our trees can only be edited under an action space that only permits small mutations. For two trees, $z_{A}$ and $z_{B}$, we linearly compare the syntax structure. For changes that are already $\leq\sigma_{\text{small}}$, we add that to our mutation list. For changes that are $>\sigma_{\text{small}}$, we find the first mutation that reduces the distance between the two trees. Therefore, for any two programs, $z_{A}$ and $z_{B}$, we can compute the first step of the mutation path in $O(|z_{A}|+|z_{B}|)$ time.

### 3.3 Value Network & Search

We additionally train a value network, $v_{\phi}(x_{A},x_{B})$, which takes as input two rendered images, $x_{A}$ and $x_{B}$, and predicts the edit distance between the underlying programs that generated those images. Since we have already computed edit paths between trees during training, we have direct access to the ground-truth program edit distance for any pair of rendered images, allowing us to train this value network in a supervised manner.

Using our policy, $q_{\phi}(z_{t-1}|z_{t},x_{t};x_{0})$, and our value, $v_{\phi}(x_{t_{A}},x_{t_{B}})$, we can perform beam-search for a given target image, $x_{0}$, and a randomly initialized program $z_{t}$. At each iteration, we maintain a collection of nodes in our search tree with the most promising values and only expand those nodes.

### 3.4 Architecture

Figure: Figure 3: We train $q_{\phi}(z_{t-1}|z_{t},x_{t};x_{0})$ as a decoder only vision-language transformer following Tsimpoukelli et al. []. We use an NF-ResNet as the image encoder, which is a normalizer-free convolutional architecture proposed by Brock et al. []. The image encoder encodes the current image, $x_{t}$, and the target images, $x_{0}$. The current program is tokenized according to the vocabulary in our context-free grammar. The decoder first predicts an edit location in the current program, and then tokens that replace what the edit location should be replaced by. We constrain the autoregressive decoding by our context-free grammar by masking only the valid token logits.
Refer to caption: extracted/5633576/figures/arch.png

Figure 3 shows an overview of our neural architecture. We use a vision-language model described by Tsimpoukelli et al. [] as our denoising model, $q_{\phi}(z_{t-1}|z_{t},x_{t};x_{0})$. We use an off-the-shelf implementation [] of NF-ResNet-26 as our image encoder, which is a normalizer-free convolutional architecture proposed by Brock et al. [] to avoid test time instabilities with Batch-Norm []. We implement a custom tokenizer, using the terminals of our CFG as tokens. The rest of the edit model is a small decoder-only transformer [, ].

We add two additional types of tokens: an <EDIT> token, which serves as a start-of-sentence token for the model; and <POS x> tokens, which allow the model to reference positions within its context. Given a current image, a target image, and a current tokenized program, we train this transformer model to predict the edit position and the replacement text autoregressively. While making predictions, the decoding is constrained under the grammar. We mask out the prediction logits to only include edit positions that represent nodes in the syntax tree, and only produce replacements that are syntactically valid for the selected edit position.

We set $\sigma_{\text{small}}=2$, which means the network is only allowed to produce edits with fewer than two primitives. For training data, we sample an infinite stream of random expressions from the CFG. We choose a random number of noise steps, $s\in[1,5]$, to produce a mutated expression. For some percentage of the examples, $\rho$, we instead sample a completely random new expression as our mutated expression. We trained for $3$ days for the environments we tested on a single Nvidia A6000 GPU.

## 4 Experiments

### 4.1 Environments

We conduct experiments on four domain-specific graphics languages, with complete grammar specifications provided in Appendix B.

##### CSG2D

A 2D constructive solid geometry language where primitive shapes are added and subtracted to create more complex forms, as explored in our baseline methods [, ]. We also create CSG2D-Sketch, which has an added observation model that simulates hand-drawn sketches using the algorithm from Wood et al. [].

##### TinySVG

A language featuring primitive shapes with color, along with Arrange commands for horizontal and vertical alignment, and Move commands for shape offsetting. Figure 1 portrays an example program. Unlike the compositional nature of CSG2D, TinySVG is hierarchical: sub-expressions can be combined into compound objects for high-level manipulation. We also create, Rainbow, a simplified version of TinySVG without Move commands for ablation studies due to its reduced computational demands.

We implemented these languages using the Lark [] and Iceberg [] Python libraries, with our tree-diffusion implementation designed to be generic and adaptable to any context-free grammar and observation model.

### 4.2 Baselines

We use two prior works, Ellis et al. [] and CSGNet [] as baseline methods.

##### CSGNet

Sharma et al. [] employed a convolutional and recurrent neural network to generate program statements from an input image. For a fair comparison, we re-implemented CSGNet using the same vision-language transformer architecture as our method, representing the modern autoregressive approach to code generation. We use rejection sampling, repeatedly generating programs until a match is found.

##### REPL Flow

Ellis et al. [] proposed a method to build programs one primitive at a time until all primitives have been placed. They also give a policy network access to a REPL, i.e., the ability to execute code and see outputs. Notably, this current image is rendered from the current partial program. As such, we require a custom partial compiler. This is straightforward for CSG2D since it is a compositional language. We simply render the shapes placed so far. For TinySVG, it is not immediately obvious how this partial compiler should be written. This is because the rendering happens bottom-up. Primitives get arranged, and those arrangements get arranged again (see Figure 1). Therefore, we only use this baseline method with CSG2D. Due to its similarities with Generative Flow Networks [], we refer to our modified method as “REPL Flow”.

##### Test tasks

For TinySVG we used a held-out test set of randomly generated expressions and their images. For the CSG2D task, we noticed that all methods were at ceiling performance on an in-distribution held-out test set. In Ellis et al. [], the authors created a harder test set with more objects. However, simply adding more objects in an environment like CSG2D resulted in simpler final scenes, since sampling a large object that subtracts a large part of the scene becomes more likely. Instead, to generate a hard test set, we filtered for images at the $95th$ percentile or more on incompressibility with the LZ4 [, ] compression algorithm.

##### Evaluation

In CSG2D, we accepted a predicted program as matching the specification if it achieved an intersection-over-union (IoU) of $0.99$ or more. In TinySVG, we accepted an image if $99\%$ of the pixels were within $0.005\approx\frac{1}{256}$.

All methods were trained with supervised learning and were not fine-tuned with reinforcement learning. All methods used the grammar-based constrained decoding method described in Section 3.4, which ensured syntactically correct outputs. While testing, we measured performance based on the number of compilations needed for a method to complete the task.

Figure: Figure 4: Performance of our approach in comparison to baseline methods in CSG2D and TinySVG languages. We give the methods $n=256$ images from the test set and measure the number of nodes expanded to find a solution. The auto-regressive baseline was queried with rejection sampling. Our policy outperforms previous methods, and our policy combined with search helps boost performance further. Error bars show standard deviation across 5 random seeds.
Refer to caption: x1.png

Figure 4 shows the performance of our method compared to the baseline methods. In both the CSG2D and TinySVG environments, our tree diffusion policy rollouts significantly outperform the policies of previous methods. Our policy combined with beam search further improves performance, solving problems with fewer calls to the renderer than all other methods. Figure 6 shows successful qualitative examples of our system alongside outputs of baseline methods. We note that our system can fix smaller issues that other methods miss. Figure 7 shows some examples of recovered programs from sketches in the CSG2D-Sketch language, showing how the observation model does not necessarily need to be a deterministic rendering; it can also consist of stochastic hand-drawn images.

### 4.3 Ablations

Figure: (a)
Refer to caption: x2.png

To understand the impact of our design decisions, we performed ablation studies on the simplified Rainbow environment using a smaller transformer model.

First, we examined the effect of removing the current image (no REPL) from the policy network’s input. As shown in Figure 5(a), this drastically hindered performance, confirming the importance of a REPL-like interface observed by Ellis et al. [].

Next, we investigated the necessity of our reverse mutation path algorithm. While training on the last mutation step alone provides a valid path, it introduces noise by potentially targeting suboptimal intermediate states. Figure 5(a) demonstrates that utilizing the reverse mutation path significantly improves performance, particularly in finding solutions with fewer steps. However, both methods eventually reach similar performance levels, suggesting that a noisy path, while less efficient, can still lead to a solution.

Finally, we explored whether the incremental noise process is crucial, given our tree edit path algorithm. Couldn’t we directly sample two random expressions, calculate the path, and train the network to imitate it? We varied the training data composition between pure forward diffusion ($\rho=0.0$) and pure random initialization ($\rho=1.0$) as shown in Figure 5(b). We found that a small proportion ($\rho=0.2$) of pure random initializations combined with forward diffusion yielded the best results. This suggests that forward diffusion provides a richer training distribution around target points, while random initialization teaches the model to navigate the program space more broadly. The emphasis on fine-grained edits from forward diffusion proves beneficial for achieving exact pixel matches in our evaluations.

Figure: Figure 6: Qualitative examples of our method and baselines on two inverse graphics languages, CSG2D (top two rows) and TinySVG (bottom two rows). The leftmost column shows the ground-truth rendered programs from our test set. The next columns show rendered programs from various methods. Our methods are able to finely adjust and match the ground-truth programs more closely.
Refer to caption: extracted/5633576/figures/examples_labels.png

Figure: Figure 7: Examples of programs recovered for input sketches in the CSG2D-Sketch language. The input sketches are from our observation model that simulates hand-drawn sketches (top-row). The output programs rendered (bottom row) are able to match the input sketches by adding and subtracting basic shapes. Video results for these sketches can be found at [https://tree-diffusion.github.io](https://tree-diffusion.github.io).
Refer to caption: extracted/5633576/figures/sketches.png

## 5 Conclusion

In this work, we proposed a neural diffusion model that operates on syntax trees for program synthesis. We implemented our approach for inverse graphics tasks, where our task is to find programs that would render a given image. Unlike previous work, our model can construct programs, view their output, and in turn edit these programs, allowing it to fix its mistakes in a feedback loop. We quantitatively showed how our approach outperforms our baselines at these inverse graphics tasks. We further studied the effects of key design decisions via ablation experiments.

##### Limitations

There are several significant limitations to this work. First, we operate on expressions with no variable binding, loops, strings, continuous parameters, etc. While we think our approach can be extended to support these, it needs more work and careful design. Current large-language models can write complicated programs in many domains, while we focus on a very narrow task. Additionally, the task of inverse graphics might just be particularly suited for inverse graphics where small mutations make informative changes in the image output.

##### Future Work

In the future, we hope to be able to leverage large-scale internet data on programs to train our system, making small mutations to their syntax tree and learning to invert them. We would also like to study this approach in domains other than inverse graphics. Additionally, we would like to extend this approach to work with both the discrete syntax structure and continuous floating-point constants.

##### Impact

Given the narrow scope of the implementation, we don’t think there is a direct societal impact, other than to inform future research direction in machine-assisted programming. We hope future directions of this work, specifically in inverse graphics, help artists, engineering CAD modelers, and programmers with a tool to convert ideas to precise programs for downstream use quickly.

## Acknowledgments and Disclosure of Funding

We would like to thank Kathy Jang, David Wu, Cam Allen, Sam Toyer, Eli Bronstein, Koushik Sen, and Pieter Abbeel for discussions, feedback, and technical support. Shreyas was supported in part by the AI2050 program at Schmidt Futures (Grant G-22-63471). Erik was supported by fellowships from the Future of Life Institute and Open Philanthropy.

## Appendix

## Appendix A Mutation Algorithm

Figure: Figure 8: An example expression from CSG2D represented as a tree to help illustrate the mutation algorithm. The green nodes are candidate nodes with primitives count $\sigma(z)\leq 2$. Our mutation algorithm only mutates these nodes.
Refer to caption: extracted/5633576/figures/example_tree.png

Here we provide additional details on how we sample small mutations for tree diffusion. We will first repeat the algorithm mentioned in Section 3 in more detail.

Our goal is to take some syntax tree and apply a small random mutation. The only type of mutation we consider is a replacement mutation. We first collect a set of candidate nodes that we are allowed to replace. If we select a node too high up in the tree, we end up replacing a very large part of the tree. To make sure we only change a small part of the tree we only select nodes with $\leq\sigma_{\text{small}}$ primitives. In Figure 8, if we set $\sigma_{\text{small}}=2$, we get all the green nodes. We sample a node, $m$, uniformly from this green set. We know the production rule for $m$ from the CFG. For instance, if we selected node $15$, the only replacements allowed are $+$ or $-$. If we selected node $46$, we can only replace it with an angle. If we selected node $11$, we can replace it with any subexpression. When we sample a replacement, we ensure that the replacement is $\leq\sigma_{\text{small}}$, and that it is different than $m$. Here we show $4$ random mutation steps on a small expression,

During our experiments we realized that this style of random mutations biases expression to get longer on average, since there are many more leaves than parents of leaves. This made the network better at going from very long expressions to target expressions, but not very good at editing shorter expressions into longer ones. This also made our model’s context window run out frequently when expressions got too long. To make the mutation length effects more uniform, we add a slight modification to the algorithm mentioned above and in Section 3.

For each of the candidate nodes, we find the set of production rules for the candidates. We then select a random production rule, $r$, and then select a node from the candidates with the production rule $r$, as follows,

$$ $\displaystyle C$ $\displaystyle=\{n\in\text{SyntaxTree}(z)\mid\sigma(n)\leq\sigma_{\text{small}}\}$ $\displaystyle R$ $\displaystyle=\{\text{ProductionRule}(n)\mid n\in C\}$ $\displaystyle r$ $\displaystyle\sim\text{Uniform}[R]$ $\displaystyle M$ $\displaystyle=\{n\in C\mid\text{ProductionRule}(n)=r\}$ $\displaystyle m$ $\displaystyle\sim\text{Uniform}[M]$ $$

For CSG2D, this approach empirically biased our method to make expressions shorter $30.8\%$, equal $49.2\%$, and longer $20.0\%$ of the times ($n=10,000$).

## Appendix B Context-Free Grammars

Here we provide the exact context-free grammars of the languages used in this work.

Figure: (a)
Refer to caption: extracted/5633576/figures/csg_example.png

### B.1 CSG2D

### B.2 TinySVG

## Appendix C Sketch Simulation

Figure: Figure 10: Examples of the same scene being called multiple times by our sketch observation model.
Refer to caption: extracted/5633576/figures/sketch_sim.png

As mentioned in the main text, we implement the CSG2D-Sketch environment, which is the same as CSG2D with a hand-drawn sketch observation model. We do this to primarily show how this sort of a generative model can possibly be applied to a real-world task, and that observations do not need to be deterministic. Our sketch algorithm can be found in our codebase, and is based off the approach described in Wood et al. [].

Our compiler uses Iceberg [] and Google’s 2D Skia library to perform boolean operations on primitive paths. The resulting path consists of line and cubic bézier commands. We post-process these commands to generate sketches. For each command, we first add Gaussian noise to all points stated in those commands. For each line, we randomly pick a point near the $50\%$ and $75\%$ of the line, add Gaussian noise, and fit a Catmull-Rom spline []. For all curves, we sample random points at uniform intervals and fit Catmull-Rom splines. We have a special condition for circles, where we ensure that the start and end points are randomized to create the effect of the pen lifting off. Additionally we randomize the stroke thickness.

Figure 10 shows the same program rendered multiple times using our randomized sketch simulator.

## Appendix D Complexity Filtering

Figure: (a)
Refer to caption: extracted/5633576/figures/complex.png

As mentioned in Section 4, while testing our method alongside baseline methods, we reached ceiling performance for all our methods. Ellis et al. [] got around this by creating a “hard” test case by sampling more objects. For us, when we increased the number of objects to increase complexity, we saw that it increased the probability that a large object would be sampled and subtract from the whole scene, resulting in simpler scenes. This is shown by Figure 11(b), which is our training distribution. Even though we sample a large number of objects, the scenes don’t look visually interesting. When we studied the implementation details of Ellis et al. [], we noticed that during random generation of expressions, they ensured that each shape did not change more that $60\%$ or less than $10\%$ of the pixels in the scene. Instead of modifying our tree sampling method, we instead chose to rejection sample based on the compressibility of the final rendered image.

## Appendix E Tree Path Algorithm

Figure: Algorithm 1 TreeDiff: Find the first set of mutations to turn one tree to another.

Algorithm 1 shows the high-level pseudocode for how we find the first step of mutations to transform tree $A$ into tree $B$. We linearly walk down both trees until we find a node that is different. If the target node is small, i.e., its $\sigma(z)\leq\sigma_{\text{small}}$, then we can simply mutate the source to the target. If the target node is larger, we sample a random small expression with the correct production rule, and compute the path from this small expression to the target. This gives us the first step to convert the source node to the target node. Repeatedly using Algorithm 1 gives us the full path to convert one expression to another. We note that this path is not necessarily the optimal path, but a valid path that is less noisy than the path we would get by simply chasing the last random mutation.

Figure 12 conceptually shows why computing this tree path might be necessary. The circle represents the space of programs. Consider a starting program $z_{0}$. Each of the black arrows represents a random mutation that kicks the program to a slightly different program, so $z_{0}\to z_{1}$, then $z_{2}\to z_{3}\ldots$. If we provide the neural network the supervised target to go from $z_{5}$ to $z_{4}$, we are teaching the network to take an inefficient path to $z_{0}$. The green path is the direct path from $z_{5}\to z_{0}$.

Figure: Figure 12: A conceptual illustration of why we need tree path-finding. The red path represents the naive target for the neural network. The green path represents the path-finding algorithm’s target.
Refer to caption: extracted/5633576/figures/why_path.png

## Appendix F Implementation Details

We implement our architecture in PyTorch []. For our image encoder we use the NF-ResNet26 [] implementation from the open-sourced library by Wightman []. Images are of size $128\times 128\times 1$ for CSG2D and $128\times 128\times 3$ for TinySVG. We pass the current and target images as a stack of image planes into the image encoder. Additionally, we provide the absolute difference between current and target image as additional planes.

Our decoder-only transformer [, ] uses $8$ layers, $16$ heads, with an embedding size of $256$. We use batch size $32$ and optimize with Adam [] with a learning rate of $3\times 10^{-4}$. The image embeddings are of the same size as the transformer embeddings. We use $4$ prefix tokens for the image embeddings. We used a maximum context size of $512$ tokens. For both environments, we sampled expressions with at most $8$ primitives. Our method and all baseline methods used this architecture. We did not do any hyperparameter sweeps or tuning.

For the autoregressive (CSGNet) baseline, we trained the model to output ground-truth programs from target images, and provided a blank current image. For tree diffusion methods, we initialized the search and rollouts using the output of the autoregressive model, which counted as a single node expansion. For our re-implementation of Ellis et al. [], we flattened the CSG2D tree into shapes being added from left to right. We then randomly sampled a position in this shape array, compiled the output up until the sampled position, and trained the model to output the next shape using constrained grammar decoding.

This is a departure from the pointer network architecture in their work. We think that the lack of prior shaping, departure from a graphics specific pointer network, and not using reinforcement learning to fine-tune leads to a performance difference between their results and our re-implementation. We note that our method does not require any of these additional features, and thus the comparison is fairer. For tree diffusion search, we used a beam size of $64$, with a maximum node expansion budget of $5000$ nodes.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 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?

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

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

How many steps did you train for? I don't think this is covered in the appendix.
soundness: 3
strengths: The method is simple but non-obvious
The more general problem of program synthesis conditioned on desired outputs is very relevant
The authors use randomly generated programs as a dataset which sidesteps dataset curation in favor of just a specification of the language
The paper is well-written, easy to understand, and has nice and (mostly) clear figures
summary: The authors propose a program synthesis method based on "tree diffusion". They randomly corrupt programs (with some constraints) and learn to invert the corruptions, conditioned on the output of the corrupted program and the target output (an image rendering in their case).
weaknesses: The paper is somewhat limited in scope (simple problem setup) in ways that make it not entirely obvious how the method "scales" to more complex relevant tasks like code generation.

Some minor things covered in Questions

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - 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?
- What is the cost of training the value network that predicts the edit distance?
- Given the recent advances in vision-language models, how does your approach compare against contemporary models like VILA or LLaMA? The current baselines only include older models (4+ years old), and evaluating against recent state-of-the-art would provide a more helpful comparison.
soundness: 4
strengths: - Innovative Approach: The paper presents a novel combination of autoregressive, diffusion, and search methodologies, which, despite being applied to a specific domain, holds potential for broader applications. The reverse mutation path algorithm also provides an efficient way to generate training targets. 
- Clarity and Replicability: The manuscript is well-written and easy to follow, providing sufficient detail to enable replication of the experiments.
- Comprehensive Ablation Studies: The authors conduct thorough ablation studies on key hyperparameters and the impact of integrating search, enhancing the understanding of their method's efficacy.
summary: The paper introduces an innovative approach to inverse graphics tasks by combining diffusion models with transformers. The authors present the first application of diffusion to program synthesis using explicit syntax tree updates, validating their method on CSG2D and TinySVG environments.
weaknesses: - Literature Coverage: 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. 
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: Here are some of the main points I would like to make:

1. The edit distance introduced by Pawlik & Augsten is narrowed to allow only small changes. If there are big changes, the change that reduces the distance between the trees the most is chosen. While this can be used as a training signal, it is assumed and was mentioned in section 3.3.2 that access to the ground-truth mutations is available. However, access to this ground truth may not be readily available in most cases. How is this ground-truth data obtained? If not automated, the authors must comment on the limited scalability of the approach.

2. In section 4.2 under evaluation, it is not clear what the criteria for considering a match between the synthesized and true plan is. Re: “In TinySVG, we accepted an image if 99% of the pixels were within 0.005 ≈ $\frac{1}{256}$.” Is this “within 0.005” a tolerance on the 8-bit pixel color intensity? If so, please state explicitly and explain how this metric is applied to the RGB images in TinySVG.

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

4. 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. It is unclear as to just how much improvement the proposed approach affords.
soundness: 3
strengths: 1. The main strength of this paper is the design of a neurosymbolic framework to evaluate the automated (i.e. diffusion-based) conversion of images into context-free grammar. This formal evaluation ensures that the desired specifications are met through iterative observation of the execution results and verification.

2. The authors extend the approach to accept hand-drawn sketches and illustrate examples in the appendix confirming the applicability of the approach in several real-world settings.

3. The supplementary videos illustrate the overall problem that the authors are attempting to solve and showcases the "edits" made by the framework.
summary: This paper proposed a program synthesis framework using mutations on syntax trees via a neural diffusion model for inverse graphic tasks. These tasks aim to convert images that can contain free hand-sketches of shapes or a set of computer-generated colored shapes into images depicting a computer-generated rendering matching the input. The authors defend the claim that the approach presents the ability to edit trees generated using a base model as opposed to incrementally autoregressive approaches that fail to narrow down the search space. The authors apply their method to inverse graphics tasks and present results in two settings (CSG2D and TinySVG) and show improved performance in the number of problems solved compared to baselines (REPL VLM and VLM) along with an ablation study to investigate the individual contributions of constituent components.
weaknesses: There are three main weaknesses I would like to bring up. The authors are encouraged to rebut and provide legitimate explanations, if any, against these and the review decision may be adjusted accordingly.

1. A claim made by the author states that the proposed method focuses on editing the program synthesized from the image, unlike prior works that autoregressively generate programs that are incrementally better. In doing so, the authors propose adding random noise to modify a base syntax tree generated from CSG2D. Despite 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. 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). An analysis and example to demonstrate this is lacking and should be included.

2. The overall architecture presented in Figure 3 is difficult to understand at first glance. In addition, the descriptions provided in section 3.4 (the model architecture) do not present enough detail to understand Figure 3.  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 (i.e. are there low-level implementations in play that ensure that entire blocks from “(“ to “)” are parsed out during replacing? How are varying input lengths handled? ). Replacing "(" with "(Quad 8..." seems to break the pairing of parenthesis. In addition, it is not clear what the purpose of "EOS" is in this context.

3. 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. The control reaches the same performance at about 50 nodes. Is this a "significant" efficiency gain when the maximum node expansion budget was two orders of magnitude higher (i.e. 5000)? There are no computational or time-related metrics presented which help put this into context.

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

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. 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?
2. If so, why not use that (training on all denoising steps) as a baseline?
3. Can you define the "Rollout" method, and the differences with the "Search" one?
4. What's the relationship between "Number of nodes expanded" and the "number of compilations needed"?
soundness: 3
strengths: Originality
-------------
The paper takes inspiration from existing ideas and benchmarks, but they are clearly cited, and the novel aspects are well described. For instance, a backward edit path that's better than reversing the corruption path, removing the need of a partial renderer, and relying on beam search rather than full-fledged reinforcement learning.

Quality
----------

Experiments demonstrate the advantages of the proposed approach, and properly ablate the different aspects and contributions.

Clarity
---------
The paper is overall clear and straightforward to follow. With the additional details of the appendix, the approach should be re-implementable by a different team.

Significance
-----------------
Using ML models to directly manipulate and modify programs, rather than either generate a whole program autoregressively, or emit edition instructions (which could be invalid or result in an invalid program) could make iterative program generation better or easier.
The fact that no reinforcement learning is required, but observation of the output of intermediate programs can simply be combined with beam search is also an interesting result.
summary: The paper presents an algorithm to train a diffusion model on the abstract syntax tree (AST) of programs written in a simple procedural language. That language produces 2D images by combining geometric shape primitives, and the resulting image is used to guide the denoising process.

Using an additional network estimating the distance from a rendered image to the target image, beam search is applied to generate a sequence of AST edits (node replacements) that produce a predicted program, approximating the target image. That approach transfers to generating geometric images from noisy sketches.
weaknesses: Originality
-------------
No major weakness here, the work is in the continuation of previous cited work.

Quality
----------
Comparison with baselines might have been more extensive, specifically the RL-based algorithms from previous work, which could have better shown how "brittle" they were.

Clarity
---------
A few things were not clearly defined in the experiments and ablation sections (see "questions" below).

Significance
----------------
Overall, the CSG2D and TinySVG languages are a small-scale benchmark, but it's unclear whether the proposed approach would scale to large, structured programs in general purpose languages.
For instance, it might not be possible to find a sequence of valid programs created by short mutations between two relatively close programs. For instance, going from recursion to a loop, from an implicit lamda to a declared function, or from a for loop to a list comprehension. Even splitting a function into smaller pieces may require either large edits, or intermediate unparseable states.

After discussion
----------------------
The authors provided clarifications, other reviewers raised a few additional concerns, overall I'm maintaining my score.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
questions: See discussions above.
soundness: 3
strengths: - This paper proposed a novel solution to the reverse CG field, to synthesis programs for visual symbolic reasoning. The proposed method address the hard task through the unique lens of syntax tree, and achieves notably better results.

- The idea of permuting on syntax tree allows for more efficient model, with better performance.

- The efforts to make demo video makes the paper easier to understoand and spread.
summary: This paper addresses the edge of visual symbolic reasoning and code generation. It addresses the important task of generating code (symbolic sequence) to depict images with visual feedbacks. It applies diffusion model-like approaches to permute the program syntax tree and guarantee the correctness of the generated code. Through iterations, the model is able to recover the image with high fidelity using discrete preset symbols.
weaknesses: This is a good paper, with minor weakness points below.

- 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.

- 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.

- Two highly related work should be cited and discussed:

"Outline, Then Details: Syntactically Guided Coarse-To-Fine Code Generation" in ICML 2023, which explore the possibility of syntax tree to generate code, and via coarse-to-fine multi-round generation approach.

"Symbolic Visual Reinforcement Learning: A Scalable Framework with Object-Level Abstraction and Differentiable Expression Search" in TPAMI, which also learns visual symbolic programs, not to depict the image but to interact with the environments. Rainbow environment is also leveraged in their experiments.

Meta-review: This paper introduces a novel diffusion-based approach to program synthesis, leveraging syntax trees to iteratively refine programs while preserving syntactic validity. The authors position their method as an alternative to traditional autoregressive models, which are limited by sequential errors and a lack of execution feedback. By modeling program edits as a diffusion process, the approach effectively integrates syntax-aware transformations with execution-guided search, making it particularly suitable for tasks like inverse graphics. The integration of execution outputs into the synthesis pipeline further enhances the model's ability to refine programs to meet desired specifications.

The paper demonstrates strong empirical results on controlled benchmarks such as CSG2D and TinySVG, showing improvements over competitive baselines. Comprehensive ablation studies highlight the contributions of key components, including the use of tree-based diffusion and search-guided optimization. The method is presented clearly, with supplementary videos and an interactive demo that significantly aid understanding and replicability. While the scope of the current work is limited to simpler grammars, the authors acknowledge this as a future direction, and the framework’s scalability to more complex programming tasks holds significant promise. Overall, this paper offers a compelling contribution to neurosymbolic AI and program synthesis and is well-suited for acceptance.

Decision: Accept (Spotlight)