ProlificDreamer（VSD） 論文閱讀筆記

這是一篇 text to 3D 方向的突破性的文章，效果確實非常棒，據說一作的朋友圈中也說“他們將這個領域從20分提升到了70分的水平”，預測之後會有許多基於該方法的優秀文章與產品出現，畢竟之前 SDS follow 的文章也有很大一批。本閱讀筆記就簡要去記錄一下這篇論文的主要方法和思想。

Variational Score Distillation (VSD) 是將之前提出的 Score Distillation Sampling (SDS)方法進行了泛化，換言之文章證明了 SDS 是 VSD 的一個特殊情況。目前所有 2D to 3D 做生成的基本都是基於 SDS 方法（DreamFusion），也有基於 SJC 的工作，但二者方法大體上基本一致，取得了一定的效果。

背景

論文首先簡要回顧一下 Diffusion 和 SDS 方法。

Diffusion Models

對於前向過程 \(\{q_t\}_{t \in [0, 1]}\) 會向數據點 \(x_0\) 逐漸添加噪聲 \(x_0 \sim q_0(x_0)\) ，反向過程 \(\{p_t\}_{t \in [0, 1]}\) 生成數據。前向滿足 \(q_t(x_t|x_0) := \mathcal N(\alpha_t x_0, \sigma_t^2 I)\) 且反向滿足 \(q_t(x_t) := \int q_t(x_t | x_0) q_0(x_0) \mathrm d x_0\) （相關推導可以看原版Diffusion Model 的內容，網上很多）。這裏的 \(\alpha_t\) 和 \(\sigma_t\) 爲超參，滿足 \(\alpha_0 \approx 1, \sigma_0 \approx 0, \alpha_1 \approx 0, \sigma_1 \approx 1\) （即 \(\lim_{t \to 1} q_t(x_t | x_0) \approx \mathcal N(0, I)\) ）

爲了學習這個反向的過程，也就是生成的過程，這裏需要訓練一個降噪的網絡 \(\epsilon_\phi (x_t, t)\) ，用如下經典 loss 去訓練：

\[\mathcal L_{\mathrm{Diff}} (\phi) := \mathbb E_{x_0 \sim q_0(x_0), t \sim \mathcal U(0, 1), \epsilon \sim \mathcal N(0, I)} [ \omega(t) || \epsilon_\phi(\alpha_t x_0 + \sigma_t \epsilon) - \epsilon||_2^2] \]

\(\omega(t)\) 是關於 \(t\) 的一個權重函數（這個 MSE loss 意義是讓網絡預測的噪聲 \(\epsilon_\phi(\alpha_t x_0 + \sigma_t \epsilon)\) 儘量貼近真實噪聲 \(\epsilon\) ）。這樣保證了在 \(p_t \approx q_t\) 情況下滿足可以採樣出 \(p_0 \approx q_0\) 。這裏可以用 SDE 那篇文章中的 score function 表示 \(\nabla_{x_t} \log q_t(x_t) \approx \nabla_{x_t} \log p_t(x_t) \approx - \epsilon_{\phi} (x_t, t) / \sigma_t\) 。

對於常見的 text-to-image 生成過程，用的是 \(\epsilon_\phi(x_t, t, y)\) 其中 \(y\) 爲 text prompt，classifier-free guidance（CFG）是常用的方法，用於平衡多樣性和質量。用 \(\hat \epsilon_\phi(x_t, t, y) := (1+s)\epsilon_\phi(x_t, t, y) - s \epsilon_\phi(x_t, t, \empty)\) 來表示，其中 \(s\) 爲 guidance scale。

Text-to-3D 生成 score distillation sampling (SDS)

SDS方法顧名思義就是用來蒸餾預訓練好的擴散模型，目前已經成了 2D-to-3D 路線中唯一的一個方法。即如果有預訓練好的 2D Diffusion 模型 \(p_t(x_t|y)\) 與網絡 \(\epsilon_{\mathrm{pretrain}} (x_t, t, y)\) ，那麼 SDS 會優化在一個參數 \(\theta \in \Theta\) 下的 3D 表達。

引入一些符號：相機參數 \(c\) （內參+外參）、分佈 \(p(c)\) 以及可微渲染函數 \(g(\cdot, c): \Theta \to \mathbb R^d\) 。\(q_t^\theta(x_t | c)\) 爲渲染圖 \(g(\theta, c)\) 在時刻 \(t\) 前向擴散過程中的分佈（這裏分佈指的是 2D Diffusion）。

那麼對於 SDS 他們需要優化的是參數 \(\theta\) ，即最小化 loss：

\[\min_{\theta \in \Theta} \mathcal L_{\mathrm{SDS}}(\theta) := \mathbb E_{t, c} [(\sigma_t / \alpha_t) \omega(t) D_{\mathrm{KL}} (q_t^{\theta}(x_t | c) || p_t(x_t | y))]\\ \mathrm{where~~} t \sim \mathcal U(0.02, 0.98), \epsilon \sim \mathcal N(0, I), x_t = \alpha_t g(\theta, c) + \sigma_t \epsilon \]

即讓可微渲染圖在第 \(t\) 步擴散與預訓練模型的分佈近似，從而讓最後整體在 \(t \to 0\) 的分佈進行近似，用這裏 KL 散度保證這個事情。

我們要優化這個得分函數，只需要求得這個梯度，可以近似爲

\[\begin{equation} \nabla_\theta \mathcal L_{\mathrm{SDS}}(\theta) \approx \mathbb E_{t, \epsilon, c} [\omega(t)(\underbrace{\epsilon_{\mathrm{pretrain}}(x_t, t, y) - \epsilon}_{\mathrm{Noise ~Residual}}) \underbrace{\frac{\partial g(\theta, c)}{\partial \theta}}_{\mathrm{Generator~Jacobian}}] \end{equation} \]

這就是基於 SDS 的核心邏輯 \(\theta_{i+1} = \mathrm{opt.step}(\theta_i, \nabla_\theta \mathcal L(x_i))\) ，不斷去優化下一個對應獲得參數。相關推導以及理論可以參考 DreamFusion 及其引用。

3D 表示

這裏用經典的 NeRF（Neural Radiance Fields）和 DMTet（Deep Marching Tetrahedra）兩種主流方法作爲 3D 表示。

• NeRF 用 MLP 表示3D，利用 \((x, y, z)\) 座標輸入得出 \((r,g,b,\sigma)\) 即顏色與密度表示。這裏 \(\theta\) 爲 MLP 的參數，那麼給出相機位姿 \(c\) ，就有渲染可微渲染函數 \(g(\theta, c)\) ，意義爲從每個像素髮射射線，計算沿着每條射線採樣點的顏色的加權和，合成每個像素的顏色。NeRF在優化方面非常靈活，並能夠表示極其複雜的場景。

• DMTet 這種方法是需要拉出三角網格，用網格點的顏色表示紋理。這裏的 \(\theta\) 爲三角網格座標和紋理的參數，那麼有可微渲染函數 \(g(\theta, c)\) ，意義爲從像素髮射射線並計算射線與網格表面的交點來獲得每個像素的顏色。這個方法消耗顯存小，並且渲染非常快，所以可以支持高像素，但是對於複雜場景支持一般。

Variational Score Distillation（VSD）

這篇文章提出的核心就是 VSD ，通過在 3D場景的分佈中採樣，用3D參數 particles 去表示 3D 分佈，基於 Wasserstein 梯度流推導出一種基於梯度的 particles 更新規則。

從3D場景中採樣進行變分推斷

和前面一致，若有 text prompt \(y\) ，那麼就存在一個對於所有 3D 表示的隨機概率分佈，在一個 3D 參數 \(\theta\) 表示下，這個分佈我們可以記成一個概率密度函數 \(\mu(\theta | y)\) 。並且記 \(q_0^\mu (x_0 | c, y)\) 爲渲染圖 \(x_0 := g(\theta, c)\) 的隱式分佈（我的理解是這個爲基於網絡參數對渲染圖的分佈表達），那麼有 \(q_0^\mu(x_0 | y) := \int q_0^\mu (x_0 | c, y) p(c) \mathrm{d} c\) 爲這個的邊緣分佈（對於原來聯合分佈相對於相機分佈 \(p(c)\) 的部分）。並且我們這裏記 \(p_0(x_0 | y)\) 爲 \(t = 0\) 時預訓練的 2D Diffusion 的邊緣分佈。

那麼需要獲得高質量的 3D 表示的話，作者在這裏提出了用 KL Loss 擬合這個邊緣分佈的方法：

\[\min_\mu D_{\mathrm{KL}} (q_0^\mu (x_0 | y) || p_0(x_0 | y)) \]

這是一個經典的變分推斷問題，用變分分佈 \(q_0^\mu (x_0 | y)\) 去近似目標分佈 \(p_0(x_0 | y)\) ，前者的邊緣分佈意義可以理解爲沒有相機渲染視角指導下的圖片分佈，其應該對應上 2D Diffusion 從 web scale 圖片下獲取的分佈。

直接最優化這個函數是比較複雜的，因爲 \(p_0\) 是相當複雜的一個分佈，並且在高維空間中這個分佈的高密度區域是很稀疏的。有一種解釋是如果這個分佈是支撐在高維空間中的低維流形時，他們重疊部分的測度爲 \(0\) 的概率幾乎是 \(1\) ，這樣直接去近似的 Loss 是很難下降的。

但仍可以借鑑擴散模型的思路去轉化爲優化每一個 \(t\) 的擴散分佈，這個之前有許多成功的先例，形如 SDS ，這裏就使用變分分數蒸餾（VSD）方法去優化這個 loss：

\[\begin{aligned} \mu^{*} &:= \mathop{\arg \min}_{\mu} \mathbb E_{t} [(\sigma_t / \alpha_t) \omega(t) D_{\mathrm{KL}} (q_t^\mu(x_t | y) || p_t(x_t | y))]\\ &= \mathop{\arg \min}_{\mu} \mathbb E_{t, c} [(\sigma_t / \alpha_t) \omega(t) D_{\mathrm{KL}}(q_t^\mu(x_t | c, y) || p_t(x_t | y))] \end{aligned} \]

這裏大部分定義和前面 SDS 可以類比的，其中分佈定義有 \(q_t^\mu(x_t | c, y) := \int q_0^\mu (x_0 | c, y) p_{t0}(x_t | x_0) \mathrm d x_0\) ，類似地 \(q_t^\mu(x_t, y) := \int q_0^\mu (x_0 | y) p_{t0}(x_t | x_0) \mathrm d x_0\) 爲前者的邊緣分佈，並且 \(p_t(x_t | y) = \int p_0(x_0 | y) p_{t0} (x_t | x_0) \mathrm d x_0\) 爲時間 \(t\) 對應的加噪分佈，根據Diffusion Model 的推導有 \(p_{t0} (x_t | x_0) = \mathcal N(x_t | \alpha_t x_0, \sigma_t^2 I)\) ，這裏 \(\omega(t)\) 與前面的定義類似爲時間權重函數。等式成立的條件是因爲聯合分佈和邊緣分佈在 KL 散度下的一致性。

這裏就可以發現該方法與 SDS 的區別了，SDS是爲了3D表示的參數 \(\theta\) 也就是單個數據點，但 VSD 是爲了得出整體 3D的概率分佈空間 \(\mu\) ，從整個空間中我們再採樣出數據點 \(\theta\) 。

而且作者在附錄中證明了優化這個額外 KL loss 就可以達到最初近似分佈的目的，即對於任意 \(t>0\) 都有

\[D_{\mathrm{KL}} (q_t^\mu (x_t | y) || p_t(x_t | y)) = 0 \iff q_0^{\mu}(x_0 | y) = p_0(x_0 | y) \]

VSD 的更新方法

爲了解決 \(\mu^* = \mathop{\arg \min}_{\mu} \mathbb E_{t, c} [(\sigma_t / \alpha_t) \omega(t) D_{\mathrm{KL}}(q_t^\mu(x_t | c, y) || p_t(x_t | y))]\) ，最簡單的方法是訓練另一個以 \(\mu\) 爲參數的生成模型，但這會帶來許多計算上的消耗。這裏介紹了 particle-based 的變分推斷方法，即使用 \(n=4\) 個（平衡效果和計算消耗）3D 參數 \(\{\theta\}_{i = 1}^n\) 作爲 particles 並且遵循一個更新法則。直觀地說，用這 \(n\) 個 particles 來代表當前的分佈 \(\mu\) ，並且 \(\theta^{(i)}\) 會從優化後的分佈 \(\mu^*\) 中進行採樣（如果收斂的話）。

這樣的優化過程可以被寫成模擬一個關於 \(\theta\) 的 ODE 的過程，下面進行描述：

• 對於 VSD 的 Wasserstein 梯度流：

  Wasserstein梯度流：基於Wasserstein距離的梯度信息，定義了一個概率分佈隨時間演化的微分方程。在Wasserstein梯度流中，考慮一個初始概率分佈 \(P_0\) 和目標概率分佈\(P_T\) ，目標是找到一個時間依賴的概率分佈 \(P(t)\) ，其中 \(t\) 在 \([0, T]\) 範圍內變化，使得 \(P(t)\) 在每個時間點 \(t\) 的 Wasserstein 距離與目標分佈 \(P_T\) 最小化。

  剩餘相關證明和推導可以見原文附錄，下面僅作定理敘述。

  從初始分佈 \(\mu_0\) 開始，對於任意時刻 \(\tau \ge 0\)，記在分佈空間最優化 \(\mu^*\) 過程的 Wasserstein 梯度流爲 \(\{\mu_\tau\}_{\tau \ge 0}\) 其中 \(\mu_\infty = \mu^*\) ，那麼我們可以從 \(\mu_\tau\) 中採樣 \(\theta_\tau\) ，首先採樣 \(\theta_0 \sim \mu_0 (\theta_0 | y)\) 然後之後的用如下 ODE 進行模擬：

  \[\frac{\mathrm d \theta_\tau}{\mathrm d \tau} = - \mathbb E_{t, \epsilon, c} [\omega(t)(\underbrace{-\sigma_t \nabla_{x_t} \log p_t(x_t | y)}_{\mathrm{score~of~noisy~real~images}} - \underbrace{(-\sigma_t \nabla_{x_t} \log q_t^{\mu_\tau}(x_t | c, y) ))}_{\mathrm{score~of~noisy~rendered~images}} \frac{\partial g(\theta_\tau, c)}{\partial \theta_\tau}] \]

只需要求出足夠大的 \(\tau\) ，就能近似到希望求出的分佈 \(\mu^*\) 。

右式類似於 SDS 的梯度求法，其中對於 score of noisy real images \(-\sigma_t \nabla_{x_t} \log p_t(x_t | y)\) 可以直接用預訓練好的 2D Diffusion 降噪網絡 \(\epsilon_{\mathrm{pretrain}} (x_t, t, y)\) 進行近似，對於後者 \(- \sigma_t \nabla_{x_t} \log q_t^{\mu_\tau} (x_t | c, y)\) 我們這裏採用另外一個降噪網絡 \(\epsilon_\phi (x_t, t, c, y)\) 進行近似，這個我們用 \(\{\theta^{(i)}\}_{i = 1}^n\) 渲染好的圖片以及標準的Diff Loss 進行訓練：

\[\min_{\phi} \sum_{i = 1}^{n} \mathbb E_{t \sim \mathcal U(0, 1), \epsilon \sim \mathcal N(0, I), c \sim p(c)} [ || \epsilon_\phi(\alpha_t g(\theta^{(i)}, c) + \sigma_t \epsilon, t, c, y) - \epsilon||_2^2] \]

即用這個 MSE Loss 去擬合這個標準噪聲。

這篇文章用了一個 U-Net 或者對於原來預訓練模型 \(\epsilon_{\mathrm{pretrain}} (x_t, t, y)\) 的 LoRA 去表示這個參數 \(\epsilon_\phi\) ，並且添加了相機的位姿 \(c\) 作爲 condition embedding 到網絡中。作者根據實驗說 LoRA 在 few-shot 下的保真度非常好。

注意到對於每個 ODE 的時刻 \(\tau\) ，需要去確認 \(\epsilon_\phi\) 符合當前的 3D分佈 \(q_t^{\mu_\tau}\) ，這樣需要輪流優化 \(\epsilon_\phi\) 和 \(\theta^{(i)}\) 。對於 particle 我們用 \(\theta^{(i)} \leftarrow \theta^{(i)} - \eta \nabla_\theta \mathcal L_{\mathrm{VSD}} (\theta^{(i)})\) 進行優化（其中 \(\eta > 0\) 爲 learning rate），由之前的式子，這個梯度可以被表示爲

\[\nabla_{\theta} \mathcal L_{\mathrm{VSD}} (\theta) \triangleq \mathbb E_{t, \epsilon, c} [\omega(t) (\epsilon_{\mathrm{pretrain}}(x_t, t, y) - \epsilon_\phi (x_t, t, c, y) \frac{\partial g(\theta, c)}{\partial \theta}] \]

在這裏有 \(x_t = \alpha_t g(\theta, c) + \sigma_t \epsilon\) 即渲染圖在時刻 \(t\) 加噪後的情況。

僞代碼如下：

網絡結構

整體網絡結構與 SDS 方法是類似的，結合上面的僞代碼即可看懂。

• 首先隨機從當前分佈中採樣網絡的3D參數 \(\theta\) ，以及相機位姿 \(c\)
• 然後利用可微渲染方法，渲染出對應的 2D 圖片 \(x_0 = g(\theta, c)\)
• 而後更新3D表達的參數 \(\theta\) 爲 \(\theta - \eta_1 E_{t, \epsilon, c} [\omega(t) (\epsilon_{\mathrm{pretrain}}(x_t, t, y) - \epsilon_\phi (x_t, t, c, y) \frac{\partial g(\theta, c)}{\partial \theta}\)
• 最後更新 LoRA 中蘊含的潛在 3D 分佈 \(\phi\) 爲 \(\phi - \eta_2 \nabla_\phi \mathbb E_{t, \epsilon} ||\epsilon_\phi(x_t, t, c, y) - \epsilon||_2^2\) ，即更加貼近原本 \(q_t^{\mu_\tau}\) 的分佈，如此往復直至收斂。