編譯 | 洪雨欣

編輯 | 陳彩嫻

3 月 4 日，MIT 計(jì)算機(jī)系的教授 Peter Holderrieth 和 Ezra Erives 在 YouTube 上傳了其最新課程系列“Generative AI with Stochastic Differential Equations”，從數(shù)學(xué)的角度探討了當(dāng)前 AIGC 領(lǐng)域應(yīng)用最廣泛的兩種算法：去噪擴(kuò)散模型和流匹配。

生成模型將噪聲逐步轉(zhuǎn)換為數(shù)據(jù)來生成對(duì)象，這一系列的演化過程可以通過模擬常微分方程（ODE）或隨機(jī)微分方程（SDE）實(shí)現(xiàn)。此外，Holderrieth 和 Erives 還討論了如何利用深度神經(jīng)網(wǎng)絡(luò)大規(guī)模構(gòu)建、訓(xùn)練和模擬這些微分方程。

該課程系列一共有 6 個(gè)章節(jié)，AI 科技評(píng)論針對(duì)原視頻作了不改原意的整理如下：

一、流模型和擴(kuò)散模型

先讓我們定義一下“generate”。

比如我們想要生成一張狗的圖片，在這一過程中，存在著一系列“可能是狗”的圖片。

在機(jī)器學(xué)習(xí)中，通常將這種“可能合適”圖片的多樣性視為一個(gè)概率分布，我們稱之為數(shù)據(jù)分布，這個(gè)分布會(huì)賦予那些看起來更像狗的圖片更高的可能性。所以，圖片的準(zhǔn)確度會(huì)被其在數(shù)據(jù)分布下的可能性所取代。

因此，我們可以從數(shù)學(xué)上將生成任務(wù)表達(dá)為從數(shù)據(jù)分布中進(jìn)行采樣 。

在機(jī)器學(xué)習(xí)中，我們需要數(shù)據(jù)來訓(xùn)練模型，包括互聯(lián)網(wǎng)上公開可用的圖像、YouTube視頻、或者是蛋白質(zhì)數(shù)據(jù)庫(kù)。在生成建模中，一個(gè)數(shù)據(jù)集由來自數(shù)據(jù)分布的有限數(shù)量的樣本組成，我們將其表示為集合 1 到集合 n。

在許多情況下，我們會(huì)希望基于特定條件來生成一個(gè)對(duì)象。例如，我們希望生成“一只狗在覆蓋著雪的山坡上奔跑，背景是山脈”的圖像 。我們可以從條件數(shù)據(jù)分布中進(jìn)行采樣，其中 y 是一個(gè)條件變量 ，它可能是狗，可能是貓，也可能是風(fēng)景。條件生成模型通?？梢詫?duì)任意變量進(jìn)行條件設(shè)定，我們希望用不同的文本提示進(jìn)行條件設(shè)定，因此，我們要尋求一個(gè)能夠針對(duì)任意變量進(jìn)行條件設(shè)定的單一模型。

所以，生成模型可以總結(jié)為：從數(shù)據(jù)分布中生成樣本。

那么我們?cè)撊绾潍@取樣本呢？

假設(shè)我們可以從一個(gè)初始分布中獲取樣本，比如高斯分布，那么生成模型的目標(biāo)就是將初始分布中采樣得到的樣本轉(zhuǎn)換為數(shù)據(jù)分布的樣本。在大多數(shù)應(yīng)用中，我們都會(huì)將初始分布設(shè)定為一個(gè)簡(jiǎn)單的高斯分布 。

接下來，我們將描述如何通過微分方程模擬來獲得流模型和擴(kuò)散模型。

首先，我們可以通過常微分方程（ODE）構(gòu)建一個(gè)流模型，我們的目標(biāo)是將一個(gè)初始分布轉(zhuǎn)換為數(shù)據(jù)分布。常微分方程的存在性與唯一性定理是微分方程理論中的一個(gè)基本結(jié)果，所以在機(jī)器學(xué)習(xí)的實(shí)際過程中，常微分方程的解不僅存在，而且是唯一的。

在算法1中，我們總結(jié)了從流模型中進(jìn)行采樣的過程：

向量場(chǎng)是帶有參數(shù)的神經(jīng)網(wǎng)絡(luò)。目前，我們將向量場(chǎng)看作是一個(gè)通用的神經(jīng)網(wǎng)絡(luò)，即一個(gè)帶有參數(shù)的連續(xù)函數(shù)。要注意的一點(diǎn)是，神經(jīng)網(wǎng)絡(luò)是對(duì)向量場(chǎng)進(jìn)行參數(shù)化，而非對(duì)流進(jìn)行參數(shù)化。

接下來，我們可以通過隨機(jī)微分方程（SDE），用相同的方式來構(gòu)建一個(gè)生成模型。SDE是通過布朗運(yùn)動(dòng)構(gòu)建而成的，布朗運(yùn)動(dòng)是一項(xiàng)源于對(duì)物理擴(kuò)散過程研究的基本隨機(jī)過程，相當(dāng)于一種連續(xù)的隨機(jī)游走。與處理常微分方程類似，我們使用隨機(jī)初始分布來模擬隨機(jī)微分方程。為了參數(shù)化這個(gè)SDE，我們可以簡(jiǎn)單地參數(shù)化其核心組成部分———向量場(chǎng)。

在算法2中，我們總結(jié)了從SDE中進(jìn)行采樣的過程：

一個(gè)擴(kuò)散模型由一個(gè)神經(jīng)網(wǎng)絡(luò)和固定的擴(kuò)散系數(shù)組成，其中神經(jīng)網(wǎng)絡(luò)的參數(shù)可以用于參數(shù)化一個(gè)向量場(chǎng) 。當(dāng)擴(kuò)散系數(shù)為0時(shí)，擴(kuò)散模型就成為了一個(gè)流模型 。

二、搭建訓(xùn)練目標(biāo)

我們將通過最小化均方誤差來訓(xùn)練目標(biāo)。首先，我們需要設(shè)置一個(gè)訓(xùn)練目標(biāo)。像神經(jīng)網(wǎng)絡(luò)一樣，訓(xùn)練目標(biāo)本身也應(yīng)該是一個(gè)向量場(chǎng)。除此之外，它應(yīng)該做到我們期望神經(jīng)網(wǎng)絡(luò)能做到的事情：將噪聲轉(zhuǎn)換為數(shù)據(jù)。

構(gòu)建訓(xùn)練目標(biāo)的第一步是指定一條概率路徑。直觀來講，概率路徑規(guī)定了從初始噪聲分布到數(shù)據(jù)分布的逐步插值過程。

條件概率路徑是一組關(guān)于數(shù)據(jù)點(diǎn)的隨機(jī)分布，是針對(duì)單個(gè)數(shù)據(jù)點(diǎn)的性質(zhì)或行為。換句話說，條件概率路徑會(huì)逐漸將單個(gè)數(shù)據(jù)點(diǎn)轉(zhuǎn)換為初始分布。而邊緣概率路徑則是跨越整個(gè)數(shù)據(jù)點(diǎn)分布的性質(zhì)或行為，你可以將概率路徑看成分布空間中的一條軌跡，每個(gè)條件概率路徑都會(huì)誘導(dǎo)出一個(gè)邊緣概率路徑。

去噪擴(kuò)散模型所使用的是高斯概率路徑，條件概率路徑會(huì)對(duì)高斯分布和單個(gè)數(shù)據(jù)點(diǎn)對(duì)應(yīng)的隨機(jī)分布進(jìn)行插值 ，邊緣概率路徑會(huì)對(duì)整個(gè)數(shù)據(jù)點(diǎn)的分布進(jìn)行插值。

現(xiàn)在我們通過使用定義的概率路徑來構(gòu)建流模型的訓(xùn)練目標(biāo)。

每個(gè)數(shù)據(jù)點(diǎn)的訓(xùn)練目標(biāo)表示一個(gè)條件向量場(chǎng)，定義為ODE產(chǎn)生條件概率路徑。我們利用ODE模擬概率路徑，藍(lán)色背景為數(shù)據(jù)分布，紅色背景為高斯分布。上排是條件路徑，下排是邊緣概率路徑。可以看出，條件向量場(chǎng)沿著條件概率路徑，邊緣向量場(chǎng)沿著邊緣概率路徑。

我們可以用連續(xù)方程來證明這一點(diǎn)。讓我們考慮一個(gè)具有向量場(chǎng)訓(xùn)練目標(biāo)的流模型，其中初始隨機(jī)變量服從初始分布。連續(xù)性方程的公式為：在點(diǎn) x 處概率質(zhì)量的變化等于向量場(chǎng)u處概率質(zhì)量的流出減去流入。

我們剛剛成功地為流模型構(gòu)建了一個(gè)訓(xùn)練目標(biāo)，我們可以通過Fokker - Planck方程將連續(xù)性方程從ODE拓展到SDE。

如示意圖所示，我們使用ODE來模擬概率路徑。藍(lán)色背景為數(shù)據(jù)分布，紅色背景為高斯分布。上排是條件路徑，下排是邊緣概率路徑?？梢钥闯觯琒DE將初始分布的樣本傳輸?shù)綏l件路徑的樣本和邊緣路徑的樣本。

對(duì)于擴(kuò)散系數(shù)，我們可以構(gòu)造一個(gè)遵循相同概率路徑的SDE。在這個(gè)場(chǎng)景下，用條件概率路徑和條件向量場(chǎng)替換邊緣概率和邊緣向量場(chǎng)，相同的結(jié)論仍然成立，所以我們可以借助條件得分函數(shù)來表示邊緣得分函數(shù)，這樣該隨機(jī)微分方程就如所期望的那樣“將噪聲轉(zhuǎn)換為數(shù)據(jù)”。

三、訓(xùn)練流模型和擴(kuò)散模型

像之前的流模型一樣，我們希望神經(jīng)網(wǎng)絡(luò)等于邊際向量場(chǎng)。換句話說，我們希望減小神經(jīng)網(wǎng)絡(luò)和邊際向量場(chǎng)之間的均方誤差。

首先，抽取一個(gè)隨機(jī)時(shí)間。其次，我們從數(shù)據(jù)集中抽取一個(gè)隨機(jī)點(diǎn)，從邊際概率路徑中進(jìn)行采樣，例如添加一些噪聲，并計(jì)算神經(jīng)網(wǎng)絡(luò)。最后，計(jì)算神經(jīng)網(wǎng)絡(luò)的輸出與邊際向量場(chǎng)之間的均方誤差。我們將利用條件速度場(chǎng)的可處理性定義條件流匹配損失。在這里，我們使用條件向量場(chǎng)而不是邊際向量場(chǎng)，這是因?yàn)槲覀冇脳l件向量場(chǎng)的解析公式最小化上述損失。

一旦神經(jīng)網(wǎng)絡(luò)被訓(xùn)練好了，我們就能夠?qū)α髂Ｐ瓦M(jìn)行模擬以此來得到采樣，這套流程就是流匹配。

為了更靠近邊緣得分，我們可以使用一個(gè)稱之為得分網(wǎng)絡(luò)的神經(jīng)網(wǎng)絡(luò)。我們可以設(shè)計(jì)一個(gè)得分匹配損失和一個(gè)條件得分匹配損失。理想情況下我們希望最小化得分匹配損失，由于我們不知道邊緣得分，我們可以使用條件得分匹配損失來實(shí)現(xiàn)這一點(diǎn)。

訓(xùn)練完成后，我們可以選擇任意擴(kuò)散系數(shù)，然后模擬該隨機(jī)微分方程以生成樣本。理論上，在完美訓(xùn)練的情況下，每個(gè)擴(kuò)散系數(shù)都應(yīng)該給出服從真實(shí)數(shù)據(jù)分布的樣本。但在實(shí)際中，我們會(huì)遇到兩類錯(cuò)誤：第一，由于對(duì)SDE的模擬不完善而產(chǎn)生的數(shù)值誤差；第二，訓(xùn)練誤差，即訓(xùn)練模型并不完全等同于目標(biāo)模型。因此，存在一個(gè)最優(yōu)但未知的噪聲水平系數(shù)——這可以通過實(shí)際測(cè)試不同的值來憑經(jīng)驗(yàn)確定。

我們可以通過高斯概率路徑來訓(xùn)練得分匹配模型，條件得分匹配損失也被稱為去噪得分匹配，它是用于學(xué)習(xí)擴(kuò)散模型的最早程序之一。我們可以看到，當(dāng)beta接近0時(shí)，條件得分損失在數(shù)值上是不穩(wěn)定的，也就是說，只有添加足夠多的噪聲，去噪得分匹配才有效。

雖然流匹配僅允許通過ODE進(jìn)行確定性的模擬過程，但去噪擴(kuò)散模型允許進(jìn)行確定性（概率流常微分方程）或者隨機(jī)性（隨機(jī)微分方程采樣）的模擬。然而，不同于流匹配或隨機(jī)插值，后者能夠通過任意的概率路徑將初始分布轉(zhuǎn)換為數(shù)據(jù)分布。

去噪擴(kuò)散模型僅適用于高斯初始分布和高斯概率路徑。對(duì)于高斯概率路徑來說，我們無(wú)需分別訓(xùn)練向量場(chǎng)和得分網(wǎng)絡(luò)，它們可以在訓(xùn)練后相互轉(zhuǎn)換。

文獻(xiàn)中流行的擴(kuò)散模型的替代公式有：

離散時(shí)間：通常使用通過離散時(shí)間得出SDE的近似值。

倒置時(shí)間約定：一種流行的做法是采用倒置時(shí)間約定，其中時(shí)間對(duì)應(yīng)數(shù)據(jù)分布。

正向過程：正向過程（或加噪過程）是構(gòu)建高斯概率路徑的方法。

通過時(shí)間反轉(zhuǎn)構(gòu)建訓(xùn)練目標(biāo)：也可以通過對(duì)SDE進(jìn)行時(shí)間反轉(zhuǎn)來構(gòu)建訓(xùn)練目標(biāo)。

四、搭建圖像生成器

假如要構(gòu)建Stable Diffusion 3和Meta Movie Gen Video這樣的圖像生成模型，我們首先需要構(gòu)建條件生成引導(dǎo)機(jī)制。我們還將了解無(wú)分類器引導(dǎo)，這是一種用于提高條件生成質(zhì)量的流行技術(shù)。

我們選擇使用“引導(dǎo)”一詞來代替“條件”，以指代基于y進(jìn)行條件化的操作。引導(dǎo)式生成建模的目標(biāo)是能夠針對(duì)任意 y 從數(shù)據(jù)分布中進(jìn)行采樣。

在流匹配和得分匹配的語(yǔ)言體系中，我們定義一個(gè)引導(dǎo)擴(kuò)散模型是由一個(gè)引導(dǎo)向量場(chǎng)和一個(gè)時(shí)間相關(guān)的擴(kuò)散系數(shù)共同組成的。如果我們?cè)O(shè)想固定標(biāo)簽y的取值，那么我們就回到了無(wú)引導(dǎo)的生成問題。我們可以相應(yīng)地使用條件流匹配目標(biāo)來構(gòu)建一個(gè)生成模型。

請(qǐng)注意，因?yàn)閥不影響條件概率路徑或條件向量場(chǎng)，所以在所有 y 以及所有時(shí)間里面都可以得到一個(gè)引導(dǎo)條件流匹配目標(biāo)。引導(dǎo)目標(biāo)與無(wú)引導(dǎo)目標(biāo)的主要區(qū)別之一在于，引導(dǎo)目標(biāo)是從數(shù)據(jù)分布中采樣z和y，而不僅僅是采樣z。原因在于，我們的數(shù)據(jù)分布現(xiàn)在原則上是一個(gè)關(guān)于圖像和文本提示的聯(lián)合分布，這樣會(huì)讓生成過程更準(zhǔn)確。

雖然上述條件訓(xùn)練程序在理論上是有效的，但這種程序得到的圖像樣本與期望標(biāo)簽的擬合度不夠好。由于人為強(qiáng)化引導(dǎo)變量 y 的影響可以提高感知質(zhì)量，這一見解被提煉成一種無(wú)分類器引導(dǎo)的技術(shù)，該技術(shù)在最先進(jìn)的擴(kuò)散模型領(lǐng)域得到了廣泛應(yīng)用。

為簡(jiǎn)單起見，我們?cè)诖藘H關(guān)注高斯概率路徑的情況。在一個(gè)模型中同時(shí)訓(xùn)練條件模型和無(wú)條件模型，這被稱為無(wú)分類器引導(dǎo)。算法8展示了如何將無(wú)分類器引導(dǎo)構(gòu)造擴(kuò)展到擴(kuò)散模型的環(huán)境中：

在擴(kuò)散模型的開發(fā)中，我們需要介紹一種特定類型的卷積神經(jīng)網(wǎng)絡(luò)U-Net。它最初是為圖像分割而設(shè)計(jì)的，其關(guān)鍵特征在于輸入和輸出都具有圖像的形狀，這使得它非常適合將向量場(chǎng)參數(shù)化。

U-Net由一系列編碼器Ei、相應(yīng)的一系列解碼器Di以及位于它們之間的一個(gè)潛在處理塊組成，我們將這個(gè)潛在處理塊稱為中編碼器。隨著輸入通過編碼器，其表示形式中的通道數(shù)量增加，而圖像的高度和寬度減小。編碼器和解碼器通常都由一系列卷積層（其間包含激活函數(shù)、池化操作等）組成。輸入在到達(dá)第一個(gè)編碼器塊之前，通常會(huì)先被送入一個(gè)初始預(yù)編碼塊以增加通道數(shù)量。

U-Net的一種替代方案是擴(kuò)散Transformer（DiTs），它摒棄了卷積操作，純粹使用注意力機(jī)制。擴(kuò)散Transformer基于視覺Transformer（ViTs），其主要思想本質(zhì)上是將圖像分割成多個(gè)圖像塊，對(duì)每個(gè)圖像塊進(jìn)行嵌入，然后在圖像塊之間進(jìn)行注意力計(jì)算。例如，Stable Diffusion 3就采用條件流匹配進(jìn)行訓(xùn)練，它將速度場(chǎng)參數(shù)化為一種改進(jìn)的DiT。

大規(guī)模應(yīng)用的一個(gè)常見問題是數(shù)據(jù)維度極高，導(dǎo)致消耗過多內(nèi)存。例如，我們想要生成分辨率為1000×10000像素的高分辨率圖像，這會(huì)產(chǎn)生100萬(wàn)個(gè)維度。

為了減少內(nèi)存使用量，一種常見的設(shè)計(jì)模式是在潛在空間中進(jìn)行操作，該潛在空間可被視為分辨率較低的數(shù)據(jù)的壓縮版本。

具體而言，通常的方法是將流模型或擴(kuò)散模型與自編碼器相結(jié)合。首先通過自編碼器將訓(xùn)練數(shù)據(jù)集編碼到潛在空間中，然后在潛在空間中訓(xùn)練流模型或擴(kuò)散模型。

采樣時(shí)，首先使用訓(xùn)練好的流模型或擴(kuò)散模型在潛在空間中進(jìn)行采樣，然后通過解碼器對(duì)輸出進(jìn)行解碼。

直觀地說，一個(gè)訓(xùn)練良好的自編碼器能夠過濾掉語(yǔ)義上無(wú)意義的細(xì)節(jié)，從而使生成模型能夠聚焦于重要的、感知上相關(guān)的特征。

到目前為止，幾乎所有用于圖像和視頻生成最先進(jìn)的方法都涉及在自編碼器的潛在空間中訓(xùn)練流模型或擴(kuò)散模型——這就是所謂的潛在擴(kuò)散模型。

然而，需要注意的是，在訓(xùn)練擴(kuò)散模型之前也需要訓(xùn)練自編碼器，模型的性能也取決于自編碼器將圖像壓縮到潛在空間以及恢復(fù)美觀圖像的能力。

Stable Diffusion 3 使用了我們?cè)谶@項(xiàng)研究中的條件流匹配目標(biāo)。正如他們的論文所概述的，他們對(duì)各種流和擴(kuò)散替代方案進(jìn)行了廣泛的測(cè)試，并發(fā)現(xiàn)流匹配表現(xiàn)最佳。在訓(xùn)練方面，它使用了無(wú)分類器指導(dǎo)訓(xùn)練。

為了增強(qiáng)文本條件作用，Stable Diffusion 3 利用了三種不同類型的文本嵌入，其中包括 CLIP 嵌入和谷歌 T5-XXL 編碼器預(yù)訓(xùn)練實(shí)例產(chǎn)生的序列輸出。

CLIP 嵌入提供了輸入文本的粗略、總體嵌入，而 T5 嵌入提供了更細(xì)粒度的上下文層次，使模型能夠關(guān)注條件文本的特定元素。為了適應(yīng)這些序列上下文嵌入，擴(kuò)散 Transformer不僅要關(guān)注圖像，還要關(guān)注文本嵌入，從而將條件能力從最初為 DiT 提出的方案擴(kuò)展到序列上下文嵌入。

這種修改后的 DiT 被稱為多模態(tài) DiT（MM-DiT）。他們最大的模型擁有 80 億個(gè)參數(shù)。

在采樣方面，他們使用 50 步法則（評(píng)估網(wǎng)絡(luò)50次），采用歐拉模擬方案，并使用 2.0-5.0 之間的無(wú)分類器指導(dǎo)權(quán)重。

Movie Gen Video利用具有相同CondOT路徑的條件流匹配目標(biāo)。與Stable Diffusion 3一樣，Movie Gen Video也在凍結(jié)的預(yù)訓(xùn)練自編碼器的潛在空間中運(yùn)行。

值得關(guān)注的是，為了減少內(nèi)存消耗，自編碼器對(duì)視頻來說比圖像更加重要——這就是為什么目前大多數(shù)視頻生成器在生成視頻的長(zhǎng)度方面相當(dāng)受限的原因。通過引入一個(gè)時(shí)間自編碼器（TAE）來處理增加的時(shí)間維度，該自編碼器將原始視頻映射到潛在空間。

為了適應(yīng)長(zhǎng)視頻，一種時(shí)間平鋪程序程序會(huì)將視頻切分成片段，每個(gè)片段分別編碼后拼接在一起。模型本身由一個(gè)類似DiT的主干網(wǎng)絡(luò)給出，xt沿時(shí)間和空間維度被分塊，然后將圖像塊傳遞給一個(gè)Transformer，Transformer會(huì)采用圖像塊之間的自注意力以及與語(yǔ)言模型嵌入的交叉注意力。

對(duì)于文本條件，Movie Gen Video采用了三種類型的文本嵌入：UL2嵌入，用于細(xì)粒度的基于文本的推理；ByT5嵌入，用于關(guān)注字符級(jí)細(xì)節(jié)；以及MetaCLIP嵌入，在共享的文本 - 圖像嵌入空間中進(jìn)行訓(xùn)練。他們最大的模型有300億個(gè)參數(shù)。

五、應(yīng)用于機(jī)器人的擴(kuò)散模型

擴(kuò)散模型會(huì)根據(jù)機(jī)器人的觀察結(jié)果進(jìn)行條件化，當(dāng)對(duì)未來軌跡上的航點(diǎn)進(jìn)行去噪時(shí)，這個(gè)航點(diǎn)上的每個(gè)點(diǎn)都是 10 赫茲的command，一旦完成了下一組推理，就會(huì)切換到新的組，所以這個(gè)擴(kuò)散過程是循環(huán)的。當(dāng)開始噪聲處理時(shí)，軌跡是隨機(jī)噪聲，擴(kuò)散模型會(huì)將其擴(kuò)散成連貫的東西，你可以用這個(gè)基本配方做很多事情。

我們使用RIS相機(jī)時(shí)，有些東西無(wú)法從靜態(tài)相機(jī)中輕易看到，但使用腕戴式相機(jī)時(shí)，它可以和場(chǎng)景相機(jī)之間共享信息，實(shí)際上腕戴式相機(jī)充當(dāng)了圖像隨機(jī)化的功能，它總是在移動(dòng)而且多樣性很高，這在很大程度上解決了靜態(tài)場(chǎng)景相機(jī)的問題。

我們通過安全層發(fā)送指令，然后傳輸給機(jī)器人，過程中會(huì)有一個(gè)正在填充的低級(jí)命令的緩沖區(qū)，未來的 16 或 32 個(gè)動(dòng)作將會(huì)轉(zhuǎn)儲(chǔ)到緩沖區(qū)，然后立即再次開始推理。如果在推理運(yùn)行之前執(zhí)行了緩沖區(qū)中一半的動(dòng)作，我們會(huì)覆蓋緩沖區(qū)中的內(nèi)容并繼續(xù)。如果機(jī)器上可能還有其他程序在運(yùn)行，可以在刷新緩沖區(qū)之前先執(zhí)行更遠(yuǎn)的命令。

機(jī)器人會(huì)把“手移動(dòng)到某個(gè)特定點(diǎn)”轉(zhuǎn)換成一個(gè)命令，命令的力度越大，施加的力度就越大，這里的命令聽起來像位置的移動(dòng)，但實(shí)際上卻是力度的感知。我們給它提供了一點(diǎn)歷史記錄，它可以查看之前命令了什么，在哪里命令的，兩者之間的增量也可以讓它推斷出力度。

我們?cè)谶@些觸覺傳感器上做了相當(dāng)多的工作，我們使用了幾種不同的版本，一種是用一個(gè)凝膠或充氣膜，后面有一個(gè)攝像頭或其他類似的傳感器。例如擰緊瓶蓋，實(shí)際上機(jī)器人很難感知到何時(shí)才算完全擰緊，裝上這些傳感器后，當(dāng)感覺到瓶子變緊時(shí)，它自然就會(huì)停止。

擴(kuò)散模型非常巧妙，因?yàn)樗梢詫?duì)多模態(tài)分布進(jìn)行建模，所以在圖像里，這些模型通常不會(huì)模糊模式之間的界限，你會(huì)得到一個(gè)清晰的樣本。目前我們面對(duì)的問題是長(zhǎng)視界多模態(tài)（Long Horizon Modality），這些較長(zhǎng)的步驟可能具有不同的順序，擴(kuò)散模型在這方面表現(xiàn)不佳，因?yàn)樗鼈儧]有足夠的歷史記錄來了解它們應(yīng)該采用哪種模式。

如果機(jī)器人操縱的物體不同，可以在訓(xùn)練數(shù)據(jù)中顯示這種變化，這樣就不會(huì)每次教一些新東西時(shí)都需要重新建立它的認(rèn)知，我們?cè)谠椭幸部梢允謩?dòng)添加更多數(shù)據(jù)去修補(bǔ)錯(cuò)誤，但這樣做會(huì)比較麻煩。

想提高穩(wěn)健性和性能，就需要使用大量數(shù)據(jù)進(jìn)行多任務(wù)學(xué)習(xí)，這是一個(gè)非常有效的方法，唯一的問題機(jī)器人數(shù)據(jù)的缺乏，我們得從多樣化的來源獲取更多的數(shù)據(jù)。當(dāng)訓(xùn)練數(shù)據(jù)足夠大，特征空間足夠豐富時(shí)，我們就能做出很多精彩的東西。

我們還做了一些流匹配的內(nèi)部工作，確保匹配過程不會(huì)改變機(jī)器人的行為。

六、應(yīng)用于蛋白質(zhì)設(shè)計(jì)的擴(kuò)散模型

我將討論蛋白質(zhì)生成的擴(kuò)散模型和蛋白質(zhì)設(shè)計(jì)的下游模型。

為什么我們需要AI來生成蛋白質(zhì)？

打個(gè)比方，我們要設(shè)計(jì)一種叫做抗體的蛋白質(zhì)，這些蛋白質(zhì)會(huì)攻擊入侵身體的病毒。如果我們能使用AI立即產(chǎn)生抗體，事情便會(huì)簡(jiǎn)單很多。數(shù)據(jù)顯示，研發(fā)一種藥物大約需要10年時(shí)間，約26億美元，AI也許可以減少藥物研發(fā)的時(shí)間和精力。

一個(gè)生成蛋白質(zhì)的算法是從結(jié)構(gòu)到序列的逆向而行，這是一種基于神經(jīng)網(wǎng)絡(luò)的模型。如果要真正生產(chǎn)和合成蛋白質(zhì)，需要?jiǎng)?chuàng)建一個(gè)結(jié)構(gòu)并篩選出一個(gè)序列，擴(kuò)散模型實(shí)現(xiàn)了這一步。

我們希望人工智能可以快速產(chǎn)出新的分子庫(kù)，并在實(shí)驗(yàn)中進(jìn)行高效的微調(diào) 。如果能做到這一點(diǎn)，我們就可以真正將生物學(xué)與人工智能聯(lián)系起來，并擁有這個(gè)迭代循環(huán)的過程。一旦我們擁有結(jié)構(gòu)生成的基本算法，我們就可以利用它們進(jìn)行下游蛋白質(zhì)的設(shè)計(jì)任務(wù)。

那么我們?nèi)绾螛?gòu)建蛋白質(zhì)結(jié)構(gòu)模型呢？我們有三個(gè)選擇。第一，給每個(gè)原子構(gòu)建一個(gè)3D坐標(biāo)，這樣做的好處是我們可以精準(zhǔn)控制每個(gè)原子的位置。壞處是，鍵是不固定的，導(dǎo)致模型難以擴(kuò)展，性能不佳。

第二種選擇是只構(gòu)建扭轉(zhuǎn)角，對(duì)角度進(jìn)行擴(kuò)散，這就是所謂的杠桿臂效應(yīng)。這樣做可以固定鍵，但是問題又來了，我們沒辦法控制原子的位置。

現(xiàn)在來看第三種選擇，即沿著鏈建?？蚣堋＿@樣4個(gè)鍵當(dāng)中可以固定3個(gè)，同時(shí)還能精確控制框架的位置。

我們可以用 SE（3）來表示固定鍵，這是一個(gè)特殊的unclean group，它既有平移又有旋轉(zhuǎn)的功能。我們可以通過將兩個(gè)向量從中心坐標(biāo)移到其他兩個(gè)坐標(biāo)。這樣既保持了鍵的固定，又可以隨意移動(dòng)原子的位置。

但是，當(dāng)你將原子組合在一起時(shí)，你只是保存了一個(gè)常數(shù)因子，所以都會(huì)產(chǎn)生擴(kuò)展方面的問題。蛋白質(zhì)框架的擴(kuò)散，就好比我們將元素組合在一起子，從噪聲開始，然后是旋轉(zhuǎn)和平移元素，再到噪聲和數(shù)據(jù)空間之間移動(dòng)。這四個(gè)過程由微分方程控制，并有反向采樣的過程。

這里有一點(diǎn)要強(qiáng)調(diào)的是，我們?cè)诶杪餍紊线M(jìn)行擴(kuò)散時(shí)必須弄清楚如何進(jìn)行旋轉(zhuǎn)，比如在 SE（3）上旋轉(zhuǎn)。當(dāng)你給蛋白質(zhì)添加噪音時(shí)，會(huì)產(chǎn)生了一個(gè)獨(dú)特的分布——SE（3）公式上的各向同性作用。當(dāng)你向它添加更多的噪音時(shí)，它會(huì)開始在黎曼流形周圍采樣。

所謂的幀擴(kuò)散就是以這種方式參數(shù)化蛋白質(zhì)，我們將用噪音來破壞它們，然后訓(xùn)練一個(gè)模型來消除這種噪音，使其恢復(fù)到原始狀態(tài)，最后從純?cè)胍糸_始，使用神經(jīng)網(wǎng)絡(luò)迭代地采樣數(shù)據(jù)。

在模型架構(gòu)中，我們不僅要關(guān)心標(biāo)記，還要關(guān)心每個(gè)殘基和每幀之間是如何互相作用的。當(dāng)從擴(kuò)散模型中采樣時(shí)，想要保持SE（3）的不變性，每個(gè)擴(kuò)散模型都得是一個(gè)等變分?jǐn)?shù)模型，這樣你的分?jǐn)?shù)就會(huì)隨著蛋白質(zhì)的旋轉(zhuǎn)而旋轉(zhuǎn)。如果我從一些噪聲里采樣，最終將其解碼成蛋白質(zhì)，我可能會(huì)旋轉(zhuǎn)這個(gè)噪音，而最終的蛋白質(zhì)也可能會(huì)旋轉(zhuǎn)。

從純?cè)胍糸_始生成一個(gè)隨機(jī)蛋白質(zhì)，模型會(huì)解碼產(chǎn)生一組不同的蛋白質(zhì)。

如何檢查這些蛋白質(zhì)是否良好？

我們可以通過實(shí)際性檢查來看人工智能生成結(jié)構(gòu)的序列。假設(shè)這個(gè)新蛋白質(zhì)是由擴(kuò)散模型生成的，我們將為它計(jì)算一個(gè)序列，我們會(huì)使用神經(jīng)網(wǎng)絡(luò)來預(yù)測(cè)實(shí)際結(jié)構(gòu)，這是我們目前評(píng)估 AI 生成結(jié)構(gòu)的主要方法。

如果不想每次都只對(duì)同一種蛋白質(zhì)進(jìn)行采樣，有一種方法是執(zhí)行聚類算法，只需取兩個(gè)蛋白質(zhì)計(jì)算，將其標(biāo)準(zhǔn)化為某個(gè)長(zhǎng)度，然后執(zhí)行算法。

當(dāng)你得到更長(zhǎng)的蛋白質(zhì)時(shí)，性能可能會(huì)變得更差，但模型實(shí)際上能夠超越訓(xùn)練集進(jìn)行推廣，通過學(xué)習(xí)等變量擴(kuò)散模型，從而對(duì)訓(xùn)練內(nèi)容非常不同的東西進(jìn)行采樣，這是一個(gè)非常好的跡象。

在初始化時(shí)，你必須確定模型要使用多少個(gè)token，然后給它的位置編碼，比如從1 到 200 或 1 到 400 的位置編碼，模型會(huì)查看位置編碼并決定蛋白質(zhì)生成的長(zhǎng)度。位置編碼和注意力機(jī)制也非常重要，因?yàn)閠oken會(huì)在擴(kuò)散過程中進(jìn)行自我排序。

接下來是用于蛋白質(zhì)生成的RF擴(kuò)散模型。我們都聽說過預(yù)訓(xùn)練對(duì)語(yǔ)言模型的重要性，所以我們采用預(yù)訓(xùn)練的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)神經(jīng)網(wǎng)絡(luò)，然后用擴(kuò)散對(duì)其進(jìn)行微調(diào)。

預(yù)訓(xùn)練實(shí)際上非常關(guān)鍵，因?yàn)樗黾恿松傻拈L(zhǎng)度。之前只能達(dá)到 300 的長(zhǎng)度，通過有預(yù)訓(xùn)練的RF2模型可以達(dá)到800 或甚至 1000 的長(zhǎng)度，并生成大量蛋白質(zhì)。這個(gè)模型有一個(gè)0.5的臨界點(diǎn)，如果低于這個(gè)值，那么生成的結(jié)果會(huì)與訓(xùn)練集中的任何東西都不太相似。最重要的一點(diǎn)是，RF2是針對(duì)條件生成進(jìn)行訓(xùn)練的。

現(xiàn)在我們來講對(duì)稱生成，如果添加對(duì)稱噪聲，模型將會(huì)生成對(duì)稱蛋白質(zhì)。

比如，這里有一個(gè)病毒，它會(huì)生成一個(gè)與它結(jié)合的蛋白質(zhì)。你也可以給它一個(gè)二維指令讓其生成一個(gè)特定的蛋白質(zhì)?；蛘呖梢宰鲆粋€(gè)修復(fù)任務(wù)，讓模型只生成一部分蛋白質(zhì)。你還可以將多個(gè)條件結(jié)合在一起，添加一些額外的術(shù)語(yǔ)，讓生成偏離它所訓(xùn)練的內(nèi)容，轉(zhuǎn)而偏向更理想的蛋白質(zhì)。

但在過程中，模型可能會(huì)生成一種蛋白質(zhì)，這種蛋白質(zhì)可能會(huì)與小分子發(fā)生沖突。但通過對(duì)稱約束，你可以讓模型與分子完美結(jié)合。我們可以告訴模型，讓它能夠按照指令生成對(duì)稱的蛋白質(zhì)。

一個(gè)現(xiàn)有的蛋白質(zhì)部分，模型會(huì)遵循對(duì)稱約束來填充這個(gè)蛋白質(zhì)的其余部分。重要的一點(diǎn)是我們必須給它提供條件，模型將遵循這些指令在特定位置結(jié)合，否則模型將結(jié)合在某個(gè)隨機(jī)位置。

我們采用對(duì)稱生成用噪聲生成一種蛋白質(zhì)，可以看到蛋白質(zhì)的 3D 體積與RF擴(kuò)散的結(jié)果非常接近。

結(jié)合劑生產(chǎn)蛋白質(zhì)，之前實(shí)驗(yàn)的成功率幾乎只有0.1% 左右，但是通過擴(kuò)散模型，我們能夠達(dá)到30%的成功率。我們?cè)O(shè)計(jì)了一個(gè)流感病毒蛋白質(zhì)的結(jié)合劑，與AI 算法得出的結(jié)果非常接近。

那么，我應(yīng)該先做序列還是先做結(jié)構(gòu)呢？這就是我們?cè)噲D回答的問題。我們的研究結(jié)果是多流算法，它是一種平移旋轉(zhuǎn)的流匹配算法，也是一種離散空間的流匹配算法。

我對(duì)未來一兩年的預(yù)測(cè)是，擴(kuò)散模型的應(yīng)用將會(huì)超越蛋白質(zhì)生成，捕獲所有的分子。我們不僅要研究蛋白質(zhì)生成，還要研究小分子和其他類型的非蛋白質(zhì)物質(zhì)的生成。

還有一件重要的事情是，我們希望能夠從實(shí)驗(yàn)中迭代改進(jìn)，并制定一個(gè)隨著時(shí)間的推移會(huì)越來越好的模型，我們希望經(jīng)驗(yàn)?zāi)Ｐ团c實(shí)驗(yàn)數(shù)據(jù)保持一致，這可能會(huì)對(duì)新生物學(xué)產(chǎn)生重大影響。

原文鏈接：

https://www.youtube.com/channel/UCDcrzjRkgrhKocTgLwBZpMA

https://diffusion.csail.mit.edu/

雷峰網(wǎng)(公眾號(hào)：雷峰網(wǎng))

雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。