ECCV 2024 | 探索離散Token視覺生成中的自適應推理策略

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

本論文第一作者倪讚林是清華大學自動化系 2022 級直博生，師從黃高副教授，主要研究方向為高效深度學習與圖像生成。他曾在 ICCV、CVPR、ECCV、ICLR 等國際會議上發表多篇學術論文。

近年來， AIGC 領域發展十分迅猛。在計算機視覺中，擴散模型已經成為一種有效且常用的模型。相反，在自然語言處理領域，內容的生成通常通過使用 Transformer 去生成離散的 token。受到這種差異的啟發，越來越多的研究（以 VQGAN 作為典型代表）開始探索這種基於離散 token 的生成範式在視覺合成中的應用。與擴散模型不同，這些方法採用類似於語言模型的離散 token 作為生成的基本單元。

為什麼要探索基於離散 token 的生成模型？我們認為主要有以下幾點原因：

1) 由於與語言模型範式相同，它們可以直接利用語言模型中已經成熟的訓練、推理技術

2) 有助於開發更先進的、具有共享 token 空間的、scalable 的多模態基礎模型

3) 有助於建構統一視覺理解與生成能力的通用視覺基礎模型

在基於離散 token 的生成里，近幾年非自回歸 Transformer (Non-autoregressive Transformer, NAT) 展現出了顯著的計算效率與生成質量方面的潛力，不同於傳統的自回歸生成範式，NAT 能夠在僅 4 到 8 步內生成質量不錯的圖像。它的生成過程如下圖所示：

簡單來說，這類模型從一個完全 mask 的 token map 開始，每步並行解碼多個 token，直到所有的 token 被解碼，然後再通過一個預訓練的解碼器把 token 空間變換到圖像空間，得到生成結果。

雖然 「並行解碼」 的機制讓模型生成過程更加地高效和靈活。但它也引入了許多複雜的設計挑戰，例如每一步應該解碼多少 token、應該選擇哪些 token、以及采樣時應該有多大的隨機性等。

現有的工作通過構建一套包含多個人工設計的調度函數（統一記為）的生成策略來緩解這一問題。如下圖所示：

NAT 的預設生成策略，其中 T 為總生成推理步數，t 為當前推理步，均為超參數。每一行策略的具體含義可以參見原文

然而，這種人為設計的方式不僅需要大量的專業知識和人力成本，最終得到的策略函數仍然可能並非最佳。與此同時，我們認為不同樣本都有其獨特的特性，一個應用於所有樣本的、全局共享的生成策略可能難以靈活應對樣本之間的差異性。

基於上述觀察，我們提出 AdaNAT，核心思想是引入一個可學習的策略網絡，自動根據每個樣本自適應地配置生成策略：

現有工作與 AdaNAT 的對比，這裏

• 論文標題：AdaNAT: Exploring Adaptive Policy for Token-Based Image Generation

• 論文地址：https://arxiv.org/abs/2409.00342

• 代碼與預訓練模型已經開源：https://github.com/LeapLabTHU/AdaNAT

方法介紹

然而，訓練一個自適應、自動配置生成策略的網絡面臨一個直接的挑戰：基於離散 token 的生成過程是不可微的，因此我們無法直接採用標準的端到端優化技術如梯度下降的方法來優化這個網絡。為應對這一挑戰，我們將生成策略的優化問題形式化為馬爾可夫決策過程（MDP），在此基礎上，策略網絡可以自然地被定義為一個「觀察生成狀態，自適應地配置策略以最大化生成質量」的 agent，也因此可以通過強化學習算法（如策略梯度）進行訓練。

另一個值得注意的點是，在我們的問題中，設計合適的獎勵函數對於有效訓練策略網絡至關重要。為瞭解決這一問題，我們首先考慮了兩種現成的設計選擇：

• 標準評估指標如 Fréchet Inception Distance (FID) 

• 預訓練的圖像獎勵模型，如 ImageReward 模型

然而，我們的實驗結果表明，儘管這些設計能夠有效地最大化獎勵函數，但最終的生成模型往往無法生成足夠高質量或足夠多樣化的圖像，如下圖所示：

以 FID 作為獎勵函數的圖像生成效果 (FID=2.56)

以預訓練的獎勵模型作為獎勵函數的圖像生成效果 (FID=33.1)

換句話說，策略網絡傾向於 「過擬合」 這些預先設定的、靜態的獎勵函數。受到這一現象的啟發，我們的核心思路是在策略網絡學習的同時，動態更新獎勵函數以防止策略網絡過擬合，讓二者 「相互對抗，共同進步」。這種思想自然地讓我們聯想到生成對抗網絡 (GAN) 的想法，因此，我們提出了一個對抗獎勵模型，該模型類似於 GAN 中的判別器，以生成樣本為真實圖片的概率作為獎勵。當策略網絡學習最大化獎勵時，我們同時優化獎勵模型，以更好地區分真實樣本和生成樣本。總結來看，AdaNAT 的方法示意圖如下：

AdaNAT 示意圖

實驗結果

我們在多個基準數據集上驗證了 AdaNAT 的有效性，包括 ImageNet 的 class-conditional 生成以及 MSCOCO 和 CC3M 數據集的文到圖生成。

得益於 NAT 生成範式中並行解碼的優勢，AdaNAT 在 ImageNet-256 和 ImageNet-512 數據集上，相較於主流的擴散模型，在低開銷場景下，推理開銷至少減少了 2-3 倍，同時生成效果更佳:

ImageNet-256 class-conditional 生成結果

ImageNet-512 class-conditional 生成結果

此外，在文到圖生成方面，AdaNAT 也有著不錯的表現：

MSCOCO 文到圖生成結果

CC3M 文到圖生成結果

在模型的優化過程中，我們也能明顯看到生成質量隨著策略網絡的學習而提升，同時 FID 指標也有相應的下降：

AdaNAT 的優化過程可視化

消融實驗發現，引入可學習、自適應的策略都對提升 NAT 的生成質量有所幫助：

消融實驗

最後，我們也可視化了 AdaNAT 模型生成的圖片，總的來看，AdaNAT 的生成樣本同時具有良好的生成質量與多樣性：

AdaNAT 生成結果可視化