蘋果文生圖大模型亮相：俄羅斯套娃式擴散，支援1024x1024解析度

習慣了 Stable Diffusion，如今終於又迎來一個俄羅斯套娃式（Matryoshka）Diffusion 模型，還是蘋果做的。

原文來源：機器之心

圖片來源：由無界 AI生成

在生成式 AI 時代，擴散模型已經成為圖像、視頻、3D、音訊和文本生成等生成式 AI 應用的流行工具。 然而將擴散模型拓展到高解析度領域仍然面臨巨大挑戰，這是因為模型必須在每個步驟重新編碼所有的高解析度輸入。 解決這些挑戰需要使用帶有注意力塊的深層架構，這使得優化更困難，消耗的算力和記憶體也更多。

怎麼辦呢？ 最近的一些工作專注於研究用於高解析度圖像的高效網路架構。 但是現有方法都沒有展示出超過 512×512 解析度的效果，並且生成質量落後於主流的級聯或 latent 方法。

我們以 OpenAI DALL-E 2、谷歌 IMAGEN 和英偉達 eDiffI 為例，它們通過學習一個低解析度模型和多個超解析度擴散模型來節省算力，其中每個元件都單獨訓練。 另一方面，latent擴散模型（LDM）僅學習低解析度擴散模型，並依賴單獨訓練的高解析度自編碼器。 對於這兩種方案，多階段式 pipeline 使訓練與推理複雜化，從而往往需要精心調整或進行超參。

本文中，研究者提出了俄羅斯套娃式擴散模型（Matryoshka Diffusion Models，MDM）它是用於端到端高解析度圖像生成的全新擴散模型。 代碼很快將釋出。

論文位址：

該研究提出的主要觀點是將低解析度擴散過程作為高解析度生成的一部分，通過使用嵌套 UNet 架構在多個解析度上執行聯合擴散過程。

該研究發現：MDM 與嵌套 UNet 架構一起實現了 1）多解析度損失：大大提高了高解析度輸入去噪的收斂速度; 2）高效的漸進式訓練計劃，從訓練低解析度擴散模型開始，按照計劃逐步添加高解析度輸入和輸出。 實驗結果表明，多解析度損失與漸進式訓練相結合可以讓訓練成本和模型質量獲得更好的平衡。

該研究在類條件圖像生成以及文本條件圖像和視頻生成方面評估了 MDM。 MDM 讓訓練高解析度模型無需使用級聯或潛在擴散（latent diffusion）。 消融研究表明，多解析度損失和漸進訓練都極大地提高了訓練效率和品質。

我們來欣賞以下 MDM 生成的圖片和視頻。

方法概覽

研究者介紹稱，MDM 擴散模型在高解析度中進行端到端訓練，同時利用層級結構的數據形成。 MDM 首先在擴散空間中泛化了標準擴散模型，然後提出了專用的嵌套架構和訓練流程。

首先來看如何在擴展空間對標準擴散模型進行泛化。

與級聯或 latent 方法的不同之處在於，MDM 通過在一個擴展空間中引入多解析度擴散過程，學得了具有層級結構的單個擴散過程。 具體如下圖 2 所示。

具體來講，給定一個數據點 x ∈ R^N，研究者定義了與時間相關的隱變數 z_t = z_t^1 ， . . . ， z_t^R ∈ R^N_1+... NR。

研究者表示，在擴展空間中進行擴散建模有以下兩點優點。 其一，我們在推理期間通常關心全解析度輸出 z_t^R，那麼所有其他中等解析度被看作是額外的隱變數 z_t^r，增加了建模分佈的複雜度。 其二，多解析度依賴性為跨 z_t^r 共用權重和計算提供了機會，從而以更高效的方式重新分配計算，並實現高效訓練和推理。

接下來看嵌套架構（NestedUNet）如何工作。

與典型的擴散模型類似，研究者使用UNet網路結構來實現MDM，其中並行使用殘差連接和計算塊以保留細粒度的輸入資訊。 這裏的計算塊包含多層捲積和自注意力層。 NestedUNet 與標準 UNet 的代碼分別如下。

除了相較於其他層級方法的簡單性，NestedUNet 允許以最高效的方式對計算進行分配。 如下圖 3 所示，研究者早期探索發現，當以最低解析度分配大部分參數和計算時，MDM 實現了明顯更好的擴展性。

最後是學習。

研究者使用常規去噪目標在多個解析度下訓練 MDM，如下公式 （3） 所示。

這裡用到了漸進式訓練。 研究者按照上述公式 （3） 直接對 MDM 進行端到端訓練，並展示出了比原始基線方法更好的收斂性。 他們發現，使用類似於 GAN 論文中提出的簡單漸進式訓練方法，極大地加速了高解析度模型的訓練。

這一訓練方法從一開始就避免了高成本的高解析度訓練，加速了整體收斂。 不僅如此，他們還合併了混合解析度訓練，該訓練方法在單個 batch 中同時訓練具有不同最終解析度的樣本。

實驗及結果

MDM 是一種通用技術，適用於可以逐步壓縮輸入維度的任何問題。 MDM 與基線方法的比較如下圖 4 所示。

表 1 給出了在 ImageNet（FID-50K）和 COCO（FID-30K）上的比較結果。

下圖 5、6、7 展示了 MDM 在圖像生成（圖 5）、文字到圖像（圖 6）和文本到視頻（圖 7）方面的結果。 儘管是在相對較小的數據集上進行訓練的，但 MDM 仍顯示出生成高解析度圖像和視頻的強大零樣本（zero-shot）能力。