Transformer 已經(jīng)為多種自然語言任務帶來了突飛猛進的進步，并且最近也已經(jīng)開始向計算機視覺領域滲透，開始在一些之前由 CNN 主導的任務上暫露頭角。近日，加州大學圣迭戈分校與 Google Research 的一項研究提出了使用視覺 Transformer 來訓練 GAN。為了有效應用該方法，研究者還提出了多項改進技巧，使新方法在一些指標上可比肩前沿 CNN 模型。
卷積神經(jīng)網(wǎng)絡（CNN）在卷積（權重共享和局部連接）和池化（平移等變）方面的強大能力，讓其已經(jīng)成為了現(xiàn)今計算機視覺領域的主導技術。但最近，Transformer 架構已經(jīng)開始在圖像和視頻識別任務上與 CNN 比肩。其中尤其值得一提的是視覺 Transformer（ViT）。這種技術會將圖像作為 token 序列（類似于自然語言中的詞）來解讀。Dosovitskiy et al. 的研究表明，ViT 在 ImageNet 基準上能以更低的計算成本取得相當?shù)姆诸悳蚀_度。不同于 CNN 中的局部連接性，ViT 依賴于在全局背景中考慮的表征，其中每個 patch 都必須與同一圖像的所有 patch 都關聯(lián)處理。 ViT 及其變體盡管還處于早期階段，但已有研究展現(xiàn)了其在建模非局部上下文依賴方面的優(yōu)秀前景，并且也讓人看到了其出色的效率和可擴展性。自 ViT 在前段時間誕生以來，其已經(jīng)被用在了目標檢測、視頻識別、多任務預訓練等多種不同任務中。 近日，加州大學圣迭戈分校與 Google Research 的一項研究提出了使用視覺 Transformer 來訓練 GAN。這篇論文的研究議題是：不使用卷積或池化，能否使用視覺 Transformer 來完成圖像生成任務？更具體而言：能否使用 ViT 來訓練生成對抗網(wǎng)絡（GAN）并使之達到與已被廣泛研究過的基于 CNN 的 GAN 相媲美的質(zhì)量？
論文鏈接：https:///pdf/2107.04589.pdf 為此，研究者遵照最本原的 ViT 設計，使用純粹基本的 ViT（如圖 2(A)）訓練了 GAN。其中的難點在于，GAN 的訓練過程在與 ViT 耦合之后會變得非常不穩(wěn)定，并且對抗訓練常常會在判別器訓練的后期受到高方差梯度（或尖峰梯度）的阻礙。此外，梯度懲罰、譜歸一化等傳統(tǒng)的正則化方法雖然能有效地用于基于 CNN 的 GAN 模型（如圖 4），但這些正則化方法卻無法解決上述不穩(wěn)定問題。使用了適當?shù)恼齽t化方法后，基于 CNN 的 GAN 訓練不穩(wěn)定的情況并不常見，因此對基于 ViT 的 GAN 而言，這是一個獨有的挑戰(zhàn)。 針對這些問題，為了實現(xiàn)訓練動態(tài)的穩(wěn)定以及促進基于 ViT 的 GAN 的收斂，這篇論文提出了多項必需的修改。 在判別器中，研究者重新審視了自注意力的 Lipschitz 性質(zhì)，在此基礎上他們設計了一種加強了 Lipschitz 連續(xù)性的譜歸一化。不同于難以應付不穩(wěn)定情況的傳統(tǒng)譜歸一化方法，這些技術能非常有效地穩(wěn)定基于 ViT 的判別器的訓練動態(tài)。此外，為了驗證新提出的技術的作用，研究者還執(zhí)行了控制變量研究。對于基于 ViT 的生成器，研究者嘗試了多種不同的架構設計并發(fā)現(xiàn)了對層歸一化和輸出映射層的兩項關鍵性修改。實驗表明，不管使用的判別器是基于 ViT 還是基于 CNN，基于修改版 ViT 的生成器都能更好地促進對抗訓練。 為了更具說服力，研究者在三個標準的圖像合成基準上進行了實驗。結(jié)果表明，新提出的模型 ViTGAN 極大優(yōu)于之前的基于 Transformer 的 GAN 模型，并且在沒有使用卷積和池化時也取得了與 StyleGAN2 等領先的基于 CNN 的 GAN 相媲美的表現(xiàn)。作者表示，新提出的 ViTGAN 算得上是在 GAN 中使用視覺 Transformer 的最早嘗試之一，更重要的是，這項研究首次表明 Transformer 能在 CIFAR、CelebA 和 LSUN 臥室數(shù)據(jù)集等標準圖像生成基準上超過當前最佳的卷積架構。 方法 圖 1 展示了新提出的 ViTGAN 架構，其由一個 ViT 判別器和一個基于 ViT 的生成器構成。研究者發(fā)現(xiàn)，直接使用 ViT 作為判別器會讓訓練不穩(wěn)定。為了穩(wěn)定訓練動態(tài)和促進收斂，研究者為生成器和判別器都引入了新技術：(1) ViT 判別器上的正則化和 (2) 新的生成器架構。 圖 1：新提出的 ViTGAN 框架示意圖。生成器和判別器都是基于視覺 Transformer（ViT）設計的。判別器分數(shù)是從分類嵌入推導得到的（圖中記為 *）；生成器是基于 patch 嵌入逐個 patch 生成像素。 增強 Transformer 判別器的 Lipschitz 性質(zhì)。在 GAN 判別器中，Lipschitz 連續(xù)性發(fā)揮著重要的作用。人們最早注意到它的時候是將其用作近似 WGAN 中 Wasserstein 距離的一個條件，之后其又在使用 Wasserstein 損失之外的其它 GAN 設置中得到了確認。其中，尤其值得關注的是 ICML 2019 論文《Lipschitz generative adversarial nets》，該研究證明 Lipschitz 判別器能確保存在最優(yōu)的判別函數(shù)以及唯一的納什均衡。但是，ICML 2021 的一篇論文《The lipschitz constant of self-attention》表明標準點積自注意力層的 Lipschitz 常數(shù)可以是無界的，這就會破壞 ViT 中的 Lipschitz 連續(xù)性。為了加強 ViT 判別器的 Lipschitz 性質(zhì)，研究者采用了上述論文中提出的 L2 注意力。如等式 7 所示，點積相似度被替換成了歐幾里得距離，并且還關聯(lián)了投影矩陣的權重，以用于自注意力中的查詢和鍵（key）。這項改進能提升用于 GAN 判別器的 Transformer 的穩(wěn)定性。 經(jīng)過改進的譜歸一化。為了進一步強化 Lipschitz 連續(xù)性，研究者還在判別器訓練中使用了譜歸一化。標準譜歸一化是使用冪迭代來估計每層神經(jīng)網(wǎng)絡的投影矩陣的譜范數(shù)，然后再使用估計得到的譜范數(shù)來除權重矩陣，這樣所得到的投影矩陣的 Lipschitz 常量就等于 1。研究者發(fā)現(xiàn)，Transformer 模塊對 Lipschitz 常數(shù)的大小很敏感，當使用了譜歸一化時，訓練速度會非常慢。類似地，研究者還發(fā)現(xiàn)當使用了基于 ViT 的判別器時，R1 梯度懲罰項會有損 GAN 訓練。另有研究發(fā)現(xiàn)，如果 MLP 模塊的 Lipschitz 常數(shù)較小，則可能導致 Transformer 的輸出坍縮為秩為 1 的矩陣。為了解決這個問題，研究者提出增大投影矩陣的譜范數(shù)。 他們發(fā)現(xiàn)，只需在初始化時將譜范數(shù)與每一層的歸一化權重矩陣相乘，便足以解決這個問題。具體而言，譜歸一化的更新規(guī)則如下，其中 σ 是計算權重矩陣的標準譜范： 重疊圖像塊。由于 ViT 判別器具有過多的學習能力，因此容易過擬合。在這項研究中，判別器和生成器使用了同樣的圖像表征，其會根據(jù)一個預定義的網(wǎng)絡 P×P 來將圖像分割為由非重疊 patch 組成的序列。如果不經(jīng)過精心設計，這些任意的網(wǎng)絡劃分可能會促使判別器記住局部線索，從而無法為生成器提供有意義的損失。為了解決這個問題，研究者采用了一種簡單技巧，即讓 patch 之間有所重疊。對于 patch 的每個邊緣，都將其擴展 o 個像素，使有效 patch 尺寸變?yōu)?(P+2o)×(P+2o)。 這樣得到的序列長度與原來一樣，但對預定義網(wǎng)格的敏感度更低。這也有可能讓 Transformer 更好地了解當前 patch 的鄰近 patch 是哪些，由此更好地理解局部特性。 生成器設計 基于 ViT 架構設計生成器并非易事，其中一大難題是將 ViT 的功能從預測一組類別標簽轉(zhuǎn)向在一個空間區(qū)域生成像素。 圖 2：生成器架構。左圖是研究者研究過的三種生成器架構：(A) 為每個位置嵌入添加中間隱藏嵌入 w，(B) 將 w 預置到序列上，(C) 使用由 w 學習到的仿射變換（圖中的 A）計算出的自調(diào)制型層范數(shù)（SLN/self-modulated layernorm）替換歸一化。右圖是用在 Transformer 模塊中的自調(diào)制運算的細節(jié)。 研究者先研究了多種生成器架構，發(fā)現(xiàn)它們都比不上基于 CNN 的生成器。于是他們遵循 ViT 的設計原理提出了一種全新的生成器。圖 2(c) 展示了這種 ViTGAN 生成器，其包含兩大組件：Transformer 模塊和輸出映射層。 為了促進訓練過程，研究者為新提出的生成器做出了兩項改進：
自調(diào)制型層范數(shù)（SLN）。新的做法不是將噪聲向量 z 作為輸入發(fā)送給 ViT，而是使用 z 來調(diào)制層范數(shù)運算。之所以稱這樣的操作為自調(diào)制，是因為該過程無需外部信息；
用于圖塊生成的隱式神經(jīng)表征。為了學習從 patch 嵌入到 patch 像素值的連續(xù)映射，研究者使用了隱式神經(jīng)表征。當結(jié)合傅里葉特征或正弦激活函數(shù)一起使用時，隱式表征可將所生成的樣本空間約束到平滑變化的自然信號空間。研究發(fā)現(xiàn)，在使用基于 ViT 的生成器訓練 GAN 時，隱式表征的作用尤其大。
需要指出，由于生成器和判別器的圖像網(wǎng)格不同，因此序列長度也不一樣。進一步的研究發(fā)現(xiàn)，當需要將模型擴展用于更高分辨率的圖像時，只需增大判別器的序列長度或特征維度就足夠了。 實驗結(jié)果 表 1：幾種代表性 GAN 架構在無條件圖像生成基準的結(jié)果比較。Conv 和 Pool 各自代表卷積和池化?！?表示越低越好；↑ 表示越高越好。 表 1 給出了在圖像合成的三個標準基準上的主要結(jié)果。本論文提出的新方法能與以下基準架構比肩。TransGAN 是目前唯一完全不使用卷積的 GAN，其完全基于 Transformer 構建。這里比較的是其最佳的變體版本 TransGAN-XL。Vanilla-ViT 是一種基于 ViT 的 GAN，其使用了圖 2(A) 的生成器和純凈版 ViT 判別器，但未使用本論文提出的改進技術。 表 3a 中分別比較了圖 2(B) 所示的生成器架構。此外，BigGAN 和 StyleGAN2 作為基于 CNN 的 GAN 的最佳模型也被納入了比較。 圖 3：定性比較。在 CIFAR-10 32 × 32、CelebA 64 × 64 和 LSUN Bedroom 64 × 64 數(shù)據(jù)集上，ViTGAN 與 StyleGAN2、Transformer 最佳基準、純凈版生成器和判別器的 ViT 的結(jié)果比較。 圖 4：(a-c) ViT 判別器的梯度幅度（在所有參數(shù)上的 L2 范數(shù)），(d-f) FID 分數(shù)（越低越好）隨訓練迭代的變化情況。 可以看到，新提出方法的表現(xiàn)與使用 R1 懲罰項和譜范數(shù)的兩個純凈版 ViT 判別器基準相當。其余架構對所有方法來說都一樣?？梢娦路椒芸朔荻确鹊募夥宀崿F(xiàn)顯著更低的 FID（在 CIFAR 和 CelebA 上）或相近的 FID（在 LSUN 上）。
