英偉達公布StyleGAN-XL：參數量3倍於StyleGAN3，計算時間僅為1/5

科技 06-06

選自arXiv

作者：AXEL SAUER

機器之心編譯

編輯：蛋醬

StyleGAN-XL 首次在 ImageNet 上實現了 1024^2 解析度圖像合成。

近年來，計算機圖形學領域的研究者一直致力於生成高解析度的模擬圖像，並經歷了一波以數據為中心的真實可控內容創作浪潮。其中英偉達的 StyleGAN 在圖像質量和可控性方面為生成模型樹立了新的標杆。

但是，當用 ImageNet 這樣的大型非結構化數據集進行訓練時，StyleGAN 還不能取得令人滿意的結果。另一個存在的問題是，當需要更大的模型時，或擴展到更高的解析度時，這些方法的成本會高得令人望而卻步。

比如，英偉達的 StyleGAN3 項目消耗了令人難以想像的資源和電力。研究者在論文中表示，整個項目在 NVIDIA V100 內部集群上消耗了 92 個 GPU year（即單個 GPU 一年的計算）和 225 兆瓦時（Mwh）的電力。有人說，這相當於整個核反應堆運行大約 15 分鐘。

最初，StyleGAN 的提出是為了明確區分變數因素，實現更好的控制和插值質量。但它的體系架構比標準的生成器網路更具限制性，這些限制似乎會在諸如 ImageNet 這種複雜和多樣化的數據集上訓練時帶來相應代價。

此前有研究者嘗試將 StyleGAN 和 StyleGAN2 擴展到 ImageNet [Grigoryev et al. 2022; Gwern 2020]，導致結果欠佳。這讓人們更加相信，對於高度多樣化的數據集來說，StyleGAN 可能會從根本上受到限制。

受益於更大的 batch 和模型尺寸，BigGAN [Brock et al. 2019] 是 ImageNet 上的圖像合成 SOTA 模型。最近，BigGAN 的性能表現正在被擴散模型 [Dhariwal and Nichol 2021] 超越。也有研究發現，擴散模型能比 GAN 實現更多樣化的圖像合成，但是在推理過程中速度明顯減慢，以前的基於 GAN 的編輯工作不能直接應用。

此前在擴展 StyleGAN 上的失敗嘗試引出了這樣一個問題：架構約束是否從根本上限制了基於 Style 的生成器，或者 missing piece 是否是正確的訓練策略。最近的一項工作 [Sauer et al. 2021] 引入了 Projected GAN，將生成和實際的樣本投射到一個固定的、預訓練的特徵空間。重組 GAN 設置這種方式顯著改進了訓練穩定性、訓練時間和數據效率。然而，Projected GAN 的優勢只是部分地延伸到了這項研究的單模態數據集上的 StyleGAN。

為了解決上述種種問題，英偉達的研究者近日提出了一種新的架構變化，並根據最新的 StyleGAN3 設計了漸進式生長的策略。研究者將改進後的模型稱為 StyleGAN-XL，該研究目前已經入選了 SIGGRAPH 2022。

論文地址：https://arxiv.org/pdf/2202.00273.pdf
代碼地址：https://github.com/autonomousvision/stylegan_xl

這些變化結合了 Projected GAN 方法，超越了此前在 ImageNet 上訓練 StyleGAN 的表現。為了進一步改進結果，研究者分析了 Projected GAN 的預訓練特徵網路，發現當計算機視覺的兩種標準神經結構 CNN 和 ViT [ Dosovitskiy et al. 2021] 聯合使用時，性能顯著提高。最後，研究者利用了分類器引導這種最初為擴散模型引入的技術，用以注入額外的類信息。

總體來說，這篇論文的貢獻在於推動模型性能超越現有的 GAN 和擴散模型，實現了大規模圖像合成 SOTA。論文展示了 ImageNet 類的反演和編輯，發現了一個強大的新反演範式 Pivotal Tuning Inversion (PTI)[ Roich et al. 2021] ，這一範式能夠與模型很好地結合，甚至平滑地嵌入域外圖像到學習到的潛在空間。高效的訓練策略使得標準 StyleGAN3 的參數能夠增加三倍，同時僅用一小部分訓練時間就達到擴散模型的 SOTA 性能。

這使得 StyleGAN-XL 能夠成為第一個在 ImageNet-scale 上演示 1024^2 解析度圖像合成的模型。

將 StyleGAN 擴展到 ImageNet

實驗表明，即使是最新的 StyleGAN3 也不能很好地擴展到 ImageNet 上，如圖 1 所示。特別是在高解析度時，訓練會變得不穩定。因此，研究者的第一個目標是在 ImageNet 上成功地訓練一個 StyleGAN3 生成器。成功的定義取決於主要通過初始評分 (IS)[Salimans et al. 2016] 衡量的樣本質量和 Fréchet 初始距離 (FID)[Heusel et al. 2017] 衡量的多樣性。

在論文中，研究者也介紹了 StyleGAN3 baseline 進行的改動，所帶來的提升如下表 1 所示：

研究者首先修改了生成器及其正則化損失，調整了潛在空間以適應 Projected GAN (Config-B) 和類條件設置 (Config-C)；然後重新討論了漸進式增長，以提高訓練速度和性能 (Config-D)；接下來研究了用於 Projected GAN 訓練的特徵網路，以找到一個非常適合的配置 (Config-E)；最後，研究者提出了分類器引導，以便 GAN 通過一個預訓練的分類器 (Config-F) 提供類信息。

這樣一來，就能夠訓練一個比以前大得多的模型，同時需要比現有技術更少的計算量。StyleGAN-XL 在深度和參數計數方面比標準的 StyleGAN3 大三倍。然而，為了在 512^2 像素的解析度下匹配 ADM [Dhariwal and Nichol 2021] 先進的性能，在一台 NVIDIA Tesla V100 上訓練模型需要 400 天，而以前需要 1914 天。(圖 2)。

實驗結果

在實驗中，研究者首先將 StyleGAN-XL 與 ImageNet 上的 SOTA 圖像合成方法進行比較。然後對 StyleGAN-XL 的反演和編輯性能進行了評價。研究者將模型擴展到了 1024^2 像素的解析度，這是之前在 ImageNet 上沒有嘗試過的。在 ImageNet 中，大多數圖像的解析度較低，因此研究者用超解析度網路 [Liang et al. 2021] 對數據進行了預處理。

圖像合成

如表 2 所示，研究者在 ImageNet 上對比了 StyleGAN-XL 和現有最強大的 GAN 模型及擴散模型的圖像合成性能。