跳到主要內容
paul.notebook

№ 007互動

一團雜訊怎麼長成一顆蘋果：在瀏覽器裡訓練 3D 擴散模型

選兩種水果，按下訓練。一個從零寫起的擴散模型會在你的瀏覽器裡，學會把一團彩色的 3D 雜訊點，一步一步收成那兩顆水果，連顏色一起。

發布
更新
閱讀時間
9 分鐘

畫圖的 AI 是從一張全是雜訊的圖開始，一點一點把雜訊擦掉，擦到最後剩下一張圖。這句話大家都聽過，但「擦掉雜訊」到底是什麼意思，光看成品看不出來。

所以這裡做一個小到可以在你面前現場訓練的版本。它不畫圖，它捏水果：三千多個彩色的點，從一團雜訊開始，收成你選的兩種水果。 為了讓你同時看到兩種，畫面上一次長三團：兩種水果各一團，再加上一團兩者的混合。

選好水果，按下「開始訓練」。

fig 01/diffusion / train

還沒訓練。現在它對水果一無所知，所以雜訊走完 40 步還是一團雜訊。

左邊長成
右邊長成
訓練步數
0
誤差
–
步／秒
–
畫面上有三團點雲，名字寫在各自的下方：兩邊是你選的兩種水果，標著「×」的那一團是同一個模型被要求「兩種各一半」時長出來的東西（後面會講）。三團每隔四五秒就從全新的雜訊重新開始，用模型當下的程度走 40 步。

一開始什麼都不會發生：沒訓練過的模型不知道水果是什麼，雜訊走完 40 步還是雜訊。按下訓練之後，先出現的是顏色：大約五秒，左邊那團就偏向一種顏色、右邊偏向另一種。十秒左右看得出形狀，半分鐘之後連蘋果的葉子都有了。

先講清楚這是什麼、不是什麼。這個模型不是「會畫水果」，它是把你選的這兩顆水果背起來了。 它這輩子只看過這兩樣東西，所以它生得出來的也只有這兩樣，以及介於兩者之間的東西。真的畫圖模型看過幾十億張圖，才有辦法畫出沒看過的組合。

但「怎麼背、怎麼從雜訊裡把它還原出來」的方法，和那些大模型用的是同一套。這篇要拆的就是那一套方法。每一次長出來的點都落在不同的位置，因為每次的起始雜訊都不同；模型記住的不是三千個點的座標，而是「水果在哪裡」。

它只學了一件很笨的事

訓練的時候，模型從頭到尾只做一種練習題：

  1. 從水果表面隨機拿一個點 x0x_0。
  2. 隨機挑一個雜訊等級 tt，往這個點上加雜訊。
  3. 把加完雜訊的點拿給模型看，問它：「我剛剛加的雜訊是什麼？」

第 2 步寫成式子只有一行：

xt=αˉt x0+1−αˉt ε,ε∼N(0,I)x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)

αˉt\bar\alpha_t 是一張事先排好的表：等級 0 的時候是 1（全部都是原本的點），等級 100 的時候幾乎是 0（全部都是雜訊）。拖動滑桿，看這一行式子對一顆水果做了什麼：

fig 02/diffusion / add noise
這裡沒有模型，只有上面那一行加法。雜訊等級 100 的時候，原本的水果一點都不剩。

注意這個式子不用一級一級加：要等級 73 的點，直接代 t=73t = 73 就有了。所以每一道練習題都很便宜，隨機抽一個點、隨機抽一個等級，馬上就能出題。

要模型學的，就是讓它的答案 εθ\varepsilon_\theta 和真正加進去的 ε\varepsilon 越接近越好：

L=Ex0, t, ε ∥ε−εθ(xt,t)∥2L = \mathbb{E}_{x_0,\, t,\, \varepsilon}\ \big\| \varepsilon - \varepsilon_\theta(x_t, t) \big\|^2

整個訓練就是這一個損失函數。沒有判別器，沒有對抗，就是一個迴歸問題。

它學到的其實是一張箭頭地圖

「猜雜訊」聽起來很抽象，但換個角度看就很具體。如果模型說「這個點裡面的雜訊是 ε\varepsilon」，那把雜訊拿掉，就是把這個點往 −ε-\varepsilon 的方向推。所以模型其實是在回答：站在空間裡的這個位置，水果在哪個方向？

下面把你訓練的模型切開來看。這是穿過第一種水果正中間的一個垂直切面，每個格點上的箭頭，就是模型會把站在那裡的點往哪邊推：

fig 03/diffusion / field

上面的模型還沒訓練過，所以箭頭是亂指的。訓練之後回來看，它們會轉向水果。

彩色的輪廓是真正的水果在這個切面上的樣子。箭頭越偏粉紅越長。格點的顏色輸入是這顆水果的平均色，再依雜訊等級變淡（和真的點在那個等級會有的顏色一致）。

外面的箭頭全部指向水果，這不意外。有意思的是把雜訊等級拉低之後（第一種水果選蘋果看得最清楚），水果裡面的箭頭會轉向，改成往外指向表面。模型知道這顆水果是空心的：點只存在於表面上，所以站在蘋果正中央的點，最近的「家」在外面。

雜訊等級高的時候則相反，所有箭頭都指向同一個地方，大約是水果的正中央。原因值得停下來想一下，因為它解釋了擴散模型幾乎所有的行為：

一個很吵的點，可能來自水果表面上的任何地方。模型不知道是哪一個，而它被訓練成讓平方誤差最小。讓平方誤差最小的答案，永遠是所有可能答案的平均，越可能的答案佔得越重。

水果表面所有點的平均，就是水果的正中央。所以在高雜訊的時候，模型能給的最好答案就是「往中間走」。它沒有學壞，這就是那道題的正確答案。

為什麼要走很多步

會回答「雜訊是什麼」之後，生成就是把這件事反過來做。從純雜訊 xTx_T 開始，每一步先問模型，然後算出模型此刻認為的成品：

x^0=xt−1−αˉt  εθ(xt,t)αˉt\hat{x}_0 = \frac{x_t - \sqrt{1 - \bar\alpha_t}\; \varepsilon_\theta(x_t, t)}{\sqrt{\bar\alpha_t}}

這只是把加雜訊那一行式子倒過來解 x0x_0。接著不要直接相信它，而是只往那個方向走一小段，走到一個比較低的雜訊等級 ss：

xs=αˉs  x^0+1−αˉs  εθ(xt,t)x_s = \sqrt{\bar\alpha_s}\; \hat{x}_0 + \sqrt{1 - \bar\alpha_s}\; \varepsilon_\theta(x_t, t)

然後在新的位置再問一次。

下面是你自己訓練的那個模型走的一次完整過程，可以前後拖。請特別試試右邊的「模型此刻猜的成品」：

fig 04/diffusion / steps
總共走幾步
畫面上顯示

…

左右各一種水果，和儀器 01 是同一個模型。「點現在的位置」是 x_t，「模型此刻猜的成品」是上面式子裡的 x̂₀。

把滑桿拉到第 4、5 步左右，切到「模型此刻猜的成品」。你會看到兩團糊糊的球，顏色大致對，形狀完全沒有。（最前面兩三步看到的不是球，而是一堆貼在方盒子邊上的點：那時 αˉt\bar\alpha_t 幾乎是 0，上面的式子等於除以零，模型只要猜錯一點點就被放大到 ±1.5 的上限。那是數值問題，不是模型的答案。）這不是壞掉，這是它此刻最好的答案，也就是上一節說的「平均」：雜訊還很大的時候，模型猜的成品就是整顆水果的平均，一團在正中央的、平均顏色的東西。

這也是為什麼不能一步到位。如果第一步就完全相信模型的猜測，得到的就是那團平均。走很多小步的意思是：每次只相信一點點，走到雜訊小一些的地方，那裡的「所有可能答案」範圍變窄了，平均起來就不那麼糊，於是下一次的猜測會更具體。形狀是這樣一輪一輪從糊變清楚的。

把「總共走幾步」調成 3 或 5，可以看到步數不夠的後果：每一步都跨得太大，細的地方（葉子、香蕉的彎）來不及長出來。

顏色也只是座標

每個點其實不是三個數字，是六個：x、y、z、紅、綠、藍。對模型來說，顏色和位置沒有差別，都是要從雜訊裡還原的數字，上面所有式子裡的 xx 都是這六個數字。

所以它學到的不是「形狀」再加上「上色」，而是一個六維空間裡的東西：「在這個位置的點，應該是這個顏色」。蘋果頂端的點是綠色的，是因為六維空間裡，那個位置配上綠色的組合才在水果上。儀器 02 把雜訊拉高的時候顏色也會跟著亂掉，就是這個原因。

它怎麼知道要長哪一種

同一個模型要捏兩種水果，靠的是多給它看兩個數字 cc：[1, 0] 代表「我要第一種」，[0, 1] 代表「我要第二種」。模型從 εθ(xt,t)\varepsilon_\theta(x_t, t) 變成 εθ(xt,t,c)\varepsilon_\theta(x_t, t, c)，其他完全不變。訓練時每個點都附上自己來自哪一種水果，取樣時你要哪一種就給哪一組。

文字生成圖片的模型做的是同一件事，只是 cc 換成了一段文字算出來的一長串數字。

那儀器 01 的第三團呢？它拿到的是 [0.5, 0.5]，一組模型在訓練時從來沒看過的數字。沒有人教過它「一半蘋果一半香蕉」長什麼樣子，它只能自己在兩個答案之間找一個折衷。我選蘋果和香蕉訓練的那一次，它長出來的是一顆蘋果形狀、香蕉顏色、頭上還留著葉子的東西。換一對水果，折衷的方式也會不一樣。

叫它「更聽話一點」

真的畫圖模型還有一個幾乎人人都在用的技巧，叫 classifier-free guidance。做法有兩半：

  • 訓練時，每十個點裡有一個故意不告訴模型它來自哪一種水果（cc 全部填 0）。於是同一個模型同時學會兩件事：「知道要哪一種的時候怎麼推」和「不知道的時候怎麼推」。
  • 取樣時，兩種都問，然後把兩個答案的差距放大 ww 倍：
ε~=εθ(xt,t,∅)+w⋅(εθ(xt,t,c)−εθ(xt,t,∅))\tilde\varepsilon = \varepsilon_\theta(x_t, t, \varnothing) + w \cdot \big( \varepsilon_\theta(x_t, t, c) - \varepsilon_\theta(x_t, t, \varnothing) \big)

「知道」和「不知道」的差，就是「這個條件造成的那部分方向」。w=1w = 1 是原樣，w>1w > 1 是把那部分方向誇大。

fig 05/diffusion / guidance

…

每次拖動滑桿，都是用同一團起始雜訊重新取樣，所以畫面上的差別只來自 w。

把 ww 拉到 0，兩團會長成同一個東西：一顆介於兩種水果之間的「平均水果」，因為模型完全沒聽你要哪一種。w=1w = 1 是正常的樣子。繼續往上拉，水果會先變得更「典型」、點更集中，再往上就開始變形：在我的那次訓練裡，w=4w = 4 的蘋果只剩一圈紅色的邊，香蕉縮成一小條。

這是真的模型也有的取捨：引導越強，結果越像「標準答案」，多樣性越低，太強就壞掉。在我的那次訓練裡，1 到 2 之間最好看；畫圖模型常見的預設值在 7 左右，因為它們的條件（一段文字）比「兩種水果選一種」模糊得多，需要更用力地聽。

這和真的畫圖模型差在哪

  • 它處理的是點，不是圖。 這裡的一筆資料是一個點、六個數字，所以一個兩萬多個參數的全連接網路就夠了。圖片模型的一筆資料是一整張圖、幾十萬個數字，而且像素之間的關係很重要，所以要用大得多的卷積或 Transformer 網路，而且通常不直接在像素上做，而是先把圖壓縮到一個小的 latent 空間裡再做擴散。
  • 它只看過兩樣東西。 開頭說過了：這是背誦，不是創作。差別來自資料量，不是方法。
  • 其他全部一樣。 加雜訊的式子、「猜雜訊」的損失函數、一步一步往回走、用條件控制、用 guidance 加強，這五件事和你在用的畫圖模型裡的是同一套。

來源