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

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 也使得圖像插值更符合人的直覺：

結語

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