習慣了Stable Diffusion，如今終於又迎來一個俄羅斯娃娃式（Matryoshka）Diffusion 模型，還是蘋果做的。
原文來源：機器之心
圖片來源：由無界AI生成
在生成式AI 時代，擴散模型已成為影像、視訊、3D、音訊和文字生成等生成式AI 應用的熱門工具。然而將擴散模型拓展到高解析度領域仍然面臨巨大挑戰，這是因為模型必須在每個步驟重新編碼所有的高解析度輸入。解決這些挑戰需要使用具有註意力塊的深層架構，這使得優化更困難，消耗的算力和記憶體也更多。
怎麼辦呢？最近的一些工作專注於研究用於高解析度影像的高效網路架構。但現有方法都沒有顯示出超過512×512 解析度的效果，且產生品質落後於主流的級聯或latent 方法。
我們以OpenAI DALL-E 2、GoogleIMAGEN 和英偉達eDiffI 為例，它們透過學習一個低解析度模型和多個超解析度擴散模型來節省算力，其中每個元件都單獨訓練。另一方面，latent 擴散模型（LDM）僅學習低解析度擴散模型，並依賴單獨訓練的高解析度自編碼器。對於這兩種方案，多階段式pipeline 使訓練與推理複雜化，從而往往需要精心調整或進行超參。
本文中，研究者提出了俄羅斯娃娃式擴散模型（Matryoshka Diffusion Models，MDM）它是用於端到端高解析度影像生成的全新擴散模型。代碼很快將釋出。
論文網址：https://arxiv.org/pdf/2310.15111.pdf
研究提出的主要觀點是將低解析度擴散過程作為高解析度產生的一部分，透過使用嵌套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）能力。
資訊來源：由0x資訊編譯自8BTC。版權所有，未經許可，不得轉載