Compose a graphics program just by looking at a hand-drawn sketch. Berkeley, California, teaches diffusion models new skills

It turns out that diffusion models can be used not only to generate images and videos, but also to synthesize new programs.

Suppose we give the model a hand-drawn "5" shape graphic, it can modify the program through continuous mutations, and finally get a program that can output the target graphic.

This model comes from a research team at the University of California, Berkeley. This new method of program synthesis proposed by them uses a neural diffusion model to directly operate syntax trees.

Paper 1 As Shreyas Kapur, a doctoral student at the school, his supervisor is Stuart Russell, professor of computer science at the school.

• Paper title: Diffusion On Syntax Trees For Program Synthesis
• Paper address: https://arxiv.org/pdf/2405.20519
• Project address: https://diffusion-diffusion. github.io/
• Code base: https://github.com/revalo/tree-diffusion

Diffusion models have previously achieved great success in the field of image generation. The team found that by utilizing the diffusion mechanism, the model can learn to iteratively optimize the program while ensuring syntactic validity. The key to this new approach is to enable an efficient debugging process by allowing the model to observe the output of the program at each step.

The ability to iterate has been demonstrated in systems such as AlphaZero, and the iterative nature of the diffusion mechanism will naturally be used in search-based program synthesis.

The team found that by training a value model at the same time as the diffusion model, they could guide the denoising process in it to a program that outputs the desired results. This allows efficient exploration of the program space and makes more informed decisions at every step of the generation process.

To implement this approach, the team chose an inverse graphics task, which assumes that a domain-specific language is used to draw images.

The team said: "The task of reverse engineering is a natural fit for our approach, because a small change in the code can lead to meaningful semantic changes in the resulting image."

For example, if A misplaced shape appears in the image and can be easily seen and positioned in program space. Figure 1 gives some examples.

The main contributions of this research include:

1. Propose a new method using diffusion on syntax trees

2. Implement the method for the inverse graphics task and find that the new method is superior than the previous method.

Method

In short, the core idea of ​​this method is: develop a denoising diffusion model for syntax trees, similar to the image diffusion model developed for vision tasks.

First, let’s look at a task example from Ellis et al.’s paper “Write, execute, assess: Program synthesis with a repl”: generating a constructed solid geometry (CSG2D) program based on an image. Using CSG2D, we can combine simple primitives like circles and quadrilaterals using Boolean operations like addition and subtraction to create more complex shapes using the following context-free grammar (CFG):

in Figure 2 , z₀ is the target program, x₀ is the rendered version of z₀.

The task is to reverse x₀ to recover z₀. First, the denoising process randomly mutates y=16 to y=10. Then transform the subtree with two shapes on the left into a new subtree with only one shape. The goal here is to train a neural network that can reverse this denoising process, starting from z₃ and x₃, based on the image x₀, to get z₀.

The following will first describe how to add "noise" to the syntax tree, then describe in detail how to train a neural network that reverses this noise, and finally describe how to use this neural network to perform a search.

Sampling small mutations

Let z_t be the program at time t. Let p_N (z_{t+1}|z_t) be the distribution on which program z_t is randomly mutated to z_{t+1}. Here we hope that the p_N mutation satisfies two points: (1) it is small, and (2) it can obtain syntactically valid z_{t+1}.

To do this, the team explored a large body of computer security literature on grammar-based fuzz testing. To ensure that mutations are small, they first define a function σ(z) that gives the "size" of program z. In the experiment, a set of endpoints (terminals) in CFG are defined as primitives.

For example, if written in their CSG2D language, the above primitives would be {Quad, Circle}. When working with this language, the team's approach is to let σ(z) = σ_primitive (z), which counts the number of primitives. σ(z) may also contain options such as depth, number of nodes, etc.

Then, randomly sample the procedure from CFG based on the exact constraint σ_min

In order to mutate a given program z, first generate a set of candidate nodes in a certain σ_small range in its syntax tree:

Then, uniformly sample a mutation node from this set :

Since it can read the entire syntax tree and CFG, it knows which generation rule can get m, and thus can ensure that syntactically valid mutations are obtained. For example, if m is a number, its replacement should also be a number. If m is a general subexpression, it can be replaced by any general subexpression. Therefore, m' can be sampled, which is a surrogate for m:

Strategy

Forward process

The team treats the program synthesis problem as an inference problem. Let p (x|z) be an observation model, where x can be any type of observation. The task is to reverse the direction of this observation model, i.e. to obtain a program z given some observation x.

First, take a program z₀ from a data set D, or in this case randomly sample a program from the CFG. That is, sample a z₀ that satisfies σ(z₀) ≤ σ_max. Noise is then added to z₀, taking s steps, where s ∼ Uniform [1, s_max] and s_max is a hyperparameter:

Then, a conditional neural network is trained that models the following distribution .

where ϕ is the parameter of the neural network, z_t is the current program, x_t is the current output of the program, and x₀ is the target output that needs to be solved.

Reverse mutation path

Due to the ability to obtain ground truth mutations, the trajectories sampled can be reversed simply through the forward process Markov chain z₀ → z₁ →... to generate The goal of training a neural network. At first glance, this might seem like a reasonable choice. However, directly training the model to reverse the last mutation risks creating a far noisier signal for the neural network.

For example, in a much larger syntax tree, the mutation path of a color is:

The color of the target image x₀ is Red, and the color of the mutated image x₂ is Green.If you simply teach the model to reverse the above Markov chain, you may train the network to turn Green into Blue, even though you can directly train the network to turn Green into Red.

Therefore, in order to create a better training signal, the editing path between the target tree and the mutation tree can be calculated. The team used a tree edit path algorithm loosely based on tree edit distance. The generalized tree edit distance problem allows insertion, deletion, and replacement of arbitrary nodes. But the setting here is different. Tree editing can only be achieved in an action space that only allows small mutations.

For two trees z_A and z_B, their syntactic structures can be compared in a linear way. For changes that satisfy ≤ σ_small, they are added to the mutation list. For changes > σ_small, look for the first mutation that reduces the distance between the two trees. Therefore, for any two programs z_A and z_B, the first step of the mutation path can be calculated in O (|z_A| + |z_B|) time.

Value Network and Search

The team also trained a value network v_ϕ (x_A, x_B), whose input is two rendered images x_A and x_B, and the prediction is to generate these two The edit distance between the underlying programs of the image. Since the edit distance between trees has been calculated during training, for any pair of rendered images, their ground truth procedural edit distance is directly obtained, which can be used to train this value in a supervised manner network.

Using the above new strategy and new value network proposed by the team, it is possible to perform beam search for any target image x₀ and a randomly initialized program z_t. At each iteration, a set of nodes in the search tree with the most promising values ​​is maintained and only these nodes are expanded.

Architecture

Figure 3 shows an overview of the newly proposed neural architecture.

Their denoising model q_ϕ (z_{t−1}|z_t, x_t; x₀) uses the visual- Language model. As for the image encoder, they used an off-the-shelf implementation of NF-ResNet-26; a normalizer-less convolutional architecture that avoids the temporal instability issues when using Batch-Norm.

The team implemented a customized tokenizer that uses their CFG’s endpoint as token. The remainder of the editing model is a small decoder-only Transformer.

They also added two other types of tokens: which serves as a sentence starting token for the model, and token which allows the model to refer to locations within its context.

Given a current image, a target image, and the current tokenized program, train the Transformer model to predict edit positions and replacement text in an autoregressive manner. When making predictions, the decoding process is constrained by the grammar. The team masked the predicted logit to include only edit positions that represent syntactic tree nodes, and only obtained substitutions that were syntactically valid for the selected edit positions.

Here we set σ_small = 2, which means that the network is only allowed to produce edits with less than two primitives. For training data, what they do is sample an infinite stream of random expressions from the CFG. For noise steps s, they randomly select one from [1, 5]. A certain proportion of the samples ρ is completed by randomly sampling new expressions as mutated expressions. They trained for three days using a single NVIDIA A6000 GPU.

Experiments

They conducted experiments on 4 domain-specific graphics languages: CSG2D, CSG2D-Sketch, TinySVG, Rainbow.

The selected benchmark methods are "Write, execute, assess: Program synthesis with a repl" proposed by Ellis et al. and "CSGNet: Neural shape parser for constructive solid geometry" proposed by Sharma et al.

Figure 4 compares the performance of the new method with the baseline method.

It can be seen that in the CSG2D and TinySVG environments, the newly proposed tree diffusion strategy is significantly better than the strategy of the previous method.The performance of this strategy can be further improved if combined with beam search, allowing the problem to be solved with fewer calls to the renderer than other methods.

The figure below gives some successful examples of the new system and the output of the benchmark method. As you can see, the new system can fix smaller problems that other methods miss.

The following video shows two examples of sketch-based recovery procedures using the CSG2D-Sketch language, which shows that the observation model does not necessarily require deterministic rendering; it can also be composed of random hand-drawn images.

To understand the impact of these new designs, the team also conducted ablation experiments using a smaller Transformer model in a simplified Rainbow environment, the results are shown in Figure 5. Overall, the results of these design choices are proven.

Please refer to the original paper for more details.

The above is the detailed content of Compose a graphics program just by looking at a hand-drawn sketch. Berkeley, California, teaches diffusion models new skills. For more information, please follow other related articles on the PHP Chinese website!