Imagen的工作原理解讀

本文詳細(xì)解讀了 Imagen 的工作原理，分析并理解其高級(jí)組件以及它們之間的關(guān)聯(lián)。

近年來(lái)，多模態(tài)學(xué)習(xí)受到重視，特別是文本 - 圖像合成和圖像 - 文本對(duì)比學(xué)習(xí)兩個(gè)方向。一些 AI 模型因在創(chuàng)意圖像生成、編輯方面的應(yīng)用引起了公眾的廣泛關(guān)注，例如 OpenAI 先后推出的文本圖像模型 DALL?E 和 DALL-E 2，以及英偉達(dá)的 GauGAN 和 GauGAN2。 谷歌也不甘落后，在 5 月底發(fā)布了自己的文本到圖像模型 Imagen，看起來(lái)進(jìn)一步拓展了字幕條件（caption-conditional）圖像生成的邊界。

僅僅給出一個(gè)場(chǎng)景的描述，Imagen 就能生成高質(zhì)量、高分辨率的圖像，無(wú)論這種場(chǎng)景在現(xiàn)實(shí)世界中是否合乎邏輯。下圖為 Imagen 文本生成圖像的幾個(gè)示例，在圖像下方顯示出了相應(yīng)的字幕。

這些令人印象深刻的生成圖像不禁讓人想了解：Imagen 到底是如何工作的呢？ 近期，開(kāi)發(fā)者講師 Ryan O'Connor 在 AssemblyAI 博客撰寫(xiě)了一篇長(zhǎng)文《How Imagen Actually Works》，詳細(xì)解讀了 Imagen 的工作原理，對(duì) Imagen 進(jìn)行了概覽介紹，分析并理解其高級(jí)組件以及它們之間的關(guān)聯(lián)。 Imagen 工作原理概覽 在這部分，作者展示了 Imagen 的整體架構(gòu)，并對(duì)其它的工作原理做了高級(jí)解讀；然后依次更透徹地剖析了 Imagen 的每個(gè)組件。如下動(dòng)圖為 Imagen 的工作流程。

首先，將字幕輸入到文本編碼器。該編碼器將文本字幕轉(zhuǎn)換成數(shù)值表示，后者將語(yǔ)義信息封裝在文本中。Imagen 中的文本編碼器是一個(gè) Transformer 編碼器，其確保文本編碼能夠理解字幕中的單詞如何彼此關(guān)聯(lián)，這里使用自注意力方法。 如果 Imagen 只關(guān)注單個(gè)單詞而不是它們之間的關(guān)聯(lián)，雖然可以獲得能夠捕獲字幕各個(gè)元素的高質(zhì)量圖像，但描述這些圖像時(shí)無(wú)法以恰當(dāng)?shù)姆绞椒从匙帜徽Z(yǔ)義。如下圖示例所示，如果不考慮單詞之間的關(guān)聯(lián)，就會(huì)產(chǎn)生截然不同的生成效果。

雖然文本編碼器為 Imagen 的字幕輸入生成了有用的表示，但仍需要設(shè)計(jì)一種方法生成使用這一表示的圖像，也即圖像生成器。為此，Imagen 使用了擴(kuò)散模型，它是一種生成模型，近年來(lái)得益于其在多項(xiàng)任務(wù)上的 SOTA 性能而廣受歡迎。 擴(kuò)散模型通過(guò)添加噪聲來(lái)破壞訓(xùn)練數(shù)據(jù)以實(shí)現(xiàn)訓(xùn)練，然后通過(guò)反轉(zhuǎn)這個(gè)噪聲過(guò)程來(lái)學(xué)習(xí)恢復(fù)數(shù)據(jù)。給定輸入圖像，擴(kuò)散模型將在一系列時(shí)間步中迭代地利用高斯噪聲破壞圖像，最終留下高斯噪聲或電視噪音靜態(tài)（TV static）。下圖為擴(kuò)散模型的迭代噪聲過(guò)程：

然后，擴(kuò)散模型將向后 work，學(xué)習(xí)如何在每個(gè)時(shí)間步上隔離和消除噪聲，抵消剛剛發(fā)生的破壞過(guò)程。訓(xùn)練完成后，模型可以一分為二。這樣可以從隨機(jī)采樣高斯噪聲開(kāi)始，使用擴(kuò)散模型逐漸去噪以生成圖像，具體如下圖所示：

總之，經(jīng)過(guò)訓(xùn)練的擴(kuò)散模型從高斯噪聲開(kāi)始，然后迭代地生成與訓(xùn)練圖像類(lèi)似的圖像。很明顯的是，無(wú)法控制圖像的實(shí)際輸出，僅僅是將高斯噪聲輸入到模型中，并且它會(huì)輸出一張看起來(lái)屬于訓(xùn)練數(shù)據(jù)集的隨機(jī)圖像。 但是，目標(biāo)是創(chuàng)建能夠?qū)⑤斎氲?Imagen 的字幕的語(yǔ)義信息封裝起來(lái)的圖像，因此需要一種將字幕合并到擴(kuò)散過(guò)程中的方法。如何做到這一點(diǎn)呢？ 上文提到文本編碼器產(chǎn)生了有代表性的字幕編碼，這種編碼實(shí)際上是向量序列。為了將這一編碼信息注入到擴(kuò)散模型中，這些向量被聚合在一起，并在它們的基礎(chǔ)上調(diào)整擴(kuò)散模型。通過(guò)調(diào)整這一向量，擴(kuò)散模型學(xué)習(xí)如何調(diào)整其去噪過(guò)程以生成與字幕匹配良好的圖像。過(guò)程可視化圖如下所示：

由于圖像生成器或基礎(chǔ)模型輸出一個(gè)小的 64x64 圖像，為了將這一模型上采樣到最終的 1024x1024 版本，使用超分辨率模型智能地對(duì)圖像進(jìn)行上采樣。 對(duì)于超分辨率模型，Imagen 再次使用了擴(kuò)散模型。整體流程與基礎(chǔ)模型基本相同，除了僅僅基于字幕編碼調(diào)整外，還以正在上采樣的更小圖像來(lái)調(diào)整。整個(gè)過(guò)程的可視化圖如下所示：

這個(gè)超分辨率模型的輸出實(shí)際上并不是最終輸出，而是一個(gè)中等大小的圖像。為了將該圖像放大到最終的 1024x1024 分辨率，又使用了另一個(gè)超分辨率模型。兩個(gè)超分辨率架構(gòu)大致相同，因此不再贅述。而第二個(gè)超分辨率模型的輸出才是 Imagen 的最終輸出。 為什么 Imagen 比 DALL-E 2 更好？ 確切地回答為什么 Imagen 比 DALL-E 2 更好是困難的。然而，性能差距中不可忽視的一部分源于字幕以及提示差異。DALL-E 2 使用對(duì)比目標(biāo)來(lái)確定文本編碼與圖像（本質(zhì)上是 CLIP）的相關(guān)程度。文本和圖像編碼器調(diào)整它們的參數(shù)，使得相似的字幕 - 圖像對(duì)的余弦相似度最大化，而不同的字幕 - 圖像對(duì)的余弦相似度最小化。 性能差距的一個(gè)顯著部分源于 Imagen 的文本編碼器比 DALL-E 2 的文本編碼器大得多，并且接受了更多數(shù)據(jù)的訓(xùn)練。作為這一假設(shè)的證據(jù)，我們可以在文本編碼器擴(kuò)展時(shí)檢查 Imagen 的性能。下面為 Imagen 性能的帕累托曲線(xiàn)：

放大文本編碼器的效果高得驚人，而放大 U-Net 的效果卻低得驚人。這一結(jié)果表明，相對(duì)簡(jiǎn)單的擴(kuò)散模型只要以強(qiáng)大的編碼為條件，就可以產(chǎn)生高質(zhì)量的結(jié)果。 鑒于 T5 文本編碼器比 CLIP 文本編碼器大得多，再加上自然語(yǔ)言訓(xùn)練數(shù)據(jù)必然比圖像 - 字幕對(duì)更豐富這一事實(shí)，大部分性能差距可能歸因于這種差異。 除此以外，作者還列出了 Imagen 的幾個(gè)關(guān)鍵要點(diǎn)，包括以下內(nèi)容：

擴(kuò)展文本編碼器是非常有效的；

擴(kuò)展文本編碼器比擴(kuò)展 U-Net 大小更重要；

動(dòng)態(tài)閾值至關(guān)重要；

噪聲條件增強(qiáng)在超分辨率模型中至關(guān)重要；

將交叉注意用于文本條件反射至關(guān)重要；

高效的 U-Net 至關(guān)重要。

這些見(jiàn)解為正在研究擴(kuò)散模型的研究人員提供了有價(jià)值的方向，而不是只在文本到圖像的子領(lǐng)域有用。

審核編輯 ：李倩