Transformer也能生成圖像，新型ViTGAN性能比肩基于CNN的GAN

Transformer 已經(jīng)為多種自然語(yǔ)言任務(wù)帶來(lái)了突飛猛進(jìn)的進(jìn)步，并且最近也已經(jīng)開(kāi)始向計算機視覺(jué)領(lǐng)域滲透，開(kāi)始在一些之前由 CNN 主導的任務(wù)上暫露頭角。近日，加州大學(xué)圣迭戈分校與 Google Research 的一項研究提出了使用視覺(jué) Transformer 來(lái)訓練 GAN。為了有效應用該方法，研究者還提出了多項改進(jìn)技巧，使新方法在一些指標上可比肩前沿 CNN 模型。

卷積神經(jīng)網(wǎng)絡(luò )（CNN）在卷積（權重共享和局部連接）和池化（平移等變）方面的強大能力，讓其已經(jīng)成為了現今計算機視覺(jué)領(lǐng)域的主導技術(shù)。但最近，Transformer 架構已經(jīng)開(kāi)始在圖像和視頻識別任務(wù)上與 CNN 比肩。其中尤其值得一提的是視覺(jué) Transformer（ViT）。這種技術(shù)會(huì )將圖像作為 token 序列（類(lèi)似于自然語(yǔ)言中的詞）來(lái)解讀。Dosovitskiy et al. 的研究表明，ViT 在 ImageNet 基準上能以更低的計算成本取得相當的分類(lèi)準確度。不同于 CNN 中的局部連接性，ViT 依賴(lài)于在全局背景中考慮的表征，其中每個(gè) patch 都必須與同一圖像的所有 patch 都關(guān)聯(lián)處理。

ViT 及其變體盡管還處于早期階段，但已有研究展現了其在建模非局部上下文依賴(lài)方面的優(yōu)秀前景，并且也讓人看到了其出色的效率和可擴展性。自 ViT 在前段時(shí)間誕生以來(lái)，其已經(jīng)被用在了目標檢測、視頻識別、多任務(wù)預訓練等多種不同任務(wù)中。

近日，加州大學(xué)圣迭戈分校與 Google Research 的一項研究提出了使用視覺(jué) Transformer 來(lái)訓練 GAN。這篇論文的研究議題是：不使用卷積或池化，能否使用視覺(jué) Transformer 來(lái)完成圖像生成任務(wù)？更具體而言：能否使用 ViT 來(lái)訓練生成對抗網(wǎng)絡(luò )（GAN）并使之達到與已被廣泛研究過(guò)的基于 CNN 的 GAN 相媲美的質(zhì)量？

論文鏈接：https://arxiv.org/pdf/2107.04589.pdf

為此，研究者遵照最本原的 ViT 設計，使用純粹基本的 ViT（如圖 2(A)）訓練了 GAN。其中的難點(diǎn)在于，GAN 的訓練過(guò)程在與 ViT 耦合之后會(huì )變得非常不穩定，并且對抗訓練常常會(huì )在判別器訓練的后期受到高方差梯度（或尖峰梯度）的阻礙。此外，梯度懲罰、譜歸一化等傳統的正則化方法雖然能有效地用于基于 CNN 的 GAN 模型（如圖 4），但這些正則化方法卻無(wú)法解決上述不穩定問(wèn)題。使用了適當的正則化方法后，基于 CNN 的 GAN 訓練不穩定的情況并不常見(jiàn)，因此對基于 ViT 的 GAN 而言，這是一個(gè)獨有的挑戰。

針對這些問(wèn)題，為了實(shí)現訓練動(dòng)態(tài)的穩定以及促進(jìn)基于 ViT 的 GAN 的收斂，這篇論文提出了多項必需的修改。

在判別器中，研究者重新審視了自注意力的 Lipschitz 性質(zhì)，在此基礎上他們設計了一種加強了 Lipschitz 連續性的譜歸一化。不同于難以應付不穩定情況的傳統譜歸一化方法，這些技術(shù)能非常有效地穩定基于 ViT 的判別器的訓練動(dòng)態(tài)。此外，為了驗證新提出的技術(shù)的作用，研究者還執行了控制變量研究。對于基于 ViT 的生成器，研究者嘗試了多種不同的架構設計并發(fā)現了對層歸一化和輸出映射層的兩項關(guān)鍵性修改。實(shí)驗表明，不管使用的判別器是基于 ViT 還是基于 CNN，基于修改版 ViT 的生成器都能更好地促進(jìn)對抗訓練。

為了更具說(shuō)服力，研究者在三個(gè)標準的圖像合成基準上進(jìn)行了實(shí)驗。結果表明，新提出的模型 ViTGAN 極大優(yōu)于之前的基于 Transformer 的 GAN 模型，并且在沒(méi)有使用卷積和池化時(shí)也取得了與 StyleGAN2 等領(lǐng)先的基于 CNN 的 GAN 相媲美的表現。作者表示，新提出的 ViTGAN 算得上是在 GAN 中使用視覺(jué) Transformer 的最早嘗試之一，更重要的是，這項研究首次表明 Transformer 能在 CIFAR、CelebA 和 LSUN 臥室數據集等標準圖像生成基準上超過(guò)當前最佳的卷積架構。

方法

圖 1 展示了新提出的 ViTGAN 架構，其由一個(gè) ViT 判別器和一個(gè)基于 ViT 的生成器構成。研究者發(fā)現，直接使用 ViT 作為判別器會(huì )讓訓練不穩定。為了穩定訓練動(dòng)態(tài)和促進(jìn)收斂，研究者為生成器和判別器都引入了新技術(shù)：(1) ViT 判別器上的正則化和 (2) 新的生成器架構。

圖 1：新提出的 ViTGAN 框架示意圖。生成器和判別器都是基于視覺(jué) Transformer（ViT）設計的。判別器分數是從分類(lèi)嵌入推導得到的（圖中記為 *）；生成器是基于 patch 嵌入逐個(gè) patch 生成像素。

增強 Transformer 判別器的 Lipschitz 性質(zhì)。在 GAN 判別器中，Lipschitz 連續性發(fā)揮著(zhù)重要的作用。人們最早注意到它的時(shí)候是將其用作近似 WGAN 中 Wasserstein 距離的一個(gè)條件，之后其又在使用 Wasserstein 損失之外的其它 GAN 設置中得到了確認。其中，尤其值得關(guān)注的是 ICML 2019 論文《Lipschitz generative adversarial nets》，該研究證明 Lipschitz 判別器能確保存在最優(yōu)的判別函數以及唯一的納什均衡。但是，ICML 2021 的一篇論文《The lipschitz constant of self-attention》表明標準點(diǎn)積自注意力層的 Lipschitz 常數可以是****的，這就會(huì )破壞 ViT 中的 Lipschitz 連續性。為了加強 ViT 判別器的 Lipschitz 性質(zhì)，研究者采用了上述論文中提出的 L2 注意力。如等式 7 所示，點(diǎn)積相似度被替換成了歐幾里得距離，并且還關(guān)聯(lián)了投影矩陣的權重，以用于自注意力中的查詢(xún)和鍵（key）。這項改進(jìn)能提升用于 GAN 判別器的 Transformer 的穩定性。

經(jīng)過(guò)改進(jìn)的譜歸一化。為了進(jìn)一步強化 Lipschitz 連續性，研究者還在判別器訓練中使用了譜歸一化。標準譜歸一化是使用冪迭代來(lái)估計每層神經(jīng)網(wǎng)絡(luò )的投影矩陣的譜范數，然后再使用估計得到的譜范數來(lái)除權重矩陣，這樣所得到的投影矩陣的 Lipschitz 常量就等于 1。研究者發(fā)現，Transformer 模塊對 Lipschitz 常數的大小很敏感，當使用了譜歸一化時(shí)，訓練速度會(huì )非常慢。類(lèi)似地，研究者還發(fā)現當使用了基于 ViT 的判別器時(shí)，R1 梯度懲罰項會(huì )有損 GAN 訓練。另有研究發(fā)現，如果 MLP 模塊的 Lipschitz 常數較小，則可能導致 Transformer 的輸出坍縮為秩為 1 的矩陣。為了解決這個(gè)問(wèn)題，研究者提出增大投影矩陣的譜范數。

他們發(fā)現，只需在初始化時(shí)將譜范數與每一層的歸一化權重矩陣相乘，便足以解決這個(gè)問(wèn)題。具體而言，譜歸一化的更新規則如下，其中 σ 是計算權重矩陣的標準譜范：

重疊圖像塊。由于 ViT 判別器具有過(guò)多的學(xué)習能力，因此容易過(guò)擬合。在這項研究中，判別器和生成器使用了同樣的圖像表征，其會(huì )根據一個(gè)預定義的網(wǎng)絡(luò ) P×P 來(lái)將圖像分割為由非重疊 patch 組成的序列。如果不經(jīng)過(guò)精心設計，這些任意的網(wǎng)絡(luò )劃分可能會(huì )促使判別器記住局部線(xiàn)索，從而無(wú)法為生成器提供有意義的損失。為了解決這個(gè)問(wèn)題，研究者采用了一種簡(jiǎn)單技巧，即讓 patch 之間有所重疊。對于 patch 的每個(gè)邊緣，都將其擴展 o 個(gè)像素，使有效 patch 尺寸變?yōu)?(P+2o)×(P+2o)。

這樣得到的序列長(cháng)度與原來(lái)一樣，但對預定義網(wǎng)格的敏感度更低。這也有可能讓 Transformer 更好地了解當前 patch 的鄰近 patch 是哪些，由此更好地理解局部特性。

生成器設計

基于 ViT 架構設計生成器并非易事，其中一大難題是將 ViT 的功能從預測一組類(lèi)別標簽轉向在一個(gè)空間區域生成像素。

圖 2：生成器架構。左圖是研究者研究過(guò)的三種生成器架構：(A) 為每個(gè)位置嵌入添加中間隱藏嵌入 w，(B) 將 w 預置到序列上，(C) 使用由 w 學(xué)習到的仿射變換（圖中的 A）計算出的自調制型層范數（SLN/self-modulated layernorm）替換歸一化。右圖是用在 Transformer 模塊中的自調制運算的細節。

研究者先研究了多種生成器架構，發(fā)現它們都比不上基于 CNN 的生成器。于是他們遵循 ViT 的設計原理提出了一種全新的生成器。圖 2(c) 展示了這種 ViTGAN 生成器，其包含兩大組件：Transformer 模塊和輸出映射層。

為了促進(jìn)訓練過(guò)程，研究者為新提出的生成器做出了兩項改進(jìn)：

自調制型層范數（SLN）。新的做法不是將噪聲向量 z 作為輸入發(fā)送給 ViT，而是使用 z 來(lái)調制層范數運算。之所以稱(chēng)這樣的操作為自調制，是因為該過(guò)程無(wú)需外部信息；

用于圖塊生成的隱式神經(jīng)表征。為了學(xué)習從 patch 嵌入到 patch 像素值的連續映射，研究者使用了隱式神經(jīng)表征。當結合傅里葉特征或正弦激活函數一起使用時(shí)，隱式表征可將所生成的樣本空間約束到平滑變化的自然信號空間。研究發(fā)現，在使用基于 ViT 的生成器訓練 GAN 時(shí)，隱式表征的作用尤其大。

需要指出，由于生成器和判別器的圖像網(wǎng)格不同，因此序列長(cháng)度也不一樣。進(jìn)一步的研究發(fā)現，當需要將模型擴展用于更高分辨率的圖像時(shí)，只需增大判別器的序列長(cháng)度或特征維度就足夠了。

實(shí)驗結果

表 1：幾種代表性 GAN 架構在無(wú)條件圖像生成基準的結果比較。Conv 和 Pool 各自代表卷積和池化?！?表示越低越好；↑ 表示越高越好。

表 1 給出了在圖像合成的三個(gè)標準基準上的主要結果。本論文提出的新方法能與以下基準架構比肩。TransGAN 是目前唯一完全不使用卷積的 GAN，其完全基于 Transformer 構建。這里比較的是其最佳的變體版本 TransGAN-XL。Vanilla-ViT 是一種基于 ViT 的 GAN，其使用了圖 2(A) 的生成器和純凈版 ViT 判別器，但未使用本論文提出的改進(jìn)技術(shù)。

表 3a 中分別比較了圖 2(B) 所示的生成器架構。此外，BigGAN 和 StyleGAN2 作為基于 CNN 的 GAN 的最佳模型也被納入了比較。

圖 3：定性比較。在 CIFAR-10 32 × 32、CelebA 64 × 64 和 LSUN Bedroom 64 × 64 數據集上，ViTGAN 與 StyleGAN2、Transformer 最佳基準、純凈版生成器和判別器的 ViT 的結果比較。

圖 4：(a-c) ViT 判別器的梯度幅度（在所有參數上的 L2 范數），(d-f) FID 分數（越低越好）隨訓練迭代的變化情況。

可以看到，新提出方法的表現與使用 R1 懲罰項和譜范數的兩個(gè)純凈版 ViT 判別器基準相當。其余架構對所有方法來(lái)說(shuō)都一樣?？梢?jiàn)新方法能克服梯度幅度的尖峰并實(shí)現顯著(zhù)更低的 FID（在 CIFAR 和 CelebA 上）或相近的 FID（在 LSUN 上）。

表 3：在 CIFAR-10 數據集上對 ViTGAN 執行的控制變量研究。左圖：對生成器架構的控制變量研究。右圖：對判別器架構的控制變量研究。