文本到圖像生成是 2022 年最火的 AIGC 方向之一，被《science》評(píng)選為 2022 年度十大科學(xué)突破。最近，谷歌的一篇文本到圖像生成新論文《Muse: Text-To-Image Generation via Masked Generative Transformers》又引起高度關(guān)注。

• 論文地址：https://arxiv.org/pdf/2301.00704v1.pdf
• 項(xiàng)目地址：https://muse-model.github.io/

該研究提出了一種使用掩碼圖像建模方法進(jìn)行文本到圖像合成的新模型，其中的圖像解碼器架構(gòu)以來自預(yù)訓(xùn)練和 frozen T5-XXL 大型語言模型 (LLM) 編碼器的嵌入為條件。

與谷歌先前的 Imagen 模型類似，該研究發(fā)現(xiàn)基于預(yù)訓(xùn)練 LLM 進(jìn)行調(diào)整對(duì)于逼真、高質(zhì)量的圖像生成至關(guān)重要。Muse 模型是建立在 Transformer (Vaswani et al., 2017) 架構(gòu)之上。

與建立在級(jí)聯(lián)像素空間（pixel-space）擴(kuò)散模型上的 Imagen (Saharia et al., 2022) 或 Dall-E2 (Ramesh et al., 2022) 相比，Muse 由于使用了離散 token，效率顯著提升。與 SOTA 自回歸模型 Parti (Yu et al., 2022) 相比，Muse 因使用并行解碼而效率更高。

基于在 TPU-v4 上的實(shí)驗(yàn)結(jié)果，研究者估計(jì) Muse 在推理速度上比 Imagen-3B 或 Parti-3B 模型快 10 倍以上，比 Stable Diffusion v1.4 (Rombach et al., 2022) 快 2 倍。研究者認(rèn)為：Muse 比 Stable Diffusion 推理速度更快是因?yàn)?Stable Diffusion v1.4 中使用了擴(kuò)散模型，在推理時(shí)明顯需要更多次迭代。

另一方面，Muse 效率的提升沒有造成生成圖像質(zhì)量下降、模型對(duì)輸入文本 prompt 的語義理解能力降低的問題。該研究根據(jù)多個(gè)標(biāo)準(zhǔn)評(píng)估了 Muse 的生成結(jié)果，包括 CLIP 評(píng)分 (Radford et al., 2021) 和 FID (Heusel et al., 2017)。Muse-3B 模型在 COCO (Lin et al., 2014) 零樣本驗(yàn)證基準(zhǔn)上取得了 0.32 的 CLIP 分?jǐn)?shù)和 7.88 的 FID 分?jǐn)?shù)。

下面我們看看 Muse 生成效果：

文本 - 圖像生成：Muse 模型從文本提示快速生成高質(zhì)量的圖像（在 TPUv4 上，對(duì)于 512x512 分辨率的圖像需要時(shí)間為 1.3 秒，生成 256x256 分辨率的圖像需要時(shí)間為 0.5 秒）。例如生成「一只熊騎著自行車，一只鳥棲息在車把上」：

Muse 模型通過對(duì)文本提示條件下的圖像 token 進(jìn)行迭代重新采樣，為用戶提供了零樣本、無掩碼編輯（mask-free editing）。

Muse 還提供了基于掩碼的編輯，例如「在美麗的秋葉映照下，有一座涼亭在湖上」。

模型簡(jiǎn)介

Muse 建立在許多組件之上，圖 3 提供了模型體系架構(gòu)概述。

具體而言所包含的組件有：

預(yù)訓(xùn)練文本編碼器：該研究發(fā)現(xiàn)利用預(yù)訓(xùn)練大型語言模型（LLM）可以提高圖像生成質(zhì)量。他們假設(shè)，Muse 模型學(xué)會(huì)了將 LLM 嵌入中的豐富視覺和語義概念映射到生成的圖像。給定一個(gè)輸入文本字幕，該研究將其通過凍結(jié)的 T5-XXL 編碼器，得到一個(gè) 4096 維語言嵌入向量序列。這些嵌入向量線性投影到 Transformer 模型。

使用 VQGAN 進(jìn)行語義 Tokenization：該模型的核心組件是使用從 VQGAN 模型獲得的語義 token。其中，VQGAN 由一個(gè)編碼器和一個(gè)解碼器組成，一個(gè)量化層將輸入圖像映射到一個(gè)學(xué)習(xí)碼本中的 token 序列。該研究全部使用卷積層構(gòu)建編碼器和解碼器，以支持對(duì)不同分辨率圖像進(jìn)行編碼。

基礎(chǔ)模型：基礎(chǔ)模型是一個(gè)掩碼 transformer，其中輸入是投影到 T5 的嵌入和圖像 token。該研究保留所有的文本嵌入（unmasked），隨機(jī)掩碼不同比例的圖像 token，并用一個(gè)特殊的 [mask] token 替換它們。

超分辨率模型：該研究發(fā)現(xiàn)使用級(jí)聯(lián)模型是有益的：首先是生成 16 × 16 潛在映射（對(duì)應(yīng)于 256 × 256 圖像）的基礎(chǔ)模型，然后是將基礎(chǔ)的潛在映射上采樣到的超分辨率模型，也就是 64 × 64 的潛在映射（對(duì)應(yīng)于一個(gè) 512 × 512 的圖像）。

解碼器微調(diào)：為了進(jìn)一步提高模型生成精細(xì)細(xì)節(jié)的能力，該研究通過添加更多的殘差層和通道來增加 VQGAN 解碼器的容量，同時(shí)保持編碼器容量不變。然后微調(diào)新的解碼器層，同時(shí)凍結(jié) VQGAN 編碼器權(quán)重、碼本和 transformer（即基礎(chǔ)模型和超分辨率模型）。

除了以上組件外，Muse 還包含可變掩碼比率組件、在推理時(shí)迭代并行解碼組件等。

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

如下表所示，與其他模型相比，Muse 縮短了推理時(shí)間。

下表為不同模型在 zero-shot COCO 上測(cè)量的 FID 和 CLIP 得分：

如下表所示，Muse（632M (base)+268M (super-res) 參數(shù)模型）在 CC3M 數(shù)據(jù)集上訓(xùn)練和評(píng)估時(shí)得到了 6.06 的 SOTA FID 分?jǐn)?shù)。

下圖是 Muse 與 Imagen、DALL-E 2 在相同 prompt 下生成結(jié)果的例子。

感興趣的讀者可以閱讀論文原文，了解更多研究細(xì)節(jié)。