扩散模型:不仅能生成图像和视频,还能合成新程序

事实证明,扩散模型不仅能用于生成图像和视频,还能用于合成新程序。假设我们给模型一张手绘的「5」状图形,它就能通过不断突变来修改程序,最终得到能输出目标图形的程序。该模型来自加州大学伯克利分校的一个研究团队,他们提出的这种程序合成新方法使用了神经扩散模型来直接操作句法树。

论文标题:Diffusion On Syntax Trees For Program Synthesis

论文地址:https://arxiv.org/pdf/2405.20519

项目地址:https://tree-diffusion.github.io/

代码库:https://github.com/revalo/tree-diffusion

扩散模型之前已经在图像生成领域取得了巨大成功。而该团队发现,通过利用扩散机制,可让模型学会迭代地优化程序,同时确保句法有效性(syntactic validity)。这种新方法的关键在于可让模型观察程序在每一步的输出,从而实现有效的调试流程。

迭代的能力已经在 AlphaZero 等系统中得到了体现,而扩散机制的迭代性质也就很自然地会被用于基于搜索的程序合成。该团队发现,通过在训练扩散模型的同时训练一个价值模型,就可以引导其中的去噪过程得到能输出所需结果的程序。这样一来,便能高效地探索程序空间,在生成过程中的每一步都做出更明智的决策。

为了实现该方法,该团队选择了逆向图形任务(inverse graphics task),即假定使用特定领域的语言来绘制图像。该团队表示:「逆向图形任务天然就适合我们的方法,因为代码中的一点微小改变就能导致所得图像出现有意义的语义变化。」举个例子,如果图像中出现了一个放置错误的图形,那就能被轻松看见并定位到程序空间中。图 1 给出了一些示例。

这项研究的主要贡献包括:

1. 提出一种在句法树上使用扩散的全新方法

2. 针对逆向图形任务实现了该方法,并发现新方法优于之前的方法。

方法:简而言之,该方法的核心思想是:为句法树开发一个去噪扩散模型,类似于为视觉任务开发的图像扩散模型。

以下是重构后的内容:

首先,我们来看一个来自 Ellis et al. 论文《Write, execute, assess: Program synthesis with a repl》中的任务示例。该任务的目标是根据图像生成一个构造实体几何(CSG2D)程序。使用 CSG2D,我们可以使用加法和减法等布尔运算将圆和四边形等简单的原语组合起来,从而能使用下面的上下文无关语法(CFG)创建出更复杂的形状。

在图 2 中,z0 是目标程序,x0 是 z0 渲染过的版本。任务就是逆转 x0,从而恢复得到 z0。首先,去噪过程将 y=16 随机突变为 y=10。然后再将左侧有两个形状的子树转变成只有一个形状的新子树。这里的目标是基于图像 x0,从 z3 和 x3 开始,训练一个能逆向该去噪过程的神经网络,从而得到 z0。

为了实现这个任务,研究者们采用了以下方法:

1. 将「噪声」添加到句法树中:首先定义了一个函数 σ(z),其能给出程序 z 的「大小」。在实验中,则是将 CFG 中的一组端点(terminal)定义为原语。举个例子,如果用他们的 CSG2D 语言来写,上述原语就为 {Quad, Circle}。使用该语言时,该团队的做法是让 σ(z) = σ_primitive (z),这能统计原语的数量。σ(z) 还可能包含深度、节点数量等选项。

2. 基于精确的约束条件 σ_min < σ(z) ≤ σ_max 从 CFG 随机采样程序。该函数被命名为 ConstrainedSample (σ_min, σ_max)。为 σ_max 设定一个较小的值就能随机采样小程序。在生成小突变时,设定 σ_max = σ_small。

3. 为了给定程序 z 发生突变,首先可在其句法树中生成一个在某个 σ_small 范围内的候选节点集合。然后,从该集合中均匀采样一个突变节点。

4. 最后,使用该神经网络来执行搜索,以找到从 z3 和 x3 开始的逆向去噪过程的路径。

原文重构内容如下:

策略:

前向过程:

将程序合成问题视为一个推理问题,令p(x|z)为观察模型,其中x可以是任意类型的观察。任务是逆转这个观察模型的方向,即给定某个观察x得到一个程序z。该团队采用的方法是从CFG中随机采样一个程序作为初始程序z0,然后通过下式描述的过程向z0添加噪声,执行s步,其中s∼Uniform[1,s_max],而s_max是一个超参数。最后,训练一个建模以下分布的条件神经网络φ,其中z_t是当前的程序,x_t是该程序当前的输出,x0是需要求解的目标输出。

逆向突变路径:

由于能够获取基本真值突变,因此可以简单地通过前向过程马尔可夫链z0→z1→...来逆向采样的轨迹,从而生成用以训练神经网络的目标。然而,直接训练模型逆向最后一次突变有可能为神经网络创造出远远更有噪声的信号。为了解决这个问题,可以计算目标树与突变树之间的编辑路径。该团队使用了一种大体上基于树编辑距离(tree edit distance)的树编辑路径算法。广义的树编辑距离问题允许插入、删除和替换任意节点,但这里的设定不同,对树的编辑仅能在一个只允许小突变的动作空间中实现。

举个例子,在一个大得多的句法树中,一种颜色的突变路径为:目标图像x0的颜色是Red,而突变后图像x2的颜色是Green。如果只是简单教模型逆向上述马尔可夫链,则可能会训练网络将Green变成Blue,即便可以直接训练网络将Green变成Red。因此,为了创建更好的训练信号,可以计算目标树与突变树之间的编辑路径。

对于两个树 z_A 和 z_B,可以通过线性方式比较它们的句法结构。对于已经满足 ≤ σ_small 的改变,就将其加入到突变列表中。对于 > σ_small 的改变,则寻找能降低两个树之间距离的首个突变。因此,对于任意两个程序 z_A 和 z_B 而言,可以在 O (|z_A| + |z_B|) 时间内计算出突变路径的第一步。

价值网络与搜索

该团队还训练了一个价值网络 v_φ (x_A, x_B),其输入为两张经过渲染的图像 x_A 和 x_B,预测的是生成这两张图像的底层程序之间的编辑距离。由于在训练期间已经计算出了树之间的编辑距离,因此对于任意一对经过渲染的图像,都能直接获得它们的基本真值程序编辑距离,这就可以用于以监督式方法训练该价值网络。

使用该团队提出的上述新策略和新价值网络,就可以为任意目标图像 x0 和随机初始化的程序 z_t 执行波束搜索。在每一次迭代中,都要维护搜索树中一组有最有希望值的节点,并且仅扩展这些节点。

架构

图 3 展示了新提出的神经架构的概况。他们的去噪模型 q_φ (z_{t−1}|z_t, x_t; x0) 使用的是 Tsimpoukelli et al. 在论文《Multimodal few-shot learning with frozen language models》中描述的视觉 - 语言模型。至于图像编码器,他们使用的是现成可用的 NF-ResNet-26 的实现;这是一种无归一化器的卷积架构,可避免使用 Batch-Norm 时的时间不稳定问题。

该团队实现了一种定制化的 token 化器,其中使用了他们的 CFG 的端点为 token。该编辑模型的其余部分是一个小型的仅解码器 Transformer。他们还添加了另外两种类型的 token:用作该模型的句子起点 token 的以及允许模型在其上下文中引用位置的 token。给定一张当前图像,一张目标图像,和当前一个已 token 化的程序,训练该 Transformer 模型使之能以自回归方式预测编辑位置和替换文本。在进行预测时,解码过程受到语法的约束。该团队对预测 logit 进行了掩蔽,使之仅包含能表示句法树节点的编辑位置,以及仅得到对于所选编辑位置句法上有效的替换。

您好,我了解到您想了解关于加州伯克利分校的研究团队提出的程序合成新方法。这个方法使用了神经扩散模型来直接操作句法树,可以通过不断突变来修改程序,最终得到能输出目标图形的程序。该模型在 4 种特定领域的图形语言上进行了实验:CSG2D、CSG2D-Sketch、TinySVG、Rainbow。所选用的基准方法为 Ellis et al. 提出的《Write, execute, assess: Program synthesis with a repl》以及 Sharma et al. 提出的《CSGNet: Neural shape parser for constructive solid geometry》。

如果您还有其他问题或需要更多信息,请告诉我。