港中文博士提出首個基於Transformer的條件GAN：成像質量不如CNN

編輯：LRS

【新智元導讀】Transformer在影象分類任務上經過充分訓練已經足以完全超越CNN模型，但GAN仍然是Transformer無法踏足的領域。最近港中文博士提出首個基於Transformer的條件GAN模型STransGAN，緩解了Transformer的部分問題，但成像質量仍不如CNN。

Transformer不僅在自然語言表達方面表現出色，在計算機視覺方面的潛力也被挖掘出來，不斷稱霸各大CV榜單。

Transformer的成功主要歸功於注意力層的表示學習能力，這種能力也能夠支援Transformer模型在其他領域上的應用。

最近，有研究人員將生成對抗網路（GAN）中常用的CNN主幹網路替換為Transformer用於影象合成。一些前期研究結果表明，將Transformer直接用於GAN中並不是一件很容易的事，特別是之前被設計用於分類的GAN中經常會導致CNN影象合成效能的下降。

例如與基於CNN的StyleGAN2實現的3。16 FID相比，ViT作為主幹網路的GAN模型僅在64×64 Celeba資料集中實現了8。92的FID。Frechet Inception Distance （FID）計算結果為真實影象計算的特徵向量與生成的影象之間的距離。

此外，這種Transformer結構也會使GAN網路的訓練變得更加不穩定，嚴重依賴於手工調整超引數。

針對這個問題，來自香港中文大學的研究人員發表了一篇論文，旨在瞭解GAN 模型中Transformer 的內在行為，以縮小基於Transformer的GAN模型與基於CNN主幹的GAN模型之間的效能差距。文中不僅研究了無條件的影象合成，而且還研究瞭如何更少地探索的條件設定。

這項研究也是首次在條件設計（conditional setting）下成功使用基於Transformer的GAN模型。

文章的第一作者是徐瑞，目前是香港中文大學多媒體實驗室四年級博士生，導師是湯曉鷗教授，本科畢業於清華大學電子工程系。主要研究方向為深入學習及其在計算機視覺中的應用，正在進行的工作包括影象/影片繪畫和影象合成。他在影象/影片分割、檢測和例項分割方面有豐富的研究經驗。

研究結果主要為Transformer 在GAN 中的使用提供了三點實踐和設計原則：

1、區域性性（locality）很重要。

特徵提取的區域性性對於影象分類中Transformer的效率和效能至關重要，在GAN 影象生成實驗中也可以觀察到同樣的結果。特別是，在現有基於Transformer 的GANs 中實施的全域性自注意力操作會降低影象合成效能，並且在計算上無法應用於高解析度影象生成。在這些方法中，Swin Layer被證明是提供區域性感應偏差的最有效的模組。

文中提出了一種新的基於Transformer的GAN網路架構設計，生成器部分稱為STrans-G。

模型首先從一個簡單的基線結構Trans-G開始，它由標準的視覺Transformer 塊組成。然而，Trans-G生成的樣本通常包含嚴重的人工痕跡（artifacts）和違和的細節，從而導致在視覺上看質量比較差。

透過分析注意力層的內在行為，可以發現全域性注意力總是打破影象資料的區域性性，特別是在合成高解析度特徵時。

這一發現也促使研究人員探索各種區域性注意機制在生成真實高解析度影象中的作用。在仔細比較了不同的區域性注意機制之後，最終選擇了Swin架構作為模型架構塊，構建了一個無CNN的生成器STrans-G。對注意力距離（attention distance）的進一步分析清楚地顯示了全域性注意和區域性注意之間的差異。

2、留意判別器（discriminator）中的殘差連線。

Transformer 在每個自注意力層的子層和點全連線層（pointwise fully connected layer）上使用殘差連線。透過對範數比的詳細分析，可以發現在基於Transformer 的判別器中，殘差連線往往主導資訊流。在判別器中執行自注意力和全連線操作的子層在訓練過程中被無意中繞過，從而導致影象合成質量低下和收斂緩慢。研究人員透過將每個殘差連線替換為跳躍投影層來解決這個問題，這樣可以更好地保持殘差塊中的資訊流。

判別器部分稱為STrans-D，包括了一些經驗上的策略。

首先採用輕量級卷積塊將原始輸入取樣降低4倍，並將影象張量投影到任意維度，而不是像大多數視覺Transformer那樣從embedding模組開始。卷積token 抽取器與patch embedding相比，採用了重疊的patch，保留了更多的細節資訊。

第二，研究人員在所有注意力模組和MLP中採用均衡學習率（equalized learning rate）。這是由判別器中的Transformer塊在使用小學習率來穩定其訓練時的緩慢和不滿意的收斂所做出的改變。透過在整個判別器中設定一個更大的學習率可以解決這一問題，並引入了一種特殊的sclaer，在執行時將Transformer 塊的可學習引數相乘。

此外，研究人員將GeLU替換為LeakyRelu，並在注意和MLP模組的末尾新增非線性啟用函式。

3、為Transformer 單獨設計策略，而非條件正則化。

研究人員發現，傳統的注入條件資訊的方法對基於Transformer的條件GAN沒有很好的效果。罪魁禍首是透過Transformer生成器中的殘差連線在大資訊流中。如果將條件資訊注入主分支，那基本上都被忽略掉了，對最終輸出幾乎沒有任何影響。研究人員提出了一種在主幹中採用條件歸一化層的可行方法，有助於在整個Transformer生成器中保留條件資訊。

一個嘗試是在Transformer塊中直接採用AdaLN，但沒有取得效果。特別是研究人員發現FID在前期的訓練不再下降了。為了找到故障的發生處，研究人員又繪製了此基線配置的標準比率。存在多個具有高範數比的塊，表明主分支中的一些AdaNorm層對中間特徵的貢獻很小，導致條件資訊丟失。

為了保證條件資訊的注入，一個簡單的解決方案是將AdaNorm應用於trunk。透過這種方式，保證了shortcut 和MLP分支的特性都包含類別的資訊。

在實驗部分，為了降低計算成本，研究人員將Transformer塊中MLP模組的通道擴充套件率設定為2。輸入token維度為512，預設採用四個注意力header。選擇Adam最佳化器（β1=0，β2=0。99）來訓練生成器和判別器。Strans-G和Strans-D分別以0。0001和0。002的學習率進行最佳化。

在無條件的生成中，Strans-G在Celeba 64x64中顯著地超出了之前的所有方法。它在FFHQ 256x256的高解析度設定方面也取得了相當的效能。