看張手繪草圖就能合成圖形程序，加州伯克利讓擴散模型掌握新技能
假設我們給模型一張手繪的「5」狀圖形，它就能通過不斷突變來修改程序，最終得到能輸出目標圖形的程序。
該模型來自加州大學伯克利分校的一個研究團隊，他們提出的這種程序合成新方法使用了神經擴散模型來直接操作句法樹。
論文一作為該校博士生 Shreyas Kapur，其導師為該校計算機科學教授 Stuart Russell。
- 論文標題：Diffusion On Syntax Trees For Program Synthesis
- 論文地址：https://arxiv.org/pdf/2405.20519
- 項目地址：https://tree-diffusion.github.io/
- 代碼庫：https://github.com/revalo/tree-diffusion?
擴散模型之前已經在圖像生成領域取得了巨大成功。而該團隊發(fā)現，通過利用擴散機制，可讓模型學會迭代地優(yōu)化程序，同時確保句法有效性（syntactic validity）。這種新方法的關鍵在于可讓模型觀察程序在每一步的輸出，從而實現有效的調試流程。
迭代的能力已經在 AlphaZero 等系統(tǒng)中得到了體現，而擴散機制的迭代性質也就很自然地會被用于基于搜索的程序合成。
該團隊發(fā)現，通過在訓練擴散模型的同時訓練一個價值模型，就可以引導其中的去噪過程得到能輸出所需結果的程序。這樣一來，便能高效地探索程序空間，在生成過程中的每一步都做出更明智的決策。
為了實現該方法，該團隊選擇了逆向圖形任務（inverse graphics task），即假定使用特定領域的語言來繪制圖像。
該團隊表示：「逆向圖形任務天然就適合我們的方法，因為代碼中的一點微小改變就能導致所得圖像出現有意義的語義變化?！?/p>
舉個例子，如果圖像中出現了一個放置錯誤的圖形，那就能被輕松看見并定位到程序空間中。圖 1 給出了一些示例。
這項研究的主要貢獻包括：
1. 提出一種在句法樹上使用擴散的全新方法
2. 針對逆向圖形任務實現了該方法，并發(fā)現新方法優(yōu)于之前的方法。
方法
簡而言之，該方法的核心思想是：為句法樹開發(fā)一個去噪擴散模型，類似于為視覺任務開發(fā)的圖像擴散模型。
首先來看一個來自 Ellis et al. 論文《Write, execute, assess: Program synthesis with a repl》中的任務示例：根據圖像生成一個構造實體幾何（CSG2D）程序。使用 CSG2D，我們可以使用加法和減法等布爾運算將圓和四邊形等簡單的原語組合起來，從而能使用下面的上下文無關語法（CFG）創(chuàng)建出更復雜的形狀：
在圖 2 中，z? 是目標程序，x? 是 z? 渲染過的版本。
任務就是逆轉 x?，從而恢復得到 z?。首先，去噪過程將 y=16 隨機突變?yōu)?y=10。然后再將左側有兩個形狀的子樹轉變成只有一個形狀的新子樹。這里的目標是基于圖像 x?，從 z? 和 x? 開始，訓練一個能逆向該去噪過程的神經網絡，從而得到 z?。
下面首先將介紹如何將「噪聲」添加到句法樹中，然后將詳細描述如何訓練逆向該噪聲的神經網絡，最后將描述如何使用該神經網絡來執(zhí)行搜索。
采樣微小突變
令 z_t 為時間 t 時的程序。令 p_N (z_{t+1}|z_t) 為將程序 z_t 隨機突變成 z_{t+1} 所基于的分布。這里希望 p_N 突變滿足兩點：(1) 很小，(2) 能得到句法有效的 z_{t+1}。
為此，該團隊探究了大量有關基于語法的模糊測試的計算機安全文獻。為了確保突變很小，他們首先定義了一個函數 σ(z)，其能給出程序 z 的「大小」。在實驗中，則是將 CFG 中的一組端點（terminal）定義為原語。
舉個例子，如果用他們的 CSG2D 語言來寫，上述原語就為 {Quad, Circle}。使用該語言時，該團隊的做法是讓 σ(z) = σ_primitive (z)，這能統(tǒng)計原語的數量。σ(z) 還可能包含深度、節(jié)點數量等選項。
然后，基于精確的約束條件 σ_min < σ(z) ≤ σ_max 從 CFG 隨機采樣程序。該函數被命名為 ConstrainedSample (σ_min, σ_max)。為 σ_max 設定一個較小的值就能隨機采樣小程序。在生成小突變時，設定 σ_max = σ_small。
為了讓給定程序 z 發(fā)生突變，首先可在其句法樹中生成一個在某個 σ_small 范圍內的候選節(jié)點集合：
然后，從該集合中均勻采樣一個突變節(jié)點：
由于能讀取整個句法樹和 CFG，因此知道了哪條生成規(guī)則能得到 m，并由此可以確保得到句法有效的突變。舉個例子，如果 m 是一個數值，那么替代它的也應該是一個數值。如果 m 是一個一般子表達式，那就可以把它替換成任何一般子表達式。因此，可采樣出 m'，這是 m 的替代：
策略
前向過程
該團隊將程序合成問題視為一個推理問題。令 p (x|z) 為觀察模型，其中 x 可以是任意類型的觀察。任務是逆轉這個觀察模型的方向，即給定某個觀察 x 得到一個程序 z。
首先，從一個數據集 D 取出某個程序 z?，或者這里的做法是從 CFG 隨機采樣一個程序。即采樣一個滿足 σ(z?) ≤ σ_max 的 z?。然后通過下式描述的過程向 z? 添加噪聲，執(zhí)行 s 步，其中 s ～ Uniform [1, s_max]，而 s_max 是一個超參數：
然后，訓練一個建模以下分布的條件神經網絡。
其中 ? 是該神經網絡的參數，z_t 是當前的程序，x_t 是該程序當前的輸出，x? 是需要求解的目標輸出。
逆向突變路徑
由于能夠獲取基本真值突變，因此可以簡單地通過前向過程馬爾可夫鏈 z? → z? →... 來逆向采樣的軌跡，從而生成用以訓練神經網絡的目標。乍一看，這或許是個合理選擇。但是，直接訓練模型逆向最后一次突變有可能為神經網絡創(chuàng)造出遠遠更有噪聲的信號。
舉個例子，在一個大得多的句法樹中，一種顏色的突變路徑為：
目標圖像 x? 的顏色是 Red，而突變后圖像 x? 的顏色是 Green。如果只是簡單教模型逆向上述馬爾可夫鏈，則可能會訓練網絡將 Green 變成 Blue，即便可以直接訓練網絡將 Green 變成 Red。
因此，為了創(chuàng)建更好的訓練信號，可以計算目標樹與突變樹之間的編輯路徑。該團隊使用了一種大體上基于樹編輯距離（tree edit distance）的樹編輯路徑算法。廣義的樹編輯距離問題允許插入、刪除和替換任意節(jié)點。但這里的設定不同，對樹的編輯僅能在一個只允許小突變的動作空間中實現。
對于兩個樹 z_A 和 z_B，可以線性方式比較它們的句法結構。對于已經滿足 ≤ σ_small 的改變，就將其加入到突變列表中。對于 > σ_small 的改變，則尋找能降低兩個樹之間距離的首個突變。因此，對于任意兩個程序 z_A 和 z_B 而言，可以在 O (|z_A| + |z_B|) 時間內計算出突變路徑的第一步。
價值網絡與搜索
該團隊另外還訓練了一個價值網絡 v_? (x_A, x_B)，其輸入為兩張經過渲染的圖像 x_A 和 x_B，預測的是生成這兩張圖像的底層程序之間的編輯距離。由于在訓練期間已經計算出了樹之間的編輯距離，因此對于任意一對經過渲染的圖像，都能直接獲得它們的基本真值程序編輯距離，這就可以用于以監(jiān)督式方法訓練該價值網絡。
使用該團隊提出的上述新策略和新價值網絡，就可以為任意目標圖像 x? 和隨機初始化的程序 z_t 執(zhí)行波束搜索。在每一次迭代中，都要維護搜索樹中一組有最有希望值的節(jié)點，并且僅擴展這些節(jié)點。
架構
圖 3 展示了新提出的神經架構的概況。
他們的去噪模型 q_? (z_{t?1}|z_t, x_t; x?) 使用的是 Tsimpoukelli et al. 在論文《Multimodal few-shot learning with frozen language models》中描述的視覺 - 語言模型。至于圖像編碼器，他們使用的是現成可用的 NF-ResNet-26 的實現；這是一種無歸一化器的卷積架構，可避免使用 Batch-Norm 時的時間不穩(wěn)定問題。
該團隊實現了一種定制化的 token 化器，其中使用了他們的 CFG 的端點為 token。該編輯模型的其余部分是一個小型的僅解碼器 Transformer。
他們還添加了另外兩種類型的 token：用作該模型的句子起點 token 的 <EDIT> 以及允許模型在其上下文中引用位置的 token <POS x>。
給定一張當前圖像，一張目標圖像，和當前一個已 token 化的程序，訓練該 Transformer 模型使之能以自回歸方式預測編輯位置和替換文本。在進行預測時，解碼過程受到語法的約束。該團隊對預測 logit 進行了掩蔽，使之僅包含能表示句法樹節(jié)點的編輯位置，以及僅得到對于所選編輯位置句法上有效的替換。
這里設置 σ_small = 2，這意味著只允許網絡產生少于兩個原語的編輯。對于訓練數據，他們的做法是從 CFG 采樣一個無限的隨機表達式流。對于噪聲步數 s，他們是從 [1, 5] 中隨機選取一個。樣本中有一定的比例 ρ 是完成隨機采樣新表達式作為突變表達式。他們使用單臺英偉達 A6000 GPU 訓練了三天時間。
實驗
他們在 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》。
圖 4 比較了新方法與基準方法的性能。
可以看到，在 CSG2D 和 TinySVG 環(huán)境中，新提出的樹擴散策略明顯優(yōu)于之前方法的策略。如果組合使用波束搜索，該策略的性能還能進一步提升，在解決問題時相比其它方法可以更少地調用渲染器。
下圖給出了新系統(tǒng)的一些成功示例以及基準方法的輸出?？梢钥吹剑孪到y(tǒng)能修復其它方法遺漏的較小問題。
如下視頻展示了兩個使用 CSG2D-Sketch 語言基于草圖恢復程序的示例，其表明觀察模型并不一定需要確定性的渲染；它也可由隨機的手繪圖像構成。
為了理解這些新設計的影響，該團隊也在簡化的 Rainbow 環(huán)境中使用一個更小的 Transformer 模型進行了消融實驗，結果見圖 5?？傮w而言，這些設計選擇的效果得到了證明。
更多細節(jié)內容請參考原論文。
本文轉自 機器之心 ，作者：機器之心
