揭秘擴散模型背後的數學理念。

作者丨洪雨欣

編輯丨陳彩嫻

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是針對條件生成進行訓練的。

現在我們來講對稱生成，如果新增對稱噪聲，模型將會生成對稱蛋白質。

比如，這裡有一個病毒，它會生成一個與它結合的蛋白質。你也可以給它一個二維指令讓其生成一個特定的蛋白質。或者可以做一個修復任務，讓模型只生成一部分蛋白質。你還可以將多個條件結合在一起，新增一些額外的術語，讓生成偏離它所訓練的內容，轉而偏向更理想的蛋白質。

但在過程中，模型可能會生成一種蛋白質，這種蛋白質可能會與小分子發生衝突。但透過對稱約束，你可以讓模型與分子完美結合。我們可以告訴模型，讓它能夠按照指令生成對稱的蛋白質。

一個現有的蛋白質部分，模型會遵循對稱約束來填充這個蛋白質的其餘部分。重要的一點是我們必須給它提供條件，模型將遵循這些指令在特定位置結合，否則模型將結合在某個隨機位置。