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

選自arXiv

作者:Zhuowen Tu、Ce Liu等

機器之心編譯

編輯:Panda

Transformer 已經為多種自然語言任務帶來了突飛猛進的進步,並且最近也已經開始向計算機視覺領域滲透,開始在一些之前由 CNN 主導的任務上暫露頭角。近日,加州大學聖迭戈分校與 Google Research 的一項研究提出了使用視覺 Transformer 來訓練 GAN。為了有效應用該方法,研究者還提出了多項改進技巧,使新方法在一些指標上可比肩前沿 CNN 模型。

卷積神經網路(CNN)在卷積(權重共享和局部連接)和池化(平移等變)方面的強大能力,讓其已經成為了現今計算機視覺領域的主導技術。但最近,Transformer 架構已經開始在圖像和視訊識別任務上與 CNN 比肩。其中尤其值得一提的是視覺 Transformer(ViT)。這種技術會將圖像作為 token 序列(類似於自然語言中的詞)來解讀。Dosovitskiy et al. 的研究表明,ViT 在 ImageNet 基準上能以更低的計算成本取得相當的分類準確度。不同於 CNN 中的局部連接性,ViT 依賴於在全局背景中考慮的表徵,其中每個 patch 都必須與同一圖像的所有 patch 都關聯處理。

ViT 及其變體儘管還處於早期階段,但已有研究展現了其在建模非局部上下文依賴方面的優秀前景,並且也讓人看到了其出色的效率和可擴展性。自 ViT 在前段時間誕生以來,其已經被用在了目標檢測、視訊識別、多工預訓練等多種不同任務中。

近日,加州大學聖迭戈分校與 Google Research 的一項研究提出了使用視覺 Transformer 來訓練 GAN。這篇論文的研究議題是:不使用卷積或池化,能否使用視覺 Transformer 來完成圖像生成任務?更具體而言:能否使用 ViT 來訓練生成對抗網路(GAN)並使之達到與已被廣泛研究過的基於 CNN 的 GAN 相媲美的質量?

論文連結:https://arxiv.org/pdf/2107.04589.pdf

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

針對這些問題,為了實現訓練動態的穩定以及促進基於 ViT 的 GAN 的收斂,這篇論文提出了多項必需的修改。

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

為了更具說服力,研究者在三個標準的圖像合成基準上進行了實驗。結果表明,新提出的模型 ViTGAN 極大優於之前的基於 Transformer 的 GAN 模型,並且在沒有使用卷積和池化時也取得了與 StyleGAN2 等領先的基於 CNN 的 GAN 相媲美的表現。作者表示,新提出的 ViTGAN 算得上是在 GAN 中使用視覺 Transformer 的最早嘗試之一,更重要的是,這項研究首次表明 Transformer 能在 CIFAR、CelebA 和 LSUN 臥室資料集等標準圖像生成基準上超過當前最佳的卷積架構。

方法

圖 1 展示了新提出的 ViTGAN 架構,其由一個 ViT 判別器和一個基於 ViT 的生成器構成。研究者發現,直接使用 ViT 作為判別器會讓訓練不穩定。為了穩定訓練動態和促進收斂,研究者為生成器和判別器都引入了新技術:(1) ViT 判別器上的正則化和 (2) 新的生成器架構。

圖 1:新提出的 ViTGAN 框架示意圖。生成器和判別器都是基於視覺 Transformer(ViT)設計的。判別器分數是從分類嵌入推導得到的(圖中記為 *);生成器是基於 patch 嵌入逐個 patch 生成像素。

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

經過改進的譜歸一化。為了進一步強化 Lipschitz 連續性,研究者還在判別器訓練中使用了譜歸一化。標準譜歸一化是使用冪迭代來估計每層神經網路的投影矩陣的譜範數,然後再使用估計得到的譜範數來除權重矩陣,這樣所得到的投影矩陣的 Lipschitz 常量就等於 1。研究者發現,Transformer 模組對 Lipschitz 常數的大小很敏感,當使用了譜歸一化時,訓練速度會非常慢。類似地,研究者還發現當使用了基於 ViT 的判別器時,R1 梯度懲罰項會有損 GAN 訓練。另有研究發現,如果 MLP 模組的 Lipschitz 常數較小,則可能導致 Transformer 的輸出坍縮為秩為 1 的矩陣。為了解決這個問題,研究者提出增大投影矩陣的譜範數。

他們發現,只需在初始化時將譜範數與每一層的歸一化權重矩陣相乘,便足以解決這個問題。具體而言,譜歸一化的更新規則如下,其中 σ 是計算權重矩陣的標準譜範:

重疊圖像塊。由於 ViT 判別器具有過多的學習能力,因此容易過擬合。在這項研究中,判別器和生成器使用了同樣的圖像表徵,其會根據一個預定義的網路 P×P 來將圖像分割為由非重疊 patch 組成的序列。如果不經過精心設計,這些任意的網路劃分可能會促使判別器記住局部線索,從而無法為生成器提供有意義的損失。為了解決這個問題,研究者採用了一種簡單技巧,即讓 patch 之間有所重疊。對於 patch 的每個邊緣,都將其擴展 o 個像素,使有效 patch 尺寸變為 (P+2o)×(P+2o)。

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

生成器設計

基於 ViT 架構設計生成器並非易事,其中一大難題是將 ViT 的功能從預測一組類別標籤轉向在一個空間區域生成像素。

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

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

為了促進訓練過程,研究者為新提出的生成器做出了兩項改進:

  • 自調製型層範數(SLN)。新的做法不是將噪聲向量 z 作為輸入發送給 ViT,而是使用 z 來調製層範數運算。之所以稱這樣的操作為自調製,是因為該過程無需外部資訊;

  • 用於圖塊生成的隱式神經表徵。為了學習從 patch 嵌入到 patch 像素值的連續對映,研究者使用了隱式神經表徵。當結合傅立葉特徵或正弦激活函數一起使用時,隱式表徵可將所生成的樣本空間約束到平滑變化的自然信號空間。研究發現,在使用基於 ViT 的生成器訓練 GAN 時,隱式表徵的作用尤其大。

需要指出,由於生成器和判別器的圖像網格不同,因此序列長度也不一樣。進一步的研究發現,當需要將模型擴展用於更高解析度的圖像時,只需增大判別器的序列長度或特徵維度就足夠了。

實驗結果

表 1:幾種代表性 GAN 架構在無條件圖像生成基準的結果比較。Conv 和 Pool 各自代表卷積和池化。↓ 表示越低越好;↑ 表示越高越好。

表 1 給出了在圖像合成的三個標準基準上的主要結果。本論文提出的新方法能與以下基準架構比肩。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 資料集上,ViTGAN 與 StyleGAN2、Transformer 最佳基準、純淨版生成器和判別器的 ViT 的結果比較。

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

可以看到,新提出方法的表現與使用 R1 懲罰項和譜範數的兩個純淨版 ViT 判別器基準相當。其餘架構對所有方法來說都一樣。可見新方法能克服梯度幅度的尖峰並實現顯著更低的 FID(在 CIFAR 和 CelebA 上)或相近的 FID(在 LSUN 上)。

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

亞馬遜雲科技中國峰會

2021亞馬遜雲科技中國峰會將在中國上海、北京、深圳三大城市舉辦。本次峰會以「構建新格局,重塑雲時代」為題,並攜手眾多業內領先的技術踐行者們一起同你分享「雲時代的構建故事與重塑經驗」。

7月21日-22日,2021亞馬遜雲科技中國峰會上海站將有多位重磅業內專家及行業大咖在現場分享獨到的行業見解。

面向開發者,本次峰會將專門設立開發者專區,並聯合 Apache 等各類開源社區,以及多位開源牛人,為開發者們帶來腦洞大開的內容分享!

面向行業,近百位來自業內各領域的合作伙伴、客戶及亞馬遜雲科技技術專家,共同組成的強大嘉賓陣容,為你帶來行業最佳實踐分享及領先技術成果發佈解讀!

識別下方小程序,立即報名。

識別下方小程序,立即報名