0
編譯 | 洪雨欣
編輯 | 陳彩嫻
3 月 4 日，MIT 計算機系的教授 Peter Holderrieth 和 Ezra Erives 在 YouTube 上傳了其最新課程系列“Generative AI with Stochastic Differential Equations”，從數學的角度探討了當前 AIGC 領域應用最廣泛的兩種算法：去噪擴散模型和流匹配。
生成模型將噪聲逐步轉換為數據來生成對象，這一系列的演化過程可以通過模擬常微分方程（ODE）或隨機微分方程（SDE）實現。此外，Holderrieth 和 Erives 還討論了如何利用深度神經網絡大規模構建、訓練和模擬這些微分方程。
該課程系列一共有 6 個章節，AI 科技評論針對原視頻作了不改原意的整理如下：
一、流模型和擴散模型
先讓我們定義一下“generate”。
比如我們想要生成一張狗的圖片，在這一過程中，存在著一系列“可能是狗”的圖片。
在機器學習中，通常將這種“可能合適”圖片的多樣性視為一個概率分布，我們稱之為數據分布，這個分布會賦予那些看起來更像狗的圖片更高的可能性。所以，圖片的準確度會被其在數據分布下的可能性所取代。
因此，我們可以從數學上將生成任務表達為從數據分布中進行采樣 。
在機器學習中，我們需要數據來訓練模型，包括互聯網上公開可用的圖像、YouTube視頻、或者是蛋白質數據庫。在生成建模中，一個數據集由來自數據分布的有限數量的樣本組成，我們將其表示為集合 1 到集合 n。
在許多情況下，我們會希望基于特定條件來生成一個對象。例如，我們希望生成“一只狗在覆蓋著雪的山坡上奔跑，背景是山脈”的圖像 。我們可以從條件數據分布中進行采樣，其中 y 是一個條件變量 ，它可能是狗，可能是貓，也可能是風景。條件生成模型通常可以對任意變量進行條件設定，我們希望用不同的文本提示進行條件設定，因此，我們要尋求一個能夠針對任意變量進行條件設定的單一模型。
所以，生成模型可以總結為：從數據分布中生成樣本。
那么我們該如何獲取樣本呢？
假設我們可以從一個初始分布中獲取樣本，比如高斯分布，那么生成模型的目標就是將初始分布中采樣得到的樣本轉換為數據分布的樣本。在大多數應用中，我們都會將初始分布設定為一個簡單的高斯分布 。
接下來，我們將描述如何通過微分方程模擬來獲得流模型和擴散模型。
首先，我們可以通過常微分方程（ODE）構建一個流模型，我們的目標是將一個初始分布轉換為數據分布。常微分方程的存在性與唯一性定理是微分方程理論中的一個基本結果，所以在機器學習的實際過程中，常微分方程的解不僅存在，而且是唯一的。
在算法1中，我們總結了從流模型中進行采樣的過程：
向量場是帶有參數的神經網絡。目前，我們將向量場看作是一個通用的神經網絡，即一個帶有參數的連續函數。要注意的一點是，神經網絡是對向量場進行參數化，而非對流進行參數化。
接下來，我們可以通過隨機微分方程（SDE），用相同的方式來構建一個生成模型。SDE是通過布朗運動構建而成的，布朗運動是一項源于對物理擴散過程研究的基本隨機過程，相當于一種連續的隨機游走。與處理常微分方程類似，我們使用隨機初始分布來模擬隨機微分方程。為了參數化這個SDE，我們可以簡單地參數化其核心組成部分———向量場。
在算法2中，我們總結了從SDE中進行采樣的過程：
一個擴散模型由一個神經網絡和固定的擴散系數組成，其中神經網絡的參數可以用于參數化一個向量場 。當擴散系數為0時，擴散模型就成為了一個流模型 。
二、搭建訓練目標
我們將通過最小化均方誤差來訓練目標。首先，我們需要設置一個訓練目標。像神經網絡一樣，訓練目標本身也應該是一個向量場。除此之外，它應該做到我們期望神經網絡能做到的事情：將噪聲轉換為數據。
構建訓練目標的第一步是指定一條概率路徑。直觀來講，概率路徑規定了從初始噪聲分布到數據分布的逐步插值過程。
條件概率路徑是一組關于數據點的隨機分布，是針對單個數據點的性質或行為。換句話說，條件概率路徑會逐漸將單個數據點轉換為初始分布。而邊緣概率路徑則是跨越整個數據點分布的性質或行為，你可以將概率路徑看成分布空間中的一條軌跡，每個條件概率路徑都會誘導出一個邊緣概率路徑。
去噪擴散模型所使用的是高斯概率路徑，條件概率路徑會對高斯分布和單個數據點對應的隨機分布進行插值 ，邊緣概率路徑會對整個數據點的分布進行插值。
現在我們通過使用定義的概率路徑來構建流模型的訓練目標。
每個數據點的訓練目標表示一個條件向量場，定義為ODE產生條件概率路徑。我們利用ODE模擬概率路徑，藍色背景為數據分布，紅色背景為高斯分布。上排是條件路徑，下排是邊緣概率路徑。可以看出，條件向量場沿著條件概率路徑，邊緣向量場沿著邊緣概率路徑。
我們可以用連續方程來證明這一點。讓我們考慮一個具有向量場訓練目標的流模型，其中初始隨機變量服從初始分布。連續性方程的公式為：在點 x 處概率質量的變化等于向量場u處概率質量的流出減去流入。
我們剛剛成功地為流模型構建了一個訓練目標，我們可以通過Fokker - Planck方程將連續性方程從ODE拓展到SDE。
如示意圖所示，我們使用ODE來模擬概率路徑。藍色背景為數據分布，紅色背景為高斯分布。上排是條件路徑，下排是邊緣概率路徑。可以看出，SDE將初始分布的樣本傳輸到條件路徑的樣本和邊緣路徑的樣本。
對于擴散系數，我們可以構造一個遵循相同概率路徑的SDE。在這個場景下，用條件概率路徑和條件向量場替換邊緣概率和邊緣向量場，相同的結論仍然成立，所以我們可以借助條件得分函數來表示邊緣得分函數，這樣該隨機微分方程就如所期望的那樣“將噪聲轉換為數據”。
三、訓練流模型和擴散模型
像之前的流模型一樣，我們希望神經網絡等于邊際向量場。換句話說，我們希望減小神經網絡和邊際向量場之間的均方誤差。
首先，抽取一個隨機時間。其次，我們從數據集中抽取一個隨機點，從邊際概率路徑中進行采樣，例如添加一些噪聲，并計算神經網絡。最后，計算神經網絡的輸出與邊際向量場之間的均方誤差。我們將利用條件速度場的可處理性定義條件流匹配損失。在這里，我們使用條件向量場而不是邊際向量場，這是因為我們用條件向量場的解析公式最小化上述損失。
一旦神經網絡被訓練好了，我們就能夠對流模型進行模擬以此來得到采樣，這套流程就是流匹配。
為了更靠近邊緣得分，我們可以使用一個稱之為得分網絡的神經網絡。我們可以設計一個得分匹配損失和一個條件得分匹配損失。理想情況下我們希望最小化得分匹配損失，由于我們不知道邊緣得分，我們可以使用條件得分匹配損失來實現這一點。
訓練完成后，我們可以選擇任意擴散系數，然后模擬該隨機微分方程以生成樣本。理論上，在完美訓練的情況下，每個擴散系數都應該給出服從真實數據分布的樣本。但在實際中，我們會遇到兩類錯誤：第一，由于對SDE的模擬不完善而產生的數值誤差；第二，訓練誤差，即訓練模型并不完全等同于目標模型。因此，存在一個最優但未知的噪聲水平系數——這可以通過實際測試不同的值來憑經驗確定。
我們可以通過高斯概率路徑來訓練得分匹配模型，條件得分匹配損失也被稱為去噪得分匹配，它是用于學習擴散模型的最早程序之一。我們可以看到，當beta接近0時，條件得分損失在數值上是不穩定的，也就是說，只有添加足夠多的噪聲，去噪得分匹配才有效。
雖然流匹配僅允許通過ODE進行確定性的模擬過程，但去噪擴散模型允許進行確定性（概率流常微分方程）或者隨機性（隨機微分方程采樣）的模擬。然而，不同于流匹配或隨機插值，后者能夠通過任意的概率路徑將初始分布轉換為數據分布。
去噪擴散模型僅適用于高斯初始分布和高斯概率路徑。對于高斯概率路徑來說，我們無需分別訓練向量場和得分網絡，它們可以在訓練后相互轉換。
文獻中流行的擴散模型的替代公式有：
離散時間：通常使用通過離散時間得出SDE的近似值。
倒置時間約定：一種流行的做法是采用倒置時間約定，其中時間對應數據分布。
正向過程：正向過程（或加噪過程）是構建高斯概率路徑的方法。
通過時間反轉構建訓練目標：也可以通過對SDE進行時間反轉來構建訓練目標。
四、搭建圖像生成器
假如要構建Stable Diffusion 3和Meta Movie Gen Video這樣的圖像生成模型，我們首先需要構建條件生成引導機制。我們還將了解無分類器引導，這是一種用于提高條件生成質量的流行技術。
我們選擇使用“引導”一詞來代替“條件”，以指代基于y進行條件化的操作。引導式生成建模的目標是能夠針對任意 y 從數據分布中進行采樣。
在流匹配和得分匹配的語言體系中，我們定義一個引導擴散模型是由一個引導向量場和一個時間相關的擴散系數共同組成的。如果我們設想固定標簽y的取值，那么我們就回到了無引導的生成問題。我們可以相應地使用條件流匹配目標來構建一個生成模型。
請注意，因為y不影響條件概率路徑或條件向量場，所以在所有 y 以及所有時間里面都可以得到一個引導條件流匹配目標。引導目標與無引導目標的主要區別之一在于，引導目標是從數據分布中采樣z和y，而不僅僅是采樣z。原因在于，我們的數據分布現在原則上是一個關于圖像和文本提示的聯合分布，這樣會讓生成過程更準確。
雖然上述條件訓練程序在理論上是有效的，但這種程序得到的圖像樣本與期望標簽的擬合度不夠好。由于人為強化引導變量 y 的影響可以提高感知質量，這一見解被提煉成一種無分類器引導的技術，該技術在最先進的擴散模型領域得到了廣泛應用。
為簡單起見，我們在此僅關注高斯概率路徑的情況。在一個模型中同時訓練條件模型和無條件模型，這被稱為無分類器引導。算法8展示了如何將無分類器引導構造擴展到擴散模型的環境中：
在擴散模型的開發中，我們需要介紹一種特定類型的卷積神經網絡U-Net。它最初是為圖像分割而設計的，其關鍵特征在于輸入和輸出都具有圖像的形狀，這使得它非常適合將向量場參數化。
U-Net由一系列編碼器Ei、相應的一系列解碼器Di以及位于它們之間的一個潛在處理塊組成，我們將這個潛在處理塊稱為中編碼器。隨著輸入通過編碼器，其表示形式中的通道數量增加，而圖像的高度和寬度減小。編碼器和解碼器通常都由一系列卷積層（其間包含激活函數、池化操作等）組成。輸入在到達第一個編碼器塊之前，通常會先被送入一個初始預編碼塊以增加通道數量。
U-Net的一種替代方案是擴散Transformer（DiTs），它摒棄了卷積操作，純粹使用注意力機制。擴散Transformer基于視覺Transformer（ViTs），其主要思想本質上是將圖像分割成多個圖像塊，對每個圖像塊進行嵌入，然后在圖像塊之間進行注意力計算。例如，Stable Diffusion 3就采用條件流匹配進行訓練，它將速度場參數化為一種改進的DiT。
大規模應用的一個常見問題是數據維度極高，導致消耗過多內存。例如，我們想要生成分辨率為1000×10000像素的高分辨率圖像，這會產生100萬個維度。
為了減少內存使用量，一種常見的設計模式是在潛在空間中進行操作，該潛在空間可被視為分辨率較低的數據的壓縮版本。
具體而言，通常的方法是將流模型或擴散模型與自編碼器相結合。首先通過自編碼器將訓練數據集編碼到潛在空間中，然后在潛在空間中訓練流模型或擴散模型。
采樣時，首先使用訓練好的流模型或擴散模型在潛在空間中進行采樣，然后通過解碼器對輸出進行解碼。
直觀地說，一個訓練良好的自編碼器能夠過濾掉語義上無意義的細節，從而使生成模型能夠聚焦于重要的、感知上相關的特征。
到目前為止，幾乎所有用于圖像和視頻生成最先進的方法都涉及在自編碼器的潛在空間中訓練流模型或擴散模型——這就是所謂的潛在擴散模型。
然而，需要注意的是，在訓練擴散模型之前也需要訓練自編碼器，模型的性能也取決于自編碼器將圖像壓縮到潛在空間以及恢復美觀圖像的能力。
Stable Diffusion 3 使用了我們在這項研究中的條件流匹配目標。正如他們的論文所概述的，他們對各種流和擴散替代方案進行了廣泛的測試，并發現流匹配表現最佳。在訓練方面，它使用了無分類器指導訓練。
為了增強文本條件作用，Stable Diffusion 3 利用了三種不同類型的文本嵌入，其中包括 CLIP 嵌入和谷歌 T5-XXL 編碼器預訓練實例產生的序列輸出。
CLIP 嵌入提供了輸入文本的粗略、總體嵌入，而 T5 嵌入提供了更細粒度的上下文層次，使模型能夠關注條件文本的特定元素。為了適應這些序列上下文嵌入，擴散 Transformer不僅要關注圖像，還要關注文本嵌入，從而將條件能力從最初為 DiT 提出的方案擴展到序列上下文嵌入。
這種修改后的 DiT 被稱為多模態 DiT（MM-DiT）。他們最大的模型擁有 80 億個參數。
在采樣方面，他們使用 50 步法則（評估網絡50次），采用歐拉模擬方案，并使用 2.0-5.0 之間的無分類器指導權重。
Movie Gen Video利用具有相同CondOT路徑的條件流匹配目標。與Stable Diffusion 3一樣，Movie Gen Video也在凍結的預訓練自編碼器的潛在空間中運行。
值得關注的是，為了減少內存消耗，自編碼器對視頻來說比圖像更加重要——這就是為什么目前大多數視頻生成器在生成視頻的長度方面相當受限的原因。通過引入一個時間自編碼器（TAE）來處理增加的時間維度，該自編碼器將原始視頻映射到潛在空間。
為了適應長視頻，一種時間平鋪程序程序會將視頻切分成片段，每個片段分別編碼后拼接在一起。模型本身由一個類似DiT的主干網絡給出，xt沿時間和空間維度被分塊，然后將圖像塊傳遞給一個Transformer，Transformer會采用圖像塊之間的自注意力以及與語言模型嵌入的交叉注意力。
對于文本條件，Movie Gen Video采用了三種類型的文本嵌入：UL2嵌入，用于細粒度的基于文本的推理；ByT5嵌入，用于關注字符級細節；以及MetaCLIP嵌入，在共享的文本 - 圖像嵌入空間中進行訓練。他們最大的模型有300億個參數。
五、應用于機器人的擴散模型
擴散模型會根據機器人的觀察結果進行條件化，當對未來軌跡上的航點進行去噪時，這個航點上的每個點都是 10 赫茲的command，一旦完成了下一組推理，就會切換到新的組，所以這個擴散過程是循環的。當開始噪聲處理時，軌跡是隨機噪聲，擴散模型會將其擴散成連貫的東西，你可以用這個基本配方做很多事情。
我們使用RIS相機時，有些東西無法從靜態相機中輕易看到，但使用腕戴式相機時，它可以和場景相機之間共享信息，實際上腕戴式相機充當了圖像隨機化的功能，它總是在移動而且多樣性很高，這在很大程度上解決了靜態場景相機的問題。
我們通過安全層發送指令，然后傳輸給機器人，過程中會有一個正在填充的低級命令的緩沖區，未來的 16 或 32 個動作將會轉儲到緩沖區，然后立即再次開始推理。如果在推理運行之前執行了緩沖區中一半的動作，我們會覆蓋緩沖區中的內容并繼續。如果機器上可能還有其他程序在運行，可以在刷新緩沖區之前先執行更遠的命令。
機器人會把“手移動到某個特定點”轉換成一個命令，命令的力度越大，施加的力度就越大，這里的命令聽起來像位置的移動，但實際上卻是力度的感知。我們給它提供了一點歷史記錄，它可以查看之前命令了什么，在哪里命令的，兩者之間的增量也可以讓它推斷出力度。
我們在這些觸覺傳感器上做了相當多的工作，我們使用了幾種不同的版本，一種是用一個凝膠或充氣膜，后面有一個攝像頭或其他類似的傳感器。例如擰緊瓶蓋，實際上機器人很難感知到何時才算完全擰緊，裝上這些傳感器后，當感覺到瓶子變緊時，它自然就會停止。
擴散模型非常巧妙，因為它可以對多模態分布進行建模，所以在圖像里，這些模型通常不會模糊模式之間的界限，你會得到一個清晰的樣本。目前我們面對的問題是長視界多模態（Long Horizon Modality），這些較長的步驟可能具有不同的順序，擴散模型在這方面表現不佳，因為它們沒有足夠的歷史記錄來了解它們應該采用哪種模式。
如果機器人操縱的物體不同，可以在訓練數據中顯示這種變化，這樣就不會每次教一些新東西時都需要重新建立它的認知，我們在原型中也可以手動添加更多數據去修補錯誤，但這樣做會比較麻煩。
想提高穩健性和性能，就需要使用大量數據進行多任務學習，這是一個非常有效的方法，唯一的問題機器人數據的缺乏，我們得從多樣化的來源獲取更多的數據。當訓練數據足夠大，特征空間足夠豐富時，我們就能做出很多精彩的東西。
我們還做了一些流匹配的內部工作，確保匹配過程不會改變機器人的行為。
六、應用于蛋白質設計的擴散模型
我將討論蛋白質生成的擴散模型和蛋白質設計的下游模型。
為什么我們需要AI來生成蛋白質？
打個比方，我們要設計一種叫做抗體的蛋白質，這些蛋白質會攻擊入侵身體的病毒。如果我們能使用AI立即產生抗體，事情便會簡單很多。數據顯示，研發一種藥物大約需要10年時間，約26億美元，AI也許可以減少藥物研發的時間和精力。
一個生成蛋白質的算法是從結構到序列的逆向而行，這是一種基于神經網絡的模型。如果要真正生產和合成蛋白質，需要創建一個結構并篩選出一個序列，擴散模型實現了這一步。
我們希望人工智能可以快速產出新的分子庫，并在實驗中進行高效的微調 。如果能做到這一點，我們就可以真正將生物學與人工智能聯系起來，并擁有這個迭代循環的過程。一旦我們擁有結構生成的基本算法，我們就可以利用它們進行下游蛋白質的設計任務。
那么我們如何構建蛋白質結構模型呢？我們有三個選擇。第一，給每個原子構建一個3D坐標，這樣做的好處是我們可以精準控制每個原子的位置。壞處是，鍵是不固定的，導致模型難以擴展，性能不佳。
第二種選擇是只構建扭轉角，對角度進行擴散，這就是所謂的杠桿臂效應。這樣做可以固定鍵，但是問題又來了，我們沒辦法控制原子的位置。
現在來看第三種選擇，即沿著鏈建模框架。這樣4個鍵當中可以固定3個，同時還能精確控制框架的位置。
我們可以用 SE（3）來表示固定鍵，這是一個特殊的unclean group，它既有平移又有旋轉的功能。我們可以通過將兩個向量從中心坐標移到其他兩個坐標。這樣既保持了鍵的固定，又可以隨意移動原子的位置。
但是，當你將原子組合在一起時，你只是保存了一個常數因子，所以都會產生擴展方面的問題。蛋白質框架的擴散，就好比我們將元素組合在一起子，從噪聲開始，然后是旋轉和平移元素，再到噪聲和數據空間之間移動。這四個過程由微分方程控制，并有反向采樣的過程。
這里有一點要強調的是，我們在黎曼流形上進行擴散時必須弄清楚如何進行旋轉，比如在 SE（3）上旋轉。當你給蛋白質添加噪音時，會產生了一個獨特的分布——SE（3）公式上的各向同性作用。當你向它添加更多的噪音時，它會開始在黎曼流形周圍采樣。
所謂的幀擴散就是以這種方式參數化蛋白質，我們將用噪音來破壞它們，然后訓練一個模型來消除這種噪音，使其恢復到原始狀態，最后從純噪音開始，使用神經網絡迭代地采樣數據。
在模型架構中，我們不僅要關心標記，還要關心每個殘基和每幀之間是如何互相作用的。當從擴散模型中采樣時，想要保持SE（3）的不變性，每個擴散模型都得是一個等變分數模型，這樣你的分數就會隨著蛋白質的旋轉而旋轉。如果我從一些噪聲里采樣，最終將其解碼成蛋白質，我可能會旋轉這個噪音，而最終的蛋白質也可能會旋轉。
從純噪音開始生成一個隨機蛋白質，模型會解碼產生一組不同的蛋白質。
如何檢查這些蛋白質是否良好？
我們可以通過實際性檢查來看人工智能生成結構的序列。假設這個新蛋白質是由擴散模型生成的，我們將為它計算一個序列，我們會使用神經網絡來預測實際結構，這是我們目前評估 AI 生成結構的主要方法。
如果不想每次都只對同一種蛋白質進行采樣，有一種方法是執行聚類算法，只需取兩個蛋白質計算，將其標準化為某個長度，然后執行算法。
當你得到更長的蛋白質時，性能可能會變得更差，但模型實際上能夠超越訓練集進行推廣，通過學習等變量擴散模型，從而對訓練內容非常不同的東西進行采樣，這是一個非常好的跡象。
在初始化時，你必須確定模型要使用多少個token，然后給它的位置編碼，比如從1 到 200 或 1 到 400 的位置編碼，模型會查看位置編碼并決定蛋白質生成的長度。位置編碼和注意力機制也非常重要，因為token會在擴散過程中進行自我排序。
接下來是用于蛋白質生成的RF擴散模型。我們都聽說過預訓練對語言模型的重要性，所以我們采用預訓練的蛋白質結構預測神經網絡，然后用擴散對其進行微調。
預訓練實際上非常關鍵，因為它增加了生成的長度。之前只能達到 300 的長度，通過有預訓練的RF2模型可以達到800 或甚至 1000 的長度，并生成大量蛋白質。這個模型有一個0.5的臨界點，如果低于這個值，那么生成的結果會與訓練集中的任何東西都不太相似。最重要的一點是，RF2是針對條件生成進行訓練的。
現在我們來講對稱生成，如果添加對稱噪聲，模型將會生成對稱蛋白質。
比如，這里有一個病毒，它會生成一個與它結合的蛋白質。你也可以給它一個二維指令讓其生成一個特定的蛋白質。或者可以做一個修復任務，讓模型只生成一部分蛋白質。你還可以將多個條件結合在一起，添加一些額外的術語，讓生成偏離它所訓練的內容，轉而偏向更理想的蛋白質。
但在過程中，模型可能會生成一種蛋白質，這種蛋白質可能會與小分子發生沖突。但通過對稱約束，你可以讓模型與分子完美結合。我們可以告訴模型，讓它能夠按照指令生成對稱的蛋白質。
一個現有的蛋白質部分，模型會遵循對稱約束來填充這個蛋白質的其余部分。重要的一點是我們必須給它提供條件，模型將遵循這些指令在特定位置結合，否則模型將結合在某個隨機位置。
我們采用對稱生成用噪聲生成一種蛋白質，可以看到蛋白質的 3D 體積與RF擴散的結果非常接近。
結合劑生產蛋白質，之前實驗的成功率幾乎只有0.1% 左右，但是通過擴散模型，我們能夠達到30%的成功率。我們設計了一個流感病毒蛋白質的結合劑，與AI 算法得出的結果非常接近。
那么，我應該先做序列還是先做結構呢？這就是我們試圖回答的問題。我們的研究結果是多流算法，它是一種平移旋轉的流匹配算法，也是一種離散空間的流匹配算法。
我對未來一兩年的預測是，擴散模型的應用將會超越蛋白質生成，捕獲所有的分子。我們不僅要研究蛋白質生成，還要研究小分子和其他類型的非蛋白質物質的生成。
還有一件重要的事情是，我們希望能夠從實驗中迭代改進，并制定一個隨著時間的推移會越來越好的模型，我們希望經驗模型與實驗數據保持一致，這可能會對新生物學產生重大影響。
原文鏈接：
https://www.youtube.com/channel/UCDcrzjRkgrhKocTgLwBZpMA
https://diffusion.csail.mit.edu/
雷峰網(公眾號：雷峰網)
雷峰網原創文章，未經授權禁止轉載。詳情見轉載須知。
