谷歌開始卷自己,AI架構Pathways加持,推出200億生成模型

語言: CN / TW / HK

機器之心報道

機器之心編輯部

繼 Imagen 後,谷歌又推出了文字 - 影象生成模型 Parti。

你見過一隻小狗破殼而出嗎?或者用飛艇俯瞰蒸汽朋克中的城市?又或者兩個機器人在電影院像人類一樣看電影…… 這些聽起來可能有些天馬行空,但一種名為「文字到影象生成」的新型機器學習技術使這些成為可能。

谷歌研究院的科學家和工程師一直致力於探索使用各種 AI 技術生成文字到影象的方法。

今年 5 月底,谷歌推出 AI 創作神器Imagen,它結合了 Transformer 語言模型和高保真擴散模型的強大功能,在文字到影象的合成中提供前所未有的逼真度和語言理解能力。與僅使用影象 - 文字資料進行模型訓練的先前工作相比,Imagen 的關鍵突破在於:谷歌的研究者發現在純文字語料庫上預訓練的大型 LM 的文字嵌入對文字到影象的合成顯著有效。Imagen 的文字到影象生成可謂天馬行空,能生成多種奇幻卻逼真的有趣影象。

Imagen 生成效果是這樣的,比如正在戶外享受騎行的柴犬(下圖左)以及狗狗照鏡子發現自己是隻貓(下圖右):

時隔沒多久,谷歌又推出了 Parti(Pathways Autoregressive Text-to-Image),該模型最高可擴充套件至 200 億引數,並且隨著可使用引數數量的增長,其輸出的影象也能夠更加逼真。

值得一提的是,這是谷歌大牛 Jeff Dean 提出的多工 AI 大模型藍圖 Pathways 的一部分。

我們先來看下 Parti 效果,袋熊在瀑布旁,揹著書包,拄著柺杖眺望著遠方:

埃及阿努比斯肖像,在洛杉磯背景下,戴著飛行員護目鏡,穿著白色 t 恤和黑色皮夾克:

一隻熊貓戴著一頂巫師帽騎在馬上:

下面我們介紹一下 Parti 的實現原理。

Parti 模型

與 DALL-E、CogView 和 Make-A-Scene 類似,Parti 是一個兩階段模型,由影象 tokenizer 和自迴歸模型組成,如下圖 3 所示。第一階段訓練一個 tokenizer,該 tokenizer 可以將影象轉換為一系列離散的視覺 token,用於訓練並在推理時重建影象。第二階段訓練從文字 token 生成影象 token 的自迴歸序列到序列模型。

影象 Tokenizer

首先,該研究訓練了一個 ViT-VQGAN-Small 模型(8 個塊,8 個頭,模型維度 512,隱藏維度 2048,總引數約為 30M),並且學習了 8192 張影象 token 類別用於程式碼本。

為了進一步提高第二階段編碼器 - 解碼器訓練後重建影象的視覺靈敏度,該研究凍結了 tokenizer 的編碼器和程式碼本,並微調更大尺寸的 tokenizer 解碼器(32 個塊,16 個頭,模型維度 1280,隱藏維度 5120, 總引數約 600M)。影象 tokenizer 的輸入和輸出使用 256×256 解析度。

最後,雖然解析度為 256×256 的影象捕獲了大部分內容、結構和紋理,但更高解析度的影象具有更大的視覺衝擊力。為此,該研究在影象 tokenizer 上採用了一個簡單的超解析度模組,如下圖 4 所示。

文字到影象生成的編碼器 - 解碼器架構

如上圖 3 所示,該研究第二階段訓練了一種標準的編碼器 - 解碼器 Transformer 模型,將文字到影象視為序列到序列建模問題。該模型將文字作為輸入,並使用從第一階段影象 tokenizer 生成的光柵化影象潛在程式碼的下一個 token 預測進行訓練。對於文字編碼,該研究構建了一個 sentence-piece 模型,詞彙量為 16000。在推理時,模型對影象 token 進行自迴歸取樣,隨後使用 ViT-VQGAN 解碼器將其解碼為畫素。

該研究使用的文字 token 最大為 128,影象 token 的長度固定為 1024。所有模型都使用 conv-shaped 掩碼稀疏注意力。該研究訓練了四種變體,引數量從 3.5 億到 200 億不等,如下表 1 所示。

以下為對 Parti 模型四種大小比較結果,可以觀察到:模型效能和輸出影象質量在持續地提高;20B 模型尤其擅長於那些抽象的、需要世界知識的、特定視角的、或符號渲染的 prompt。

在悉尼歌劇院前的草地上,一隻袋鼠穿著橙色衛衣,戴著藍色墨鏡,胸前掛著「歡迎朋友」的牌子。

松鼠把蘋果送給了小鳥。

文字編碼器預訓練

該研究在兩個資料集上預訓練文字編碼器:具有 BERT [36] 預訓練目標的 Colossal Clean Crawled Corpus (C4) [35],以及具有對比學習目標的影象文字資料。預訓練後,該研究繼續訓練編碼器和解碼器,在 8192 個離散影象 token 的詞彙表上使用 softmax 交叉熵損失生成文字到影象。

預訓練後的文字編碼器在 GLUE 上的效能與 BERT 相當;然而,在文字到影象生成的完整編碼器 - 解碼器訓練過程之後,文字編碼器會降級。

擴充套件

該研究在 Lingvo 上來實現模型,並在 CloudTPUv4 硬體上使用 GSPMD 進行擴充套件,以用於訓練和推理。GSPMD 是一個基於 XLA 編譯器的模型分佈系統,它允許將 TPU 叢集視為單個虛擬裝置,並在幾個張量上使用 sharding annotations 來指示編譯器自動分發資料並在數千個裝置上進行計算。

該研究用資料並行性訓練 350M 和 750M 模型。對於 3B 模型,該研究使用 4 路內層模型並行(參見下圖 5)和 128 路資料並行。

下圖 6 為分散式訓練策略整體架構圖:

實驗

下表 5 給出了自動影象質量評估的主要結果。與基於擴散的 Imagen 模型相比,Parti 獲得了相媲美的零樣本 FID 分數。

下表 6 為 Parti 字幕評估結果(captioner evaluation [55]),Parti 優於其他模型:

下圖 8 顯示,儘管 Parti 沒有接受過 MS-COCO 字幕或影象方面的訓練,但表現更好。

下圖 9 總結了 MS-COCO 零樣本 FID 分數:

更多內容,請參考原論文。

參考連結:https://blog.google/technology/research/how-ai-creates-photorealistic-images-from-text/

THE END

轉載請聯絡本公眾號獲得授權

投稿或尋求報道:[email protected]