StyleGAN-XL 首次在 ImageNet 上實(shí)現(xiàn)了 1024^2 分辨率圖像合成。

近年來，計(jì)算機(jī)圖形學(xué)領(lǐng)域的研究者一直致力于生成高分辨率的仿真圖像，并經(jīng)歷了一波以數(shù)據(jù)為中心的真實(shí)可控內(nèi)容創(chuàng)作浪潮。其中英偉達(dá)的 StyleGAN 在圖像質(zhì)量和可控性方面為生成模型樹立了新的標(biāo)桿。

但是，當(dāng)用 ImageNet 這樣的大型非結(jié)構(gòu)化數(shù)據(jù)集進(jìn)行訓(xùn)練時(shí)，StyleGAN 還不能取得令人滿意的結(jié)果。另一個(gè)存在的問題是，當(dāng)需要更大的模型時(shí)，或擴(kuò)展到更高的分辨率時(shí)，這些方法的成本會(huì)高得令人望而卻步。

比如，英偉達(dá)的 StyleGAN3 項(xiàng)目消耗了令人難以想象的資源和電力。研究者在論文中表示，整個(gè)項(xiàng)目在 NVIDIA V100 內(nèi)部集群上消耗了 92 個(gè) GPU year（即單個(gè) GPU 一年的計(jì)算）和 225 兆瓦時(shí)（Mwh）的電力。有人說，這相當(dāng)于整個(gè)核反應(yīng)堆運(yùn)行大約 15 分鐘。

最初，StyleGAN 的提出是為了明確區(qū)分變量因素，實(shí)現(xiàn)更好的控制和插值質(zhì)量。但它的體系架構(gòu)比標(biāo)準(zhǔn)的生成器網(wǎng)絡(luò)更具限制性，這些限制似乎會(huì)在諸如 ImageNet 這種復(fù)雜和多樣化的數(shù)據(jù)集上訓(xùn)練時(shí)帶來相應(yīng)代價(jià)。

此前有研究者嘗試將 StyleGAN 和 StyleGAN2 擴(kuò)展到 ImageNet ［Grigoryev et al. 2022; Gwern 2020］，導(dǎo)致結(jié)果欠佳。這讓人們更加相信，對(duì)于高度多樣化的數(shù)據(jù)集來說，StyleGAN 可能會(huì)從根本上受到限制。

受益于更大的 batch 和模型尺寸，BigGAN ［Brock et al. 2019］ 是 ImageNet 上的圖像合成 SOTA 模型。最近，BigGAN 的性能表現(xiàn)正在被擴(kuò)散模型 ［Dhariwal and Nichol 2021］ 超越。也有研究發(fā)現(xiàn)，擴(kuò)散模型能比 GAN 實(shí)現(xiàn)更多樣化的圖像合成，但是在推理過程中速度明顯減慢，以前的基于 GAN 的編輯工作不能直接應(yīng)用。

此前在擴(kuò)展 StyleGAN 上的失敗嘗試引出了這樣一個(gè)問題：架構(gòu)約束是否從根本上限制了基于 Style 的生成器，或者 missing piece 是否是正確的訓(xùn)練策略。最近的一項(xiàng)工作 ［Sauer et al. 2021］ 引入了 Projected GAN，將生成和實(shí)際的樣本投射到一個(gè)固定的、預(yù)訓(xùn)練的特征空間。重組 GAN 設(shè)置這種方式顯著改進(jìn)了訓(xùn)練穩(wěn)定性、訓(xùn)練時(shí)間和數(shù)據(jù)效率。然而，Projected GAN 的優(yōu)勢(shì)只是部分地延伸到了這項(xiàng)研究的單模態(tài)數(shù)據(jù)集上的 StyleGAN。

為了解決上述種種問題，英偉達(dá)的研究者近日提出了一種新的架構(gòu)變化，并根據(jù)最新的 StyleGAN3 設(shè)計(jì)了漸進(jìn)式生長的策略。研究者將改進(jìn)后的模型稱為 StyleGAN-XL，該研究目前已經(jīng)入選了 SIGGRAPH 2022。

論文地址：https://arxiv.org/pdf/2202.00273.pdf

代碼地址：https://github.com/autonomousvision/stylegan_xl

這些變化結(jié)合了 Projected GAN 方法，超越了此前在 ImageNet 上訓(xùn)練 StyleGAN 的表現(xiàn)。為了進(jìn)一步改進(jìn)結(jié)果，研究者分析了 Projected GAN 的預(yù)訓(xùn)練特征網(wǎng)絡(luò)，發(fā)現(xiàn)當(dāng)計(jì)算機(jī)視覺的兩種標(biāo)準(zhǔn)神經(jīng)結(jié)構(gòu) CNN 和 ViT ［ Dosovitskiy et al. 2021］ 聯(lián)合使用時(shí)，性能顯著提高。最后，研究者利用了分類器引導(dǎo)這種最初為擴(kuò)散模型引入的技術(shù)，用以注入額外的類信息。

總體來說，這篇論文的貢獻(xiàn)在于推動(dòng)模型性能超越現(xiàn)有的 GAN 和擴(kuò)散模型，實(shí)現(xiàn)了大規(guī)模圖像合成 SOTA。論文展示了 ImageNet 類的反演和編輯，發(fā)現(xiàn)了一個(gè)強(qiáng)大的新反演范式 Pivotal Tuning Inversion （PTI）［ Roich et al. 2021］ ，這一范式能夠與模型很好地結(jié)合，甚至平滑地嵌入域外圖像到學(xué)習(xí)到的潛在空間。高效的訓(xùn)練策略使得標(biāo)準(zhǔn) StyleGAN3 的參數(shù)能夠增加三倍，同時(shí)僅用一小部分訓(xùn)練時(shí)間就達(dá)到擴(kuò)散模型的 SOTA 性能。

這使得 StyleGAN-XL 能夠成為第一個(gè)在 ImageNet-scale 上演示 1024^2 分辨率圖像合成的模型。

將 StyleGAN 擴(kuò)展到 ImageNet

實(shí)驗(yàn)表明，即使是最新的 StyleGAN3 也不能很好地?cái)U(kuò)展到 ImageNet 上，如圖 1 所示。特別是在高分辨率時(shí)，訓(xùn)練會(huì)變得不穩(wěn)定。因此，研究者的第一個(gè)目標(biāo)是在 ImageNet 上成功地訓(xùn)練一個(gè) StyleGAN3 生成器。成功的定義取決于主要通過初始評(píng)分 （IS）［Salimans et al. 2016］ 衡量的樣本質(zhì)量和 Fréchet 初始距離 （FID）［Heusel et al. 2017］ 衡量的多樣性。

在論文中，研究者也介紹了 StyleGAN3 baseline 進(jìn)行的改動(dòng)，所帶來的提升如下表 1 所示：

研究者首先修改了生成器及其正則化損失，調(diào)整了潛在空間以適應(yīng) Projected GAN （Config-B） 和類條件設(shè)置 （Config-C）；然后重新討論了漸進(jìn)式增長，以提高訓(xùn)練速度和性能 （Config-D）；接下來研究了用于 Projected GAN 訓(xùn)練的特征網(wǎng)絡(luò)，以找到一個(gè)非常適合的配置 （Config-E）；最后，研究者提出了分類器引導(dǎo)，以便 GAN 通過一個(gè)預(yù)訓(xùn)練的分類器 （Config-F） 提供類信息。

這樣一來，就能夠訓(xùn)練一個(gè)比以前大得多的模型，同時(shí)需要比現(xiàn)有技術(shù)更少的計(jì)算量。StyleGAN-XL 在深度和參數(shù)計(jì)數(shù)方面比標(biāo)準(zhǔn)的 StyleGAN3 大三倍。然而，為了在 512^2 像素的分辨率下匹配 ADM ［Dhariwal and Nichol 2021］ 先進(jìn)的性能，在一臺(tái) NVIDIA Tesla V100 上訓(xùn)練模型需要 400 天，而以前需要 1914 天。（圖 2）。

實(shí)驗(yàn)結(jié)果

在實(shí)驗(yàn)中，研究者首先將 StyleGAN-XL 與 ImageNet 上的 SOTA 圖像合成方法進(jìn)行比較。然后對(duì) StyleGAN-XL 的反演和編輯性能進(jìn)行了評(píng)價(jià)。研究者將模型擴(kuò)展到了 1024^2 像素的分辨率，這是之前在 ImageNet 上沒有嘗試過的。在 ImageNet 中，大多數(shù)圖像的分辨率較低，因此研究者用超分辨率網(wǎng)絡(luò) ［Liang et al. 2021］ 對(duì)數(shù)據(jù)進(jìn)行了預(yù)處理。

圖像合成

如表 2 所示，研究者在 ImageNet 上對(duì)比了 StyleGAN-XL 和現(xiàn)有最強(qiáng)大的 GAN 模型及擴(kuò)散模型的圖像合成性能。

有趣的是，StyleGAN-XL 在所有分辨率下都實(shí)現(xiàn)了高度的多樣性，這可以歸功于漸進(jìn)式生長策略。此外，這種策略使擴(kuò)大到百萬像素分辨率的合成變成可能。

在 1024^2 這一分辨率下，StyleGAN-XL 沒有與 baseline 進(jìn)行比較，因?yàn)槭艿劫Y源限制，且它們的訓(xùn)練成本高得令人望而卻步。

圖 3 展示了分辨率提高后的生成樣本可視化結(jié)果。

反演和操縱

同時(shí)，還可以進(jìn)一步細(xì)化所得到的重構(gòu)結(jié)果。將 PTI ［Roich et al. 2021］ 和 StyleGAN-XL 相結(jié)合，幾乎可以精確地反演域內(nèi) （ImageNet 驗(yàn)證集） 和域外圖像。同時(shí)生成器的輸出保持平滑，如下圖 4 所示。

圖 5、圖 6 展示了 StyleGAN-XL 在圖像操縱方面的性能：

審核編輯 ：李倩