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

机器之心發表於2024-09-18

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 步內生成質量不錯的影像。它的生成過程如下圖所示：

^{NAT 生成過程的示意圖}

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

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

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

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

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

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

現有工作與 AdaNAT 的對比，這裡是當前生成樣本在時刻時的 token map

論文標題：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 的方法示意圖如下：