VQ-Diffusion

向量量化擴散（VQ-Diffusion）是由中國科學技術大學和微軟開發的一種條件潛在擴散模型。與大多數常見研究的擴散模型不同，VQ-Diffusion 的加噪和去噪過程在量化潛在空間中進行，即潛在空間由一組離散向量組成。離散擴散模型比其連續對應模型探索較少，並提供了與自迴歸（AR）模型有趣的比較點。

• Hugging Face 模型卡片
• Hugging Face Spaces
• 原始實現
• 論文

演示

🧨 Diffusers 允許您僅用幾行程式碼執行 VQ-Diffusion。

安裝依賴項

pip install 'diffusers[torch]' transformers ftfy

載入管道

from diffusers import VQDiffusionPipeline

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")

如果您想使用 FP16 權重

from diffusers import VQDiffusionPipeline
import torch

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16")

移至 GPU

pipe.to("cuda")

執行管道！

prompt = "A teddy bear playing in the pool."

image = pipe(prompt).images[0]

架構

VQ-VAE

影像使用 VQ-VAE 編碼器編碼成一組離散的“tokens”或嵌入向量。為此，影像被分割成小塊，然後每個小塊被替換為程式碼本中詞彙量固定大小的最接近的條目。這減少了輸入畫素空間的維度。VQ-Diffusion 使用 Taming Transformers 中的 VQGAN 變體。這篇部落格文章是更好地理解 VQ-VAE 的一個很好的資源。

VQ-Diffusion 使用了一個預訓練的 VQ-VAE，在擴散訓練過程中被凍結。

前向過程

在前向擴散過程中，每個潛在 token 可以保持不變，被重新取樣到不同的潛在向量（每個都有相等的機率），或者被遮蔽。一旦一個潛在 token 被遮蔽，它將保持遮蔽狀態。$\alpha_t$、$\beta_t$ 和 $\gamma_t$ 是控制從步驟 $t-1$ 到步驟 $t$ 的前向擴散過程的超引數。$\gamma_t$ 是未被遮蔽的 token 變為遮蔽的機率。$\alpha_t + \beta_t$ 是未被遮蔽的 token 保持不變的機率。該 token 可以以 $\beta_t$ 的機率轉換為任何單個未被遮蔽的潛在向量。換句話說，$\alpha_t + K \beta_t + \gamma_t = 1$，其中 $K$ 是未被遮蔽的潛在向量的數量。更多細節請參見論文第 4.1 節。

近似逆向過程

編碼器-解碼器 Transformer 近似未加噪的潛在變數 $x_0$ 的類別，並以提示 $y$ 為條件。編碼器是一個 CLIP 文字編碼器，其權重被凍結。解碼器 Transformer 為所有潛在畫素提供未遮蔽的全域性注意力，並輸出向量嵌入的分類分佈的對數機率。解碼器 Transformer 在一次前向傳播中預測未加噪潛在變數的整個分佈，為 $x_t$ 提供全域性自注意力。將問題框定為離散值上的條件序列到序列問題，為編碼器-解碼器 Transformer 為何如此合適提供了一些直觀的解釋。

AR 模型部分提供了 VQ-Diffusion 架構與基於 AR Transformer 模型的比較的額外背景資訊。

Taming Transformers 對將原始畫素轉換為壓縮潛在空間中的離散 token 進行了很好的討論，從而使 Transformer 在影像資料上變得計算可行。

VQ-Diffusion 在上下文中

擴散模型

當代擴散模型大多是連續的。在前向過程中，連續擴散模型迭代地新增高斯噪聲。逆向過程透過 $p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t))$ 來近似。在更簡單的 DDPM 情況下，協方差矩陣是固定的，訓練一個 U-Net 來預測 $x_t$ 中的噪聲，並且 $x_{t-1}$ 由噪聲推導。

近似逆向過程在結構上類似於離散逆向過程。然而在離散情況下，預測 $x_t$ 中的噪聲沒有明確的類比，直接預測 $x_0$ 的分佈是一個更明確的目標。

關於離散擴散模型的文獻量少於連續擴散模型。使用非平衡熱力學的深度無監督學習引入了基於二項分佈的擴散模型。Argmax Flows and Multinomial Diffusion 將離散擴散擴充套件到多項分佈，並訓練一個 Transformer 來預測用於語言建模任務的無噪聲分佈。離散狀態空間中的結構化去噪擴散模型 透過替代加噪過程（均勻、吸收、離散化高斯和 token 嵌入距離）推廣了多項擴散。替代加噪過程在連續擴散模型中也是可能的，但如論文所述，只有加性高斯噪聲受到了 significant attention。

自迴歸模型

將 VQ-Diffusion 與 AR 模型進行比較可能更有趣，因為它們更常以 Transformer 進行離散分佈預測為特徵。雖然 Transformer 在 AR 建模中取得了成功，但它們仍然面臨影像解析度增加導致推理速度線性下降、錯誤累積和方向性偏差的問題。VQ-Diffusion 在所有三個痛點上都有所改進。

AR 影像生成模型透過將影像機率因子化來表徵，使得每個畫素都以前一個畫素以光柵掃描順序（從左到右，從上到下）為條件，即 $p(x) = \prod_i p(x_i | x_{i-1}, x_{i-2}, ... x_{2}, x_{1})$。因此，可以透過直接最大化對數似然來訓練模型。此外，作用於實際畫素（非潛在）值的 AR 模型，從離散多項分佈中預測通道值，即首先從 256 路 softmax 中取樣紅色通道值，然後綠色通道預測以紅色通道值為條件。

AR 影像生成模型在架構上不斷發展，許多工作致力於使 Transformer 在計算上可行。在基於 Transformer 的模型之前，PixelRNN、PixelCNN 和 PixelCNN++ 是最先進的模型。

Image Transformer 很好地討論了非 Transformer 模型以及向基於 Transformer 模型過渡的過程（省略了引文，請參閱論文）。

訓練迴圈神經網路以順序預測即使是小影像的每個畫素在計算上也非常具有挑戰性。因此，使用卷積神經網路的可並行化模型，如 PixelCNN，最近受到了更多關注，並且現在在質量上已經超越了 PixelRNN。

與 RNN 相比，CNN 的一個缺點是其感受野通常相當有限。這會對其建模影像中常見的長距離現象（如對稱性和遮擋）的能力產生不利影響，尤其是在層數較少的情況下。已證明擴大感受野可以顯著提高質量（Salimans 等人）。然而，這樣做會顯著增加引數數量，從而影響計算效能，並使此類模型的訓練更具挑戰性。

... 自注意力可以在必然順序的 PixelRNN 的幾乎無限感受野與更具並行性的 PixelCNN 及其各種擴充套件的有限感受野之間的權衡中實現更好的平衡。

Image Transformer 透過限制區域性畫素鄰域的自注意力來使用 Transformer。

Taming Transformers 和 DALL-E 1 結合了卷積和 Transformer。兩者都訓練一個 VQ-VAE 來學習離散潛在空間，然後在壓縮的潛在空間中訓練一個 Transformer。Transformer 的上下文是全域性的但被遮蔽，因為注意力提供給所有先前預測的潛在畫素，但模型仍然是 AR，因此不能提供注意力給尚未預測的畫素。

ImageBART 結合了卷積、Transformer 和擴散過程。它學習了一個離散潛在空間，並透過一個短多項擴散過程進一步壓縮。然後訓練單獨的編碼器-解碼器 Transformer 來逆轉擴散過程中的每個步驟。編碼器 Transformer 提供 $x_t$ 上的全域性上下文，而解碼器 Transformer 自迴歸地預測 $x_{t-1}$ 中的潛在畫素。因此，每個畫素都會在噪聲更大的影像上獲得全域性交叉注意力。對於更復雜的資料集，使用了 2-5 個擴散步驟。

儘管取得了巨大的進步，AR 模型仍然面臨影像解析度增加導致推理速度線性下降、錯誤累積和方向性偏差的問題。對於同等大小的 AR Transformer 模型，只要擴散步驟的數量小於潛在畫素的數量，VQ-Diffusion 的推理大 O 複雜度就更優。對於 ITHQ 資料集，潛在解析度為 32x32，模型訓練到 100 個擴散步驟，實現了約 10 倍的大 O 改進。在實踐中，VQ-Diffusion“可以比 AR 方法快 15 倍，同時實現更好的影像質量”（更多細節請參閱論文）。此外，VQ-Diffusion 不需要教師強制，而是學習糾正錯誤預測的 token。在訓練期間，加噪影像既被遮蔽，其潛在畫素也被隨機 token 替換。VQ-Diffusion 還能夠提供 $x_t$ 上的全域性上下文，同時預測 $x_{t-1}$。

VQ-Diffusion 和 🧨 Diffusers 的後續步驟

目前，我們只移植了在 ITHQ 資料集上訓練的 VQ-Diffusion 模型。還有已釋出的 VQ-Diffusion 模型在 CUB-200、Oxford-102、MSCOCO、Conceptual Captions、LAION-400M 和 ImageNet 上進行了訓練。

VQ-Diffusion 還支援更快的推理策略。網路重引數化依賴於以未加噪影像為條件的擴散過程的後驗是可處理的。使用時間步長 $\Delta t$ 時，也會應用類似的公式，該時間步長跳過若干個逆向擴散步驟，$p_\theta (x_{t - \Delta t } | x_t, y) = \sum_{\tilde{x}_0=1}^{K}{q(x_{t - \Delta t} | x_t, \tilde{x}_0)} p_\theta(\tilde{x}_0 | x_t, y)$。

改進的向量量化擴散模型 透過離散無分類器指導和替代推理策略來解決“聯合分佈問題”，從而提高了 VQ-Diffusion 的樣本質量——更多細節請參見第 3.2 節。離散無分類器指導已合併到 Diffusers 中，但替代推理策略尚未新增。

歡迎貢獻！