本論文第一作者倪贊林是清華大學(xué)自動(dòng)化系 2022 級(jí)直博生，師從黃高副教授，主要研究方向?yàn)楦咝疃葘W(xué)習(xí)與圖像生成。他曾在 ICCV、CVPR、ECCV、ICLR 等國(guó)際會(huì)議上發(fā)表多篇學(xué)術(shù)論文。

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

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

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

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

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

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

NAT 生成過(guò)程的示意圖

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

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

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

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

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

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

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

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

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

方法介紹

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

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

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

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

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

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

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

AdaNAT 示意圖

實(shí)驗(yàn)結(jié)果

我們?cè)诙鄠€(gè)基準(zhǔn)數(shù)據(jù)集上驗(yàn)證了 AdaNAT 的有效性，包括 ImageNet 的 class-conditional 生成以及 MSCOCO 和 CC3M 數(shù)據(jù)集的文到圖生成。

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

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

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

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

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

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

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

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

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

消融實(shí)驗(yàn)

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

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

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

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