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

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

為什么要探索基于離散 token 的生成模型？我們認(rèn)為主要有以下幾點(diǎn)原因：

1) 由于與語言模型范式相同，它們可以直接利用語言模型中已經(jīng)成熟的訓(xùn)練、推理技術(shù)

2) 有助于開發(fā)更先進(jìn)的、具有共享 token 空間的、scalable 的多模態(tài)基礎(chǔ)模型

3) 有助于建構(gòu)統(tǒng)一視覺理解與生成能力的通用視覺基礎(chǔ)模型

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

NAT 生成過程的示意圖

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

雖然 “并行解碼” 的機(jī)制讓模型生成過程更加地高效和靈活。但它也引入了許多復(fù)雜的設(shè)計挑戰(zhàn)，例如每一步應(yīng)該解碼多少 token、應(yīng)該選擇哪些 token、以及采樣時應(yīng)該有多大的隨機(jī)性等。

現(xiàn)有的工作通過構(gòu)建一套包含多個人工設(shè)計的調(diào)度函數(shù)（統(tǒng)一記為）的生成策略來緩解這一問題。如下圖所示：

NAT 的默認(rèn)生成策略，其中 T 為總生成推理步數(shù)，t 為當(dāng)前推理步，均為超參數(shù)。每一行策略的具體含義可以參見原文

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

基于上述觀察，我們提出 AdaNAT，核心思想是引入一個可學(xué)習(xí)的策略網(wǎng)絡(luò)，自動根據(jù)每個樣本自適應(yīng)地配置生成策略：

現(xiàn)有工作與 AdaNAT 的對比，這里

是當(dāng)前生成樣本在時刻時的 token map

• 論文標(biāo)題：AdaNAT: Exploring Adaptive Policy for Token-Based Image Generation
• 論文地址：https://arxiv.org/abs/2409.00342
• 代碼與預(yù)訓(xùn)練模型已經(jīng)開源：https://github.com/LeapLabTHU/AdaNAT

方法介紹

然而，訓(xùn)練一個自適應(yīng)、自動配置生成策略的網(wǎng)絡(luò)面臨一個直接的挑戰(zhàn)：基于離散 token 的生成過程是不可微的，因此我們無法直接采用標(biāo)準(zhǔn)的端到端優(yōu)化技術(shù)如梯度下降的方法來優(yōu)化這個網(wǎng)絡(luò)。為應(yīng)對這一挑戰(zhàn)，我們將生成策略的優(yōu)化問題形式化為馬爾可夫決策過程（MDP），在此基礎(chǔ)上，策略網(wǎng)絡(luò)可以自然地被定義為一個「觀察生成狀態(tài)，自適應(yīng)地配置策略以最大化生成質(zhì)量」的 agent，也因此可以通過強(qiáng)化學(xué)習(xí)算法（如策略梯度）進(jìn)行訓(xùn)練。

另一個值得注意的點(diǎn)是，在我們的問題中，設(shè)計合適的獎勵函數(shù)對于有效訓(xùn)練策略網(wǎng)絡(luò)至關(guān)重要。為了解決這一問題，我們首先考慮了兩種現(xiàn)成的設(shè)計選擇：

• 標(biāo)準(zhǔn)評估指標(biāo)如 Fréchet Inception Distance (FID)
• 預(yù)訓(xùn)練的圖像獎勵模型，如 ImageReward 模型

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

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

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

換句話說，策略網(wǎng)絡(luò)傾向于 “過擬合” 這些預(yù)先設(shè)定的、靜態(tài)的獎勵函數(shù)。受到這一現(xiàn)象的啟發(fā)，我們的核心思路是在策略網(wǎng)絡(luò)學(xué)習(xí)的同時，動態(tài)更新獎勵函數(shù)以防止策略網(wǎng)絡(luò)過擬合，讓二者 “相互對抗，共同進(jìn)步”。這種思想自然地讓我們聯(lián)想到生成對抗網(wǎng)絡(luò) (GAN) 的想法，因此，我們提出了一個對抗獎勵模型，該模型類似于 GAN 中的判別器，以生成樣本為真實圖片的概率作為獎勵。當(dāng)策略網(wǎng)絡(luò)學(xué)習(xí)最大化獎勵時，我們同時優(yōu)化獎勵模型，以更好地區(qū)分真實樣本和生成樣本。總結(jié)來看，AdaNAT 的方法示意圖如下：

AdaNAT 示意圖

實驗結(jié)果

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

得益于 NAT 生成范式中并行解碼的優(yōu)勢，AdaNAT 在 ImageNet-256 和 ImageNet-512 數(shù)據(jù)集上，相較于主流的擴(kuò)散模型，在低開銷場景下，推理開銷至少減少了 2-3 倍，同時生成效果更佳:

ImageNet-256 class-conditional 生成結(jié)果

ImageNet-512 class-conditional 生成結(jié)果

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

MSCOCO 文到圖生成結(jié)果

CC3M 文到圖生成結(jié)果

在模型的優(yōu)化過程中，我們也能明顯看到生成質(zhì)量隨著策略網(wǎng)絡(luò)的學(xué)習(xí)而提升，同時 FID 指標(biāo)也有相應(yīng)的下降：

AdaNAT 的優(yōu)化過程可視化

消融實驗發(fā)現(xiàn)，引入可學(xué)習(xí)、自適應(yīng)的策略都對提升 NAT 的生成質(zhì)量有所幫助：

消融實驗

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

AdaNAT 生成結(jié)果可視化

本文轉(zhuǎn)自 機(jī)器之心  ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/wTcVBpAb97Kq1cEVbOMNbA??