通過顯式建模數據流形幾何的擴散圖像生成
研究人員提出了MIND（數據流形感知圖像擴散模型），通過將離散補丁標記化集成到連續擴散模型的得分函數中，顯式建模流形幾何。該方法在ImageNet 256×256上實現了FID 2.06（MIND-B，1.3億參數），優於LlamaGen-3B（31億參數）的FID。
圖像生成模型的核心目標是能夠從底層數據流形中採樣高質量的數據點。然而，這一任務要求模型能夠學習並解碼一個密集、低維且緊湊的參數化空間，傳統方法往往難以兼顧結構表達與生成效率。近日，來自研究團隊的一篇新論文提出了數據流形感知圖像擴散模型（MIND），為這一問題提供了全新的解決思路。
MIND是一個新穎的框架，其核心創新在於將離散補丁標記化技術無縫集成到連續擴散模型的得分函數中。通過這種方式，MIND能夠顯式地對數據流形的幾何結構進行建模，充分利用離散標記在結構量化方面的優勢，同時保持連續擴散模型並行生成的靈活性。這種結合使得模型在生成過程中能夠更準確地捕捉數據的底層結構。
為了實現端到端的可微訓練，研究團隊引入了一種新穎的軟top-k聚合機制。該機制可以高效地處理離散標記的選擇問題，同時保證整個訓練過程的梯度可以順利回傳。此外，為了應對Transformer主幹網絡在低維輸入上表現出的頻譜偏差，團隊設計了雙分支高頻特徵嵌入層。這一設計通過增強模型對高頻信號的感知能力，有效緩解了頻譜偏差帶來的性能損失。
在推理階段，MIND採用了一種多階段過渡採樣方案。該方案能夠根據當前的時間步動態調整採樣策略，從而在生成質量和效率之間取得更好的平衡。
實驗部分，研究團隊在ImageNet 256×256數據集上進行了廣泛的評估。經過80個週期的訓練，基礎模型在無引導條件下實現了22.73的FID，這一成績幾乎是原始DiT-B/2基線（43.47 FID）的一半。與DiT和SiT基線相比，MIND在平均FID上分別降低了15.95和9.06，提升幅度顯著。在有引導的圖像生成任務中，MIND展現出了更強的競爭力。僅含1.3億參數的MIND-B模型實現了2.06的FID，這一結果甚至超越了擁有31億參數的LlamaGen-3B模型。進一步地，擁有7.15億參數的MIND-XL模型將FID降低至1.95，達到了新的高度。
MIND為基於擴散的圖像生成領域帶來了全新的視角，其明確的幾何建模思路有望推動該領域的進一步發展。研究團隊表示，相關代碼將公開提供，以促進學術界的交流與創新。