今年 2 月初，Sora 的發(fā)布讓 AI 社區(qū)更加看到了基礎(chǔ)擴(kuò)散模型的潛力。連同以往出現(xiàn)的 Stable Diffusion、PixArt-α 和 PixArt-Σ，這些模型在生成真實(shí)圖像和視頻方面取得了顯著的成功。這意味著開始了從經(jīng)典 U-Net 架構(gòu)到基于 Transformer 的擴(kuò)散主干架構(gòu)的范式轉(zhuǎn)變。

值得注意的是，通過(guò)這種改進(jìn)的架構(gòu)，Sora 和 Stable Diffusion 3 可以生成任意分辨率的樣本，并表現(xiàn)出對(duì) scaling 定律的嚴(yán)格遵守，即增加參數(shù)大小可以實(shí)現(xiàn)更好的結(jié)果。 

不過(guò)，推出者們只對(duì)自家模型的設(shè)計(jì)選擇提供有限的指導(dǎo)，并且缺乏詳細(xì)的實(shí)現(xiàn)說(shuō)明和公開的預(yù)訓(xùn)練檢查點(diǎn)，限制了它們?cè)谏鐓^(qū)使用和復(fù)刻方面的效用。并且，這些方法是針對(duì)特定任務(wù)（例如圖像或視頻生成任務(wù)）量身定制的，這阻礙了潛在的跨模態(tài)適應(yīng)性。

為了彌補(bǔ)這些差距，上海 AI Lab、港中文和英偉達(dá)的研究者聯(lián)合推出了 Lumina-T2X 系列模型，通過(guò)基于流（Flow-based）的大型擴(kuò)散 Transformers（Flag-DiT）打造，旨在將噪聲轉(zhuǎn)換為圖像、視頻、多視圖 3D 對(duì)象和基于文本描述的音頻。

其中，Lumina-T2X 系列中最大的模型包括具有 70 億參數(shù)的 Flag-DiT 和一個(gè)多模態(tài)大語(yǔ)言模型 SPHINX。SPHINX 是一個(gè)文本編碼器，它具有 130 億參數(shù)，能夠處理 128K tokens。

• 論文地址：https://arxiv.org/pdf/2405.05945
• GitHub 地址：https://github.com/Alpha-VLLM/Lumina-T2X
• 模型下載地址：https://huggingface.co/Alpha-VLLM/Lumina-T2I/tree/main
• 論文標(biāo)題：Lumina-T2X: Transforming Text into Any Modality, Resolution, and Duration via Flow-based Large Diffusion Transformers

具體來(lái)講，基礎(chǔ)的文本到圖像模型 Lumina-T2I 利用流匹配框架，在精心整理的高分辨率真實(shí)圖像文本對(duì)數(shù)據(jù)集上進(jìn)行訓(xùn)練，只需要使用很少的計(jì)算資源就能取得真實(shí)感非常不錯(cuò)的結(jié)果。

如圖 1 所示，Lumina-T2I 可以生成任意分辨率和寬高比的高質(zhì)量圖像，并進(jìn)一步實(shí)現(xiàn)高級(jí)功能，包括分辨率外推、高分辨率編輯、構(gòu)圖生成和風(fēng)格一致生成，所有這些都以免訓(xùn)練的方式無(wú)縫集成到框架中。 

此外，為了增強(qiáng)跨各種模態(tài)的生成能力，Lumina-T2X 從頭開始對(duì)視頻 - 文本、多視圖 - 文本和語(yǔ)音 - 文本對(duì)進(jìn)行獨(dú)立訓(xùn)練，從而可以合成視頻、多視圖 3D 對(duì)象以及文本語(yǔ)音指示。例如，Lumina-T2V 僅用有限的資源和時(shí)間進(jìn)行訓(xùn)練，可以生成任何寬高比和時(shí)長(zhǎng)的 720p 視頻，顯著縮小了開源模型與 Sora 之間的差距。

我們先來(lái)一睹實(shí)現(xiàn)效果如何。比如生成視頻：

生成單張圖像：

3D 生成：

語(yǔ)音生成：

構(gòu)圖生成：

風(fēng)格一致性生成：

更大分辨率外推：

圖像編輯：

可以說(shuō)，Lumina-T2X 系列模型真正實(shí)現(xiàn)了圖像、視頻、3D 和語(yǔ)音的「大一統(tǒng)」。

目前，研究者已經(jīng)推出了分別使用 Flag-DiT 2B 和 Gemma 2B 作為文本編碼器的 Lumina-Next-T2I 模型，可以在 gradio 上試玩。

• 試用地址 1：http://106.14.2.150:10021/
• 試用地址 2：http://106.14.2.150:10022/

方法概覽

Flag-DiT 架構(gòu)

Flag-DiT 是 Lumina-T2X 框架的主干，它具有顯著的穩(wěn)定性、靈活性和可擴(kuò)展性。

首先是穩(wěn)定性。Flag-DiT 建立在 DiT 之上，并結(jié)合 ViT-22B 和 LLaMa 來(lái)修改，以提高訓(xùn)練穩(wěn)定性。具體來(lái)說(shuō)，F(xiàn)lag-DiT 將所有 LayerNorm 替換為 RMSNorm，使得訓(xùn)練穩(wěn)定性增強(qiáng)。

此外，F(xiàn)lag-DiT 在鍵查詢點(diǎn)積注意力計(jì)算之前結(jié)合鍵查詢歸一化（KQ-Norm）。KQ-Norm 的引入旨在通過(guò)消除注意力 logits 中極大值來(lái)防止損失發(fā)散。這種簡(jiǎn)單的修改可以防止混合精度訓(xùn)練下的發(fā)散損失，并有助于以更高的學(xué)習(xí)率進(jìn)行優(yōu)化。Flag-DiT 的詳細(xì)計(jì)算流如圖 2 所示。

其次是靈活性。DiT 僅支持具有簡(jiǎn)單標(biāo)簽條件和固定 DDPM 公式的固定分辨率圖像生成。為了解決這些問(wèn)題，研究者首先探究為什么 DiT 缺乏以任意分辨率和比例生成樣本的靈活性。他們發(fā)現(xiàn)，專為視覺識(shí)別任務(wù)而設(shè)計(jì)的 APE 很難泛化到訓(xùn)練之外未見過(guò)的分辨率和規(guī)模。

因此，受最近展現(xiàn)出強(qiáng)大上下文外推能力的 LLM 的推動(dòng)，他們用 RoPE 替換了 APE。RoPE 按照以下公式 1 和 2 以分層方式注入相對(duì)位置信息。

最后，研究者根據(jù)經(jīng)驗(yàn)，使用更大的參數(shù)和更多的訓(xùn)練樣本擴(kuò)展了 Flag-DiT。具體來(lái)講，他們探索在標(biāo)簽條件 ImageNet 生成基準(zhǔn)上將參數(shù)大小從 600M 擴(kuò)大到 7B。

Lumina-T2X 整體流程

如圖 3 所示，Lumina-T2X 在訓(xùn)練過(guò)程中主要由四個(gè)組件組成，接下來(lái)進(jìn)行一一介紹。

不同模態(tài)的逐幀編碼。在 Lumina-T2X 框架中統(tǒng)一不同模態(tài)的關(guān)鍵是將圖像、視頻、多視圖圖像和語(yǔ)音頻譜圖視為長(zhǎng)度為 T 的幀序列，然后利用特定模態(tài)的編碼器來(lái)將這些輸入轉(zhuǎn)換為形狀為 [H, W, T, C] 的潛在框架。

使用多種文本編碼器進(jìn)行文本編碼。對(duì)于文本條件生成，研究者使用預(yù)先訓(xùn)練的語(yǔ)言模型對(duì)文本提示進(jìn)行編碼。他們結(jié)合了各種大小不一的文本編碼器，其中包括 CLIP、LLaMA、SPHINX 和 Phone 編碼器，針對(duì)各種需求和模態(tài)進(jìn)行量身定制，以優(yōu)化文本調(diào)整。

輸入和目標(biāo)構(gòu)建。Lumina-T2X 在流匹配中采用線性插值方案來(lái)構(gòu)建輸入和目標(biāo)，具體如下公式 4 和 6 所示，簡(jiǎn)單靈活。并且，受到中間時(shí)間步對(duì)于擴(kuò)散模型和流模型都至關(guān)重要的觀察啟發(fā)， 研究者在訓(xùn)練期間采用時(shí)間重采樣策略從對(duì)數(shù)范數(shù)分布中采樣時(shí)間步。

網(wǎng)絡(luò)架構(gòu)和損失。研究者使用 Flag-DiT 作為去噪主干。給定噪聲輸入，F(xiàn)lag-DiT 塊通過(guò)調(diào)控機(jī)制注入添加了全局文本嵌入的擴(kuò)散時(shí)間步，并使用如下公式 9 通過(guò)零初始化注意力來(lái)進(jìn)一步集成文本調(diào)整。

Lumina-T2X 系列

Lumina-T2X 系列模型包括了 Lumina-T2I、Lumina-T2V、LuminaT2MV 和 Lumina-T2Speech。對(duì)于每種模態(tài)，Lumina-T2X 都經(jīng)過(guò)了針對(duì)不同場(chǎng)景優(yōu)化的多配置獨(dú)立訓(xùn)練，例如不同的文本編碼器、VAE 潛在空間和參數(shù)大小。具體如圖 17 所示。

Lumina-T2I 的高級(jí)應(yīng)用 

除了基本的文本生成圖像之外，文本到圖像版本的 Lumina-T2I 還支持更復(fù)雜的視覺創(chuàng)作，并作為基礎(chǔ)模型產(chǎn)生富有創(chuàng)造力的視覺效果。這包括分辨率外推、風(fēng)格一致性生成、高分辨率圖像編輯和構(gòu)圖生成。

與以往使用多種方法解決這些任務(wù)的策略不同，Lumina-T2I 可以通過(guò) token 操作統(tǒng)一解決這些問(wèn)題，如圖 4 所示。

免調(diào)整分辨率外推。RoPE 的平移不變性增強(qiáng)了 Lumina-T2X 的分辨率外推潛力，使其能夠生成域外分辨率的圖像。Lumina-T2X 分辨率最高可以外推到 2K。

風(fēng)格一致性生成?；?Transformer 的擴(kuò)散模型架構(gòu)使得 Lumina-T2I 自然地適合風(fēng)格一致性生成等自注意力操作應(yīng)用。

構(gòu)圖生成。研究者只將此操作應(yīng)用于 10 個(gè)注意力交叉層，以確保文本信息被注入到不同的區(qū)域。同時(shí)保持自注意層不變，以確保最終圖像的連貫、和諧。

高分辨率編輯。除了高分辨率生成之外，Lumina-T2I 還可以執(zhí)行圖像編輯，尤其是對(duì)于高分辨率圖像。

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

在 ImageNet 上驗(yàn)證 Flag-DiT

研究者在有標(biāo)簽條件的 256×256 和 512×512 ImageNet 上進(jìn)行實(shí)驗(yàn)，以驗(yàn)證 Flag-DiT 相對(duì)于 DiT 的優(yōu)勢(shì)。Large-DiT 是 Flag-DiT 的特化版本，采用了 DDPM 算法 ，以便與原始 DiT 進(jìn)行公平比較。研究者完全沿用了 DiT 的設(shè)置，但做了以下修改，包括混合精度訓(xùn)練、大學(xué)習(xí)率和架構(gòu)修改套件（如 QK-Norm、RoPE 和 RMSNorm）。

研究者將其與 SOTA 方法的比較，如表 2 所示，Large-DiT-7B 在不使用無(wú)分類指導(dǎo)（CFG）的情況下，在 FID 和 IS 分?jǐn)?shù)上明顯超過(guò)了所有方法，將 FID 分?jǐn)?shù)從 8.60 降至 6.09。這表明，增加擴(kuò)散模型的參數(shù)可以顯著提高樣本質(zhì)量，而無(wú)需依賴 CFG 等額外技巧。

研究者比較了 Flag-DiT、Large-DiT 和 SiT 在 ImageNet 條件生成上的性能，為了進(jìn)行公平比較，他們將參數(shù)大小固定為 600M。如圖 5 (a) 所示，在 FID 評(píng)估的所有歷時(shí)中，F(xiàn)lag-DiT 的性能始終優(yōu)于 Large-DiT。這表明，與標(biāo)準(zhǔn)擴(kuò)散設(shè)置相比，流匹配公式可以改善圖像生成。此外，與 SiT 相比，F(xiàn)lag-DiT 的 FID 分?jǐn)?shù)較低，這表明元架構(gòu)修改（包括 RMSNorm、RoPE 和 K-Q norm）不僅能穩(wěn)定訓(xùn)練，還能提高性能。

通過(guò)混合精度訓(xùn)練提高訓(xùn)練速度 。Flag-DiT 不僅能提高性能，還能提高訓(xùn)練效率和穩(wěn)定性。如表 4 所示。Flag-DiT 每秒可多處理 40% 的圖像。

ImageNet 初始化的影響 PixArt-α 利用 ImageNet 預(yù)訓(xùn)練的 DiT（學(xué)習(xí)像素依賴性）作為后續(xù) T2I 模型的初始化。為了驗(yàn)證 ImageNet 初始化的影響，研究者比較了使用 ImageNet 初始化和從頭開始訓(xùn)練的 600M 參數(shù)模型的 Lumina-T2I 速度預(yù)測(cè)損失。如圖 5 (d) 所示，從頭開始訓(xùn)練的損失水平更低，收斂速度更快。此外，從零開始可以更靈活地選擇配置和架構(gòu)，而不受預(yù)訓(xùn)練網(wǎng)絡(luò)的限制。表 1 所示的簡(jiǎn)單而快速的訓(xùn)練配方也是根據(jù)這一觀察結(jié)果設(shè)計(jì)的。

Lumina-T2I 的結(jié)果

圖 6 中展示了基本的文本到圖像生成能力。擴(kuò)散主干架構(gòu)和文本編碼器的大容量允許生成逼真的高分辨率圖像，并能準(zhǔn)確理解文本，只需使用 288 個(gè) A100 GPU 天數(shù)。

分辨率外推法不僅能帶來(lái)更大比例的圖像，還能帶來(lái)更高的圖像質(zhì)量和更強(qiáng)的細(xì)節(jié)。如圖 7 所示，當(dāng)分辨率從 1K 外推至 1.5K 時(shí)，我們可以發(fā)現(xiàn)到生成圖像的質(zhì)量和文本到圖像的對(duì)齊情況都得到了顯著提升。此外，Lumina-T2I 還能進(jìn)行外推，生成分辨率更低的圖像，如 512 分辨率，從而提供了更大的靈活性。

如圖 8 所示，通過(guò)利用一個(gè)簡(jiǎn)單的注意力共享操作，我們可以觀察到生成批次中的強(qiáng)一致性。得益于完全注意力模型架構(gòu)，研究者獲得了與參考文獻(xiàn) [58] 中相媲美的結(jié)果，而無(wú)需使用任何技巧，如自適應(yīng)實(shí)例規(guī)范化（AdaIN）。此外，研究者認(rèn)為，正如先前的研究所示，通過(guò)適當(dāng)?shù)姆崔D(zhuǎn)技術(shù)，他們可以實(shí)現(xiàn)零成本的風(fēng)格 / 概念個(gè)性化，這是未來(lái)探索的一個(gè)有前景的方向。

如圖 9 所示，研究者演示了組合生成 。在這其中可以定義任意數(shù)量的 prompt，并為每個(gè) prompt 分配任意區(qū)域。Lumina-T2I 成功生成了各種分辨率的高質(zhì)量圖像，這些圖像與復(fù)雜的輸入 prompt 相一致，同時(shí)又保持了整體的視覺一致性。這表明，Lumina-T2I 的設(shè)計(jì)選擇提供了一種靈活有效的方法，在生成復(fù)雜的高分辨率多概念圖像方面表現(xiàn)出色。

研究者對(duì)高分辨率圖像進(jìn)行風(fēng)格和主題編輯。如圖 10 所示，Lumina-T2I 可以無(wú)縫修改全局樣式或添加主題，而無(wú)需額外的訓(xùn)練。此外，他們還分析了圖像編輯中的啟動(dòng)時(shí)間和潛在特征歸一化等各種因素，如圖 11 所示。

如圖 13 (a) 所示，研究者可視化了各層和各頭部的門控值，發(fā)現(xiàn)大多數(shù)門控值接近零，只有少部分顯示出顯著的重要性。有趣的是，最關(guān)鍵的文本調(diào)節(jié)頭部主要位于中間層，這表明這些層在文本調(diào)節(jié)中起著關(guān)鍵作用。為了鞏固這一觀察結(jié)果，研究者對(duì)低于某一閾值的門控進(jìn)行了截?cái)?，發(fā)現(xiàn) 80% 的門控可以在不影響圖像生成質(zhì)量的情況下被停用，如圖 13 (b) 所示。這一觀察表明，在采樣過(guò)程中截?cái)啻蠖鄶?shù)交叉注意力操作的可能性，這可以大大減少推理時(shí)間。

Lumina-T2V 的結(jié)果

研究者觀察到，使用大批量的 Lumina-T2V 能夠收斂，而小批量則難以收斂。如圖 14 (a) 所示，將批量大小從 32 增加到 1024 會(huì)導(dǎo)致?lián)p失收斂。另一方面，與 ImageNet 實(shí)驗(yàn)中的觀察相似，增加模型參數(shù)會(huì)加速視頻生成的收斂速度。如圖 14 (b) 所示，當(dāng)參數(shù)大小從 600M 增加到 5B 時(shí)，我們能夠在相同的訓(xùn)練迭代次數(shù)下一致觀察到更低的損失。

如圖 15 所示，Lumina-T2V 的第一階段能夠生成具有場(chǎng)景動(dòng)態(tài)變化（如場(chǎng)景轉(zhuǎn)換）的短視頻，盡管生成的視頻在分辨率和持續(xù)時(shí)間上有限，總 token 數(shù)最多為 32K。經(jīng)過(guò)對(duì)更長(zhǎng)持續(xù)時(shí)間和更高分辨率視頻的第二階段訓(xùn)練后，Lumina-T2V 能夠生成 128K token 的各種分辨率和持續(xù)時(shí)間的長(zhǎng)視頻。如圖 16 所示，生成的視頻展示了時(shí)間上的一致性和更豐富的場(chǎng)景動(dòng)態(tài)，表明當(dāng)使用更多的計(jì)算資源和數(shù)據(jù)時(shí)，展現(xiàn)出有希望的擴(kuò)展趨勢(shì)。

更多詳細(xì)內(nèi)容，請(qǐng)閱讀原論文。

本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/NwwbaeRujh-02V6LRs5zMg??