假設(shè)我們給模型一張手繪的「5」?fàn)顖D形，它就能通過不斷突變來修改程序，最終得到能輸出目標(biāo)圖形的程序。

該模型來自加州大學(xué)伯克利分校的一個研究團(tuán)隊(duì)，他們提出的這種程序合成新方法使用了神經(jīng)擴(kuò)散模型來直接操作句法樹。

論文一作為該校博士生 Shreyas Kapur，其導(dǎo)師為該校計(jì)算機(jī)科學(xué)教授 Stuart Russell。

• 論文標(biāo)題：Diffusion On Syntax Trees For Program Synthesis
• 論文地址：https://arxiv.org/pdf/2405.20519
• 項(xiàng)目地址：https://tree-diffusion.github.io/
• 代碼庫：https://github.com/revalo/tree-diffusion?

擴(kuò)散模型之前已經(jīng)在圖像生成領(lǐng)域取得了巨大成功。而該團(tuán)隊(duì)發(fā)現(xiàn)，通過利用擴(kuò)散機(jī)制，可讓模型學(xué)會迭代地優(yōu)化程序，同時(shí)確保句法有效性（syntactic validity）。這種新方法的關(guān)鍵在于可讓模型觀察程序在每一步的輸出，從而實(shí)現(xiàn)有效的調(diào)試流程。

迭代的能力已經(jīng)在 AlphaZero 等系統(tǒng)中得到了體現(xiàn)，而擴(kuò)散機(jī)制的迭代性質(zhì)也就很自然地會被用于基于搜索的程序合成。

該團(tuán)隊(duì)發(fā)現(xiàn)，通過在訓(xùn)練擴(kuò)散模型的同時(shí)訓(xùn)練一個價(jià)值模型，就可以引導(dǎo)其中的去噪過程得到能輸出所需結(jié)果的程序。這樣一來，便能高效地探索程序空間，在生成過程中的每一步都做出更明智的決策。

為了實(shí)現(xiàn)該方法，該團(tuán)隊(duì)選擇了逆向圖形任務(wù)（inverse graphics task），即假定使用特定領(lǐng)域的語言來繪制圖像。

該團(tuán)隊(duì)表示：「逆向圖形任務(wù)天然就適合我們的方法，因?yàn)榇a中的一點(diǎn)微小改變就能導(dǎo)致所得圖像出現(xiàn)有意義的語義變化?！?/p>

舉個例子，如果圖像中出現(xiàn)了一個放置錯誤的圖形，那就能被輕松看見并定位到程序空間中。圖 1 給出了一些示例。

這項(xiàng)研究的主要貢獻(xiàn)包括：

1. 提出一種在句法樹上使用擴(kuò)散的全新方法

2. 針對逆向圖形任務(wù)實(shí)現(xiàn)了該方法，并發(fā)現(xiàn)新方法優(yōu)于之前的方法。

方法

簡而言之，該方法的核心思想是：為句法樹開發(fā)一個去噪擴(kuò)散模型，類似于為視覺任務(wù)開發(fā)的圖像擴(kuò)散模型。

首先來看一個來自 Ellis et al. 論文《Write, execute, assess: Program synthesis with a repl》中的任務(wù)示例：根據(jù)圖像生成一個構(gòu)造實(shí)體幾何（CSG2D）程序。使用 CSG2D，我們可以使用加法和減法等布爾運(yùn)算將圓和四邊形等簡單的原語組合起來，從而能使用下面的上下文無關(guān)語法（CFG）創(chuàng)建出更復(fù)雜的形狀：

在圖 2 中，z? 是目標(biāo)程序，x? 是 z? 渲染過的版本。

任務(wù)就是逆轉(zhuǎn) x?，從而恢復(fù)得到 z?。首先，去噪過程將 y=16 隨機(jī)突變?yōu)?y=10。然后再將左側(cè)有兩個形狀的子樹轉(zhuǎn)變成只有一個形狀的新子樹。這里的目標(biāo)是基于圖像 x?，從 z? 和 x? 開始，訓(xùn)練一個能逆向該去噪過程的神經(jīng)網(wǎng)絡(luò)，從而得到 z?。

下面首先將介紹如何將「噪聲」添加到句法樹中，然后將詳細(xì)描述如何訓(xùn)練逆向該噪聲的神經(jīng)網(wǎng)絡(luò)，最后將描述如何使用該神經(jīng)網(wǎng)絡(luò)來執(zhí)行搜索。

采樣微小突變

令 z_t 為時(shí)間 t 時(shí)的程序。令 p_N (z_{t+1}|z_t) 為將程序 z_t 隨機(jī)突變成 z_{t+1} 所基于的分布。這里希望 p_N 突變滿足兩點(diǎn)：(1) 很小，(2) 能得到句法有效的 z_{t+1}。

為此，該團(tuán)隊(duì)探究了大量有關(guān)基于語法的模糊測試的計(jì)算機(jī)安全文獻(xiàn)。為了確保突變很小，他們首先定義了一個函數(shù) σ(z)，其能給出程序 z 的「大小」。在實(shí)驗(yàn)中，則是將 CFG 中的一組端點(diǎn)（terminal）定義為原語。

舉個例子，如果用他們的 CSG2D 語言來寫，上述原語就為 {Quad, Circle}。使用該語言時(shí)，該團(tuán)隊(duì)的做法是讓 σ(z) = σ_primitive (z)，這能統(tǒng)計(jì)原語的數(shù)量。σ(z) 還可能包含深度、節(jié)點(diǎn)數(shù)量等選項(xiàng)。

然后，基于精確的約束條件 σ_min < σ(z) ≤ σ_max 從 CFG 隨機(jī)采樣程序。該函數(shù)被命名為 ConstrainedSample (σ_min, σ_max)。為 σ_max 設(shè)定一個較小的值就能隨機(jī)采樣小程序。在生成小突變時(shí)，設(shè)定 σ_max = σ_small。

為了讓給定程序 z 發(fā)生突變，首先可在其句法樹中生成一個在某個 σ_small 范圍內(nèi)的候選節(jié)點(diǎn)集合：

然后，從該集合中均勻采樣一個突變節(jié)點(diǎn)：

由于能讀取整個句法樹和 CFG，因此知道了哪條生成規(guī)則能得到 m，并由此可以確保得到句法有效的突變。舉個例子，如果 m 是一個數(shù)值，那么替代它的也應(yīng)該是一個數(shù)值。如果 m 是一個一般子表達(dá)式，那就可以把它替換成任何一般子表達(dá)式。因此，可采樣出 m'，這是 m 的替代：

策略

前向過程

該團(tuán)隊(duì)將程序合成問題視為一個推理問題。令 p (x|z) 為觀察模型，其中 x 可以是任意類型的觀察。任務(wù)是逆轉(zhuǎn)這個觀察模型的方向，即給定某個觀察 x 得到一個程序 z。

首先，從一個數(shù)據(jù)集 D 取出某個程序 z?，或者這里的做法是從 CFG 隨機(jī)采樣一個程序。即采樣一個滿足 σ(z?) ≤ σ_max 的 z?。然后通過下式描述的過程向 z? 添加噪聲，執(zhí)行 s 步，其中 s ～ Uniform [1, s_max]，而 s_max 是一個超參數(shù)：

然后，訓(xùn)練一個建模以下分布的條件神經(jīng)網(wǎng)絡(luò)。

其中 ? 是該神經(jīng)網(wǎng)絡(luò)的參數(shù)，z_t 是當(dāng)前的程序，x_t 是該程序當(dāng)前的輸出，x? 是需要求解的目標(biāo)輸出。

逆向突變路徑

由于能夠獲取基本真值突變，因此可以簡單地通過前向過程馬爾可夫鏈 z? → z? →... 來逆向采樣的軌跡，從而生成用以訓(xùn)練神經(jīng)網(wǎng)絡(luò)的目標(biāo)。乍一看，這或許是個合理選擇。但是，直接訓(xùn)練模型逆向最后一次突變有可能為神經(jīng)網(wǎng)絡(luò)創(chuàng)造出遠(yuǎn)遠(yuǎn)更有噪聲的信號。

舉個例子，在一個大得多的句法樹中，一種顏色的突變路徑為：

目標(biāo)圖像 x? 的顏色是 Red，而突變后圖像 x? 的顏色是 Green。如果只是簡單教模型逆向上述馬爾可夫鏈，則可能會訓(xùn)練網(wǎng)絡(luò)將 Green 變成 Blue，即便可以直接訓(xùn)練網(wǎng)絡(luò)將 Green 變成 Red。

因此，為了創(chuàng)建更好的訓(xùn)練信號，可以計(jì)算目標(biāo)樹與突變樹之間的編輯路徑。該團(tuán)隊(duì)使用了一種大體上基于樹編輯距離（tree edit distance）的樹編輯路徑算法。廣義的樹編輯距離問題允許插入、刪除和替換任意節(jié)點(diǎn)。但這里的設(shè)定不同，對樹的編輯僅能在一個只允許小突變的動作空間中實(shí)現(xiàn)。

對于兩個樹 z_A 和 z_B，可以線性方式比較它們的句法結(jié)構(gòu)。對于已經(jīng)滿足 ≤ σ_small 的改變，就將其加入到突變列表中。對于 > σ_small 的改變，則尋找能降低兩個樹之間距離的首個突變。因此，對于任意兩個程序 z_A 和 z_B 而言，可以在 O (|z_A| + |z_B|) 時(shí)間內(nèi)計(jì)算出突變路徑的第一步。

價(jià)值網(wǎng)絡(luò)與搜索

該團(tuán)隊(duì)另外還訓(xùn)練了一個價(jià)值網(wǎng)絡(luò) v_? (x_A, x_B)，其輸入為兩張經(jīng)過渲染的圖像 x_A 和 x_B，預(yù)測的是生成這兩張圖像的底層程序之間的編輯距離。由于在訓(xùn)練期間已經(jīng)計(jì)算出了樹之間的編輯距離，因此對于任意一對經(jīng)過渲染的圖像，都能直接獲得它們的基本真值程序編輯距離，這就可以用于以監(jiān)督式方法訓(xùn)練該價(jià)值網(wǎng)絡(luò)。

使用該團(tuán)隊(duì)提出的上述新策略和新價(jià)值網(wǎng)絡(luò)，就可以為任意目標(biāo)圖像 x? 和隨機(jī)初始化的程序 z_t 執(zhí)行波束搜索。在每一次迭代中，都要維護(hù)搜索樹中一組有最有希望值的節(jié)點(diǎn)，并且僅擴(kuò)展這些節(jié)點(diǎn)。

架構(gòu)

圖 3 展示了新提出的神經(jīng)架構(gòu)的概況。

他們的去噪模型 q_? (z_{t?1}|z_t, x_t; x?) 使用的是 Tsimpoukelli et al. 在論文《Multimodal few-shot learning with frozen language models》中描述的視覺 - 語言模型。至于圖像編碼器，他們使用的是現(xiàn)成可用的 NF-ResNet-26 的實(shí)現(xiàn)；這是一種無歸一化器的卷積架構(gòu)，可避免使用 Batch-Norm 時(shí)的時(shí)間不穩(wěn)定問題。

該團(tuán)隊(duì)實(shí)現(xiàn)了一種定制化的 token 化器，其中使用了他們的 CFG 的端點(diǎn)為 token。該編輯模型的其余部分是一個小型的僅解碼器 Transformer。

他們還添加了另外兩種類型的 token：用作該模型的句子起點(diǎn) token 的 <EDIT> 以及允許模型在其上下文中引用位置的 token <POS x>。

給定一張當(dāng)前圖像，一張目標(biāo)圖像，和當(dāng)前一個已 token 化的程序，訓(xùn)練該 Transformer 模型使之能以自回歸方式預(yù)測編輯位置和替換文本。在進(jìn)行預(yù)測時(shí)，解碼過程受到語法的約束。該團(tuán)隊(duì)對預(yù)測 logit 進(jìn)行了掩蔽，使之僅包含能表示句法樹節(jié)點(diǎn)的編輯位置，以及僅得到對于所選編輯位置句法上有效的替換。

這里設(shè)置 σ_small = 2，這意味著只允許網(wǎng)絡(luò)產(chǎn)生少于兩個原語的編輯。對于訓(xùn)練數(shù)據(jù)，他們的做法是從 CFG 采樣一個無限的隨機(jī)表達(dá)式流。對于噪聲步數(shù) s，他們是從 [1, 5] 中隨機(jī)選取一個。樣本中有一定的比例 ρ 是完成隨機(jī)采樣新表達(dá)式作為突變表達(dá)式。他們使用單臺英偉達(dá) A6000 GPU 訓(xùn)練了三天時(shí)間。

實(shí)驗(yàn)

他們在 4 種特定領(lǐng)域的圖形語言上進(jìn)行了實(shí)驗(yàn)：CSG2D、CSG2D-Sketch、TinySVG、Rainbow。

所選用的基準(zhǔn)方法為 Ellis et al. 提出的《Write, execute, assess: Program synthesis with a repl》以及 Sharma et al. 提出的《CSGNet: Neural shape parser for constructive solid geometry》。

圖 4 比較了新方法與基準(zhǔn)方法的性能。

可以看到，在 CSG2D 和 TinySVG 環(huán)境中，新提出的樹擴(kuò)散策略明顯優(yōu)于之前方法的策略。如果組合使用波束搜索，該策略的性能還能進(jìn)一步提升，在解決問題時(shí)相比其它方法可以更少地調(diào)用渲染器。

下圖給出了新系統(tǒng)的一些成功示例以及基準(zhǔn)方法的輸出?？梢钥吹?，新系統(tǒng)能修復(fù)其它方法遺漏的較小問題。

如下視頻展示了兩個使用 CSG2D-Sketch 語言基于草圖恢復(fù)程序的示例，其表明觀察模型并不一定需要確定性的渲染；它也可由隨機(jī)的手繪圖像構(gòu)成。

為了理解這些新設(shè)計(jì)的影響，該團(tuán)隊(duì)也在簡化的 Rainbow 環(huán)境中使用一個更小的 Transformer 模型進(jìn)行了消融實(shí)驗(yàn)，結(jié)果見圖 5?？傮w而言，這些設(shè)計(jì)選擇的效果得到了證明。

更多細(xì)節(jié)內(nèi)容請參考原論文。

本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/xolhLw1WYWpOXw2YGkEfEA??