NeurIPS 2024｜浙大 & 微信 & 清華：徹底解決擴散模型反演問題

机器之心發表於2024-11-02
模型

圖片

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本篇論文已經被NeurIPS 2024接收。第一作者王方懿康是微信視覺團隊實習生，同時也是浙江大學一年級碩士研究生。共同一作是微信高階研究員Hubery。通訊作者是浙江大學助理教授張超。其他作者包括來自清華大學的董玥江，來自浙江大學的朱胡旻昊，趙涵斌助理教授和錢徽教授，以及微信基礎視覺和視覺生成技術負責人李琛。

隨著擴散生成模型的發展，人工智慧步入了屬於 AIGC 的新紀元。擴散生成模型可以對初始高斯噪聲進行逐步去噪而得到高質量的取樣。當前，許多應用都涉及擴散模型的反演，即找到一個生成樣本對應的初始噪聲。當前的取樣器不能兼顧反演的準確性和取樣的質量。

為徹底解決這一問題，微信視覺團隊與浙江大學和清華大學聯手提出了基於雙向顯式線性多步法的擴散模型精確反演取樣器（BELM）這一通用演算法，並透過截斷誤差分析確定了最優的 BELM 取樣器係數。此方法在確保精確反演的同時還提升了生成樣本的質量，在影像與影片的編輯、插值等下游任務中有廣泛的應用前景。這一研究成果已被 NeurIPS 2024 會議接收。

當前，擴散模型在影像生成、文字生成、音訊生成等多個領域得到了廣泛應用，表現出了卓越的效能。擴散模型的反演操作，即找到一個生成樣本對應的初始噪聲，對若干下游任務起到關鍵的作用。傳統的 DDIM 反演會造成嚴重的不一致問題，即原始圖片加噪再去噪的結果與原圖相差甚遠。

近期，研究者們提出了多種啟發式的精確反演取樣器來解決 DDIM 反演的不一致問題。然而，這些啟發式的精確反演取樣器的理論特性尚不明確，且取樣質量常常不盡如人意，這在一定程度上限制了它們的應用。

為此，本研究引入了一種通用的精確反演取樣器正規化 —— 雙向顯式線性多步（BELM）取樣器，該正規化包含了上文提到的啟發式精確反演取樣器。該團隊在 BELM 正規化內系統地研究了區域性截斷誤差（LTE），發現現有的精確反演取樣器的 LTE 並非最優。

因此，研究團隊透過 LTE 最小化方法提出了最優的 BELM（Optimal-BELM，O-BELM）取樣器。實驗表明，O-BELM 取樣器在實現精確反演的同時，也提升了取樣的質量。
圖片

  • 論文題目：BELM: Bidirectional Explicit Linear Multi-step Sampler for Exact Inversion in Diffusion Models

  • 論文連結：https://arxiv.org/abs/2410.07273

  • 專案地址：https://github.com/zituitui/BELM

背景：DDIM 反演造成的不一致問題

由於 DDIM 的正向過程和反演過程使用的迭代式並不相同，所以 DDIM 的反演重構樣本與初始的樣本存在較大差別。

實際使用中，DDIM 的反演有顯著的不一致問題：
圖片
圖片
現有精確反演方法

Null-text-inversion

以 Null-tex-inversion 為代表的方法對 unconditional 佔位符進行 fine-tune，以達到精確反演。

問題：這類方法侷限於 text-classifier-free-guidance 場景下的擴散模型；需要額外訓練，低效。
圖片
EDICT

EDICT 是基於 DDIM 的啟發式演算法，借鑑了可逆網路的做法，有兩個相互糅合的取樣鏈。
圖片
其逆過程如下，精確可逆：
圖片
問題：需要兩倍計算量；超引數 p 不魯棒，導致取樣質量不可控。

BDIA

BDIA 改進了 EDICT，使用 x_i 的速度，x_i 和 x_{i+1} 的位置，透過下述公式實現精確可逆：
圖片
問題：超引數 gamma 不魯棒，導致取樣質量不佳。

EDICT 和 BDIA 引數的不魯棒：

EDICT 和 BDIA 超引數的意義不明，沒有理論指導如何調整，導致不同情形下超引數的選擇差別巨大。使用起來極為不便。
圖片
圖片
雙向顯式線性多步法（BELM）框架

思路起源：DDIM 的正向過程（由藍線表示）與反演過程（由紅線表示）是兩種不同的關係，這導致了 DDIM 的反演不準確。如果強制正過程與反過程使用相同關係，又會引入隱式方法，大大增加計算複雜度。如果多引入一個點，不用隱式方法也可逆（由綠線表示）。
圖片
該論文中的演算法，正向和反演過程都服從相同的關係，因此能夠精確反演。具體來說，為了系統地設計這種取樣器，首先要將擴散模型的取樣過程建模為一個 IVP（Initial Value Problem，初值問題）：
圖片
以下是 IVP 的一般形式，這實際上是一個變步長變公式線性多步方法（VSVFM）：
圖片
為了避免隱式方法的複雜計算，上式需要在正向和反向都是顯式的，該團隊稱這一性質為雙向顯性（bidirectional explicit）。
圖片
代入雙向顯性條件，可以得到一般的 k 步 BELM 取樣器：
圖片
最簡單的形式是 k=2，稱為 2-BELM，其表示式如下：
圖片
據此很容易證明，一個滿足雙向顯性性質的線性多步法取樣器擁有精確反演性質：
圖片
研究團隊還發現，前文提到的 EDICT 和 BDIA 都是 BELM 框架的特例：
圖片
這也解釋了 EDICT 和 BDIA 能夠精確反演的原因。

最優雙向顯式線性多步（O-BELM）取樣器

研究團隊在推導 BELM 框架暫時沒有給出具體的係數選擇，而啟發式的係數選擇（如 EDICT 和 BDIA）會造成取樣質量的退化。因此，他們提出使用區域性截斷誤差（LTE）來獲取最優係數。

首先分析 BELM 的區域性截斷誤差：
圖片
透過對區域性截斷誤差的最小化，我們得到了最優的 BELM 係數，我們稱此係數下的 BELM 取樣器為最優 BELM（O-BELM）：
圖片
O-BELM 的正向過程表示式如下：
圖片
O-BELM 的反演過程表示式如下：
圖片
此外，研究團隊還證明了 O-BELM 滿足穩定性和全域性收斂性：
圖片
至此，可以對比幾種不同反演取樣器的性質：
圖片
可見，O-BELM 是第一種在嚴格的理論保證下兼顧精確反演性質和取樣質量的取樣器。

實驗

重建實驗（驗證 O-BELM 精確反演性質）

latent 空間上的 O-BELM 的重建誤差為 0，這表明 O-BELM 具有精確反演的性質：
圖片
圖片
取樣實驗（驗證 O-BELM 的高質量取樣性質）

不論在無條件生成還是條件生成中，O-BELM 都表現出了高於 DDIM，EDICT 和 BDIA 的取樣質量：
圖片
下游任務 —— 影像編輯

影像編輯實驗體現了：

1. 由於 DDIM 不具有精確反演性質，編輯的結果中存在不一致問題（紅色框）；
2. 由於 EDICT 和 BDIA 具有較大的取樣誤差，編輯的結果出現了不真實區域（黃色框）；
3.O-BELM 在保持影像一致的條件下完成了高質量的編輯。
圖片
圖片
由於 O-BELM 是一個取樣方法，因此可以無縫地與 controlNet 結合，編輯效果也優於其他方法：
圖片
下游任務 —— 影像插值

由於 O-BELM 精確地建立了噪聲和生成樣本的對應關係，這個關係是 probability flow ODE 的近似，因此 O-BELM 也使得影像插值更符合人的直覺：
圖片
結語

本研究提出的雙向顯式線性多步法取樣器從理論上分析並徹底解決了現有擴散生成模型中的反演問題，進一步拓寬了擴散模型在計算機視覺領域的能力邊界。在影像和影片的編輯任務上有巨大的應用前景。

相關文章